Update model_inference.py

model_inference.py

@@ -1,30 +1,57 @@
-# model_inference.py
+from typing import List, Optional
+
+try:
+    from transformers import AutoModelForSequenceClassification, AutoTokenizer
+    import torch
+except ImportError:
+    AutoModelForSequenceClassification = None  # type: ignore
+    AutoTokenizer = None  # type: ignore
+    torch = None  # type: ignore
 
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-import torch
 
 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.
     """
 
-    def __init__(self, model_path="bert-base-chinese", device=None):
-        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 __init__(self, model_path: str = "bert-base-chinese", 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]
+
+    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 []
 
-    def predict(self, text):
-        inputs = self.tokenizer(
+        inputs = self.tokenizer(  # type: ignore[union-attr]
             text,
             return_tensors="pt",
             truncation=True,
             padding=True
-        ).to(self.device)
+        ).to(self.device)  # type: ignore[union-attr]
 
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            logits = outputs.logits
-            probs = torch.softmax(logits, dim=-1).cpu().tolist()[0]
+        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]
 
-        return probs  # list of probabilities per class
+        return probs