Generate with enter

Initialize with small Llama model

app.py

@@ -2,7 +2,8 @@ import os
 import torch
 import gradio as gr
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
 
 # Use a global variable to hold the current model and tokenizer
 current_model = None
@@ -30,16 +31,33 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
 def generate_text(prompt):
     global current_model, current_tokenizer
     if current_model is None or current_tokenizer is None:
-        return "⚠️ No model loaded yet. Please select a model first."
-    current_model.to('cuda')
+        yield "⚠️ No model loaded yet. Please select a model first."
+
+    current_model.to("cuda")
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-    outputs = current_model.generate(**inputs, max_new_tokens=256)
-    return current_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    output_ids = []
+    streamer_output = ""
+
+    def token_streamer():
+        for token_id in current_model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+            return_dict_in_generate=True,
+            output_scores=False
+        ).sequences[0]:
+            output_ids.append(token_id.item())
+            yield current_tokenizer.decode(output_ids, skip_special_tokens=True)
+
+    for partial_output in token_streamer():
+        yield partial_output
+
 
 # Model options
 model_choices = [
-    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
     "meta-llama/Llama-3.2-3B-Instruct",
+    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
     "google/gemma-7b"
 ]
 
@@ -51,14 +69,17 @@ with gr.Blocks() as demo:
     model_status = gr.Textbox(label="Model Status", interactive=False)
 
     input_text = gr.Textbox(label="Input Clinical Text")
-    output_text = gr.Textbox(label="Generated Output")
-
     generate_btn = gr.Button("Generate")
 
+    output_text = gr.Textbox(label="Generated Output")
+
     # Load model on dropdown change
     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)
+    generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text, stream=True)
+    input_text.submit(fn=generate_text, inputs=input_text, outputs=output_text, stream=True)
+
 
+load_model_on_selection("meta-llama/Llama-3.2-3B-Instruct")
 demo.launch()