app.py

import gradio as gr
import time
from transformers import pipeline

# Load the TinyLlama text generation pipeline
pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

# Define the inference function
def generate_text(prompt):
    start_time = time.time()
    results = pipe(prompt, max_length=50, 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"),
        gr.Textbox(label="Response Time")
    ],
    title="TinyLlama Text Generation"
)

# Launch the interface
iface.launch()