Upload inference.py

inference.py

@@ -0,0 +1,120 @@
+
+\"\"\"
+inference.py  — Deep Dive v2: Self-Contained Inference Module
+Model       : itsLu/mentalbert-v5-deep-dive-v2
+Task        : Binary Depression vs Suicidal re-ranker (second-tier specialist)
+
+⚠️  WARNING: This model MUST NOT be used standalone.
+    It is invoked by Quick Vibe (itsLu/mentalbert-v5-source-aware) only when:
+      (a) Quick Vibe's top1 ∈ {Depression, Suicidal} AND margin < 0.20, OR
+      (b) Quick Vibe abstains.
+    It has NO coverage for: Normal, Anxiety, Stress, Bipolar,
+    Personality Disorder, Directed Aggression.
+
+Usage:
+    from inference import DeepDiveV2
+    clf = DeepDiveV2.from_hub()
+    result = clf.predict("I don't want to be here anymore.")
+    # {'label': 'Suicidal', 'p_suicidal': 0.91, 'crisis_evidence_found': True,
+    #  'crisis_tokens_matched': ['want to die']}
+\"\"\"
+
+import json, torch, torch.nn as nn
+from typing import Any, Dict, List
+from transformers    import BertModel, BertTokenizerFast
+from huggingface_hub import hf_hub_download
+
+HF_REPO  = "itsLu/mentalbert-v5-deep-dive-v2"
+CLASSES  = ["Depression", "Suicidal"]
+
+
+class CrisisEvidenceMentalBERT(nn.Module):
+    \"\"\"BertModel + crisis-evidence token max-pooling + binary head.\"\"\"
+    def __init__(self, model_name: str, hidden_size: int = 768):
+        super().__init__()
+        self.encoder    = BertModel.from_pretrained(model_name)
+        self.dropout    = nn.Dropout(0.2)
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_size * 2, 256),
+            nn.GELU(),
+            nn.Dropout(0.2),
+            nn.Linear(256, 2),
+        )
+
+    def forward(self, input_ids, attention_mask, crisis_mask):
+        h_seq        = self.encoder(input_ids=input_ids,
+                                    attention_mask=attention_mask).last_hidden_state
+        h_cls        = h_seq[:, 0, :]
+        masked       = h_seq.masked_fill(~crisis_mask.unsqueeze(-1), float('-inf'))
+        any_crisis   = crisis_mask.any(dim=1)
+        h_crisis_raw = masked.max(dim=1).values
+        h_crisis     = torch.where(any_crisis.unsqueeze(-1), h_crisis_raw, h_cls)
+        return self.classifier(self.dropout(torch.cat([h_cls, h_crisis], dim=-1)))
+
+
+class DeepDiveV2:
+    \"\"\"High-level inference wrapper — load from Hub, call .predict(text).\"\"\"
+
+    def __init__(self, model: CrisisEvidenceMentalBERT, tokenizer: BertTokenizerFast,
+                 crisis_keywords: List[str], threshold: float, max_len: int = 256):
+        self.model           = model.eval()
+        self.tokenizer       = tokenizer
+        self.crisis_keywords = [k.lower() for k in crisis_keywords]
+        self.threshold       = threshold
+        self.max_len         = max_len
+        self.device          = next(model.parameters()).device
+
+    @classmethod
+    def from_hub(cls, repo_id: str = HF_REPO, device: str = "cpu") -> "DeepDiveV2":
+        tokenizer = BertTokenizerFast.from_pretrained(repo_id)
+        cfg_path  = hf_hub_download(repo_id=repo_id, filename="inference_config.json")
+        kw_path   = hf_hub_download(repo_id=repo_id, filename="crisis_keywords.json")
+        clf_path  = hf_hub_download(repo_id=repo_id, filename="classifier.pt")
+        with open(cfg_path) as f: cfg = json.load(f)
+        with open(kw_path)  as f: kws = json.load(f)["crisis_keywords"]
+        model = CrisisEvidenceMentalBERT(model_name=repo_id)
+        model.classifier.load_state_dict(torch.load(clf_path, map_location=device))
+        model.to(device)
+        return cls(model=model, tokenizer=tokenizer, crisis_keywords=kws,
+                   threshold=cfg["threshold"], max_len=cfg["max_len"])
+
+    def _crisis_mask(self, text: str, offsets: list) -> "torch.BoolTensor":
+        tl = text.lower()
+        spans = []
+        for kw in self.crisis_keywords:
+            s = 0
+            while (i := tl.find(kw, s)) >= 0:
+                spans.append((i, i + len(kw))); s = i + 1
+        mask = torch.zeros(self.max_len, dtype=torch.bool)
+        for ti, (s, e) in enumerate(offsets):
+            if s == e == 0: continue
+            for cs, ce in spans:
+                if s < ce and e > cs: mask[ti] = True; break
+        return mask
+
+    def predict(self, text: str) -> Dict[str, Any]:
+        \"\"\"
+        Returns:
+          label                 : 'Depression' | 'Suicidal'
+          p_suicidal            : float  — P(Suicidal)
+          crisis_evidence_found : bool   — any crisis keyword matched
+          crisis_tokens_matched : list[str]
+        \"\"\"
+        enc  = self.tokenizer(text, max_length=self.max_len, truncation=True,
+                              padding="max_length", return_offsets_mapping=True,
+                              return_tensors="pt")
+        ids  = enc["input_ids"].to(self.device)
+        amsk = enc["attention_mask"].to(self.device)
+        offs = enc["offset_mapping"].squeeze(0).tolist()
+        cmsk = self._crisis_mask(text, offs).unsqueeze(0).to(self.device)
+        matched = [kw for kw in self.crisis_keywords if kw in text.lower()]
+        with torch.no_grad():
+            probs = torch.softmax(self.model(ids, amsk, cmsk), dim=-1).squeeze(0).cpu()
+        p = float(probs[1])
+        return {"label": "Suicidal" if p >= self.threshold else "Depression",
+                "p_suicidal": round(p, 6),
+                "crisis_evidence_found": bool(cmsk.any().item()),
+                "crisis_tokens_matched": matched}
+
+    def predict_batch(self, texts: List[str]) -> List[Dict[str, Any]]:
+        return [self.predict(t) for t in texts]