Update app.py

app.py

@@ -1,38 +1,26 @@
-import gradio as gr  # Gradio for UI :contentReference[oaicite:3]{index=3}
-from transformers import pipeline, set_seed  # Transformers pipeline :contentReference[oaicite:4]{index=4}
+import gradio as gr
+from transformers import pipeline, set_seed
+from functools import lru_cache
 
-# Set random seed reproducibility
-def maybe_set_seed(seed: int):
-    if seed and seed > 0:
-        set_seed(seed)
-
-# Load the model once per Space session
-@gr.cache_resource  # caches across sessions :contentReference[oaicite:5]{index=5}
-def load_generator(model_name: str):
+# === 1. Optional: Cache the pipeline loader to avoid reloading ===
+@lru_cache(maxsize=1)
+def get_generator(model_name: str):
     return pipeline(
         "text-generation",
         model=model_name,
         trust_remote_code=True,
-        device_map="auto",      # auto GPU/CPU placement :contentReference[oaicite:6]{index=6}
+        device_map="auto"
     )
 
-# The chat function called on each submit
-def chat(
-    user_input: str,
-    history: list[tuple[str, str]],
-    model_name: str,
-    max_length: int,
-    temperature: float,
-    seed: int
-):
-    maybe_set_seed(seed)
-    generator = load_generator(model_name)
-    # Prepare instruction-wrapped prompt
+# === 2. Chat function ===
+def chat(user_input, history, model_name, max_length, temperature, seed):
+    if seed and seed > 0:
+        set_seed(seed)
+    generator = get_generator(model_name)
     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,
@@ -40,51 +28,30 @@ def chat(
         do_sample=True,
         num_return_sequences=1
     )
-    raw = outputs[0]["generated_text"]
-    # Extract only assistant’s reply
-    response = raw.split("[/INST]")[-1].strip()
-    # Append to chat history
+    response = outputs[0]["generated_text"].split("[/INST]")[-1].strip()
     history.append((user_input, response))
     return history, history
 
-# Build the Gradio Blocks interface
+# === 3. Build Gradio UI ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Mistral-7B-Instruct Chatbot (Gradio)")  # title :contentReference[oaicite:7]{index=7}
+    gr.Markdown("## 🤖 Mistral-7B-Instruct Chatbot (Gradio)")
     with gr.Row():
         with gr.Column(scale=3):
-            chatbox = gr.Chatbot(label="Chat History")
-            with gr.Row():
-                inp = gr.Textbox(
-                    placeholder="Type your message here...",
-                    show_label=False,
-                    lines=2
-                )
-                submit = gr.Button("Send")
+            chatbox = gr.Chatbot()
+            inp = gr.Textbox(show_label=False, placeholder="Type your message here...", lines=2)
+            submit = gr.Button("Send")
         with gr.Column(scale=1):
-            gr.Markdown("### Settings 🚀")
-            model_name = gr.Textbox(
-                label="Model name",
-                value="mistralai/Mistral-7B-Instruct-v0.3"
-            )
-            max_length = gr.Slider(
-                minimum=50, maximum=1024, step=50,
-                value=256, label="Max tokens"
-            )
-            temperature = gr.Slider(
-                minimum=0.0, maximum=1.0, step=0.05,
-                value=0.7, label="Temperature"
-            )
-            seed = gr.Number(
-                value=42, label="Random seed (0 disables)"
-            )
+            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")
+            temperature = gr.Slider(0.0, 1.0, 0.7, step=0.05, label="Temperature")
+            seed = gr.Number(42, label="Random seed (0 disables)")
 
-    # Wire up the event
     submit.click(
         fn=chat,
         inputs=[inp, chatbox, model_name, max_length, temperature, seed],
         outputs=[chatbox, chatbox]
     )
 
-# Launch the app; in Spaces, no need to set share=True :contentReference[oaicite:8]{index=8}
 if __name__ == "__main__":
     demo.launch()