Update model_inference.py

model_inference.py

@@ -4,54 +4,44 @@ try:
     from transformers import AutoModelForSequenceClassification, AutoTokenizer
     import torch
 except ImportError:
-    AutoModelForSequenceClassification = None  # type: ignore
-    AutoTokenizer = None  # type: ignore
-    torch = None  # type: ignore
+    AutoModelForSequenceClassification = None
+    AutoTokenizer = None
+    torch = None
 
 
 class ThreatModel:
     """
-    Wraps a transformer classifier for threat categorization.
-
-    If `transformers` or `torch` are not installed, this class will gracefully
-    degrade and simply return empty probability lists instead of crashing.
+    Transformer wrapper. If transformers is not installed,
+    falls back to dummy mode and returns empty probabilities.
     """
-
-    def __init__(self, model_path: str = "bert-base-chinese", device: Optional[str] = None):
+    def __init__(self, model_path: str, device: Optional[str] = None):
         self.available = AutoModelForSequenceClassification is not None and torch is not None
         self.model = None
         self.tokenizer = None
         self.device = "cpu"
 
         if not self.available:
-            # No transformers / torch in the environment; operate in dummy mode.
             return
 
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")  # type: ignore[attr-defined]
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)  # type: ignore[call-arg]
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)  # type: ignore[call-arg]
-        self.model.to(self.device)  # type: ignore[union-attr]
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
+        self.model.to(self.device)
 
     def predict_proba(self, text: str) -> List[float]:
-        """
-        Return a list of probabilities per class.
-
-        If the model is not available (e.g. transformers not installed),
-        returns an empty list and lets the caller decide how to handle it.
-        """
         if not self.available or self.model is None or self.tokenizer is None:
             return []
 
-        inputs = self.tokenizer(  # type: ignore[union-attr]
+        inputs = self.tokenizer(
             text,
             return_tensors="pt",
             truncation=True,
             padding=True
-        ).to(self.device)  # type: ignore[union-attr]
+        ).to(self.device)
 
-        with torch.no_grad():  # type: ignore[union-attr]
-            outputs = self.model(**inputs)  # type: ignore[operator]
-            logits = outputs.logits  # type: ignore[union-attr]
-            probs = torch.softmax(logits, dim=-1).cpu().tolist()[0]  # type: ignore[union-attr]
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=-1).cpu().tolist()[0]
 
         return probs