fix, ig

src/antispam_api/loader.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F  # noqa: N812
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from antispam_api.config import MODEL_NAME
@@ -6,12 +7,15 @@ from antispam_api.config import MODEL_NAME
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(
-    MODEL_NAME, num_labels=1,
-).to(device).eval()
+model = (
+    AutoModelForSequenceClassification
+    .from_pretrained(MODEL_NAME)
+    .to(device)
+    .eval()
+)
 
 
-def get_spam_score(text: str):
+def get_scam_score(text: str) -> float:
     encoding = tokenizer(
         text,
         padding="max_length",
@@ -24,6 +28,6 @@ def get_spam_score(text: str):
 
     with torch.no_grad():
         logits = model(input_ids, attention_mask=attention_mask).logits
-        score = torch.sigmoid(logits).cpu().numpy()[0][0]
+        probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
 
-    return float(score)
+    return float(probs[1])