Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import pipeline, set_seed
 from functools import lru_cache
 
-# === 1. Optional: Cache the pipeline loader to avoid reloading ===
+# === 1. Cache the model loader once per session ===
 @lru_cache(maxsize=1)
 def get_generator(model_name: str):
     return pipeline(
@@ -12,15 +12,20 @@ def get_generator(model_name: str):
         device_map="auto"
     )
 
-# === 2. Chat function ===
+# === 2. Chat callback ===
 def chat(user_input, history, model_name, max_length, temperature, seed):
+    # Set seed if provided
     if seed and seed > 0:
         set_seed(seed)
+    # Lazy-load the model
     generator = get_generator(model_name)
+
+    # Build prompt in Mistral’s instruction format
     prompt = (
         "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\n\n"
         f"{user_input}\n[/INST]"
     )
+    # Generate response
     outputs = generator(
         prompt,
         max_length=max_length,
@@ -29,28 +34,36 @@ def chat(user_input, history, model_name, max_length, temperature, seed):
         num_return_sequences=1
     )
     response = outputs[0]["generated_text"].split("[/INST]")[-1].strip()
-    history.append((user_input, response))
+
+    # Append to history as dicts for "messages" format
+    history.append({"role": "user",      "content": user_input})
+    history.append({"role": "assistant", "content": response})
     return history, history
 
 # === 3. Build Gradio UI ===
 with gr.Blocks() as demo:
     gr.Markdown("## 🤖 Mistral-7B-Instruct Chatbot (Gradio)")
+
+    # Chatbot and session-state
+    chatbot = gr.Chatbot(type="messages")                          # :contentReference[oaicite:3]{index=3}
+    state   = gr.State([])                                         # :contentReference[oaicite:4]{index=4}
+
     with gr.Row():
         with gr.Column(scale=3):
-            chatbox = gr.Chatbot()
-            inp = gr.Textbox(show_label=False, placeholder="Type your message here...", lines=2)
-            submit = gr.Button("Send")
+            inp    = gr.Textbox(placeholder="Type your message...", lines=2, show_label=False)
+            submit= gr.Button("Send")
         with gr.Column(scale=1):
             gr.Markdown("### Settings")
-            model_name = gr.Textbox(value="mistralai/Mistral-7B-Instruct-v0.3", label="Model name")
-            max_length = gr.Slider(50, 1024, 256, step=50, label="Max tokens")
+            model_name  = gr.Textbox(value="mistralai/Mistral-7B-Instruct-v0.3", label="Model name")
+            max_length  = gr.Slider(50, 1024, 256, step=50, label="Max tokens")
             temperature = gr.Slider(0.0, 1.0, 0.7, step=0.05, label="Temperature")
-            seed = gr.Number(42, label="Random seed (0 disables)")
+            seed        = gr.Number(42, label="Random seed (0 disables)")
 
+    # Wire the button: inputs include the gr.State; outputs update both Chatbot and state
     submit.click(
         fn=chat,
-        inputs=[inp, chatbox, model_name, max_length, temperature, seed],
-        outputs=[chatbox, chatbox]
+        inputs=[inp, state, model_name, max_length, temperature, seed],
+        outputs=[chatbot, state]
     )
 
 if __name__ == "__main__":