chest2vec CT report labeler (0.6B): self-contained AutoModel, 137-leaf ternary, CheXbert-style report F1

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,130 @@
+---
+license: cc-by-nc-sa-4.0
+language:
+- en
+base_model:
+- Qwen/Qwen3-Embedding-0.6B
+datasets:
+- chest2vec/chest2vec_labels
+pipeline_tag: text-classification
+library_name: transformers
+tags:
+- radiology
+- chest-ct
+- report-labeling
+- multi-label
+- ct-rate
+- chexbert-style-f1
+---
+
+# chest2vec CT Report Labeler (0.6B)
+
+A weakly-supervised **multi-label classifier** that reads a free-text **chest-CT report** and
+predicts a **137-leaf chest-imaging taxonomy**, with a **ternary** status per label
+(*negative / uncertain / positive*).
+
+It also provides a **CheXbert / SRR-BERT-style report-comparison F1**: label a list of
+ground-truth reports and a list of generated/predicted reports, then score them against each
+other (micro / macro / weighted F1) — useful for evaluating radiology report generation.
+
+- **Base architecture:** [`Qwen/Qwen3-Embedding-0.6B`](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B) (Apache-2.0)
+- **Adaptation:** LoRA (r=16, α=32) **merged into the weights** + last-token (EOS) pooling + L2-norm + a linear ternary head (`1024 → 137 × 3`)
+- **Self-contained:** the full model (encoder + head) ships in `model.safetensors`. Loading does **not** download Qwen3-Embedding weights — the architecture is rebuilt from the bundled config and our weights are loaded in. Tokenizer is bundled too.
+- **Params:** ~596M · weights in float32
+- **Training labels:** [`chest2vec/chest2vec_labels`](https://huggingface.co/datasets/chest2vec/chest2vec_labels) (revised CT-RATE, 137-leaf taxonomy)
+
+## Label space
+
+137 leaf labels over 9 chest-CT sections (Lungs & Airways, Pleura, Mediastinum & Hila,
+Cardiovascular, Chest Wall, Bones/Spine, Upper Abdomen, Lower Neck, Others). The exact list is
+in `config.json` (`labels`); full definitions and per-split counts are in the dataset's
+[`LABEL_HIERARCHY.md`](https://huggingface.co/datasets/chest2vec/chest2vec_labels/blob/main/LABEL_HIERARCHY.md).
+
+**Ternary head** — `softmax(logits, dim=-1)` over class indices `[0, 1, 2]`:
+
+| class index | meaning | value |
+|---:|---|---:|
+| 0 | negative | 0 |
+| 1 | uncertain | -1 |
+| 2 | positive | 1 |
+
+A label is reported **positive** when `P(class=2) ≥ threshold` (default **0.5**).
+
+## Usage
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True).eval()
+tok   = AutoTokenizer.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True)
+
+reports = ["Bibasilar atelectasis with small bilateral pleural effusions. Cardiomegaly. Coronary artery calcification."]
+
+# 1) human-readable positive labels per report
+print(model.label_reports(reports, tokenizer=tok))
+# [{'Subsegmental / linear atelectasis': 'positive', 'Pleural effusion': 'positive',
+#   'Cardiomegaly': 'positive', 'Coronary artery calcification': 'positive'}]
+
+# 2) full prediction matrices
+out = model.predict(reports, tokenizer=tok, threshold=0.5, return_ternary=True)
+out["labels"]      # list of 137 label names
+out["proba"]       # [N, 137] P(positive)
+out["positive"]    # [N, 137] in {0,1}
+out["ternary"]     # [N, 137] in {-1,0,1}
+```
+
+### CheXbert / SRR-BERT-style report comparison
+
+Label both ground-truth and predicted reports, then compute label-level F1 (GT-labels treated
+as truth):
+
+```python
+res = model.score_reports(gt_reports, pred_reports, tokenizer=tok)   # equal-length lists
+print(res["micro"]["f1"], res["macro"]["f1"], res["weighted"]["f1"])
+print(res["per_label"]["Pleural effusion"])   # {'precision':..,'recall':..,'f1':..,'support_gt':..}
+
+# or one-liner that loads the model for you:
+from modeling_chest2vec_labeler import report_f1
+report_f1(gt_reports, pred_reports, tokenizer=tok)
+```
+
+Returns `micro`, `macro`, `weighted` precision/recall/F1 over the 137 labels, plus `per_label`.
+
+## Inputs & conventions
+
+- Input is the **findings** text (the model was trained on CT-RATE findings + their refined
+  section-structured form). Reports are formatted internally as
+  `Instruct: Given the following chest CT report, extract the presence/absence of entities\nQuery: <report>`,
+  truncated to **512** tokens, with an EOS token appended and left-padding.
+- For best fidelity, run in float32 (default). bf16 is fine for throughput with negligible drift.
+
+## Evaluation
+
+Direct-paragraph evaluation (`softmax positive-class`, macro-F1 over labels with ≥30 positives —
+the stable headline; the all-labels macro is dragged down by sparse-tail labels):
+
+| Eval set | reports | macro-F1 (≥30) @0.33 | macro-F1 (≥30) @0.5 | macro-AUC (≥30) |
+|---|--:|--:|--:|--:|
+| CT-RATE revised test (public) | 1,464 | **0.875** | 0.866 | 0.989 |
+| sample1000 (private radiologist-reviewed gold) | 1,000 | **0.766** | 0.761 | 0.972 |
+
+On the public test set, a radiologist reviewed **966 reports**: **857 fully accepted, 60
+imperfect-but-acceptable, 49 failed** (94.9% acceptable). AUC barely moves public→private
+(0.989 → 0.972), i.e. label ranking transfers to real clinical reports; the F1 gap is mostly
+threshold/labeling-convention, not domain failure.
+
+## Caveats
+
+- **Weakly supervised** — trained on LLM-generated labels (not radiologist ground truth) derived
+  from report **text**, not images. Not a medical device; not for clinical use.
+- `IVC filter` is in the taxonomy for completeness but had no training positives.
+- `score_reports` measures **label agreement** between two reports as judged by this labeler;
+  like CheXbert-F1 it inherits the labeler's own error modes.
+
+## License & attribution
+
+Released under **CC-BY-NC-SA-4.0**. Built on **`Qwen/Qwen3-Embedding-0.6B`** (Apache-2.0) and
+trained using labels derived from **[CT-RATE](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE)**
+(CC-BY-NC-SA-4.0). **If you use this model, cite the CT-RATE paper** (arXiv:2403.17834) and
+acknowledge Qwen3-Embedding. See the [dataset card](https://huggingface.co/datasets/chest2vec/chest2vec_labels)
+for the full citation.

config.json

@@ -0,0 +1,273 @@
+{
+ "model_type": "chest2vec_labeler",
+ "architectures": [
+  "Chest2VecLabelerModel"
+ ],
+ "auto_map": {
+  "AutoConfig": "modeling_chest2vec_labeler.Chest2VecLabelerConfig",
+  "AutoModel": "modeling_chest2vec_labeler.Chest2VecLabelerModel"
+ },
+ "base_model": "Qwen/Qwen3-Embedding-0.6B",
+ "encoder_config": {
+  "vocab_size": 151669,
+  "max_position_embeddings": 32768,
+  "hidden_size": 1024,
+  "intermediate_size": 3072,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "use_sliding_window": false,
+  "sliding_window": null,
+  "max_window_layers": 28,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "initializer_range": 0.02,
+  "rms_norm_eps": 1e-06,
+  "use_cache": true,
+  "rope_theta": 1000000,
+  "rope_scaling": null,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "layer_types": [
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention",
+   "full_attention"
+  ],
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "architectures": [
+   "Qwen3ForCausalLM"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+   "0": "LABEL_0",
+   "1": "LABEL_1"
+  },
+  "label2id": {
+   "LABEL_0": 0,
+   "LABEL_1": 1
+  },
+  "task_specific_params": null,
+  "problem_type": null,
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": 151643,
+  "pad_token_id": null,
+  "eos_token_id": 151643,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "_name_or_path": "Qwen/Qwen3-Embedding-0.6B",
+  "transformers_version": "4.57.3",
+  "model_type": "qwen3",
+  "tf_legacy_loss": false,
+  "use_bfloat16": false,
+  "output_attentions": false
+ },
+ "hidden_size": 1024,
+ "num_classes_per_label": 3,
+ "labels": [
+  "Consolidation",
+  "Ground-glass opacity (GGO)",
+  "Crazy-paving pattern",
+  "Mosaic attenuation / air-trapping",
+  "Tree-in-bud",
+  "Centrilobular nodules / bronchiolitis pattern",
+  "Pulmonary nodule (solid / PSN / GGN)",
+  "Pulmonary mass (>3 cm)",
+  "Cavitary nodule / mass",
+  "Emphysema",
+  "Bullae / giant bulla",
+  "Pulmonary cysts / cystic lung disease",
+  "Reticulation / intralobular thickening",
+  "Interlobular septal thickening",
+  "Traction bronchiectasis / bronchiolectasis",
+  "Honeycombing",
+  "Parenchymal scarring / fibrotic band",
+  "Tracheal stenosis / malacia",
+  "Tracheal / bronchial wall thickening",
+  "Bronchiectasis",
+  "Mucoid impaction / plugging",
+  "Tracheal diverticulum",
+  "Endotracheal tube",
+  "Tracheostomy tube",
+  "Lobar / segmental atelectasis",
+  "Subsegmental / linear atelectasis",
+  "Post-lobectomy / segmentectomy",
+  "Post-pneumonectomy",
+  "Lung transplant",
+  "Lungs & Airways_others",
+  "Pleural effusion",
+  "Loculated pleural effusion",
+  "Hemothorax",
+  "Chest tube / pleural drain",
+  "Pneumothorax",
+  "Tension pneumothorax",
+  "Pleural thickening",
+  "Pleural plaques",
+  "Pleural nodule / mass",
+  "Pleura_others",
+  "Mediastinal lymphadenopathy",
+  "Hilar lymphadenopathy",
+  "Calcified mediastinal / hilar lymph nodes",
+  "Anterior mediastinal mass",
+  "Middle / posterior mediastinal mass or cyst",
+  "Thymic remnant / hyperplasia",
+  "Esophageal wall thickening / mass",
+  "Hiatal hernia",
+  "Esophageal dilation",
+  "Nasogastric / orogastric tube",
+  "Pneumomediastinum",
+  "Mediastinal hematoma / fluid collection",
+  "Mediastinum & Hila_others",
+  "Cardiomegaly",
+  "Pericardial effusion",
+  "Pericardial thickening / calcification",
+  "Coronary artery calcification",
+  "Coronary stent or bypass graft",
+  "Thoracic aortic calcification",
+  "Thoracic aortic ectasia / dilation (non-aneurysmal)",
+  "Thoracic aortic aneurysm",
+  "Aortic dissection / intramural hematoma",
+  "Main pulmonary artery enlargement",
+  "Pulmonary embolism",
+  "Aortic valve calcification",
+  "Mitral annular calcification",
+  "Pacemaker / ICD leads",
+  "Central venous catheter / PICC",
+  "LVAD / other cardiac assist device",
+  "Cardiovascular_others",
+  "Chest wall soft tissue edema / hematoma",
+  "Subcutaneous emphysema",
+  "Chest wall mass",
+  "Post-thoracotomy change",
+  "Chest wall tumor invasion",
+  "Chest Wall_others",
+  "Acute rib fracture",
+  "Non-acute / healed rib fracture",
+  "Sternal fracture",
+  "Vertebral compression fracture",
+  "Degenerative spine changes",
+  "Osteolytic bone lesion",
+  "Osteosclerotic bone lesion",
+  "Mixed osteolytic-osteosclerotic lesion",
+  "Osteopenia",
+  "Scoliosis / kyphosis",
+  "Vertebral hemangioma",
+  "Postoperative spine change / hardware",
+  "Bones / Spine_others",
+  "Hepatic steatosis",
+  "Focal liver lesion (nodule / mass)",
+  "Hepatomegaly",
+  "Liver contour irregularity / cirrhosis features",
+  "Hepatic calcification",
+  "Cholelithiasis / gallstones",
+  "Post-cholecystectomy (gallbladder operated / absent)",
+  "Gallbladder wall thickening",
+  "Hydropic gallbladder / distension",
+  "Biliary sludge",
+  "Biliary stent / catheter / drain",
+  "Splenomegaly",
+  "Accessory spleen / splenule / polysplenia",
+  "Focal splenic lesion (nodule / mass)",
+  "Pancreatic mass / focal lesion",
+  "Pancreatic lipomatosis",
+  "Adrenal nodule / mass",
+  "Adrenal thickening / hyperplasia",
+  "Adrenal calcification",
+  "Simple renal cyst",
+  "Complex renal cyst / solid renal mass",
+  "Hydronephrosis",
+  "Renal calculi / nephrolithiasis",
+  "Renal atrophy / decreased renal size",
+  "Nephrectomy (kidney absent / operated)",
+  "Ascites",
+  "Pneumoperitoneum",
+  "Bowel wall thickening / inflammation",
+  "Diverticulosis",
+  "Omental caking / peritoneal carcinomatosis",
+  "Abdominal lymphadenopathy",
+  "Abdominal aortic aneurysm (partially imaged)",
+  "Abdominal aortic calcification / atherosclerosis (partially imaged)",
+  "IVC filter",
+  "Upper Abdomen_others",
+  "Thyroid enlargement (goiter)",
+  "Thyroid nodule",
+  "Cervical / supraclavicular lymphadenopathy",
+  "Neck soft tissue mass",
+  "Lower Neck_others",
+  "Breast mass / focal asymmetry",
+  "Post-lumpectomy / post-mastectomy change",
+  "Breast implant (intact or present)",
+  "Axillary lymphadenopathy",
+  "Motion artifact / suboptimal study",
+  "Study limitation / limited evaluation (non-motion)",
+  "No significant intrathoracic abnormality",
+  "Others_others"
+ ],
+ "instruction": "Given the following chest CT report, extract the presence/absence of entities",
+ "max_len": 512,
+ "default_threshold": 0.5,
+ "torch_dtype": "float32",
+ "n_labels": 137,
+ "attn_implementation": "sdpa"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b779d82e66bc2e765627ea32c160e401bfce2cebcf986e80f7fdf3e53c6934b8
+size 2384826596

modeling_chest2vec_labeler.py

@@ -0,0 +1,241 @@
+"""
+Chest2Vec CT Report Labeler — HuggingFace `AutoModel` wrapper.
+
+A weakly-supervised multi-label classifier that maps a free-text chest-CT report to a
+137-leaf chest-imaging taxonomy with a ternary status per label
+(negative / uncertain / positive).
+
+Architecture: `Qwen/Qwen3-Embedding-0.6B` encoder (LoRA merged in) → left-padding-aware
+last-token (EOS) pooling → L2-normalization → a single linear ternary head
+(`hidden=1024 → 137 × 3`).
+
+Usage:
+
+    from transformers import AutoModel, AutoTokenizer
+    model = AutoModel.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True).eval()
+    tok   = AutoTokenizer.from_pretrained("chest2vec/chest2vec_labeler", trust_remote_code=True)
+
+    reports = ["Bibasilar atelectasis with small bilateral pleural effusions. Cardiomegaly."]
+    print(model.label_reports(reports, tokenizer=tok))          # -> [{'Pleural effusion': 'positive', ...}]
+
+    # CheXbert / SRR-BERT-style report comparison (label both, compare):
+    res = model.score_reports(gt_reports, pred_reports, tokenizer=tok)
+    print(res["micro"]["f1"], res["macro"]["f1"], res["weighted"]["f1"])
+"""
+from typing import Dict, List, Optional, Any
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel
+from transformers.modeling_outputs import ModelOutput
+from dataclasses import dataclass
+
+# class index ordering produced by the head's softmax (axis=-1)
+NEGATIVE, UNCERTAIN, POSITIVE = 0, 1, 2
+_CLASS_TO_VALUE = {NEGATIVE: 0, UNCERTAIN: -1, POSITIVE: 1}
+_CLASS_TO_NAME = {NEGATIVE: "negative", UNCERTAIN: "uncertain", POSITIVE: "positive"}
+
+
+class Chest2VecLabelerConfig(PretrainedConfig):
+    model_type = "chest2vec_labeler"
+
+    def __init__(
+        self,
+        encoder_config: Optional[dict] = None,
+        base_model: str = "Qwen/Qwen3-Embedding-0.6B",
+        hidden_size: int = 1024,
+        n_labels: int = 137,
+        num_classes_per_label: int = 3,
+        labels: Optional[List[str]] = None,
+        instruction: str = "Given the following chest CT report, extract the presence/absence of entities",
+        max_len: int = 512,
+        default_threshold: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.encoder_config = encoder_config or {}
+        self.base_model = base_model
+        self.hidden_size = hidden_size
+        self.n_labels = n_labels
+        self.num_classes_per_label = num_classes_per_label
+        self.labels = labels or []
+        self.instruction = instruction
+        self.max_len = max_len
+        self.default_threshold = default_threshold
+
+
+@dataclass
+class LabelerOutput(ModelOutput):
+    logits: torch.FloatTensor = None          # [B, num_labels, 3]
+    embedding: torch.FloatTensor = None       # [B, hidden] L2-normalized pooled
+
+
+def _build_encoder(encoder_config: dict, attn_implementation: str = "sdpa"):
+    ecfg = dict(encoder_config)
+    for k in ("architectures", "auto_map", "transformers_version", "_name_or_path", "torch_dtype"):
+        ecfg.pop(k, None)
+    model_type = ecfg.pop("model_type", "qwen3")
+    cfg = AutoConfig.for_model(model_type, **ecfg)
+    cfg.torch_dtype = "float32"
+    try:
+        cfg._attn_implementation = attn_implementation
+    except Exception:
+        pass
+    try:
+        return AutoModel.from_config(cfg, attn_implementation=attn_implementation)
+    except TypeError:
+        return AutoModel.from_config(cfg)
+
+
+def _last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    """Left-padding-aware last-token (EOS) pooling — matches the training pipeline."""
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    idx = attention_mask.sum(dim=1) - 1
+    return last_hidden_states[torch.arange(last_hidden_states.size(0), device=last_hidden_states.device), idx]
+
+
+class Chest2VecLabelerModel(PreTrainedModel):
+    config_class = Chest2VecLabelerConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: Chest2VecLabelerConfig):
+        super().__init__(config)
+        self.model = _build_encoder(config.encoder_config, getattr(config, "attn_implementation", "sdpa"))
+        self.head = nn.Linear(config.hidden_size, config.n_labels * config.num_classes_per_label)
+        self.num_labels = config.n_labels
+        self.num_classes_per_label = config.num_classes_per_label
+        self._tokenizer = None
+        self.post_init()
+
+    # ---- core forward (token tensors in, logits out) ----
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, **kwargs):
+        if position_ids is None and attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 0)
+        out = self.model(input_ids=input_ids, attention_mask=attention_mask,
+                         position_ids=position_ids, use_cache=False, return_dict=True)
+        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out.hidden_states[-1]
+        emb = _last_token_pool(h, attention_mask)
+        emb = F.normalize(emb.float(), p=2, dim=-1)
+        logits = self.head(emb).view(emb.size(0), self.num_labels, self.num_classes_per_label)
+        return LabelerOutput(logits=logits, embedding=emb)
+
+    # ---- tokenization (matches training: Instruct/Query + reserved EOS + left pad) ----
+    def _get_tokenizer(self, tokenizer=None):
+        if tokenizer is not None:
+            return tokenizer
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+            src = self.config._name_or_path or self.config.base_model
+            self._tokenizer = AutoTokenizer.from_pretrained(src, padding_side="left", trust_remote_code=True)
+            if self._tokenizer.pad_token_id is None:
+                self._tokenizer.pad_token = self._tokenizer.eos_token
+        return self._tokenizer
+
+    def _encode(self, tok, reports: List[str], max_len: int):
+        instr = self.config.instruction.strip()
+        texts = [(f"Instruct: {instr}\nQuery: {str(r).strip()}" if instr else str(r).strip()) for r in reports]
+        pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
+        eod_id = tok.convert_tokens_to_ids("<|endoftext|>")
+        if eod_id is None or eod_id < 0:
+            eod_id = pad_id
+        enc = tok(texts, add_special_tokens=False, truncation=True, max_length=max_len - 1,
+                  padding=False, return_attention_mask=False)
+        ids = [x + [eod_id] for x in enc["input_ids"]]
+        T = max((len(x) for x in ids), default=1)
+        input_ids = [[pad_id] * (T - len(x)) + x for x in ids]
+        attn = [[0] * (T - len(x)) + [1] * len(x) for x in ids]
+        return (torch.tensor(input_ids, dtype=torch.long), torch.tensor(attn, dtype=torch.long))
+
+    # ---- high-level prediction API ----
+    @torch.no_grad()
+    def predict_proba(self, reports: List[str], tokenizer=None, batch_size: int = 16,
+                      max_len: Optional[int] = None, device=None) -> torch.Tensor:
+        """Return [N, num_labels] probability of the POSITIVE class for each label."""
+        if isinstance(reports, str):
+            reports = [reports]
+        tok = self._get_tokenizer(tokenizer)
+        max_len = max_len or self.config.max_len
+        device = device or next(self.parameters()).device
+        self.eval()
+        out = []
+        for i in range(0, len(reports), batch_size):
+            ii, am = self._encode(tok, reports[i:i + batch_size], max_len)
+            logits = self(input_ids=ii.to(device), attention_mask=am.to(device)).logits
+            out.append(torch.softmax(logits.float(), dim=-1)[:, :, POSITIVE].cpu())
+        return torch.cat(out, dim=0)
+
+    @torch.no_grad()
+    def predict(self, reports: List[str], tokenizer=None, threshold: Optional[float] = None,
+                batch_size: int = 16, max_len: Optional[int] = None, device=None,
+                return_ternary: bool = False) -> Dict[str, Any]:
+        """Return {'labels': names, 'positive': [N,L] 0/1, 'proba': [N,L], ('ternary': [N,L] in {-1,0,1})}."""
+        if isinstance(reports, str):
+            reports = [reports]
+        thr = self.config.default_threshold if threshold is None else threshold
+        tok = self._get_tokenizer(tokenizer)
+        max_len = max_len or self.config.max_len
+        device = device or next(self.parameters()).device
+        self.eval()
+        proba, ternary = [], []
+        for i in range(0, len(reports), batch_size):
+            ii, am = self._encode(tok, reports[i:i + batch_size], max_len)
+            logits = self(input_ids=ii.to(device), attention_mask=am.to(device)).logits.float().cpu()
+            proba.append(torch.softmax(logits, dim=-1)[:, :, POSITIVE])
+            if return_ternary:
+                cls = logits.argmax(-1)
+                ternary.append(torch.tensor([[_CLASS_TO_VALUE[int(c)] for c in row] for row in cls]))
+        proba = torch.cat(proba, dim=0)
+        res = {"labels": list(self.config.labels), "proba": proba.numpy(),
+               "positive": (proba >= thr).int().numpy(), "threshold": thr}
+        if return_ternary:
+            res["ternary"] = torch.cat(ternary, dim=0).numpy()
+        return res
+
+    def label_reports(self, reports: List[str], tokenizer=None, threshold: Optional[float] = None,
+                      **kw) -> List[Dict[str, str]]:
+        """Return, per report, a dict {label_name: 'positive'} for labels above threshold."""
+        out = self.predict(reports, tokenizer=tokenizer, threshold=threshold, **kw)
+        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 ----
+    @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]:
+        """
+        Label both GT and predicted reports, then compute label-level F1 (CheXbert-style).
+
+        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.
+        """
+        from sklearn.metrics import precision_recall_fscore_support
+        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))}
+        return res
+
+
+def report_f1(gt_reports: List[str], pred_reports: List[str], model=None, tokenizer=None,
+              model_id: str = "chest2vec/chest2vec_labeler", **kw) -> Dict[str, Any]:
+    """Convenience wrapper: load the labeler (if not supplied) and score GT vs predicted reports."""
+    if model is None:
+        model = Chest2VecLabelerModel.from_pretrained(model_id).eval()
+    return model.score_reports(gt_reports, pred_reports, tokenizer=tokenizer, **kw)

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:def76fb086971c7867b829c23a26261e38d9d74e02139253b38aeb9df8b4b50a
+size 11423705

tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}