Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from threading import Thread
+from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
 model_id = "textcleanlm/textclean-4B"
 model = None
@@ -37,19 +38,32 @@ def load_model():
 def clean_text(text):
     model, tokenizer = load_model()
     
-    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+    inputs = tokenizer(text, return_tensors="pt", max_length=4096, truncation=True)
     inputs = {k: v.cuda() for k, v in inputs.items()}
     
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=512,
-            num_beams=4,
-            early_stopping=True
-        )
+    # Enable streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
     
-    cleaned_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return cleaned_text
+    generation_kwargs = dict(
+        **inputs,
+        max_length=4096,
+        num_beams=1,  # Set to 1 for streaming
+        do_sample=True,
+        temperature=1.0,
+        streamer=streamer,
+    )
+    
+    # Run generation in a separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # Yield text as it's generated
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+    
+    thread.join()
 
 iface = gr.Interface(
     fn=clean_text,