TextIteratorStreamer

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import torch
-from transformers import AutoProcessor, MllamaForConditionalGeneration
+from transformers import AutoProcessor, MllamaForConditionalGeneration, TextIteratorStreamer
 from PIL import Image
 import spaces
 import tempfile
@@ -91,11 +91,28 @@ def predict_text(text, url = 'https://arinsight.co/2024_FA_AEC_1200_GR1_GR2.pdf'
     # inputs = processor(image, input_text, return_tensors="pt").to(device)
     inputs = processor(text=input_text, return_tensors="pt").to("cuda")
     # Generate a response from the model
-    outputs = model.generate(**inputs, max_new_tokens=1024)
+    # outputs = model.generate(**inputs, max_new_tokens=1024)
     
-    # Decode the output to return the final response
-    response = processor.decode(outputs[1], skip_special_tokens=True, skip_prompt=True)
-    return response
+    # # Decode the output to return the final response
+    # response = processor.decode(outputs[0], skip_special_tokens=True, skip_prompt=True)
+
+
+    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
+
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+    generated_text = ""
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    
+    for new_text in streamer:
+        buffer += new_text
+        # generated_text_without_prompt = buffer
+        # # time.sleep(0.01)
+        # yield buffer
+    
+    return buffer
 
 
 # Define the Gradio interface