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

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)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

def clean_text(text):
    # Удаляем эмодзи
    text = emoji.replace_emoji(text, replace='')

    # Удаляем цифры и символы, кроме букв, пробела, точки и запятой
    text = re.sub(r'[^a-zA-Zа-яА-ЯёЁ .,]', '', text, flags=re.UNICODE)

    # Приводим текст в нижний регистр
    text = text.lower()

    # Делаем первую букву заглавной
    text = text.capitalize()

    # Убираем лишние пробелы
    text = re.sub(r'\s+', ' ', text).strip()

    return text

def classify_text(text):
    message = clean_text(text)
    encoding = tokenizer(message, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask).logits
        prediction = torch.sigmoid(outputs).cpu().numpy()[0][0]
    label = "СПАМ" if prediction >= 0.5 else "НЕ СПАМ"
    return f"{label} (вероятность: {prediction*100:.2f}%)"

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

iface.launch()