Update app.py

app.py

@@ -1,132 +1,120 @@
 import gradio as gr
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    pipeline,
-    Trainer,
-    TrainingArguments,
-    DataCollatorForLanguageModeling
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, Trainer, TrainingArguments
 from datasets import load_dataset, Dataset
 import torch
 import pandas as pd
+from transformers import DataCollatorForLanguageModeling
 from sklearn.model_selection import train_test_split
 
-
-# Configuration
 MODEL_NAME = "microsoft/DialoGPT-medium"
 DATASET_NAME = "embedding-data/Amazon-QA"
 FINETUNED_MODEL_NAME = "MujtabaShopifyChatbot"
-MAX_LENGTH = 128
-BATCH_SIZE = 8
 
 chatbot_pipe = None
-tokenizer = None
 
 def show_dataset_head(dataset, num_rows=5):
-    """Dataset preview"""
+    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])
-            print(f"\n{split} split preview:")
-            print(df[['question', 'answer']].head() if 'question' in df.columns else df.head())
+            cols = [col for col in ['query', 'pos', 'question', 'answer'] if col in df.columns]
+            if cols:
+                print("Dataset columns ", cols)
 
 def load_and_preprocess_data():
-    """Data loading with cleaning"""
-    print(f"Loading {DATASET_NAME}")
+    print("Loading dataset from ", DATASET_NAME)
     try:
         dataset = load_dataset(DATASET_NAME)
         show_dataset_head(dataset)
 
         df = pd.DataFrame(dataset['train'])
 
-        # Column normalization
         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:
-            if len(df.columns) >= 2:
-                df = df.rename(columns={df.columns[0]: 'question', df.columns[1]: 'answer'})
-            else:
-                raise ValueError("Dataset must have at least two columns for question and answer")
+            df = df.rename(columns={df.columns[0]: 'question', df.columns[1]: 'answer'})
 
-        # Cleaning
         df = df[['question', 'answer']].dropna()
-        df = df[~df['answer'].str.contains(r'\[|\^|\]', regex=True, na=False)]
-        df = df[df['answer'].str.len() > 10]
-        df = df[:10000]
+        df = df[:5000]
+
+        df['answer'] = df['answer'].astype(str).str.replace(r'\[\^|\].*', '', regex=True)
 
-        # Split
-        train_df, test_df = train_test_split(df, test_size=0.1, random_state=42)
-        return Dataset.from_pandas(train_df), Dataset.from_pandas(test_df)
+        processed_dataset = Dataset.from_pandas(df)
+        show_dataset_head(processed_dataset)
+        return processed_dataset.train_test_split(test_size=0.1)
     except Exception as e:
-        print(f"Data error: {str(e)}")
+        print("Error loading dataset ", e)
         raise
 
-def tokenize_data(train_dataset, test_dataset):
-    """Basic tokenization"""
-    global tokenizer
-    print(f"Tokenizing with {MODEL_NAME}")
+def tokenize_data(dataset):
+    print("Tokenizing data with model ", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     tokenizer.pad_token = tokenizer.eos_token
 
     def preprocess_function(examples):
-        texts = [f"{q} {tokenizer.eos_token} {a}" for q, a in zip(examples["question"], examples["answer"])]
-        return tokenizer(
-            texts,
-            max_length=MAX_LENGTH,
+        inputs = [f"question: {q} answer: {a}" for q, a in zip(examples["question"], examples["answer"])]
+
+        model_inputs = tokenizer(
+            inputs,
+            max_length=128,
             truncation=True,
-            padding="max_length",
-            return_tensors="pt"
+            padding='max_length'
         )
 
-    train_tokenized = train_dataset.map(preprocess_function, batched=True, remove_columns=['question', 'answer'])
-    test_tokenized = test_dataset.map(preprocess_function, batched=True, remove_columns=['question', 'answer'])
-    return train_tokenized, test_tokenized
+        model_inputs["labels"] = model_inputs["input_ids"].copy()
+        return model_inputs
+
+    return dataset.map(preprocess_function, batched=True)
 
-def fine_tune_model(train_data, test_data):
-    """Optimized training"""
-    print("Starting fine-tuning")
+def fine_tune_model(tokenized_dataset):
+    print("Starting fine-tuning process")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False
+    )
 
     training_args = TrainingArguments(
         output_dir="./results",
-        evaluation_strategy="steps",
-        eval_steps=500,
-        learning_rate=3e-5,
-        per_device_train_batch_size=BATCH_SIZE,
-        per_device_eval_batch_size=BATCH_SIZE,
-        num_train_epochs=4,
+        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,
-        warmup_ratio=0.1,
+        save_total_limit=3,
         fp16=torch.cuda.is_available(),
+        push_to_hub=False,
+        report_to="none",
         logging_steps=100,
-        save_steps=1000,
-        save_total_limit=2,
-        load_best_model_at_end=True,
-        report_to="none"  # Disable W&B logging
+        save_steps=500
     )
 
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=train_data,
-        eval_dataset=test_data,
-        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
+        train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["test"],
+        data_collator=data_collator
     )
 
     trainer.train()
+    print("Training completed, saving model")
     model.save_pretrained(FINETUNED_MODEL_NAME)
     tokenizer.save_pretrained(FINETUNED_MODEL_NAME)
     return model
 
 def initialize_chatbot():
-    """Initialize generation pipeline"""
-    global chatbot_pipe, tokenizer
-    print(f"Loading {FINETUNED_MODEL_NAME}")
+    global chatbot_pipe
+    print("Initializing chatbot with model ", FINETUNED_MODEL_NAME)
     try:
+        model = AutoModelForCausalLM.from_pretrained(FINETUNED_MODEL_NAME)
         tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL_NAME)
         tokenizer.pad_token = tokenizer.eos_token
-        model = AutoModelForCausalLM.from_pretrained(FINETUNED_MODEL_NAME)
 
         chatbot_pipe = pipeline(
             "text-generation",
@@ -134,48 +122,54 @@ def initialize_chatbot():
             tokenizer=tokenizer,
             device=0 if torch.cuda.is_available() else -1
         )
+        print("Chatbot initialized successfully")
     except Exception as e:
-        print(f"Initialization failed: {str(e)}")
-        raise
+        print("Error initializing chatbot ", e)
+        return None
+    return chatbot_pipe
 
 def generate_response(message, history):
-    """Direct generation without prompt engineering"""
-    if not chatbot_pipe:
-        return "System initializing..."
+    if chatbot_pipe is None:
+        print("Chatbot pipeline not initialized")
+        return "System error: Chatbot not ready"
 
     try:
+        print("Generating response for query ", message)
         response = chatbot_pipe(
-            message,
-            max_length=MAX_LENGTH,
+            f"question: {message} answer:",
+            max_length=128,
             do_sample=True,
             temperature=0.7,
-            top_k=50,
-            top_p=0.9,
-            repetition_penalty=1.2,
-            num_return_sequences=1
+            top_p=0.9
         )[0]['generated_text']
-
-        return response.split(tokenizer.eos_token)[-1].strip()
+        final_response = response.split("answer:")[-1].strip()
+        print("Generated response ", final_response)
+        return final_response
     except Exception as e:
-        print(f"Generation error: {str(e)}")
-        return "Please try again later."
+        print("Error generating response ", e)
+        return "Sorry, I encountered an error processing your request"
 
 def deploy_chatbot():
-    """Gradio interface"""
+    print("Launching chatbot interface")
     demo = gr.ChatInterface(
         fn=generate_response,
-        title="Shopify Assistant",
+        title="Mujtaba's Shopify Assistant",
+        description="Ask about products, shipping, or store policies",
         examples=[
-            "Does this work with iPhone 15?",
-            "What's the return policy?",
-            "Do you ship internationally?"
-        ]
+            "Will this work with iPhone 15?",
+            "What's the return window?",
+            "Do you ship to Lahore?"
+        ],
+        theme="soft",
+        cache_examples=False
     )
     return demo
 
 if __name__ == "__main__":
-    train_data, test_data = load_and_preprocess_data()
-    train_tokenized, test_tokenized = tokenize_data(train_data, test_data)
-    model = fine_tune_model(train_tokenized, test_tokenized)
+    dataset = load_and_preprocess_data()
+    tokenized_data = tokenize_data(dataset)
+
+    model = fine_tune_model(tokenized_data)
+
     initialize_chatbot()
     deploy_chatbot().launch()