Make it into a chatbot

app.py

@@ -4,11 +4,11 @@ import gradio as gr
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-
-# Use a global variable to hold the current model and tokenizer
+# Global model/tokenizer
 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")
@@ -27,17 +27,27 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     progress(1, desc="Model ready.")
     return f"{model_name} loaded and ready!"
 
+# Inference - yields response token-by-token
 @spaces.GPU
-def generate_text(prompt):
+def chat_with_model(history):
     global current_model, current_tokenizer
     if current_model is None or current_tokenizer is None:
-        yield "⚠️ No model loaded yet. Please select a model first."
+        yield history + [("⚠️ No model loaded.", "")]
 
     current_model.to("cuda")
-    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
 
+    # 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]"
+
+    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
     output_ids = []
-    streamer_output = ""
+
+    # Clone history to avoid mutating during yield
+    updated_history = history.copy()
+    updated_history[-1] = (history[-1][0], "")
 
     for token_id in current_model.generate(
         **inputs,
@@ -45,14 +55,17 @@ def generate_text(prompt):
         do_sample=False,
         return_dict_in_generate=True,
         output_scores=False
-        ).sequences[0]:
-
+    ).sequences[0]:
         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
 
-        yield current_tokenizer.decode(output_ids, skip_special_tokens=True)
-
+# When user submits a message
+def add_user_message(message, history):
+    return "", history + [(message, "")]
 
-# Model options
+# Model choices
 model_choices = [
     "meta-llama/Llama-3.2-3B-Instruct",
     "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
@@ -61,26 +74,30 @@ model_choices = [
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Clinical Text Testing with LLaMA, DeepSeek, and Gemma")
+    gr.Markdown("## Clinical Chatbot — LLaMA, DeepSeek, Gemma")
 
-    # State to track initial model to load
     default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
 
-    model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
-    model_status = gr.Textbox(label="Model Status", interactive=False)
+    with gr.Row():
+        model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
+        model_status = gr.Textbox(label="Model Status", interactive=False)
 
-    input_text = gr.Textbox(label="Input Clinical Text")
-    generate_btn = gr.Button("Generate")
-    output_text = gr.Textbox(label="Generated Output")
+    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")
 
-    # Auto-load default model on launch
+    # Load model on launch
     demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)
 
-    # Manual model selection
+    # Load model on dropdown selection
     model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)
 
-    # Generate with current model
-    generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)
-    input_text.submit(fn=generate_text, inputs=input_text, outputs=output_text)
+    # On message submit: update history, then stream bot reply
+    msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
+        fn=chat_with_model, inputs=chatbot, outputs=chatbot
+    )
+
+    # Clear chat
+    clear_btn.click(lambda: [], None, chatbot, queue=False)
 
 demo.launch()