ForModelTest / app.py
NeuroSpaceX's picture
Create app.py
4950ce7 verified
raw
history blame
1.05 kB
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()