Stream Tiny Aya text generation

app.py

@@ -12,7 +12,7 @@ os.environ.setdefault("GRADIO_SSR_MODE", "false")
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MODEL_ID = os.environ.get("TINY_AYA_MODEL", "CohereLabs/tiny-aya-global")
 DEFAULT_MAX_TOKENS = int(os.environ.get("TINY_AYA_MAX_TOKENS", "400"))
@@ -59,6 +59,42 @@ def generate(system: str, user: str, max_tokens: int = DEFAULT_MAX_TOKENS, tempe
     return tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
 
 
+@spaces.GPU(duration=120)
+def generate_stream(system: str, user: str, max_tokens: int = DEFAULT_MAX_TOKENS, temperature: float = 0.8):
+    if not user or not user.strip():
+        raise gr.Error("user prompt required")
+    max_tokens = max(1, min(int(max_tokens or DEFAULT_MAX_TOKENS), 1024))
+    temperature = max(0.0, min(float(temperature if temperature is not None else 0.8), 2.0))
+    inputs = tokenizer.apply_chat_template(
+        _messages(system, user),
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    def run():
+        with _lock, torch.inference_mode():
+            model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                do_sample=temperature > 0,
+                temperature=max(temperature, 1e-5),
+                top_p=0.95,
+                pad_token_id=tokenizer.eos_token_id,
+                streamer=streamer,
+            )
+
+    thread = threading.Thread(target=run, daemon=True)
+    thread.start()
+    acc = ""
+    for token in streamer:
+        acc += token
+        yield acc
+    thread.join(timeout=1)
+
+
 with gr.Blocks(title="Tiny Army Tiny Aya ZeroGPU") as demo:
     gr.Markdown("# Tiny Army Tiny Aya ZeroGPU")
     system = gr.Textbox(label="System", lines=5)
@@ -68,6 +104,7 @@ with gr.Blocks(title="Tiny Army Tiny Aya ZeroGPU") as demo:
     btn = gr.Button("Generate")
     out = gr.Textbox(label="Output", lines=10)
     btn.click(generate, inputs=[system, user, max_tokens, temperature], outputs=out, api_name="generate")
+    btn.click(generate_stream, inputs=[system, user, max_tokens, temperature], outputs=out, api_name="generate_stream")
 
 
 if __name__ == "__main__":