Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Only transformer-loadable models
 MODEL_OPTIONS = {
     "Llama-3.2-3B": "meta-llama/Llama-3.2-3B-Instruct",
     "Llama-3.2-1B": "meta-llama/Llama-3.2-1B-Instruct",
@@ -13,7 +12,7 @@ MODEL_OPTIONS = {
 }
 
 loaded = {}
-SYSTEM_PROMPT = "You are HugginGPT — a helpful assistant with memory."
+SYSTEM_PROMPT = "You are HugginGPT — a helpful assistant that remembers context."
 
 def load_model(model_key):
     model_id = MODEL_OPTIONS[model_key]
@@ -32,7 +31,6 @@ def load_model(model_key):
 def generate_response(message, history, model_choice):
     tokenizer, model = load_model(model_choice)
 
-    # build prompt with system + memory
     context = f"system: {SYSTEM_PROMPT}\n"
     if history:
         for u, a in history:
@@ -52,10 +50,17 @@ def generate_response(message, history, model_choice):
     return reply
 
 with gr.Blocks() as demo:
-    gr.ChatInterface(
-        fn=generate_response,
-        title="HugginGPT",
-        inputs=[gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="Llama-3.2-3B")]
+    gr.Markdown("## HugginGPT")
+
+    model_selector = gr.Dropdown(
+        choices=list(MODEL_OPTIONS.keys()),
+        value="Llama-3.2-3B",
+        label="Select model"
+    )
+
+    chat = gr.ChatInterface(
+        fn=lambda message, history: generate_response(message, history, model_selector.value),
+        title="HugginGPT"
     )
 
 demo.launch()