Mr report training

report-generation-training/.gitignore

@@ -0,0 +1,9 @@
+artifacts/
+logs/
+runs/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+

report-generation-training/README.md

@@ -0,0 +1,190 @@
+# MR-RATE single-writer report training
+
+This is a standalone report-generation training layer for MR-RATE. It does not
+modify FORA or an existing MR-RATE checkout.
+
+The writer target is the natural `findings` field from
+`MR-RATE-validation/reports/all_reports.csv`. The standardized
+`findings_sentences.jsonl` remains an encoder/MIL data dependency only; it is
+never used as report-generation supervision. The trainer does not infer
+abnormal/healthy labels or split reports into artificial subtasks.
+
+The two execution modes are deliberately equivalent:
+
+- `online`: a frozen MR-RATE encoder produces the complete projected visual
+  token bag during training.
+- `cached`: the same complete projected token bag is read from the upstream
+  ragged memmap cache.
+
+In both modes the same frozen 74-label `ClassifyThenAggregate` MIL head consumes
+the full bag, a trainable query resampler makes 512 language-prefix tokens, and
+one Gemma LoRA writer learns the complete findings text. The MIL probabilities
+are soft conditioning context; they are not report targets.
+
+## Required artifacts
+
+The files have distinct roles:
+
+| Input | Purpose |
+| --- | --- |
+| `all_reports.csv` | Writer supervision from `study_uid, findings` |
+| `findings_sentences.jsonl` | Reproduce the encoder/MIL cohort only |
+| labels CSV | Exact 74-class MIL schema |
+| splits CSV | Train/validation/test membership |
+| MR-RATE checkpoint | Frozen visual encoder and projection |
+| `mil_head.pt` | Frozen Classify-Then-Aggregate MIL head and provenance |
+| Gemma snapshot | Frozen language model plus trainable report LoRA |
+| exact token cache | Required only for cached mode |
+
+The real encoder checkpoint and `mil_head.pt` are intentionally placeholders
+in `configs/base.yaml`; replace them before running.
+
+## 1. Configure
+
+```bash
+cd /hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training
+cp configs/base.yaml configs/my_run.yaml
+```
+
+Edit at least:
+
+```yaml
+encoder_checkpoint: /absolute/path/to/the/mr-rate-checkpoint.pt
+mil_checkpoint: /absolute/path/to/the/corresponding/mil_head.pt
+llm_path: /absolute/path/to/gemma-3-12b-it
+
+data:
+  data_folder: /absolute/path/to/mri
+  reports_csv: /absolute/path/to/all_reports.csv
+  jsonl_file: /absolute/path/to/findings_sentences.jsonl
+  labels_file: /absolute/path/to/mrrate_labels.csv
+  splits_csv: /absolute/path/to/splits.csv
+  cached_tokens_dir: /absolute/path/to/exact_tokens
+```
+
+Do not pair an arbitrary MIL head and encoder. Preflight checks the recorded
+encoder SHA-256, architecture, label order, and cache fingerprint.
+
+## 2. Run preflight
+
+Inside the training container:
+
+```bash
+export PROJECT=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training
+export PYTHONPATH=/hnvme/workspace/b180dc51-sezgin/extra-pip:$PROJECT/src
+cd "$PROJECT"
+
+python -m mrrate_report_training.preflight \
+  --config configs/my_run.yaml \
+  --mode cached
+```
+
+Use `--mode online` for online training. Preflight rejects pooled features,
+token-capped caches, mismatched dimensions/classes, missing findings, and
+mismatched MIL/encoder provenance.
+
+## 3. Prepare an exact cache, if needed
+
+```bash
+export MRRATE_REPORT_CONFIG="$PROJECT/configs/my_run.yaml"
+bash scripts/build_exact_cache.sh train
+bash scripts/build_exact_cache.sh val
+```
+
+Run cache generation on a CUDA node. It is an expensive frozen-encoder pass,
+not part of every training epoch. The builder always uses
+`max_tokens_per_study=0`. If the verified exact cache already exists, skip
+this step.
+
+## 4. Run a small real-data smoke
+
+Use a separate output directory so the smoke cannot overwrite production
+checkpoints:
+
+```bash
+cp configs/my_run.yaml configs/my_smoke.yaml
+# Edit output_dir in configs/my_smoke.yaml, for example:
+# output_dir: runs/mrrate_single_writer_smoke
+```
+
+From an allocated GPU node, inside the container:
+
+```bash
+GPUS_PER_NODE=1 bash scripts/train_cached.sh \
+  configs/my_smoke.yaml \
+  --max-studies 8 \
+  --max-updates 2
+```
+
+For the online path:
+
+```bash
+GPUS_PER_NODE=1 bash scripts/train_online.sh \
+  configs/my_smoke.yaml \
+  --max-studies 2 \
+  --max-updates 1
+```
+
+Online mode is slower because it loads/resamples MRI volumes and runs the
+frozen encoder during every epoch. The online reader supports
+`batchXX/<study>.zip`: it extracts only the current study under node-local
+`/tmp` and removes it after encoding.
+
+## 5. Submit Slurm training
+
+The launcher defaults to two nodes with four GPUs per node. Command-line
+`sbatch` options can override the node count.
+
+```bash
+export PROJECT=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training
+export MODE=cached
+export CONFIG="$PROJECT/configs/my_run.yaml"
+export LLM_PATH=/absolute/path/to/gemma-3-12b-it
+export ENCODER_CHECKPOINT=/absolute/path/to/the/mr-rate-checkpoint.pt
+export MIL_CHECKPOINT=/absolute/path/to/the/corresponding/mil_head.pt
+export GPUS_PER_NODE=4
+export MAX_STUDIES=0
+export MAX_UPDATES=0
+
+sbatch --nodes=2 "$PROJECT/scripts/slurm_train.sh"
+```
+
+For online training, set `MODE=online`. The node launcher stages Gemma and
+both checkpoints once per node, creates node-local CUDA/Triton/Inductor
+caches, then launches one `torchrun` rank per GPU.
+
+The production settings in `configs/my_run.yaml` are:
+
+```yaml
+epochs: 1
+batch_size: 1
+gradient_accumulation: 1
+learning_rate: 0.0001
+```
+
+Every real study is seen exactly once. No study is duplicated to fill a
+distributed batch.
+
+## 6. Resume
+
+```bash
+export RESUME=/absolute/path/to/checkpoint-00000500.pt
+sbatch --nodes=2 "$PROJECT/scripts/slurm_train.sh"
+```
+
+Keep `MODE`, `CONFIG`, node count, and artifact variables the same. Checkpoints
+contain the resampler/connector, report LoRA, optimizer, scheduler, next data
+position, and per-rank RNG state.
+
+## 7. Integration tests
+
+```bash
+python tests/smoke_checks.py
+sbatch scripts/synthetic_gpu_e2e.sbatch
+```
+
+The tests cover natural findings targets, strict cache validation,
+online/cached numerical equivalence, frozen MIL behavior, single-prefix
+construction, optimizer updates, strict weight provenance, and checkpoint
+resume. The Slurm test uses a deterministic dummy dataset and does not replace
+the required real-data smoke.

report-generation-training/configs/base.yaml

@@ -0,0 +1,53 @@
+seed: 17
+mode: cached
+output_dir: runs/mrrate_single_writer_v1
+
+upstream_root: /hnvme/workspace/b180dc51-sezgin/MR-RATE-linearprobe/contrastive-pretraining
+encoder_checkpoint: /path/to/mrrate_pretraining_checkpoint.pt
+mil_checkpoint: /path/to/mil_head.pt
+llm_path: /hnvme/data/LLM-TZ/hub/models--google--gemma-3-12b-it/snapshots/96b6f1eccf38110c56df3a15bffe176da04bfd80
+
+data:
+  data_folder: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/mri
+  jsonl_file: /hnvme/workspace/b180dc51-sezgin/MR-RATE/contrastive-pretraining/data/findings_sentences.jsonl
+  reports_csv: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/reports/all_reports.csv
+  labels_file: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/mrrate_neurovfm74_labels.csv
+  splits_csv: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/splits.csv
+  cached_tokens_dir: artifacts/exact_tokens
+  use_preprocessed: false
+  preprocessed_dir: null
+  normalizer: zscore
+  space: native_space
+
+encoder:
+  name: vjepa2
+  vjepa21_checkpoint: null
+  chunk_size: 64
+  fusion_mode: late
+  pooling_strategy: simple_attn
+  extra_latent_projection: false
+  dim_latent: 512
+
+writer:
+  num_visual_queries: 512
+  resampler_depth: 2
+  resampler_heads: 8
+  lora_r: 16
+  lora_alpha: 32
+  # Covers every normal raw findings report. One corrupted 32k-token outlier
+  # (5NIUCVXWHA) is safely bounded by this ceiling.
+  max_target_tokens: 1536
+  mil_conditioning: all_classes
+  mil_proposal_dropout: 0.0
+  localization: false
+
+training:
+  epochs: 1
+  batch_size: 1
+  gradient_accumulation: 1
+  learning_rate: 0.0001
+  weight_decay: 0.01
+  warmup_ratio: 0.03
+  checkpoint_every: 500
+  shuffle: true
+  replacement_sampling: false

report-generation-training/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mrrate-report-training"
+version = "0.1.0"
+description = "Exact online and cached single-writer report training for MR-RATE"
+requires-python = ">=3.10"
+dependencies = [
+  "numpy",
+  "pyyaml",
+  "torch",
+  "transformers",
+  "peft",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+testpaths = ["tests"]

report-generation-training/scripts/build_exact_cache.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+set -euo pipefail
+project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$project_dir"
+export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}"
+exec python -m mrrate_report_training.build_cache \
+  --config "${MRRATE_REPORT_CONFIG:-configs/base.yaml}" --split "${1:?split required}"

report-generation-training/scripts/slurm_train.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=mrrate_report
+#SBATCH --partition=h200
+#SBATCH --nodes=2
+#SBATCH --gres=gpu:4
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=32
+#SBATCH --time=24:00:00
+#SBATCH --signal=B:USR1@300
+#SBATCH --output=logs/mrrate_report_%j.out
+#SBATCH --error=logs/mrrate_report_%j.err
+set -euo pipefail
+
+: "${MODE:?set MODE=online or cached}"
+: "${CONFIG:?set CONFIG to an absolute yaml path}"
+: "${LLM_PATH:?set LLM_PATH}"
+: "${ENCODER_CHECKPOINT:?set ENCODER_CHECKPOINT}"
+: "${MIL_CHECKPOINT:?set MIL_CHECKPOINT}"
+case "$MODE" in online|cached) ;; *) exit 2 ;; esac
+
+project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+sif="${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate-ib3.sif}"
+mkdir -p "$project_dir/logs"
+export MASTER_ADDR
+MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1)"
+export MASTER_PORT="${MASTER_PORT:-29541}"
+export MODE CONFIG LLM_PATH ENCODER_CHECKPOINT MIL_CHECKPOINT
+export GPUS_PER_NODE="${GPUS_PER_NODE:-4}"
+export MAX_STUDIES="${MAX_STUDIES:-0}"
+export MAX_UPDATES="${MAX_UPDATES:-0}"
+export RESUME="${RESUME:-}"
+
+exec srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \
+  --kill-on-bad-exit=1 \
+  singularity exec --nv -B /hnvme:/hnvme,/tmp:/tmp "$sif" \
+  bash "$project_dir/scripts/train_node.sh"

report-generation-training/scripts/synthetic_gpu_e2e.sbatch

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=mr_syn_e2e
+#SBATCH --partition=h200
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --time=00:30:00
+#SBATCH --output=logs/synthetic_gpu_e2e_%j.out
+#SBATCH --error=logs/synthetic_gpu_e2e_%j.err
+set -euo pipefail
+
+project=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training
+sif=/hnvme/workspace/b180dc51-sezgin/mrrate-ib3.sif
+mkdir -p "$project/logs"
+cache="/tmp/mrrate_synthetic_${SLURM_JOB_ID}"
+mkdir -p "$cache"/{cuda,triton,inductor,xdg}
+export CUDA_CACHE_PATH="$cache/cuda"
+export TRITON_CACHE_DIR="$cache/triton"
+export TORCHINDUCTOR_CACHE_DIR="$cache/inductor"
+export XDG_CACHE_HOME="$cache/xdg"
+export MRRATE_REPORT_STREAM_THRESHOLD=128
+export MRRATE_REPORT_STREAM_CHUNK=64
+
+exec singularity exec --nv -B /hnvme:/hnvme,/tmp:/tmp "$sif" bash -lc "
+  set -euo pipefail
+  export PYTHONPATH=/hnvme/workspace/b180dc51-sezgin/extra-pip:$project/src
+  export CUDA_CACHE_PATH=$CUDA_CACHE_PATH
+  export TRITON_CACHE_DIR=$TRITON_CACHE_DIR
+  export TORCHINDUCTOR_CACHE_DIR=$TORCHINDUCTOR_CACHE_DIR
+  export XDG_CACHE_HOME=$XDG_CACHE_HOME
+  export MRRATE_REPORT_STREAM_THRESHOLD=$MRRATE_REPORT_STREAM_THRESHOLD
+  export MRRATE_REPORT_STREAM_CHUNK=$MRRATE_REPORT_STREAM_CHUNK
+  cd $project
+  python tests/gpu_e2e.py
+"

report-generation-training/scripts/train_cached.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$project_dir"
+export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}"
+exec python -m torch.distributed.run --standalone --nproc-per-node="${GPUS_PER_NODE:-1}" \
+  -m mrrate_report_training.train --mode cached \
+  --config "${1:-configs/base.yaml}" "${@:2}"