Update app.py

app.py

@@ -13,7 +13,8 @@ models = {
     "ruSpamNS_small": "NeuroSpaceX/ruSpamNS_small",
     "ruSpamNS_v14": "NeuroSpaceX/ruSpamNS_v14",
     "ruSpamNS_v14_multiclass": "NeuroSpaceX/ruSpamNS_v14_multiclass",
-    "ruSpamNS_v16_multiclass": "NeuroSpaceX/ruSpamNS_v16_multiclass"
+    "ruSpamNS_v16_multiclass": "NeuroSpaceX/ruSpamNS_v16_multiclass",
+    "ruSpamNS_v17_multiclass": "NeuroSpaceX/ruSpamNS_v17_multiclass"
 }
 
 tokenizers = {name: AutoTokenizer.from_pretrained(path, use_auth_token=TOKEN) for name, path in models.items()}
@@ -40,7 +41,7 @@ def classify_text(text, model_choice):
     
     with torch.no_grad():
         outputs = model(input_ids, attention_mask=attention_mask).logits
-        if model_choice == "ruSpamNS_v14_multiclass" or model_choice == "ruSpamNS_v16_multiclass":
+        if model_choice in ["ruSpamNS_v14_multiclass", "ruSpamNS_v16_multiclass", "ruSpamNS_v17_multiclass"]:
             probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]
             labels = ["НЕ СПАМ", "СПАМ", "НЕДВИЖИМОСТЬ", "ВАКАНСИИ", "ПРОДАЖА"]
             predicted_index = probabilities.argmax()