Adding streaming

app.py

@@ -1,27 +1,27 @@
 import gradio as gr
-import time
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 
-# Load the TinyLlama text generation pipeline
-pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+# Load the model and tokenizer
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Define the inference function
+# Create a text generation pipeline
+pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
+
+# Define the inference function with streaming
 def generate_text(prompt):
-    start_time = time.time()
-    results = pipe(prompt, max_length=100, num_return_sequences=1)
-    end_time = time.time()
-    response_time = end_time - start_time
-    return results[0]['generated_text'], f"{response_time:.2f} seconds"
+    output = ""
+    for token in pipeline(prompt, max_length=100, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95):
+        output += token['generated_text']
+        yield output  # Stream the output
 
 # Create a Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
-    outputs=[
-        gr.Textbox(label="Generated Text"),
-        gr.Textbox(label="Response Time")
-    ],
-    title="TinyLlama Text Generation"
+    outputs=gr.Textbox(label="Generated Text"),
+    live=True  # Enable live streaming
 )
 
 # Launch the interface