Update app.py

app.py

@@ -2,19 +2,34 @@ import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import os
+import re
 
 MODEL_NAME = "NeuroSpaceX/ruSpamNS"
-TOKEN = os.getenv("HF_TOKEN")  # Читаем токен из переменной окружения
+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 = text.strip()
+    text = text.replace('\n', ' ')
+    text = re.sub(r'[^\w\s,.!?]', '', text, flags=re.UNICODE)
+    text = re.sub(r'[!?]', '', text)
+    return text.lower()
+
 def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    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(**inputs)
-    prediction = torch.argmax(outputs.logits, dim=1).item()
-    return "СПАМ" if prediction == 1 else "НЕ СПАМ"
+        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,
@@ -24,4 +39,4 @@ iface = gr.Interface(
     description="Введите текст, чтобы проверить, является ли он спамом."
 )
 
-iface.launch()
+iface.launch()