Update app.py

app.py

@@ -1,6 +1,6 @@
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 # Load the model and tokenizer locally
 model_name = "kz919/QwQ-0.5B-Distilled-SFT"
@@ -30,16 +30,27 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
     # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-    # Generate a response
-    outputs = model.generate(
-        inputs.input_ids,
-        max_length=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        pad_token_id=tokenizer.eos_token_id,
-        streamer = TextStreamer(tokenizer)
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # Use a thread to run the generation in parallel
+    generation_thread = threading.Thread(
+        target=model.generate,
+        kwargs=dict(
+            inputs=inputs.input_ids,
+            max_length=max_tokens,
+            streamer=streamer,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            pad_token_id=tokenizer.eos_token_id,
+        ),
     )
-    yield outputs
+    generation_thread.start()
+
+    # Stream the tokens as they are generated
+    for new_text in streamer:
+        yield new_text
 
 
 # Create the Gradio interface