Initial commit: Customer Support Chatbot with DialoGPT-medium

app.py

@@ -3,18 +3,25 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from datasets import load_dataset
 import random
+import os
 
-# Load the model and tokenizer
-model_name = "microsoft/DialoGPT-medium"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Check if fine-tuned model exists, otherwise use base model
+model_path = "./customer_support_chatbot" if os.path.exists("./customer_support_chatbot") else "microsoft/DialoGPT-medium"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path)
 
 # Load the customer support dataset
 dataset = load_dataset("Victorano/customer-support-1k")
 
 def generate_response(message, history):
-    # Encode the user message
-    input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
+    # Format the input with conversation history
+    conversation = ""
+    for user_msg, bot_msg in history:
+        conversation += f"Customer: {user_msg}\nSupport: {bot_msg}\n"
+    conversation += f"Customer: {message}\nSupport:"
+    
+    # Encode the conversation
+    input_ids = tokenizer.encode(conversation, return_tensors='pt')
     
     # Generate response
     with torch.no_grad():
@@ -31,13 +38,15 @@ def generate_response(message, history):
     
     # Decode and return the response
     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # Extract only the last response (after "Support:")
+    response = response.split("Support:")[-1].strip()
     return response
 
 # Create the Gradio interface
 with gr.Blocks(css="footer {display: none !important}") as demo:
     gr.Markdown("""
     # 🤖 Customer Support Chatbot
-    This chatbot is powered by DialoGPT-medium and trained on customer support conversations.
+    This chatbot is fine-tuned on customer support conversations using DialoGPT-medium.
     """)
     
     chatbot = gr.Chatbot(

requirements.txt

@@ -1,4 +1,7 @@
 gradio==4.19.2
 transformers==4.37.2
 torch==2.2.0
-datasets==2.17.1 
+datasets==2.17.1
+accelerate==0.27.2
+evaluate==0.4.1
+scikit-learn==1.4.0 

train.py

@@ -0,0 +1,87 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
+from datasets import load_dataset
+import numpy as np
+from typing import Dict, List
+import os
+
+def load_and_prepare_data():
+    # Load the dataset
+    dataset = load_dataset("Victorano/customer-support-1k")
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+    
+    # Function to format conversations
+    def format_conversation(example):
+        # Combine question and answer into a single conversation
+        conversation = f"Customer: {example['question']}\nSupport: {example['answer']}"
+        return {"text": conversation}
+    
+    # Apply formatting to both train and test sets
+    formatted_dataset = dataset.map(
+        format_conversation,
+        remove_columns=dataset["train"].column_names
+    )
+    
+    # Tokenize the dataset
+    def tokenize_function(examples):
+        return tokenizer(
+            examples["text"],
+            padding="max_length",
+            truncation=True,
+            max_length=512,
+            return_tensors="pt"
+        )
+    
+    tokenized_dataset = formatted_dataset.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=formatted_dataset["train"].column_names
+    )
+    
+    return tokenized_dataset, tokenizer
+
+def train_model():
+    # Load and prepare data
+    tokenized_dataset, tokenizer = load_and_prepare_data()
+    
+    # Load model
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+    
+    # Define training arguments
+    training_args = TrainingArguments(
+        output_dir="./customer_support_chatbot",
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        warmup_steps=500,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=100,
+        save_strategy="epoch",
+        evaluation_strategy="epoch",
+        load_best_model_at_end=True,
+        push_to_hub=False,
+    )
+    
+    # Initialize trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["test"],
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+    )
+    
+    # Train the model
+    trainer.train()
+    
+    # Save the model and tokenizer
+    model.save_pretrained("./customer_support_chatbot")
+    tokenizer.save_pretrained("./customer_support_chatbot")
+    
+    print("Training completed! Model saved to ./customer_support_chatbot")
+
+if __name__ == "__main__":
+    train_model()