Skip to content

Repository files navigation

CohortX Task 2 — Structuring Cohort Eligibility Criteria as Triples

Our solution for the CohortX Task 2 challenge (MICCAI 2026): extracting (subject, predicate, object) triples from free-text clinical-trial eligibility criteria.

Public leaderboard: 0.78 (FM3S with Hungarian matching).

This repository contains only the code path that produced the submitted system — training through submission — with no exploratory variants or superseded experiments.

The whole system runs on CPU, offline, within 16 GB of RAM — the challenge's hardware constraint. No cloud inference, no API calls.


Architecture

Two independent branches produce a candidate triple set per trial; a per-row selector picks between them.

                  eligibility criteria text
                             │
        ┌────────────────────┴────────────────────┐
        │                                         │
   BRANCH A                                  BRANCH B
   per-criterion hybrid                      row-level generation
        │                                         │
  rule-based splitter                    BioLORD-2023 retrieval
  → scaffold triples                     → 3 nearest training rows
        │                                    as few-shot exemplars
  flan-t5-base (fine-tuned)                       │
        │                                  Llama-3.1-8B-Instruct
  BioMistral-7B (QLoRA)                    5 samples @ temp 0.3
        │                                         │
  length-gated router                      self-consistency pick
  (≥100 chars & ≥2 triples                 (highest mean similarity
   → BioMistral, else T5)                   to the other 4)
        │                                         │
  lexical k=1 retrieval                           │
        └────────────────────┬────────────────────┘
                             │
                   cross-score both ways with
                   FM3S + Hungarian matching;
                   keep the better-explained one
                             │
                    final triple set

Why cross-scoring works as a selector. There is no gold at test time, so we score each branch's output against the other branch's output using the task metric itself, in both directions. The output that better "explains" its counterpart wins, tie-broken slightly toward the larger triple set.

Why cardinality matters more than recall. FM3S zero-pads unmatched triples, so over-generating is directly penalised. Every component is tuned to match gold triple count rather than maximise coverage — this is the single most important property of the metric and it shaped every design decision here. It is also why the two branches are selected between rather than merged: the union of both scored below either one alone.


Running the pipeline: training → submission

Install requirements-full.txt. Steps 2 and 3 need a GPU; everything else is CPU.

1. Build the splits and T5 training data

python -m src.data          # Task_2.xlsx → artifacts/{train,val,test}.jsonl
python -m src.t5_data_v36   # → artifacts/t5_v36_{train,val}.jsonl

2. Fine-tune T5 (~2 min on a single consumer GPU)

python -m src.t5_train_v36  # → artifacts/t5_model_v36/

google/flan-t5-base, 5 epochs, lr 1e-4, batch 2, seed 13. Best validation loss lands at epoch 2.

3. Fine-tune BioMistral with QLoRA, then merge and quantise

python train_qlora_biomistral.py   # → artifacts/biomistral_qlora/
python merge_lora.py               # merge adapter into BioMistral/BioMistral-7B
# then quantise the merged model to Q4_K_M GGUF with llama.cpp's convert + quantize

4-bit NF4, r=16, alpha=32, dropout 0.05 on the q/k/v/o and MLP projections; 3 epochs, lr 2e-4 cosine, paged 8-bit AdamW, seed 42. Trained on 759 examples built from all 100 annotated trials, with multi-pair criteria duplicated once. Run on a Kaggle T4 — a 7B QLoRA fine-tune does not fit in 8 GB of VRAM.

4. Generate Branch A (CPU, offline)

python run_v38a_fixed.py     # T5 val candidates       (~10 min)
python run_v38b_submit.py    # T5 test candidates      (~20 min)
python run_v50_finetuned.py  # BioMistral, val + test  (~80 min)

5. Generate Branch B and write the submission

python run_v56_multisample.py

Runs Llama-3.1 with 5 samples per row, picks the self-consistent one, then cross-scores against Branch A and writes submissions/submission_v56_ensemble.csv — the submitted file. It imports build_v50_preds and pick_best_per_row from run_v55_ensemble.py, which is why that file is present; run_v55_ensemble.py is a module here, not an entry point.

Quantised weights are expected at artifacts/biomistral_finetuned/BioMistral-7B-finetuned.Q4_K_M.gguf and artifacts/llama31_model/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf.

A note on exact reproduction

Branch A is greedy end-to-end — T5 beam search and BioMistral at temperature=0.0 — so it recomputes bit-for-bit, and the cached outputs in artifacts/ (see below) let you skip straight to step 5.

Branch B is not bit-reproducible: the multi-sample step used temperature=0.3 unseeded, so re-running it produces a similar but not identical triple set, and therefore a submission file that differs from the one originally scored. Seeding the sampler would have made the pipeline fully deterministic; not doing so was a mistake worth documenting.


Repository map

Path Purpose
src/data.py Loads Task_2.xlsx, deterministic 80/20 split (seed 13), triple (de)serialisation
src/scorer.py Our re-implementation of FM3S (WordNet path similarity + Hungarian matching)
src/splitter.py Inclusion/exclusion section parser and atomic-criterion counter
src/t5_data.py Linearised triple format shared by the T5 stages
src/t5_data_v36.py, src/t5_train_v36.py T5 training-data builder and trainer
src/retrieval3.py Lexical k=1 criterion retrieval + nearest-row deep-triple transfer
train_qlora_biomistral.py, merge_lora.py BioMistral QLoRA fine-tune, adapter merge
run_v38a_fixed.py, run_v38b_submit.py T5 per-criterion candidates (val / test)
run_v50_finetuned.py Fine-tuned BioMistral per-criterion extraction — Branch A
run_v56_multisample.py Multi-sample self-consistency, ensemble, submission — Branch B
run_v55_ensemble.py Cross-scoring selector; imported by the step above
artifacts/ Data splits and cached model outputs

The run scripts create a submissions/ directory on first use and write their CSV output there; it is not tracked in this repository.

What is in artifacts/

  • {train,val,test}.jsonl — the 80/20 split plus the 50 test trials
  • t5_v36_{train,val}.jsonl — linearised criterion→triple training pairs
  • v38_{val,test}_candidates.json — T5 per-criterion candidates (beam search, deterministic)
  • v50_biomistral_{val,test}_outputs.json — BioMistral per-criterion pairs (temperature=0.0, deterministic)

The two cache files let you skip steps 2–4: both stages are greedy, so the caches are exactly what re-running them would produce. Every run script checks for its cache and loads it instead of regenerating.

Model weights are not committed (947 MB for T5, 4.1 GB + 4.6 GB for the quantised GGUFs — all far past GitHub's file limit). Steps 2–3 regenerate them.


Environment

requirements.txt is the minimal CPU set (no deep-learning stack); requirements-full.txt adds what training and model inference need. Python 3.12.

WordNet is fetched by nltk on first use; for a fully offline run, pre-download it with python -c "import nltk; nltk.download('wordnet')".


Models and external resources

Resource Role
google/flan-t5-base Fine-tuned per-criterion triple extractor
BioMistral/BioMistral-7B QLoRA fine-tuned per-criterion extractor
meta-llama/Meta-Llama-3.1-8B-Instruct Row-level generator, few-shot
FremyCompany/BioLORD-2023 Sentence embeddings for exemplar retrieval
WordNet (via NLTK) Word similarity inside our FM3S re-implementation

No clinical ontology (UMLS, SNOMED CT, etc.) was consulted. All four checkpoints are publicly available on HuggingFace, as the challenge rules permit.


A note on the data

Task_2.xlsx and the files under artifacts/ are challenge-provided data. Please check the competition rules on redistribution before making this repository public — if redistribution is not permitted, delete Task_2.xlsx and artifacts/* and have users regenerate the splits with python -m src.data from their own copy of the spreadsheet.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages