Update app.py

app.py

@@ -47,6 +47,53 @@ class ChatModel:
             self.models[model_name] = model
             self.tokenizers[model_name] = tokenizer
 
-    # ... (keep the rest of the code the same as previous version)
+    def generate(self, message, model_name, history):
+        self.load_model(model_name)
+        config = MODEL_CONFIG[model_name]
+        
+        # Format prompt
+        prompt = config["template"].format(message=message)
+        
+        # Create pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=self.models[model_name],
+            tokenizer=self.tokenizers[model_name],
+            max_new_tokens=384,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True,
+            return_full_text=False
+        )
+        
+        response = pipe(prompt)[0]['generated_text']
+        return response.strip()
+
+model_handler = ChatModel()
+
+def chat(message, history, model_choice):
+    try:
+        response = model_handler.generate(message, model_choice, history)
+        return [(message, response)]
+    except Exception as e:
+        return [(message, f"Error: {str(e)}")]
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚀 Phi-3 vs Llama-3 Chatbot")
+    with gr.Row():
+        model_choice = gr.Dropdown(
+            choices=["phi-3", "llama3-8b"],
+            label="Select Model",
+            value="phi-3"
+        )
+    chatbot = gr.Chatbot(height=400)
+    msg = gr.Textbox(label="Message", placeholder="Type here...")
+    with gr.Row():
+        submit_btn = gr.Button("Send", variant="primary")
+        clear_btn = gr.ClearButton([msg, chatbot])
+
+    msg.submit(chat, [msg, chatbot, model_choice], chatbot)
+    submit_btn.click(chat, [msg, chatbot, model_choice], chatbot)
 
-# ... (remaining code identical to previous implementation)
+demo.launch()