Add streaming support and increase max tokens to 20000

- Implement token streaming using TextIteratorStreamer
- Increase max_new_tokens slider from 1024 to 20000
- Convert generation function to generator for real-time token output
- Add progress indicator during generation

app.py

@@ -7,7 +7,8 @@ from typing import Any, Dict, List, Tuple
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoTokenizer, pipeline, BitsAndBytesConfig
+from transformers import AutoTokenizer, pipeline, BitsAndBytesConfig, TextIteratorStreamer
+from threading import Thread
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if not HF_TOKEN:
@@ -174,7 +175,7 @@ def format_validation_message(ok: bool, issues: List[str]) -> str:
 
 
 @spaces.GPU(duration=600)
-def generate_router_plan(
+def generate_router_plan_streaming(
     user_task: str,
     context: str,
     acceptance: str,
@@ -185,12 +186,15 @@ def generate_router_plan(
     max_new_tokens: int,
     temperature: float,
     top_p: float,
-) -> Tuple[str, Dict[str, Any], str, str]:
+):
+    """Generator function for streaming token output."""
     if not user_task.strip():
-        raise gr.Error("User task is required.")
+        yield "", {}, "❌ User task is required.", ""
+        return
     
     if model_choice not in MODELS:
-        raise gr.Error(f"Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}")
+        yield "", {}, f"❌ Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}", ""
+        return
 
     try:
         prompt = build_router_prompt(
@@ -203,16 +207,43 @@ def generate_router_plan(
         )
 
         generator = load_pipeline(model_choice)
-        result = generator(
-            prompt,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-        )[0]["generated_text"]
-
-        completion = result[len(prompt) :].strip() if result.startswith(prompt) else result.strip()
-
+        
+        # Get the underlying model and tokenizer
+        model = generator.model
+        tokenizer = generator.tokenizer
+        
+        # Set up streaming
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        
+        # Prepare inputs
+        inputs = tokenizer(prompt, return_tensors="pt")
+        if hasattr(model, 'device'):
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        elif torch.cuda.is_available():
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+        
+        # Start generation in a separate thread
+        generation_kwargs = {
+            **inputs,
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "do_sample": True,
+            "streamer": streamer,
+        }
+        
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream tokens
+        completion = ""
+        for new_text in streamer:
+            completion += new_text
+            yield completion, {}, "🔄 Generating...", prompt
+        
+        # Final processing after streaming completes
+        thread.join()
+        
         try:
             json_block = extract_json_from_text(completion)
             plan = json.loads(json_block)
@@ -221,11 +252,12 @@ def generate_router_plan(
         except Exception as exc:
             plan = {}
             validation_msg = f"❌ JSON parsing failed: {exc}"
-
-        return completion, plan, validation_msg, prompt
+        
+        yield completion, plan, validation_msg, prompt
+        
     except Exception as exc:
         error_msg = f"❌ Generation failed: {str(exc)}"
-        return "", {}, error_msg, ""
+        yield "", {}, error_msg, ""
 
 
 def clear_outputs():
@@ -284,7 +316,7 @@ def build_ui():
                     placeholder="Comma-separated e.g. calculus, optimization, python",
                     value="calculus, optimization, python",
                 )
-                max_new_tokens = gr.Slider(256, 1024, value=640, step=32, label="Max New Tokens")
+                max_new_tokens = gr.Slider(256, 20000, value=640, step=32, label="Max New Tokens")
                 temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature")
                 top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
 
@@ -298,7 +330,7 @@ def build_ui():
         prompt_view = gr.Textbox(label="Full Prompt", lines=10)
 
         generate_btn.click(
-            generate_router_plan,
+            generate_router_plan_streaming,
             inputs=[
                 user_task,
                 context,
@@ -312,6 +344,7 @@ def build_ui():
                 top_p,
             ],
             outputs=[raw_output, plan_json, validation_msg, prompt_view],
+            show_progress="full",
         )
 
         clear_btn.click(fn=clear_outputs, outputs=[raw_output, plan_json, validation_msg, prompt_view])