Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import spaces
 import torch
 import os
-from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load model from secret
 model_name = os.environ.get("MODEL_ID")
@@ -17,11 +17,6 @@ model.eval()
 
 SYSTEM_PROMPT = "You are an expert biomedical assistant trained to identify randomized controlled trials (RCTs). Include RCTs and exclude non-RCTs."
 
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids, scores, **kwargs):
-        # stop generation when EOS is reached
-        return input_ids[0][-1] in [tokenizer.eos_token_id]
-
 @spaces.GPU(duration=120)
 def stream_response(title_text, abstract_text):
     user_input = f"Title: {title_text.strip()}\nAbstract: {abstract_text.strip()}"
@@ -32,8 +27,9 @@ def stream_response(title_text, abstract_text):
 
     generated = inputs["input_ids"]
     past_key_values = None
+    response_text = ""
 
-    for _ in range(1024):  # max tokens
+    for _ in range(1024):  # limit max length
         outputs = model(input_ids=generated, past_key_values=past_key_values, use_cache=True)
         next_token_logits = outputs.logits[:, -1, :]
         next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
@@ -43,31 +39,30 @@ def stream_response(title_text, abstract_text):
 
         generated = torch.cat((generated, next_token), dim=1)
         decoded_output = tokenizer.decode(generated[0], skip_special_tokens=True)
-        # Stream only the assistant's reply
+
         if "<|assistant|>" in decoded_output:
-            reply = decoded_output.split("<|assistant|>")[-1].strip()
-            yield reply
+            response_text = decoded_output.split("<|assistant|>")[-1].strip()
+            yield response_text
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 RCT Classifier Demonstration (Streaming)")
+    gr.Markdown("## 🧠 RCT Classifier Demonstration (Streaming Enabled)")
 
     chatbot = gr.Chatbot()
-
     with gr.Row():
-        title = gr.Textbox(label="Title", placeholder="Enter article title")
+        title = gr.Textbox(label="Title", placeholder="Enter title")
         abstract = gr.Textbox(label="Abstract", placeholder="Enter abstract", lines=6)
-
     submit = gr.Button("Classify")
 
-    def wrapper(title, abstract):
-        message = f"Title: {title.strip()}\nAbstract: {abstract.strip()}"
-        return [(message, stream_response(title, abstract))]  # yield response
+    def stream_chat(title_text, abstract_text):
+        user_message = f"Title: {title_text.strip()}\nAbstract: {abstract_text.strip()}"
+        yield (user_message, "")  # show user message
+        response_stream = stream_response(title_text, abstract_text)
+        collected = ""
+        for partial in response_stream:
+            collected = partial
+            yield (user_message, collected)
 
-    submit.click(
-        fn=wrapper,
-        inputs=[title, abstract],
-        outputs=chatbot
-    )
+    submit.stream(fn=stream_chat, inputs=[title, abstract], outputs=chatbot)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()