chest2vec_labeler / README.md
lukeingawesome's picture
Evaluation: replace threshold-based F1 tables with AUROC / 3-class F1 / binary F1 (uncertain→neg) on CT-RATE + SNUH, with 95% CIs; add ≥1-positive subset
f4e90fc verified
|
Raw
History Blame Contribute Delete
8.72 kB
metadata
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 (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 (revised CT-RATE, 137-leaf taxonomy)

Label space

The model head predicts 137 leaf labels. They roll up through the chest-imaging hierarchy into 38 upper/container groups and 10 anatomy sections (the label_hierarchy in config.json), so predictions and report-comparison F1 can be reported at leaf, upper, or anatomy granularity.

  • The model outputs all 137 leaves. In the training data, 136 of them have at least one positive example; the single exception is IVC filter (kept for taxonomy completeness, but it had no positives, so the model effectively never predicts it).
  • The exact label list is in config.json (labels). Full definitions and per-split counts are in the chest2vec/chest2vec_labels dataset's LABEL_HIERARCHY.md.

This model was trained and evaluated on the chest2vec/chest2vec_labels dataset (revised CT-RATE, 137-leaf taxonomy).

Ternary headsoftmax(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

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):

res = model.score_reports(gt_reports, pred_reports, tokenizer=tok)   # equal-length lists
# scores are reported at three hierarchy levels:
for level in ("leaf", "upper", "anatomy"):
    b = res[level]
    print(level, b["n_labels"], b["micro"]["f1"], b["macro"]["f1"], b["weighted"]["f1"])
print(res["leaf"]["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)

Each level (leaf = 137 labels, upper = 38 container groups, anatomy = 10 sections) returns micro / macro / weighted precision/recall/F1 plus per_label. Upper/anatomy scores are the max-over-children roll-up of the leaf predictions (model.aggregate_hierarchy(...)). Coarser levels are easier to match, so upper/anatomy F1 are typically higher than leaf.

Per-label best F1 (threshold tuning)

The default decision threshold is a single global value, but the F1-optimal threshold differs per label. To get the best achievable F1 per label (and the threshold that achieves it) against a ground-truth label set:

# gt: a DataFrame with the 137 label columns (ternary; positive == 1), or a binary array
res = model.per_label_best_f1(reports, gt, tokenizer=tok, level="leaf", min_pos=30)
res["macro_best_f1_min_pos"]                 # macro best-F1 over labels with >= min_pos positives
res["per_label"]["Pleural effusion"]         # {'best_f1':.., 'best_threshold':.., 'n_pos':..}

Per-label threshold tuning lifts macro-F1 by ~4–6 points over the fixed-0.5 threshold (see below).

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

Reported as macro-averages over per-label leaf metrics, with 95% bootstrap percentile confidence intervals. Compare only within a row.

  • AUROC — macro one-versus-rest AUROC on the positive-class probability. Threshold-free.
  • Three-class F1 — macro F1 from the ternary head (negative / uncertain / positive). Uses argmax; no threshold involved.
  • Binary F1 — macro F1 with uncertain mapped to negative. Per-label decision thresholds are selected on the CT-RATE tuning partition (not on the reported eval set).

Two label subsets are reported per eval set:

  • ≥30 positives — leaves with at least 30 positive test instances (headline).
  • ≥1 positive — all evaluated leaves (of the 137, 129 had ≥1 positive in CT-RATE and 120 in SNUH).
Label set AUROC Three-class F1 Binary F1
CT-RATE (in distribution)
53 (≥30 pos) 0.989 (0.988, 0.991) 0.686 (0.667, 0.692) 0.888 (0.879, 0.896)
129 (≥1 pos) 0.988 (0.985, 0.991) 0.610 (0.592, 0.612) 0.769 (0.731, 0.777)
SNUH (external)
29 (≥30 pos) 0.971 (0.967, 0.975) 0.590 (0.579, 0.598) 0.766 (0.749, 0.780)
120 (≥1 pos) 0.972 (0.968, 0.975) 0.541 (0.523, 0.543) 0.652 (0.601, 0.660)

Leaf macro-AUROC barely moves in distribution → external (0.989 → 0.971): label ranking transfers to unseen data; the residual F1 gap reflects labeling-convention drift, not a domain failure. A radiologist spot-checked 966 reports of the public test labels (857 fully accepted / 60 imperfect-but-acceptable / 49 failed; see the dataset card).

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 (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 for the full citation.