fine-tuning, infering, app codes added

app.py

@@ -1,22 +1,92 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import gradio as gr
 
-app = FastAPI()
+class CustomerSupportBot:
+    def __init__(self, model_path="models/customer_support_gpt"):
+        """
+        Initialize the customer support bot with the fine-tuned model.
+        
+        Args:
+            model_path (str): Path to the saved model and tokenizer
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path)
+        
+        # Move model to GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = self.model.to(self.device)
+        
+    def generate_response(self, instruction, max_length=100, temperature=0.7):
+        """
+        Generate a response for a given customer support instruction/query.
+        
+        Args:
+            instruction (str): Customer's query or instruction
+            max_length (int): Maximum length of the generated response
+            temperature (float): Controls randomness in generation (higher = more random)
+            
+        Returns:
+            str: Generated response
+        """
+        # Format input text the same way as during training
+        input_text = f"Instruction: {instruction}\nResponse:"
+        
+        # Tokenize input
+        inputs = self.tokenizer(input_text, return_tensors="pt")
+        inputs = inputs.to(self.device)
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=50,
+                temperature=temperature,
+                num_return_sequences=1,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                do_sample=True,
+                top_p=0.95,
+                top_k=50
+            )
+        
+        # Decode and format response
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the response part
+        response = response.split("Response:")[-1].strip()
+        
+        return response
 
-model = AutoModelForCausalLM.from_pretrained("models/customer_support_gpt")
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
+# Initialize the chatbot
+bot = CustomerSupportBot()
 
-class Query(BaseModel):
-    query: str
+# Define the Gradio interface function
+def chatbot_response(message, history):
+    """
+    Generate bot response for the Gradio interface.
+    
+    Args:
+        message (str): User's input message
+        history (list): Chat history
+    """
+    bot_response = bot.generate_response(message)
+    history.append((bot_response))
+    return history
 
-@app.post("/predict/")
-def predict(query: Query):
-    inputs = tokenizer(query.query, return_tensors="pt", truncation=True)
-    outputs = model.generate(inputs['input_ids'], max_length=150)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": response}
+# Create the Gradio interface
+iface = gr.ChatInterface(
+    fn=chatbot_response,
+    title="Customer Support Chatbot",
+    description="Ask your questions to the customer support bot!",
+    examples=["How do I reset my password?", 
+             "What are your shipping policies?", 
+             "I want to return a product."],
+    # retry_btn=None,
+    # undo_btn="Remove Last",
+    # clear_btn="Clear",
+)
 
+# Launch the interface
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    iface.launch(share=False)  # Set share=True if you want to create a public link

src/infer.py

@@ -0,0 +1,89 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+class CustomerSupportBot:
+    def __init__(self, model_path="models/customer_support_gpt"):
+        """
+        Initialize the customer support bot with the fine-tuned model.
+        
+        Args:
+            model_path (str): Path to the saved model and tokenizer
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path)
+        
+        # Move model to GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = self.model.to(self.device)
+        
+    def generate_response(self, instruction, max_length=100, temperature=0.7):
+        """
+        Generate a response for a given customer support instruction/query.
+        
+        Args:
+            instruction (str): Customer's query or instruction
+            max_length (int): Maximum length of the generated response
+            temperature (float): Controls randomness in generation (higher = more random)
+            
+        Returns:
+            str: Generated response
+        """
+        # Format input text the same way as during training
+        input_text = f"Instruction: {instruction}\nResponse:"
+        
+        # Tokenize input
+        inputs = self.tokenizer(input_text, return_tensors="pt")
+        inputs = inputs.to(self.device)
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=max_length,
+                temperature=temperature,
+                num_return_sequences=1,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                do_sample=True,
+                top_p=0.95,
+                top_k=50
+            )
+        
+        # Decode and format response
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the response part
+        response = response.split("Response:")[-1].strip()
+        
+        return response
+
+def main():
+    # Initialize the bot
+    bot = CustomerSupportBot()
+    
+    # Example queries
+    example_queries = [
+        "How do I reset my password?",
+        "What are your shipping policies?",
+        "I want to return a product.",
+    ]
+    
+    # Generate and print responses
+    print("Customer Support Bot Demo:\n")
+    for query in example_queries:
+        print(f"Customer: {query}")
+        response = bot.generate_response(query)
+        print(f"Bot: {response}\n")
+
+    # Interactive mode
+    print("Enter your questions (type 'quit' to exit):")
+    while True:
+        query = input("\nYour question: ")
+        if query.lower() == 'quit':
+            break
+            
+        response = bot.generate_response(query)
+        print(f"Bot: {response}")
+
+if __name__ == "__main__":
+    main()

src/train.py

@@ -0,0 +1,107 @@
+import mlflow
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer,
+    Trainer, 
+    TrainingArguments,
+    DataCollatorForLanguageModeling
+)
+from datasets import load_dataset
+
+def prepare_data(tokenizer, dataset):
+    """Tokenize and format the dataset."""
+    def tokenize_function(examples):
+        # Combine instruction and response with a separator
+        text = [f"Instruction: {instr}\nResponse: {resp}" 
+               for instr, resp in zip(examples['instruction'], examples['response'])]
+        
+        return tokenizer(
+            text,
+            truncation=True,
+            max_length=256,
+            padding='max_length'
+        )
+
+    tokenized_datasets = dataset.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=dataset['train'].column_names
+    )
+    
+    return tokenized_datasets
+
+def fine_tune_model():
+    """
+    Fine-tune GPT-Neo on customer support data using instructions and responses.
+    """
+    # Load dataset
+    dataset = load_dataset('csv', data_files='data/raw/customer_support.csv')
+    dataset = dataset['train'].train_test_split(test_size=0.2, seed=42)
+
+    # Load model and tokenizer
+    model_name = "EleutherAI/gpt-neo-125M"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    # Add padding token if it doesn't exist
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        model.config.pad_token_id = model.config.eos_token_id
+
+    # Prepare the dataset
+    tokenized_datasets = prepare_data(tokenizer, dataset)
+
+    # Create data collator
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False  # We're not doing masked language modeling
+    )
+
+    mlflow.start_run()
+
+    # Log hyperparameters
+    mlflow.log_param("model_name", model_name)
+    mlflow.log_param("epochs", 3)
+    mlflow.log_param("batch_size", 4)
+    mlflow.log_param("learning_rate", 2e-5)
+
+    training_args = TrainingArguments(
+        output_dir="models/",
+        evaluation_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        num_train_epochs=3,
+        weight_decay=0.01,
+        save_strategy="epoch",
+        save_total_limit=2,
+        load_best_model_at_end=True,
+        report_to="mlflow"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets['train'],
+        eval_dataset=tokenized_datasets['test'],
+        data_collator=data_collator,
+    )
+
+    trainer.train()
+
+    # Save the model and tokenizer
+    model_path = "models/customer_support_gpt"
+    model.save_pretrained(model_path)
+    tokenizer.save_pretrained(model_path)
+    
+    # Log model artifacts
+    mlflow.log_artifact(model_path)
+
+    # Log evaluation metrics
+    metrics = trainer.evaluate()
+    mlflow.log_metrics(metrics)
+
+    mlflow.end_run()
+
+if __name__ == "__main__":
+    fine_tune_model()