Upload modeling_chest2vec_labeler.py with huggingface_hub

modeling_chest2vec_labeler.py

@@ -50,6 +50,7 @@ class Chest2VecLabelerConfig(PretrainedConfig):
         instruction: str = "Given the following chest CT report, extract the presence/absence of entities",
         max_len: int = 512,
         default_threshold: float = 0.5,
+        label_hierarchy: Optional[dict] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -62,6 +63,7 @@ class Chest2VecLabelerConfig(PretrainedConfig):
         self.instruction = instruction
         self.max_len = max_len
         self.default_threshold = default_threshold
+        self.label_hierarchy = label_hierarchy or {}
 
 
 @dataclass
@@ -201,35 +203,82 @@ class Chest2VecLabelerModel(PreTrainedModel):
         names = out["labels"]
         return [{names[j]: "positive" for j in range(len(names)) if row[j]} for row in out["positive"]]
 
-    # ---- CheXbert / SRR-BERT-style report-comparison F1 ----
+    # ---- hierarchy roll-up (leaf -> upper -> anatomy), max over children ----
+    def aggregate_hierarchy(self, leaf_prob):
+        """Roll leaf positive-probabilities up to upper and anatomy levels (max over children).
+
+        Mirrors the training-time evaluation: each upper group's score is the max over its
+        child-leaf probabilities; each anatomy score is the max over its upper groups plus the
+        section's `*_others` leaf. Returns (upper_prob, upper_names, anatomy_prob, anatomy_names).
+        """
+        import numpy as np
+        leaf_prob = np.asarray(leaf_prob, dtype=np.float32)
+        H = self.config.label_hierarchy or {}
+        idx = {n: i for i, n in enumerate(self.config.labels)}
+        N = leaf_prob.shape[0]
+        u_names, u_cols, a_names, a_cols = [], [], [], []
+        for anat, groups in H.items():
+            a_names.append(anat)
+            ac = np.full(N, -1.0, dtype=np.float32)
+            for up, leaves in groups.items():
+                u_names.append(up)
+                cols = [idx[l] for l in leaves if l in idx]
+                uc = leaf_prob[:, cols].max(axis=1) if cols else np.zeros(N, dtype=np.float32)
+                u_cols.append(uc)
+                ac = np.maximum(ac, uc)
+            okey = f"{anat}_others"
+            if okey in idx:
+                ac = np.maximum(ac, leaf_prob[:, idx[okey]])
+            a_cols.append(np.maximum(ac, 0.0))
+        import numpy as _np
+        up = _np.column_stack(u_cols) if u_cols else _np.zeros((N, 0), dtype=_np.float32)
+        an = _np.column_stack(a_cols) if a_cols else _np.zeros((N, 0), dtype=_np.float32)
+        return up, u_names, an, a_names
+
+    # ---- CheXbert / SRR-BERT-style report-comparison F1 (leaf / upper / anatomy) ----
     @torch.no_grad()
     def score_reports(self, gt_reports: List[str], pred_reports: List[str], tokenizer=None,
                       threshold: Optional[float] = None, batch_size: int = 16,
-                      max_len: Optional[int] = None, device=None) -> Dict[str, Any]:
+                      max_len: Optional[int] = None, device=None,
+                      levels=("leaf", "upper", "anatomy")) -> Dict[str, Any]:
         """
-        Label both GT and predicted reports, then compute label-level F1 (CheXbert-style).
+        Label both GT and predicted reports, then compute label-agreement F1 (CheXbert-style)
+        at the requested hierarchy levels.
 
-        Treats the labels extracted from `gt_reports` as ground truth and those from
-        `pred_reports` as the prediction. Returns micro / macro / weighted P,R,F1 over the
-        137 labels, plus per-label scores.
+        `gt_reports` labels are treated as truth, `pred_reports` as prediction. For each level
+        in `levels` ("leaf" = 137 labels, "upper" = container groups, "anatomy" = sections),
+        returns micro / macro / weighted precision-recall-F1 plus per-label scores.
         """
         from sklearn.metrics import precision_recall_fscore_support
+        import numpy as np
         if len(gt_reports) != len(pred_reports):
             raise ValueError("gt_reports and pred_reports must have the same length")
-        kw = dict(tokenizer=tokenizer, threshold=threshold, batch_size=batch_size, max_len=max_len, device=device)
-        y_true = self.predict(gt_reports, **kw)["positive"]
-        y_pred = self.predict(pred_reports, **kw)["positive"]
-        names = list(self.config.labels)
-
-        res: Dict[str, Any] = {"n_reports": len(gt_reports),
-                               "threshold": self.config.default_threshold if threshold is None else threshold}
-        for avg in ("micro", "macro", "weighted"):
-            p, r, f, _ = precision_recall_fscore_support(y_true, y_pred, average=avg, zero_division=0)
-            res[avg] = {"precision": float(p), "recall": float(r), "f1": float(f)}
-        p, r, f, s = precision_recall_fscore_support(y_true, y_pred, average=None,
-                                                     labels=list(range(len(names))), zero_division=0)
-        res["per_label"] = {names[j]: {"precision": float(p[j]), "recall": float(r[j]),
-                                       "f1": float(f[j]), "support_gt": int(s[j])} for j in range(len(names))}
+        thr = self.config.default_threshold if threshold is None else threshold
+        kw = dict(tokenizer=tokenizer, batch_size=batch_size, max_len=max_len, device=device)
+        gt_leaf = self.predict_proba(gt_reports, **kw).numpy()
+        pr_leaf = self.predict_proba(pred_reports, **kw).numpy()
+
+        level_inputs = {"leaf": (gt_leaf, pr_leaf, list(self.config.labels))}
+        if "upper" in levels or "anatomy" in levels:
+            gu, un, ga, an = self.aggregate_hierarchy(gt_leaf)
+            pu, _, pa, _ = self.aggregate_hierarchy(pr_leaf)
+            level_inputs["upper"] = (gu, pu, un)
+            level_inputs["anatomy"] = (ga, pa, an)
+
+        res: Dict[str, Any] = {"n_reports": len(gt_reports), "threshold": thr}
+        for lvl in levels:
+            gp, pp, names = level_inputs[lvl]
+            y_true = (gp >= thr).astype(int)
+            y_pred = (pp >= thr).astype(int)
+            block: Dict[str, Any] = {"n_labels": len(names)}
+            for avg in ("micro", "macro", "weighted"):
+                p, r, f, _ = precision_recall_fscore_support(y_true, y_pred, average=avg, zero_division=0)
+                block[avg] = {"precision": float(p), "recall": float(r), "f1": float(f)}
+            p, r, f, s = precision_recall_fscore_support(y_true, y_pred, average=None,
+                                                         labels=list(range(len(names))), zero_division=0)
+            block["per_label"] = {names[j]: {"precision": float(p[j]), "recall": float(r[j]),
+                                             "f1": float(f[j]), "support_gt": int(s[j])} for j in range(len(names))}
+            res[lvl] = block
         return res