import gradio as gr
from transformers import pipeline

# Load the GPT-2 News Generation Model
model_name = "AventIQ-AI/gpt2-news-article-generation"
generator = pipeline("text-generation", model=model_name)

# Predefined headline suggestions
headline_suggestions = [
    "Breaking: Stock Market Hits Record High",
    "Scientists Discover New Treatment for Alzheimer's",
    "Tech Giants Compete in AI Race",
    "Severe Weather Warnings Issued Across the Country",
    "New Law Passed to Improve Cybersecurity Standards"
]

def generate_news_article(headline, max_length=250):
    """Generate a news article based on the given headline."""
    response = generator(headline, max_length=max_length, num_return_sequences=1)
    return response[0]["generated_text"]

# Create Gradio UI
with gr.Blocks(theme="default") as demo:
    gr.Markdown("## 📰 GPT-2 News Article Generator")
    gr.Markdown("🔍 Enter a **news headline**, and the model will generate a news article based on it.")

    headline_input = gr.Textbox(placeholder="Enter a news headline...", label="News Headline")
    suggestion_dropdown = gr.Dropdown(choices=headline_suggestions, label="💡 Select a Sample Headline (Optional)")
    generate_button = gr.Button("📝 Generate Article")
    output_box = gr.Textbox(label="Generated News Article", interactive=False)

    # Function to update input field with selected suggestion
    def update_headline(suggestion):
        return suggestion

    suggestion_dropdown.change(update_headline, inputs=[suggestion_dropdown], outputs=[headline_input])
    generate_button.click(generate_news_article, inputs=[headline_input], outputs=[output_box])

demo.launch()