import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import os

MODEL_NAME = "NeuroSpaceX/ruSpamNS"
TOKEN = os.getenv("HF_TOKEN")  # Читаем токен из переменной окружения

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)

def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=1).item()
    return "СПАМ" if prediction == 1 else "НЕ СПАМ"

iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=3, placeholder="Введите текст..."),
    outputs="text",
    title="ruSpamNS - Проверка на спам",
    description="Введите текст, чтобы проверить, является ли он спамом."
)

iface.launch()