Update app.py

app.py

@@ -1,72 +1,103 @@
-# app.py - Complete Chatbot with Fine-tuning and Deployment
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, Trainer, TrainingArguments
 from datasets import load_dataset, Dataset
 import torch
 import pandas as pd
-from huggingface_hub import notebook_login, Repository
+from huggingface_hub import notebook_login
+from transformers import DataCollatorForSeq2Seq
 
-# Configuration
-MODEL_NAME = "t5-small"  # Lightweight model good for chatbots
-DATASET_NAME = "AmazonQA"
+MODEL_NAME = "microsoft/DialoGPT-small"
+DATASET_NAME = "embedding-data/amazon-QA"
 FINETUNED_MODEL_NAME = "MujtabaShopifyChatbot"
-HF_TOKEN = "your_huggingface_token"  # Replace with your token
+HF_TOKEN = "your_huggingface_token"
+
+chatbot_pipe = None
+
+def show_dataset_head(dataset, num_rows=5):
+    print("Displaying dataset preview ", dataset)
+    if isinstance(dataset, dict):
+        for split in dataset.keys():
+            print("Current split ", split)
+            df = pd.DataFrame(dataset[split][:num_rows])
+            cols = [col for col in ['query', 'pos', 'question', 'answer'] if col in df.columns]
+            if cols:
+                print("Dataset columns ", cols)
 
-# --- Step 1: Load and Prepare Dataset ---
 def load_and_preprocess_data():
-    print("Loading AmazonQA dataset...")
+    print("Loading dataset from ", DATASET_NAME)
     dataset = load_dataset(DATASET_NAME)
+    show_dataset_head(dataset)
     
-    # Convert to pandas for easier processing
     df = pd.DataFrame(dataset['train'])
     
-    # Preprocessing - create consistent Q&A pairs
+    if 'query' in df.columns and 'pos' in df.columns:
+        df = df.rename(columns={'query': 'question', 'pos': 'answer'})
+    elif 'question' not in df.columns or 'answer' not in df.columns:
+        df = df.rename(columns={df.columns[0]: 'question', df.columns[1]: 'answer'})
+    
     df = df[['question', 'answer']].dropna()
-    df = df[:5000]  # Use subset for faster training
+    df = df[:5000]
     
-    # Convert back to Hugging Face Dataset
-    processed_dataset = Dataset.from_pandas(df)
+    df['answer'] = df['answer'].astype(str).str.replace(r'\[\^|\].*', '', regex=True)
     
-    # Split into train and eval
-    split_dataset = processed_dataset.train_test_split(test_size=0.1)
-    return split_dataset
+    processed_dataset = Dataset.from_pandas(df)
+    show_dataset_head(processed_dataset)
+    return processed_dataset.train_test_split(test_size=0.1)
 
-# --- Step 2: Tokenization ---
 def tokenize_data(dataset):
-    print("Tokenizing data...")
+    print("Tokenizing data with model ", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     
     def preprocess_function(examples):
         inputs = [f"question: {q} answer:" for q in examples["question"]]
-        targets = examples["answer"]
+        targets = [str(a) for a in examples["answer"]]
         
-        model_inputs = tokenizer(inputs, max_length=128, truncation=True)
-        labels = tokenizer(targets, max_length=128, truncation=True)
+        model_inputs = tokenizer(
+            inputs,
+            max_length=128,
+            truncation=True,
+            padding='max_length'
+        )
+        labels = tokenizer(
+            targets,
+            max_length=128,
+            truncation=True,
+            padding='max_length'
+        )
         
         model_inputs["labels"] = labels["input_ids"]
         return model_inputs
 
-    tokenized_dataset = dataset.map(preprocess_function, batched=True)
-    return tokenized_dataset
+    return dataset.map(preprocess_function, batched=True)
 
-# --- Step 3: Fine-tuning ---
 def fine_tune_model(tokenized_dataset):
-    print("Fine-tuning model...")
+    print("Starting fine-tuning process")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
     
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer,
+        model=model,
+        padding='longest',
+        max_length=128,
+        pad_to_multiple_of=8
+    )
+    
     training_args = TrainingArguments(
         output_dir="./results",
-        evaluation_strategy="epoch",
-        learning_rate=2e-5,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
+        eval_strategy="epoch",
+        learning_rate=5e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
         num_train_epochs=3,
         weight_decay=0.01,
         save_total_limit=3,
         fp16=torch.cuda.is_available(),
-        push_to_hub=True,
-        hub_model_id=FINETUNED_MODEL_NAME,
-        hub_token=HF_TOKEN,
+        push_to_hub=False,
+        report_to="none",
+        logging_steps=100,
+        save_steps=500,
+        gradient_accumulation_steps=1
     )
     
     trainer = Trainer(
@@ -74,70 +105,77 @@ def fine_tune_model(tokenized_dataset):
         args=training_args,
         train_dataset=tokenized_dataset["train"],
         eval_dataset=tokenized_dataset["test"],
+        data_collator=data_collator,
+        tokenizer=tokenizer
     )
     
     trainer.train()
-    trainer.push_to_hub()
+    print("Training completed, saving model")
+    model.save_pretrained(FINETUNED_MODEL_NAME)
+    tokenizer.save_pretrained(FINETUNED_MODEL_NAME)
     return model
 
-# --- Step 4: Chatbot Interface ---
 def initialize_chatbot():
-    print("Loading chatbot...")
+    global chatbot_pipe
+    print("Initializing chatbot with model ", FINETUNED_MODEL_NAME)
     try:
-        # Try loading fine-tuned model first
         model = AutoModelForSeq2SeqLM.from_pretrained(FINETUNED_MODEL_NAME)
         tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL_NAME)
-    except:
-        # Fallback to pre-trained model
-        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    
-    chatbot_pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+        chatbot_pipe = pipeline(
+            "text2text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device=0 if torch.cuda.is_available() else -1
+        )
+        print("Chatbot initialized successfully")
+    except Exception as e:
+        print("Error initializing chatbot ", e)
+        return None
     return chatbot_pipe
 
 def generate_response(message, history):
-    # Format the input for the model
-    input_text = f"question: {message} answer:"
-    
-    # Generate response
-    response = chatbot_pipe(input_text, max_length=128, do_sample=True)[0]['generated_text']
+    if chatbot_pipe is None:
+        print("Chatbot pipeline not initialized")
+        return "System error: Chatbot not ready"
     
-    # Clean up the response
-    if "answer:" in response:
-        response = response.split("answer:")[-1].strip()
-    return response
+    try:
+        print("Generating response for query ", message)
+        response = chatbot_pipe(
+            f"question: {message} answer:",
+            max_length=128,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )[0]['generated_text']
+        final_response = response.split("answer:")[-1].strip()
+        print("Generated response ", final_response)
+        return final_response
+    except Exception as e:
+        print("Error generating response ", e)
+        return "Sorry, I encountered an error processing your request"
 
-# --- Step 5: Deployment ---
 def deploy_chatbot():
-    print("Launching chatbot interface...")
+    print("Launching chatbot interface")
     demo = gr.ChatInterface(
         fn=generate_response,
-        title="Mujtaba's Shopify Chatbot",
-        description="Ask me anything about products, shipping, or returns!",
+        title="Mujtaba's Shopify Assistant",
+        description="Ask about products, shipping, or store policies",
         examples=[
-            "What's the return policy?",
-            "How long does shipping take to Karachi?",
-            "Do you have size charts for kurtas?"
+            "Will this work with iPhone 15?",
+            "What's the return window?",
+            "Do you ship to Lahore?"
         ],
-        theme="soft"
+        theme="soft",
+        cache_examples=False
     )
     return demo
 
-# --- Main Execution ---
 if __name__ == "__main__":
-    # Login to Hugging Face Hub
     notebook_login()
-    
-    # Dataset preparation
     dataset = load_and_preprocess_data()
-    tokenized_dataset = tokenize_data(dataset)
-    
-    # Fine-tuning (uncomment to run)
-    # fine_tune_model(tokenized_dataset)
+    tokenized_data = tokenize_data(dataset)
     
-    # Initialize chatbot
-    chatbot_pipe = initialize_chatbot()
+    model = fine_tune_model(tokenized_data)
     
-    # Launch interface
-    demo = deploy_chatbot()
-    demo.launch(share=True)
+    initialize_chatbot()
+    deploy_chatbot().launch()