Update app.py

app.py

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