Free to choose a model

app.py

@@ -146,75 +146,50 @@ def format_prompt(messages):
 def add_user_message(user_input, history):
     return "", history + [{"role": "user", "content": user_input}]
 
-# Available models
+# Curated models
 model_choices = [
     "meta-llama/Llama-3.2-3B-Instruct",
     "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
-    "google/gemma-7b"
+    "google/gemma-7b",
+    "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
 ]
 
-# UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Clinical Chatbot (Streaming) — LLaMA, DeepSeek, Gemma")
-
-    default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
-
-    # @spaces.GPU
-    # def chat_with_model(messages):
-    #     global current_model, current_tokenizer
-    #     if current_model is None or current_tokenizer is None:
-    #         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
-    #         return
-
-    #     current_model = current_model.to("cuda").half()
-
-    #     prompt = format_prompt(messages)
-    #     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-
-    #     output_ids = []
-    #     messages = messages.copy()
-    #     messages.append({"role": "assistant", "content": ""})
-
-    #     for token_id in current_model.generate(
-    #         **inputs,
-    #         max_new_tokens=256,
-    #         do_sample=True,
-    #         return_dict_in_generate=True,
-    #         output_scores=False
-    #     ).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=False)
-    #         if output_ids[-1] == current_tokenizer.eos_token_id:
-    #             current_model.to("cpu")
-    #             torch.cuda.empty_cache()
-    #             return
-    #         messages[-1]["content"] = decoded
-    #         yield messages
-
-    #     current_model.to("cpu")
-    #     torch.cuda.empty_cache()
-    #     return
+    gr.Markdown("## Clinical Chatbot (Streaming)")
+
+    default_model = gr.State(model_choices[0])
 
     with gr.Row():
-        model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
-        model_status = gr.Textbox(label="Model Status", interactive=False)
+        mode = gr.Radio(["Choose from list", "Enter custom model"], value="Choose from list", label="Model Input Mode")
+        model_selector = gr.Dropdown(choices=model_choices, label="Select Predefined Model")
+        model_textbox = gr.Textbox(label="Or Enter HF Model Name")
 
+    model_status = gr.Textbox(label="Model Status", interactive=False)
     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 default model on startup
+    def resolve_model_choice(mode, dropdown_value, textbox_value):
+        return textbox_value.strip() if mode == "Enter custom model" else dropdown_value
+
+    # Load on launch
     demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)
 
-    # Load selected model manually
-    model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)
+    # Load on user selection
+    mode.select(fn=resolve_model_choice, inputs=[mode, model_selector, model_textbox], outputs=default_model).then(
+        load_model_on_selection, inputs=default_model, outputs=model_status
+    )
+    model_selector.change(fn=resolve_model_choice, inputs=[mode, model_selector, model_textbox], outputs=default_model).then(
+        load_model_on_selection, inputs=default_model, outputs=model_status
+    )
+    model_textbox.submit(fn=resolve_model_choice, inputs=[mode, model_selector, model_textbox], outputs=default_model).then(
+        load_model_on_selection, inputs=default_model, outputs=model_status
+    )
 
-    # Submit message + stream model response
     msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
         chat_with_model, chatbot, chatbot
     )
-
-    # Clear chat
     clear.click(lambda: [], None, chatbot, queue=False)
 
+
 demo.launch()