Reverting back to original version

app.py

@@ -1,30 +1,27 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
-import asyncio
+import time
+from transformers import pipeline
 
-# 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)
+# Load the TinyLlama text generation pipeline
+pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
-# Create a text generation pipeline
-pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
-
-# Define the inference function with streaming
-async def generate_text(prompt):
-    loop = asyncio.get_event_loop()
-    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
-        await loop.run_in_executor(None, lambda: None)  # Yield control to the event loop
+# Define the inference function
+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"
 
 # 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"),
-    live=True  # Enable live streaming
+    outputs=[
+        gr.Textbox(label="Generated Text"),
+        gr.Textbox(label="Response Time")
+    ],
+    title="TinyLlama Text Generation"
 )
 
 # Launch the interface