Update app.py

app.py

@@ -1,15 +1,16 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
+# Model configurations
 MODEL_CONFIG = {
-    "tinyllama": {
-        "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-        "template": "<|user|>\n{message}\n<|assistant|>"
+    "llama": {
+        "model_name": "meta-llama/Llama-2-7b-chat-hf",
+        "template": "[INST] {message} [/INST]"
     },
-    "phi-1.5": {
-        "model_name": "microsoft/phi-1_5",
-        "template": "Instruct: {message}\nOutput:"
+    "phi": {
+        "model_name": "microsoft/phi-2",
+        "template": "{message}"
     }
 }
 
@@ -26,32 +27,39 @@ class ChatModel:
             self.model = AutoModelForCausalLM.from_pretrained(
                 config["model_name"],
                 torch_dtype=torch.float16,
-                device_map="auto",
-                load_in_4bit=True  # Quantization to reduce memory usage
+                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)
-        config = MODEL_CONFIG[model_name]
+        formatted_message = self.format_message(message, model_name)
         
-        # Format prompt
-        prompt = config["template"].format(message=message)
-        
-        # Create pipeline
+        # Create pipeline for text generation
         pipe = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
-            max_new_tokens=150,
-            temperature=0.3,
-            do_sample=True,
             device_map="auto"
         )
         
-        response = pipe(prompt)[0]['generated_text']
-        return response.split(prompt)[-1].strip()
+        # 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):
@@ -59,19 +67,31 @@ def chat(message, history, model_choice):
     return [(message, response)]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Free-Tier LLM Chatbot")
+    gr.Markdown("# 🤖 Local LLM Chatbot\nSelect a model and start chatting!")
+    
     with gr.Row():
         model_choice = gr.Dropdown(
-            choices=["tinyllama", "phi-1.5"],
+            choices=["llama", "phi"],
             label="Select Model",
-            value="tinyllama"
+            value="phi"
         )
-    chatbot = gr.Chatbot(height=300)
-    msg = gr.Textbox(label="Message", placeholder="Type here...")
-    submit_btn = gr.Button("Send", variant="primary")
-    clear_btn = gr.ClearButton([msg, chatbot])
+    
+    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(chat, [msg, chatbot, model_choice], chatbot)
-    submit_btn.click(chat, [msg, chatbot, model_choice], 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()