import gradio as gr
from transformers import pipeline

# Load a lightweight model (GPT-2)
generator = pipeline("text-generation", model="gpt2")

def generate_completion(prompt):
    outputs = generator(prompt, max_length=50, do_sample=True, temperature=0.9)
    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()