app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")

def generate_text(prompt, max_length=100, min_length=20, temperature=1.0):
    # Tokenize the prompt
    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    # Generate text
    output = model.generate(
        input_ids,
        max_length=max_length,
        min_length=min_length,
        num_return_sequences=1,
        temperature=temperature
    )

    # Decode the generated output
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

    return generated_text

# Create the Gradio interface
iface = gr.Interface(
    generate_text,
    [
        gr.Textbox(lines=5, label="Prompt"),
        gr.Slider(0, 2048, 100, label="Max Length"),
        gr.Slider(0, 2048, 20, label="Min Length"),
        gr.Slider(0.1, 2.0, 1.0, label="Temperature"),
    ],
    "textbox",
    title="TinyLlama Text Generator",
    live=True,  # Enable live updates as the user types
)

# Launch the Gradio app
iface.launch()