import gradio as gr
from transformers import pipeline
import torch

def load_model():
    classifier = pipeline(
        "text-classification",
        model="suronek/toxic-detector-roberta",
        device=0 if torch.cuda.is_available() else -1
    )
    return classifier

print("Загружаем модель...")
toxic_classifier = load_model()
print("Модель загружена!")

def detect_toxicity(text):
    if not text.strip():
        return "Пожалуйста, введите текст для анализа."
    
    try:
        # Получаем предсказание от модели
        results = toxic_classifier(text)
        
        # Форматируем результат
        if isinstance(results, list):
            result = results[0]
        else:
            result = results
            
        label = result['label']
        confidence = result['score']
        
        # Интерпретируем результат в зависимости от модели
        if label.upper() in ['TOXIC', 'TOXICITY', "LABEL_0"]:
            status = "🔴 ТОКСИЧНЫЙ"
            explanation = "Текст содержит токсичный контент"
        elif label.upper() in ['NON_TOXIC', 'NON-TOXIC', "LABEL_1"]:
            status = "🟢 БЕЗОПАСНЫЙ"
            explanation = "Текст не содержит токсичного контента"
        else:
            status = f"📊 {label}"
            explanation = "Результат анализа"
        
        return f"""
        **Результат анализа:**
        
        **Статус:** {status}
        **Уверенность:** {confidence:.2%}
        **Объяснение:** {explanation}
        
        **Исходный текст:** "{text}"
        """
        
    except Exception as e:
        return f"❌ Ошибка при анализе: {str(e)}"

# Создаем интерфейс Gradio
with gr.Blocks(title="RoBERTa Toxic Detector", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛡️ Детектор токсичности на основе RoBERTa")
    gr.Markdown("Введите текст для анализа на токсичность с помощью нейронной сети RoBERTa")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Введите текст для анализа",
                placeholder="Напишите что-нибудь...",
                lines=3,
                max_lines=10
            )
            
            analyze_btn = gr.Button("🔍 Анализировать", variant="primary")
            
        with gr.Column():
            output = gr.Markdown(label="Результат анализа")
    
    # Примеры для тестирования
    gr.Examples(
        examples=[
            ["Привет! Как дела? Хорошего дня!"],
            ["Ты идиот и дурак!"],
            ["Мне нравится этот фильм, очень интересный сюжет"],
            ["Я ненавижу всех вокруг"],
            ["Спасибо за помощь, вы очень добры"]
        ],
        inputs=text_input,
        outputs=output,
        fn=detect_toxicity,
        cache_examples=True
    )
    
    # Обработчик события
    analyze_btn.click(
        fn=detect_toxicity,
        inputs=text_input,
        outputs=output
    )
    
    # Также можно анализировать по нажатию Enter
    text_input.submit(
        fn=detect_toxicity,
        inputs=text_input,
        outputs=output
    )

    demo.launch()