Try streaming again

main.py

@@ -98,7 +98,7 @@ def clean_output(raw_prediction: str) -> str:
 
     return predicted_code + "\n" + "\n".join(udt.declaration("") for udt in predicted_udts)
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline #, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextIteratorStreamer #, BitsAndBytesConfig
 from peft import PeftModel
 
 # Define your model and adapter identifiers
@@ -169,29 +169,53 @@ def _predict_sync(inp):
 
 
 def predict(inp):
-    # Stream simple progress messages while running inference in background thread
+    # Stream tokens without moving model weights; only move input tensors to model device
     import threading, time
-    container = {}
-    def target():
+    initial_tokens = len(tokenizer.tokenize(inp))
+    new_tokens = math.floor(initial_tokens * 1.5)
+
+    inp2 = canonicalize_output(inp)
+    # Prepare inputs on CPU then move to model device when generating
+    inputs = tokenizer(inp2, return_tensors="pt")
+
+    # Use TextIteratorStreamer to stream tokens as they're generated
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, timeout=10.0)
+
+    def _generate():
         try:
-            container['result'] = _predict_sync(inp)
+            device = next(model.parameters()).device
+            inputs_device = {k: v.to(device) for k, v in inputs.items()}
+            # Call model.generate which will use streamer to yield tokens; model weights stay on their device
+            model.generate(**inputs_device, max_new_tokens=new_tokens, streamer=streamer)
         except Exception as e:
-            container['exc'] = e
+            try:
+                streamer.end()
+            except:
+                pass
+            # Attach exception to streamer by yielding nothing and storing in container
+            streamer.exception = e
 
-    t = threading.Thread(target=target)
+    t = threading.Thread(target=_generate, daemon=True)
     t.start()
 
-    dots = 0
-    # yield periodic progress messages until inference is done
-    while t.is_alive():
-        dots = (dots + 1) % 4
-        yield "Running inference" + ("." * dots)
-        time.sleep(1)
+    aggregated = ""
+    # Yield incremental aggregated text as tokenizer emits chunks
+    for chunk in streamer:
+        aggregated += chunk
+        yield aggregated
 
     t.join()
-    if 'exc' in container:
-        raise container['exc']
-    yield container.get('result', '')
+
+    # Propagate any generation exception
+    if hasattr(streamer, 'exception'):
+        raise streamer.exception
+
+    # Final cleaning of the aggregated output
+    try:
+        cleaned = clean_output(aggregated)
+        yield cleaned
+    except Exception:
+        yield aggregated
 
 def run():