GoBlueInformatics @ #SMM4H-HeaRD 2026, Task 6: TNM Stage Prediction

System submission for #SMM4H-HeaRD 2026 Task 6 — Pathological TNM Stage Prediction from TCGA Reports. Predicts T (1–4), N (0–3), and M (0–1) labels for each pathology report.

The submission combines long-context discriminative encoders and LoRA-tuned biomedical LLMs:

For the official results see the paper (paper/smm4h_task6_tnm.tex).

Repository layout

.
├── README.md
├── requirements.txt
├── TCGA_cancer_types_binary.csv     # cancer-type one-hot metadata (bundled, ~1.3 MB)
│
├── configs/                          # one YAML per system in the table above
│   ├── config_v6.yaml
│   ├── config_v13.yaml
│   ├── config_v14.yaml
│   ├── config_v15.yaml
│   ├── config_v16.yaml
│   └── config_v17.yaml
│
├── splits/                           # 85/15 held-out split of the fully-labeled subset
│   ├── train_split.csv               # 2,055 rows used to train v14–v17 + held-out v6/v13
│   └── val_split.csv                 # 363 rows used as the internal held-out test
│
├── results/                          # outputs of the new held-out reruns (no checkpoints)
│   ├── v6_heldout/
│   └── v13_heldout/
│
├── paper/
│   ├── smm4h_task6_tnm.tex
│   └── custom.bib
│
├── dataset.py / llm_dataset.py       # encoder + LLM dataset classes
├── losses.py / metrics.py            # multi-task loss and macro-F1 / exact-match metrics
├── model.py                          # encoder model (chunked + long-context variants)
├── preprocess.py / utils.py          # text cleaning + Config / seeding / logging
│
├── train.py                          # encoder trainer, single train/val split
├── train_kfold.py                    # encoder 5-fold CV trainer (used for v6 / v13)
├── train_heldout.py                  # v6 / v13 retrained on the 2,055-row training portion
├── train_bbten.py                    # BB-TEN fine-tune (v17)
├── train_llm.py                      # OpenBioLLM-8B LoRA, full fully-labeled set (v14)
├── train_llm_v15.py                  # DeepSeek-R1-Distill-Llama-8B LoRA (v15)
├── train_llm_v16.py                  # OpenBioLLM-8B LoRA on the 2,055-row split (v16)
│
├── infer.py                          # encoder inference (single checkpoint or k-fold ensemble)
├── infer_oof.py                      # encoder k-fold out-of-fold predictions
├── infer_bbten.py                    # BB-TEN inference (stock and fine-tuned variants)
├── infer_llm.py                      # LLM inference, hand-rolled Llama-3 prompt (v14, v16)
├── infer_llm_v15.py                  # LLM inference, R1-Distill chat template (v15)
│
└── score_heldout.py                  # score any submission CSV against splits/val_split.csv

Data you need to provide

Two files are intentionally not bundled and must be obtained from the shared-task organizers (or derived from the public TCGA pathology report resource of Kefeli & Tatonetti, 2024):

Place these files in the project root (the existing configs reference them by relative path).

splits/train_split.csv and splits/val_split.csv (already bundled) are the deterministic 2,055 / 363 stratified-on-T split of the fully-labeled subset (seed 42) used throughout the paper for the held-out internal comparison.

Installation

conda create -n smm4h python=3.11
conda activate smm4h
pip install -r requirements.txt

Tested on Python 3.11, PyTorch 2.x with CUDA 12.x, NVIDIA RTX 5090 (32 GB).

Running the systems

All commands assume the project root as cwd.

v6 — BioClinicalBERT chunked ensemble

# original 5-fold ensemble (used to produce the OOF predictions on the full training file)
python train_kfold.py --config configs/config_v6.yaml

# held-out version (aligned with LLM / BB-TEN comparison): trains on the 2,055-row split
python train_heldout.py --config configs/config_v6.yaml --output_dir output_v6_heldout

v13 — BioClinical-ModernBERT-large ensemble

# original 5-fold ensemble (Test set 1 submission was made from this)
python train_kfold.py --config configs/config_v13.yaml

# held-out version
python train_heldout.py --config configs/config_v13.yaml --output_dir output_v13_heldout

# inference on a test file using the 5-fold ensemble
python infer.py --config configs/config_v13.yaml \
                --ckpts output_v13/fold_1/best_model/model.pt output_v13/fold_2/best_model/model.pt \
                        output_v13/fold_3/best_model/model.pt output_v13/fold_4/best_model/model.pt \
                        output_v13/fold_5/best_model/model.pt \
                --test_file Task6_test.csv \
                --output_dir test1_predictions

v14 — OpenBioLLM-8B LoRA (Test set 2 submission)

Trained on all 2,418 fully-labeled rows:

python train_llm.py --config configs/config_v14.yaml
python infer_llm.py --adapter output_v14/lora \
                    --test_file Task6_test_hard.csv \
                    --output output_v14/task6_hard_predictions.csv

v15 — R1-Distill-8B LoRA (split-controlled ablation)

Trained on the 2,055-row training portion of the held-out split:

python train_llm_v15.py --config configs/config_v15.yaml
python infer_llm_v15.py --adapter output_v15/lora \
                        --test_file splits/val_split.csv \
                        --output output_v15/val_predictions.csv

v16 — OpenBioLLM-8B LoRA (split-controlled ablation)

Same prompt and recipe as v14, but trained on the 2,055-row training portion only:

python train_llm_v16.py --config configs/config_v16.yaml
python infer_llm.py --adapter output_v16/lora \
                    --test_file splits/val_split.csv \
                    --output output_v16/val_predictions.csv

v17 — BB-TEN fine-tuned

Three single-task Clinical-BigBird classifiers (T, N, M), each fine-tuned on the 2,055-row training portion:

python train_bbten.py --config configs/config_v17.yaml

BB-TEN stock — inference only on the held-out split

python infer_bbten.py --test_file splits/val_split.csv --output_dir output_bbten_stock

Scoring against the held-out split

score_heldout.py scores any submission-format CSV (patient_filename,t,n,m with T in 0–3) against splits/val_split.csv:

python score_heldout.py --submit output_v16/val_predictions_submit.csv --label v16

It prints per-stage macro-F1, mean macro-F1, and exact-match accuracy, and optionally writes the metrics JSON via --out_json.

Held-out internal results (paper Table 1a)

The numbers in paper/smm4h_task6_tnm.tex Table 1(a) are reproducible from these scripts. The bundled results/v6_heldout/heldout_summary.json and results/v13_heldout/heldout_summary.json are the actual outputs of the v6 and v13 held-out reruns. The other rows of Table 1(a) are produced by training the v14–v17 systems above and scoring with score_heldout.py.

Hardware notes

LLM fine-tuning targets a single 32 GB GPU (RTX 5090). The configs use bf16, gradient checkpointing, and LoRA rank 16. On smaller GPUs you'll likely need to reduce max_length or batch size in the LLM configs.

Citation

If you use this code, please cite the system paper. The reference list in paper/custom.bib includes the upstream backbones used: BioClinicalBERT, BioClinical-ModernBERT, BB-TEN (Kefeli et al. 2024 Nature Communications), TCGA-Reports (Kefeli & Tatonetti 2024 Patterns), OpenBioLLM, DeepSeek-R1, Big Bird, and LoRA.