import gradio as gr
from transformers import pipeline

# Загружаем модель прямо из вашего нового репозитория на Hugging Face
MODEL_NAME = "ENTUM-AI/distilbert-clickbait-classifier"

try:
    classifier = pipeline("text-classification", model=MODEL_NAME)
except Exception as e:
    print(f"Error loading model: {e}")
    classifier = None

def predict(text):
    if not text.strip():
        return "Please enter a headline."
    
    if classifier is None:
        return "Model has not loaded yet or an error occurred."
        
    result = classifier(text)[0]
    label = result['label']
    score = result['score']
    
    # Форматируем красивый вывод
    if label == "Clickbait":
        return f"🚨 CLICKBAIT! (Confidence: {score:.1%})"
    else:
        return f"📰 NORMAL NEWS (Confidence: {score:.1%})"

# Настраиваем красивый интерфейс Gradio
theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
)

with gr.Blocks(theme=theme, title="Clickbait Detector 🎣") as demo:
    gr.Markdown(
        """
        # 🎣 Clickbait Headline Detector
        
        This model, based on **DistilBERT**, predicts whether a news headline or article title is "clickbait".
        It was trained on tens of thousands of real media headlines.
        
        *Enter any English headline below to check it!*
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                label="Enter headline", 
                placeholder="Example: 10 Bizarre Facts About Apples...",
                lines=3
            )
            submit_btn = gr.Button("Check Headline 🔍", variant="primary")
            
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Model Verdict", 
                lines=3,
                interactive=False
            )
            
    # Примеры для быстрого тестирования
    gr.Examples(
        examples=[
            ["10 Bizarre Facts About Apples That Will BLOW YOUR MIND! 🍎🤯"],
            ["Apple releases new quarterly earnings report showing 5% growth."],
            ["You'll Never Guess What Happened Next..."],
            ["Federal Reserve announces increase in interest rates by 0.25%"]
        ],
        inputs=input_text
    )
    
    submit_btn.click(fn=predict, inputs=input_text, outputs=output_text)

# Запускаем приложение
demo.launch()