Change chatbot

app.py

@@ -1,14 +1,14 @@
 import os
 import torch
+import time
 import gradio as gr
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Global model/tokenizer
+# Globals
 current_model = None
 current_tokenizer = None
 
-# Load model when selected
 def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     global current_model, current_tokenizer
     token = os.getenv("HF_TOKEN")
@@ -20,34 +20,40 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     current_model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,
-        device_map="cpu",
+        device_map="cpu",  # loaded to CPU initially
         use_auth_token=token
     )
 
     progress(1, desc="Model ready.")
     return f"{model_name} loaded and ready!"
 
-# Inference - yields response token-by-token
+# Format conversation as plain text
+def format_prompt(messages):
+    prompt = ""
+    for msg in messages:
+        role = msg["role"]
+        if role == "user":
+            prompt += f"User: {msg['content'].strip()}\n"
+        elif role == "assistant":
+            prompt += f"Assistant: {msg['content'].strip()}\n"
+    prompt += "Assistant:"
+    return prompt
+
 @spaces.GPU
-def chat_with_model(history):
+def chat_with_model(messages):
     global current_model, current_tokenizer
     if current_model is None or current_tokenizer is None:
-        yield history + [("⚠️ No model loaded.", "")]
+        yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
+        return
 
     current_model.to("cuda")
 
-    # Combine conversation history into prompt
-    prompt = ""
-    for user_msg, bot_msg in history:
-        prompt += f"[INST] {user_msg.strip()} [/INST] {bot_msg.strip()} "
-    prompt += f"[INST] {history[-1][0]} [/INST]"
-
+    prompt = format_prompt(messages)
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-    output_ids = []
 
-    # Clone history to avoid mutating during yield
-    updated_history = history.copy()
-    updated_history[-1] = (history[-1][0], "")
+    output_ids = []
+    messages = messages.copy()
+    messages.append({"role": "assistant", "content": ""})
 
     for token_id in current_model.generate(
         **inputs,
@@ -55,26 +61,25 @@ def chat_with_model(history):
         do_sample=False,
         return_dict_in_generate=True,
         output_scores=False
-    ).sequences[0]:
+    ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
         output_ids.append(token_id.item())
         decoded = current_tokenizer.decode(output_ids, skip_special_tokens=True)
-        updated_history[-1] = (history[-1][0], decoded)
-        yield updated_history
+        messages[-1]["content"] = decoded
+        yield messages
 
-# When user submits a message
-def add_user_message(message, history):
-    return "", history + [(message, "")]
+def add_user_message(user_input, history):
+    return "", history + [{"role": "user", "content": user_input}]
 
-# Model choices
+# Available models
 model_choices = [
     "meta-llama/Llama-3.2-3B-Instruct",
     "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
     "google/gemma-7b"
 ]
 
-# Gradio UI
+# UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Clinical Chatbot — LLaMA, DeepSeek, Gemma")
+    gr.Markdown("## Clinical Chatbot (Streaming) — LLaMA, DeepSeek, Gemma")
 
     default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
 
@@ -82,22 +87,22 @@ with gr.Blocks() as demo:
         model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
         model_status = gr.Textbox(label="Model Status", interactive=False)
 
-    chatbot = gr.Chatbot(label="Chat")
-    msg = gr.Textbox(label="Your Message", placeholder="Enter your clinical query...", show_label=False)
-    clear_btn = gr.Button("Clear Chat")
+    chatbot = gr.Chatbot(label="Chat", type="messages")
+    msg = gr.Textbox(label="Your message", placeholder="Enter clinical input...", show_label=False)
+    clear = gr.Button("Clear")
 
-    # Load model on launch
+    # Load default model on startup
     demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)
 
-    # Load model on dropdown selection
+    # Load selected model manually
     model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)
 
-    # On message submit: update history, then stream bot reply
+    # Submit message + stream model response
     msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-        fn=chat_with_model, inputs=chatbot, outputs=chatbot
+        chat_with_model, chatbot, chatbot
     )
 
     # Clear chat
-    clear_btn.click(lambda: [], None, chatbot, queue=False)
+    clear.click(lambda: [], None, chatbot, queue=False)
 
 demo.launch()