import gradio as gr
from transformers import pipeline

# Use a lighter model: distilgpt2
generator = pipeline("text-generation", model="distilgpt2")

def generate_completion(prompt):
    outputs = generator(
        prompt,
        max_length=50,
        do_sample=True,
        temperature=1.1,  # slightly higher for more randomness
        top_k=50,
        top_p=0.95
    )
    generated_text = outputs[0]['generated_text']
    # Remove the original prompt from the returned completion
    return generated_text[len(prompt):].strip()

iface = gr.Interface(
    fn=generate_completion,
    inputs=gr.Textbox(lines=1, placeholder="Enter your text here..."),
    outputs="text"
)

iface.launch()