import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load pre-trained model and tokenizer
model_name = "gpt2"  # You can also try "distilgpt2" for a smaller model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to generate text
def generate_text(prompt, max_length=100, num_return_sequences=1):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs["input_ids"],
        max_length=max_length,
        num_return_sequences=num_return_sequences,
        no_repeat_ngram_size=2,
        top_p=0.95,
        top_k=60,
        temperature=0.7,
    )
    return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]

# Create Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Enter your prompt", placeholder="Type here..."),
        gr.Slider(minimum=50, maximum=200, value=100, label="Max Length"),
        gr.Slider(minimum=1, maximum=5, value=1, label="Number of Outputs")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="GPT-2 Text Generator",
    description="Generate human-like text based on your prompt using GPT-2.",
    theme="compact",
    examples=[["Once upon a time in a land far, far away..."]]
)

iface.launch()