Update app.py

app.py

@@ -11,7 +11,8 @@ models = {
     "ruSpamNS_v13": "NeuroSpaceX/ruSpamNS_v13",
     "ruSpamNS_big": "NeuroSpaceX/ruSpamNS_big",
     "ruSpamNS_small": "NeuroSpaceX/ruSpamNS_small",
-    "ruSpamNS_v14": "NeuroSpaceX/ruSpamNS_v14"  # Добавлена новая модель
+    "ruSpamNS_v14": "NeuroSpaceX/ruSpamNS_v14",
+    "ruSpamNS_v14_multiclass": "NeuroSpaceX/ruSpamNS_v14_multiclass"
 }
 
 tokenizers = {name: AutoTokenizer.from_pretrained(path, use_auth_token=TOKEN) for name, path in models.items()}
@@ -38,16 +39,23 @@ def classify_text(text, model_choice):
     
     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}%)"
+        if model_choice == "ruSpamNS_v14_multiclass":
+            probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+            labels = ["НЕ СПАМ", "СПАМ", "НЕДВИЖИМОСТЬ", "ВАКАНСИИ"]
+            predicted_index = probabilities.argmax()
+            predicted_label = labels[predicted_index]
+            confidence = probabilities[predicted_index] * 100
+            return f"{predicted_label} (вероятность: {confidence:.2f}%)"
+        else:
+            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="Введите текст..."),
-        gr.Radio(["ruSpamNS_v13", "ruSpamNS_big", "ruSpamNS_small", "ruSpamNS_v14"], label="Выберите модель")  # Добавлена новая модель в выбор
+        gr.Radio(list(models.keys()), label="Выберите модель")
     ],
     outputs="text",
     title="ruSpamNS - Проверка на спам",