import torch
import gradio as gr
from transformers import AutoTokenizer

from model import BertClassifier

WEIGHTS = "bert_toxic_classifier.pt"
MAX_LEN = 128
# Индекс класса "токсичный" в твоём датасете (Mnwa/russian-toxic).
# Если демо путает токсичное с нетоксичным — поменяй на 0.
TOXIC_INDEX = 0

device = torch.device("cpu")
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")

model = BertClassifier(n_classes=2)
model.load_state_dict(torch.load(WEIGHTS, map_location=device))
model.to(device).eval()


@torch.no_grad()
def predict(text: str):
    if not text or not text.strip():
        return {"Введите текст": 1.0}
    enc = tokenizer(
        text,
        max_length=MAX_LEN,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    logits = model(enc["input_ids"].to(device), enc["attention_mask"].to(device))
    probs = torch.softmax(logits, dim=1)[0]
    toxic = probs[TOXIC_INDEX].item()
    return {"☣️ Токсичный": toxic, "✅ Нетоксичный": 1.0 - toxic}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Введите комментарий на русском…", label="Текст комментария"),
    outputs=gr.Label(num_top_classes=2, label="Оценка токсичности"),
    title="🛡️ Классификатор токсичности русского текста",
    description=(
        "Fine-tuning **cointegrated/rubert-tiny2** на датасете ~300k комментариев. "
        "Accuracy на отложенном тесте — 96%. "
        "Код: https://github.com/TeGPro/russian-toxic-classifier"
    ),
    examples=[
        ["Спасибо большое, очень помогло, всё заработало!"],
        ["Ты вообще читать умеешь? Бесполезный совет, только время потратил."],
        ["Отличная статья, давно искал такое объяснение."],
    ],
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch()