feat: Added Streaming

app.py

@@ -26,14 +26,30 @@ def inference(prompt: str, max_tokens: int = 256) -> str:
         tokenize = False,
         add_generation_prompt = True, # Must add for generation
     )
+
+    inputs = tokenizer(text, return_tensors="pt").to("cuda")
+    
+    # Use TextIteratorStreamer instead of TextStreamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    outputs = model.generate(
-        **tokenizer(text, return_tensors = "pt").to("cuda"),
-        max_new_tokens = 100, # Increase for longer outputs!
-        temperature = 0.7, top_p = 0.8, top_k = 20, # For non thinking
-        streamer = TextStreamer(tokenizer, skip_prompt = True),
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_tokens,
+        temperature=0.7,
+        top_p=0.8,
+        top_k=20,
+        streamer=streamer,
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Run generation in separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # Yield tokens as they come
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text  # yield cumulative text for Gradio
 
 demo = gr.Interface(
     fn=inference,