bert.py

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import streamlit as st
import torch


MODEL_PATH = "rubert2"
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model.eval()

# === Streamlit UI ===
st.set_page_config(page_title="Оценка токсичности", layout="centered")
st.title("💬 Оценка токсичности текста")

text = st.text_area("Введите сообщение", "Ты ужасный человек!")
submit = st.button("Проверить токсичность")

if submit and text.strip():
    # Токенизация
    inputs = tokenizer(text, return_tensors="pt", truncation=True)

    # Предсказание
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        score = torch.sigmoid(logits).item()  # степень токсичности

    # Вывод
    st.subheader("Результат:")
    st.write(f"**Степень токсичности:** `{score:.3f}`")

    if score > 0.8:
        st.error("⚠️ Высокая токсичность!")
    elif score > 0.4:
        st.warning("⚠️ Средняя токсичность")
    else:
        st.success("✅ Низкая токсичность")