import gradio as gr
from transformers import pipeline

# Load the GPT-2 next-line prediction model
model_name = "AventIQ-AI/gpt2-lmheadmodel-next-line-prediction-model"
generator = pipeline("text-generation", model=model_name)

def generate_next_line(prompt):
    """Generates the next line of text based on the input prompt."""
    if not prompt.strip():
        return "⚠️ Please enter a prompt."

    response = generator(prompt, max_length=len(prompt.split()) + 10, num_return_sequences=1)
    return response[0]["generated_text"]

# Example Inputs
example_prompts = [
    "Once upon a time, a young explorer discovered",
    "The AI revolutionized the way people interacted with",
    "In the distant future, humans and robots lived together in",
    "The detective examined the crime scene and found a"
]

# Create Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📝 Next-Line Prediction with GPT-2")
    gr.Markdown("Enter a sentence, and the model will predict the next line!")

    with gr.Row():
        input_text = gr.Textbox(label="✍️ Enter your text:", placeholder="Once upon a time...")

    generate_button = gr.Button("🔮 Generate Next Line")
    output_text = gr.Textbox(label="📜 Generated Text:")

    gr.Markdown("### ✨ Example Inputs")
    example_buttons = [gr.Button(example) for example in example_prompts]

    for btn in example_buttons:
        btn.click(fn=lambda text=btn.value: text, outputs=input_text)

    generate_button.click(generate_next_line, inputs=input_text, outputs=output_text)

# Launch the Gradio app
demo.launch()