Update app.py

app.py

@@ -1,175 +1,62 @@
+import os
 import gradio as gr
-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
+from groq import Groq
+from datasets import load_dataset
 
-MODEL_NAME = "microsoft/DialoGPT-medium"
+GROQ_MODEL = "llama3-70b-8192"
 DATASET_NAME = "embedding-data/Amazon-QA"
-FINETUNED_MODEL_NAME = "MujtabaShopifyChatbot"
 
-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)
-
-def load_and_preprocess_data():
-    print("Loading dataset from ", DATASET_NAME)
-    try:
-        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)
-    except Exception as e:
-        print("Error loading dataset ", e)
-        raise
-
-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):
-        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'
-        )
-
-        model_inputs["labels"] = model_inputs["input_ids"].copy()
-        return model_inputs
-
-    return dataset.map(preprocess_function, batched=True)
-
-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="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=False,
-        report_to="none",
-        logging_steps=100,
-        save_steps=500
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        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():
-    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
-
-        chatbot_pipe = pipeline(
-            "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
+def load_shopify_context():
+    dataset = load_dataset(DATASET_NAME)
+    samples = dataset['train'].select(range(3))
+    examples = []
+    for sample in samples:
+        question = sample['query']
+        if isinstance(question, list):
+            question = question[0] if len(question) > 0 else "No question"
+        question = str(question).replace('\\', '/')
+        answer = sample.get('pos', sample.get('answer', ["No answer"]))
+        if isinstance(answer, list):
+            answer = answer[0] if len(answer) > 0 else "No answer"
+        answer = str(answer).replace('\\', '/')
+        examples.append(f"Q: {question}\nA: {answer}")
+    return '\n'.join(examples)
 
 def generate_response(message, history):
-    if chatbot_pipe is None:
-        print("Chatbot pipeline not initialized")
-        return "System error: Chatbot not ready"
-
+    api_key = os.getenv("GROQ_API_KEY")
+    if not api_key:
+        return "Error: GROQ_API_KEY not set. Please add it as a secret in your Space."
+    client = Groq(api_key=api_key)
+    context = load_shopify_context()
+    conversation = []
+    for user_msg, bot_msg in history:
+        safe_user = str(user_msg).replace('\\', '/')
+        safe_bot = str(bot_msg).replace('\\', '/')
+        conversation.extend([f"User: {safe_user}", f"Assistant: {safe_bot}"])
+    safe_message = str(message).replace('\\', '/')
+    prompt = f"You are an expert Shopify support agent. Context examples:\n{context}\n{chr(10).join(conversation)}\nUser: {safe_message}\nAssistant:"
     try:
-        print("Generating response for query ", message)
-        response = chatbot_pipe(
-            f"question: {message} answer:",
-            max_length=128,
-            do_sample=True,
+        response = client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=GROQ_MODEL,
             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
+            max_tokens=256,
+            top_p=0.9,
+            stop=["<|endoftext|>"]
+        )
+        return response.choices[0].message.content
     except Exception as e:
-        print("Error generating response ", e)
-        return "Sorry, I encountered an error processing your request"
+        return f"Error: {str(e)}"
 
-def deploy_chatbot():
-    print("Launching chatbot interface")
-    demo = gr.ChatInterface(
+with gr.Blocks() as app:
+    gr.Markdown("## Shopify Q&A Assistant (Groq-powered)")
+    gr.ChatInterface(
         fn=generate_response,
-        title="Mujtaba's Shopify Assistant",
-        description="Ask about products, shipping, or store policies",
         examples=[
-            "Will this work with iPhone 15?",
-            "What's the return window?",
-            "Do you ship to Lahore?"
-        ],
-        theme="soft",
-        cache_examples=False
+            "What's your return policy?",
+            "Do you ship internationally?",
+            "Is this compatible with iPhone 15?"
+        ]
     )
-    return demo
-
-if __name__ == "__main__":
-    dataset = load_and_preprocess_data()
-    tokenized_data = tokenize_data(dataset)
-
-    model = fine_tune_model(tokenized_data)
 
-    initialize_chatbot()
-    deploy_chatbot().launch()
+app.launch()