Update app.py

app.py

@@ -1,35 +1,97 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Use Phi model (ensure to pass trust_remote_code if required)
-model_name = "microsoft/Phi-3-mini-4k-instruct"
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-
-def generate_response_phi(user_input, chat_history):
-    if chat_history is None:
-        chat_history = []
-    # Append user message to the conversation as a dict (the Phi template expects this format)
-    chat_history.append({"role": "user", "content": user_input})
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+# Model configurations
+MODEL_CONFIG = {
+    "llama": {
+        "model_name": "meta-llama/Llama-2-7b-chat-hf",
+        "template": "[INST] {message} [/INST]"
+    },
+    "phi": {
+        "model_name": "microsoft/phi-2",
+        "template": "{message}"
+    }
+}
+
+class ChatModel:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.current_model = None
+
+    def load_model(self, model_name):
+        if model_name != self.current_model:
+            config = MODEL_CONFIG[model_name]
+            self.tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
+            self.model = AutoModelForCausalLM.from_pretrained(
+                config["model_name"],
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            self.current_model = model_name
+
+    def format_message(self, message, model_name):
+        return MODEL_CONFIG[model_name]["template"].format(message=message)
+
+    def generate(self, message, model_name, history):
+        self.load_model(model_name)
+        formatted_message = self.format_message(message, model_name)
+        
+        # Create pipeline for text generation
+        pipe = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device_map="auto"
+        )
+        
+        # Generate response
+        response = pipe(
+            formatted_message,
+            max_length=200,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+        
+        return response[0]['generated_text'].replace(formatted_message, "").strip()
+
+# Initialize model handler
+model_handler = ChatModel()
+
+def chat(message, history, model_choice):
+    response = model_handler.generate(message, model_choice, history)
+    return [(message, response)]
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🤖 Local LLM Chatbot\nSelect a model and start chatting!")
     
-    # Use the tokenizer's chat template to prepare inputs
-    inputs = tokenizer.apply_chat_template(
-        chat_history, add_generation_prompt=True, return_tensors="pt"
-    )
-    # Generate response
-    output_ids = model.generate(**inputs, max_new_tokens=100)
-    generated_text = tokenizer.batch_decode(output_ids)[0]
-    # Extract assistant reply (assuming the template adds "<|assistant|>" marker)
-    answer = generated_text.split("<|assistant|>")[-1].strip()
-    chat_history.append({"role": "assistant", "content": answer})
-    return "", chat_history
-
-with gr.Blocks() as phi_demo:
-    gr.Markdown("# Phi Chatbot")
-    chatbot = gr.Chatbot()
-    state = gr.State([])
-    txt = gr.Textbox(placeholder="Enter your message")
-    txt.submit(generate_response_phi, [txt, state], [txt, chatbot])
+    with gr.Row():
+        model_choice = gr.Dropdown(
+            choices=["llama", "phi"],
+            label="Select Model",
+            value="phi"
+        )
     
-phi_demo.launch()
+    chatbot = gr.Chatbot(height=400)
+    msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")
+    
+    with gr.Row():
+        submit_btn = gr.Button("Send")
+        clear_btn = gr.ClearButton([msg, chatbot])
+
+    msg.submit(
+        fn=chat,
+        inputs=[msg, chatbot, model_choice],
+        outputs=[chatbot]
+    )
+    submit_btn.click(
+        fn=chat,
+        inputs=[msg, chatbot, model_choice],
+        outputs=[chatbot]
+    )
+
+demo.launch()