Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 import torch
 from dataclasses import dataclass
-from transformers import AutoTokenizer, PretrainedConfig, pipeline, GenerationConfig
+from transformers import AutoTokenizer, PretrainedConfig, GenerationConfig, TextIteratorStreamer
 from optimum.onnxruntime import ORTModelForCausalLM
 import onnx
 import logging
+from threading import Thread
 
 logging.basicConfig(level=logging.INFO)
 
@@ -57,7 +58,6 @@ class Sam3Config(PretrainedConfig):
         self.input_modality = input_modality
         self.head_type = head_type
         self.version = version
-        
         self.hidden_size = self.d_model
         self.num_attention_heads = self.n_heads
 
@@ -77,36 +77,49 @@ except Exception as e:
     logging.error(f"Failed to load ONNX model: {e}")
     raise e
 
-# Define a function to generate text
-def generate_text(prompt, max_length=128, temperature=0.8, top_k=60, top_p=0.9):
+# -----------------------------------------------------------------------------
+# Streaming Generation Function
+# -----------------------------------------------------------------------------
+def generate_text_stream(prompt, max_length, temperature, top_k, top_p):
+    """
+    This function acts as a generator to stream text.
+    It yields each new token as it's generated by the model.
+    """
+    # Create a streamer to iterate over the generated tokens
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # Prepare the generation inputs
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+
     # Set generation parameters within a GenerationConfig object
-    # We set use_cache=False here to bypass the onnx export issue
+    # We explicitly set use_cache=False to avoid the ONNX export bug
     gen_config = GenerationConfig(
         max_length=max_length,
         temperature=temperature,
         top_k=top_k,
         top_p=top_p,
         do_sample=True,
-        use_cache=False, 
+        use_cache=False,
     )
 
-    gen_pipeline = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=device,
+    # Create a thread to run the generation in the background
+    generation_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        generation_config=gen_config,
     )
-    
-    # Pass all generation parameters to the pipeline
-    generated_text = gen_pipeline(
-        prompt,
-        **gen_config.to_dict() 
-    )
-    return generated_text[0]["generated_text"]
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # Yield each token from the streamer as it is generated
+    for new_text in streamer:
+        yield new_text
 
-# Create and launch the Gradio interface
+# -----------------------------------------------------------------------------
+# Gradio Interface
+# -----------------------------------------------------------------------------
 demo = gr.Interface(
-    fn=generate_text,
+    fn=generate_text_stream,
     inputs=[
         gr.Textbox(label="Prompt", lines=2),
         gr.Slider(minimum=10, maximum=512, value=128, label="Max Length"),
@@ -115,8 +128,8 @@ demo = gr.Interface(
         gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top P"),
     ],
     outputs="text",
-    title="SmilyAI Sam 3.0-2 ONNX Text Generation",
-    description="A simple API and UI for text generation using the ONNX version of Sam 3.0-2."
+    title="SmilyAI Sam 3.0-2 ONNX Text Generation (Streaming)",
+    description="A simple API and UI for text generation using the ONNX version of Sam 3.0-2, with streaming output.",
 )
 
 demo.launch()