From 197975cc42e57ef7e3c88112e97752214b31bb59 Mon Sep 17 00:00:00 2001 From: Kemal Yaylali Date: Sat, 12 Sep 2026 09:13:54 +0100 Subject: [PATCH] feat(ml): train a real model, and report the number that matters rather than the flattering one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit "Variants are unscored" was accurate: nothing was ever trained, so a quarter of every rank was dead weight and the UI leaked a connection error at the reader. - scripts/make-training-set.sh derives a training table from ClinVar directly. ClinVar already carries the molecular consequence, the gene and an allele frequency, which is the feature set serving sends, so this avoids running VEP over hundreds of thousands of variants. 2-star records only. - train.py now holds out whole genes (GroupShuffleSplit). docs/data.md had said to do this since the data pass; the code was still doing a random split, which is the leak Grimm 2015 describes. - evaluate() reports missense on its own. On the last run: AUROC 0.986 over 74,239 held-out variants, but 0.872 over the 13,553 missense ones, and the docs say plainly why even that is flattered — within missense the only live feature is allele frequency, and ClinVar's benign calls often use allele frequency as evidence (ACMG BA1/BS1), so the feature partly caused the label. - the 503 now names what is missing (model@alias via tracking URI) and leaves the exception in the server log instead of the UI. - make training-set / make train; the 58 MB table is gitignored. Verified end to end: model registered as v2, the simulated NF2 case scores 0.999 on the planted variant, and it now ranks 1.00 with all four components live. Tests: api 77, ml 22, loader 16, web 32; ruff, mypy, svelte-check clean. --- .gitignore | 1 + Makefile | 10 +++++- README.md | 7 ++-- api/app/routers/cases.py | 10 ++++-- api/tests/test_scoring.py | 5 ++- docs/data.md | 33 ++++++++++++++++- ml/rarelens_ml/train.py | 70 ++++++++++++++++++++++++++++++++---- ml/tests/test_train.py | 66 +++++++++++++++++++++++++++++++++- scripts/make-training-set.sh | 47 ++++++++++++++++++++++++ 9 files changed, 235 insertions(+), 14 deletions(-) create mode 100755 scripts/make-training-set.sh diff --git a/.gitignore b/.gitignore index 3284d42..7fcc305 100644 --- a/.gitignore +++ b/.gitignore @@ -16,4 +16,5 @@ mlruns/ *.tfstate.backup .terraform/ data/*.vcf* +data/*.tsv !data/README.md diff --git a/Makefile b/Makefile index 34c74a2..b85b5d0 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,7 @@ -.PHONY: up down clean migrate test lint data hpo demo-case loader pipeline annotate images kind serverless-deploy serverless-destroy gcp-configure gcp-secrets +.PHONY: up down clean migrate test lint data hpo demo-case training-set train loader pipeline annotate images kind serverless-deploy serverless-destroy gcp-configure gcp-secrets VCF ?= data/example.vcf.gz +MLFLOW_URI ?= http://localhost:5001 TAG ?= latest # The loader container reaches docker-compose's Postgres through the host. HOST_DB_URL ?= postgresql://rarelens:rarelens@host.docker.internal:5432/rarelens @@ -36,6 +37,13 @@ hpo: ## load HPO gene-to-phenotype annotations, which the ranking matches again demo-case: ## build the simulated proband: GIAB background + one ClinVar pathogenic variant scripts/make-demo-case.sh +training-set: ## build a ClinVar training table, shaped like VEP --tab output + scripts/make-training-set.sh + +train: ## train the pathogenicity model and point the production alias at it (needs `make up`) + cd ml && MLFLOW_TRACKING_URI=$(MLFLOW_URI) uv run --extra dev \ + python -m rarelens_ml.train --tsv ../data/clinvar-training.vep.tsv --register + loader: docker build -t rarelens/loader:dev -f pipeline/loader.Dockerfile pipeline diff --git a/README.md b/README.md index f12913e..d4851fd 100644 --- a/README.md +++ b/README.md @@ -75,10 +75,13 @@ would be the container itself. `LOCAL_DATA_ROOT` is the directory a sample's `vc To train and register a model (the API scores with `models:/rarelens-pathogenicity@production`): ```bash -cd ml && MLFLOW_TRACKING_URI=http://localhost:5000 \ - uv run python -m rarelens_ml.train --tsv ../pipeline/results/.vep.tsv --register +make training-set # a ClinVar-derived training table, ~370k labelled variants +make train # fits, reports held-out metrics by gene split, moves the production alias ``` +What those metrics do and do not mean is in [docs/data.md](docs/data.md); the headline AUROC +flatters a model whose strongest feature is the consequence class. + Local Kubernetes: `make kind` builds the images, loads them into a kind cluster and applies `infra/k8s/overlays/local`. diff --git a/api/app/routers/cases.py b/api/app/routers/cases.py index a7a79e4..1afb2be 100644 --- a/api/app/routers/cases.py +++ b/api/app/routers/cases.py @@ -8,6 +8,7 @@ from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import selectinload +from app.config import settings from app.db import SessionDep from app.models import ( Case, @@ -152,9 +153,14 @@ async def score(case_id: uuid.UUID, session: SessionDep): scored, version = await score_job(job.id, session) except Exception as e: # No registry, no model behind the alias, a model that will not load: all of these are - # the environment being unready, not a bug in the request. Say so rather than throwing 500. + # the environment being unready, not a bug in the request. Name what is missing and keep + # the exception in the log, where it is useful, rather than in the UI, where it is noise. logger.exception("scoring case %s failed", case_id) - raise HTTPException(503, f"could not score with the model: {e}") from e + raise HTTPException( + 503, + f"no model available: {settings.model_name}@{settings.model_alias} " + f"via {settings.mlflow_tracking_uri}", + ) from e return ScoreOut(case_id=case_id, scored=scored, model_version=version) diff --git a/api/tests/test_scoring.py b/api/tests/test_scoring.py index 1b22921..026f448 100644 --- a/api/tests/test_scoring.py +++ b/api/tests/test_scoring.py @@ -99,4 +99,7 @@ async def test_an_unreachable_model_registry_is_explained_not_a_500( monkeypatch.setattr(scoring, "load_model", unreachable) r = await client.post(f"/api/cases/{case_id}/score") assert r.status_code == 503 - assert "connection refused" in r.json()["detail"] + detail = r.json()["detail"] + # Names what is missing; the exception itself belongs in the server log, not the UI. + assert "rarelens-pathogenicity@production" in detail + assert "Max retries" not in detail and "Traceback" not in detail diff --git a/docs/data.md b/docs/data.md index be67868..cc23999 100644 --- a/docs/data.md +++ b/docs/data.md @@ -59,6 +59,36 @@ instead of the 25 GB cache. It returns no gnomAD frequencies, so every variant l gnomAD and the rarity term stops discriminating. Fine for showing the mechanics; use the cache for anything you would quote. +## The model, and what its numbers mean + +`make training-set` builds a training table straight from ClinVar rather than running VEP over +hundreds of thousands of variants: ClinVar already carries the molecular consequence (`MC`), the +gene (`GENEINFO`) and an allele frequency (`AF_EXAC`), which is the feature set serving sends. +Only 2-star-and-above records are kept. `make train` then fits LightGBM and points the +`production` alias at the new version. + +The last run: 312,025 training and 74,239 held-out variants across 7,728 and 1,932 genes, with no +gene on both sides. + +| | AUROC | AUPRC | +|---|---|---| +| all held-out variants | 0.986 | 0.954 | +| missense only (13,553) | 0.872 | 0.725 | + +Three things to say before anyone quotes the headline number: + +1. **0.986 mostly measures how separable ClinVar's classes are by consequence.** Its pathogenic set + is largely loss of function and its benign set largely is not, so a model handed the consequence + class does well without knowing anything hard. That is why the missense row exists: missense is + where interpretation is actually difficult. +2. **Even 0.872 is flattered by circularity.** Within missense, every row has the same consequence + and impact and no CADD or AlphaMissense score, so allele frequency is doing nearly all the work + — and ClinVar's benign calls frequently *use* allele frequency as evidence (ACMG BA1/BS1). The + feature partly caused the label. +3. **It is not comparable to published CADD or AlphaMissense numbers.** Those are trained and + evaluated on different data. A fair comparison scores the same held-out rows with all three, + which needs the plugin data (see above) and is the obvious next step. + ## Evaluating the model honestly The model trains on ClinVar labels and is scored on ClinVar-labelled variants, which is exactly @@ -69,7 +99,8 @@ where published benchmarks go wrong. What to do about it: 2. **Split by gene, not by variant.** Random splits put variants from the same gene on both sides, and a model can then score a gene rather than a variant. Grimm et al. showed this inflates reported accuracy for exactly this class of tool: *Hum Mutat* 36:513–523, 2015. - [10.1002/humu.22768](https://doi.org/10.1002/humu.22768) + [10.1002/humu.22768](https://doi.org/10.1002/humu.22768) *Implemented*: + `rarelens_ml.train.split_by_gene` holds out whole genes. 3. **Prefer a time-based holdout.** Train on an older ClinVar release (monthly archives live under `vcf_GRCh38/archive_2.0/`) and test only on variants classified after that date. This is the closest thing to a prospective evaluation available without new patients. diff --git a/ml/rarelens_ml/train.py b/ml/rarelens_ml/train.py index d1a8443..a5b3682 100644 --- a/ml/rarelens_ml/train.py +++ b/ml/rarelens_ml/train.py @@ -5,6 +5,7 @@ Usage: python -m rarelens_ml.train --tsv results/clinvar.vep.tsv --register """ import argparse import re +import sys from pathlib import Path import lightgbm as lgb @@ -13,7 +14,7 @@ import pandas as pd import sklearn from mlflow import MlflowClient from sklearn.metrics import average_precision_score, roc_auc_score -from sklearn.model_selection import train_test_split +from sklearn.model_selection import GroupShuffleSplit from rarelens_ml.features import RAW_COLUMNS, build from rarelens_ml.model import PathogenicityModel @@ -66,19 +67,64 @@ def read_vep_tab(path: str | Path) -> pd.DataFrame: return df.rename(columns={"#Uploaded_variation": "Uploaded_variation"}) -def load(tsv: str) -> tuple[pd.DataFrame, pd.Series]: +def load(tsv: str) -> tuple[pd.DataFrame, pd.Series, pd.Series]: + """Returns the raw feature columns, the labels, and each row's gene for grouping.""" df = read_vep_tab(tsv).rename(columns=VEP_TO_RAW) for col in RAW_COLUMNS: # plugin columns are absent when VEP ran without CADD/AlphaMissense if col not in df: df[col] = pd.NA + if "SYMBOL" not in df: + df["SYMBOL"] = "-" y = df["CLIN_SIG"].map(label) keep = y.notna() return ( df.loc[keep, RAW_COLUMNS].reset_index(drop=True), y[keep].astype(int).reset_index(drop=True), + df.loc[keep, "SYMBOL"].reset_index(drop=True), ) +def split_by_gene( + X: pd.DataFrame, y: pd.Series, genes: pd.Series, test_size: float = 0.2, seed: int = 42 +) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series, pd.Series, pd.Series]: + """Hold out whole genes, never single variants. + + A random split puts variants of the same gene on both sides, and the model can then score the + gene rather than the variant. Grimm et al. (Hum Mutat 2015) showed this inflates reported + accuracy for exactly this class of tool; see docs/data.md. + """ + splitter = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=seed) + train_idx, test_idx = next(splitter.split(X, y, groups=genes)) + return ( + X.iloc[train_idx], X.iloc[test_idx], + y.iloc[train_idx], y.iloc[test_idx], + genes.iloc[train_idx], genes.iloc[test_idx], + ) + + +MIN_SUBSET = 50 + + +def evaluate(X: pd.DataFrame, y: pd.Series, proba) -> dict[str, float]: + """Headline metrics, plus missense on its own. + + Most of ClinVar's pathogenic set is loss of function and most of its benign set is not, so a + model given the consequence class separates them easily and the overall AUROC flatters it. + Missense is where variant interpretation is actually hard, so it gets its own number. + """ + metrics = { + "auroc": float(roc_auc_score(y, proba)), + "auprc": float(average_precision_score(y, proba)), + "test_variants": float(len(y)), + } + missense = X["consequence"].eq("missense_variant").to_numpy() + if missense.sum() >= MIN_SUBSET and len(set(y[missense])) == 2: + metrics["auroc_missense"] = float(roc_auc_score(y[missense], proba[missense])) + metrics["auprc_missense"] = float(average_precision_score(y[missense], proba[missense])) + metrics["missense_variants"] = float(missense.sum()) + return metrics + + def fit(X: pd.DataFrame, y: pd.Series) -> lgb.LGBMClassifier: return lgb.LGBMClassifier(**PARAMS).fit(build(X), y) @@ -109,16 +155,28 @@ def main() -> None: p.add_argument("--alias", default="production") a = p.parse_args() - X, y = load(a.tsv) - Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42) + X, y, genes = load(a.tsv) + Xtr, Xte, ytr, yte, train_genes, test_genes = split_by_gene(X, y, genes) + print( + f"{len(Xtr)} train / {len(Xte)} test variants; " + f"{train_genes.nunique()} / {test_genes.nunique()} genes, no gene in both", + file=sys.stderr, + ) mlflow.set_experiment(MODEL_NAME) with mlflow.start_run(): mlflow.log_params(PARAMS) clf = fit(Xtr, ytr) proba = clf.predict_proba(build(Xte))[:, 1] - mlflow.log_metrics({"auroc": roc_auc_score(yte, proba), - "auprc": average_precision_score(yte, proba)}) + metrics = evaluate(Xte, yte, proba) | {"test_genes": float(test_genes.nunique())} + mlflow.log_metrics(metrics) + summary = f"held-out AUROC {metrics['auroc']:.3f}, AUPRC {metrics['auprc']:.3f}" + if "auroc_missense" in metrics: + summary += ( + f" | missense only: AUROC {metrics['auroc_missense']:.3f}, " + f"AUPRC {metrics['auprc_missense']:.3f} over {int(metrics['missense_variants'])}" + ) + print(summary, file=sys.stderr) if a.register: version = log_and_register(clf, MODEL_NAME, a.alias) print(f"registered {MODEL_NAME} v{version} as @{a.alias}") diff --git a/ml/tests/test_train.py b/ml/tests/test_train.py index 6945f28..f8d6a78 100644 --- a/ml/tests/test_train.py +++ b/ml/tests/test_train.py @@ -68,9 +68,10 @@ def test_load_returns_raw_serving_columns_and_labels(tmp_path: Path) -> None: ["c", "22:3", "T", "intron_variant", "MODIFIER", "CHEK2", "0.3", "uncertain_significance", "1", "-"], ], ) - X, y = load(str(tsv)) + X, y, genes = load(str(tsv)) assert list(X.columns) == RAW_COLUMNS assert y.tolist() == [1, 0] # the VUS row is dropped + assert list(genes) == ["TBX1", "CHEK2"] def test_logged_model_returns_probabilities_from_raw_columns(tmp_path: Path) -> None: @@ -102,3 +103,66 @@ def test_logged_model_returns_probabilities_from_raw_columns(tmp_path: Path) -> assert scores.shape == (50,) assert ((scores >= 0) & (scores <= 1)).all() assert not set(np.unique(scores)) <= {0.0, 1.0}, "got class labels, expected probabilities" + + +def test_load_returns_the_gene_of_each_row_for_grouping(tmp_path: Path) -> None: + from rarelens_ml.train import load + + tsv = write_vep_tab( + tmp_path / "x.vep.tsv", + [ + ["a", "22:1", "G", "missense_variant", "MODERATE", "NF2", "0.0001", "pathogenic", "28", "0.9"], + ["b", "22:2", "A", "synonymous_variant", "LOW", "CHEK2", "0.12", "benign", "3", "-"], + ], + ) + _, _, genes = load(str(tsv)) + assert list(genes) == ["NF2", "CHEK2"] + + +def test_the_split_never_puts_one_gene_on_both_sides(tmp_path: Path) -> None: + """Random splits leak: a model can learn the gene instead of the variant (Grimm 2015).""" + import numpy as np + import pandas as pd + + from rarelens_ml.train import split_by_gene + + genes = pd.Series([f"GENE{i // 4}" for i in range(40)]) + X = pd.DataFrame({"impact": ["HIGH"] * 40, "consequence": ["stop_gained"] * 40, + "gnomad_af": ["0"] * 40, "cadd_phred": ["10"] * 40, "am_pathogenicity": ["-"] * 40}) + y = pd.Series(np.tile([1, 0], 20)) + + Xtr, Xte, _ytr, yte, train_genes, test_genes = split_by_gene(X, y, genes, test_size=0.3) + assert set(train_genes) & set(test_genes) == set() + assert len(Xtr) + len(Xte) == 40 + assert len(yte) > 0 + + +def test_evaluate_reports_missense_separately() -> None: + """Overall AUROC flatters a consequence-based model; missense is where the problem is.""" + import numpy as np + import pandas as pd + + from rarelens_ml.train import evaluate + + rng = np.random.default_rng(0) + consequence = ["missense_variant"] * 200 + ["stop_gained"] * 200 + X = pd.DataFrame({"consequence": consequence}) + y = pd.Series([*rng.integers(0, 2, 200), *([1] * 200)]) + proba = np.concatenate([rng.random(200), rng.uniform(0.8, 1.0, 200)]) + + metrics = evaluate(X, y, proba) + assert metrics["missense_variants"] == 200 + assert "auroc_missense" in metrics + assert metrics["auroc"] > metrics["auroc_missense"] # the easy class inflates the headline + + +def test_evaluate_omits_the_missense_metric_when_there_is_nothing_to_measure() -> None: + import numpy as np + import pandas as pd + + from rarelens_ml.train import evaluate + + X = pd.DataFrame({"consequence": ["stop_gained"] * 100}) + y = pd.Series([1] * 50 + [0] * 50) + metrics = evaluate(X, y, np.linspace(0, 1, 100)) + assert "auroc_missense" not in metrics diff --git a/scripts/make-training-set.sh b/scripts/make-training-set.sh new file mode 100755 index 0000000..6b7a418 --- /dev/null +++ b/scripts/make-training-set.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Build a training table from ClinVar, shaped like VEP --tab output so rarelens_ml.train reads it +# unchanged. +# +# ClinVar already carries the molecular consequence (MC), the gene (GENEINFO) and an allele +# frequency (AF_EXAC), which is the feature set serving sends. That avoids running VEP over +# hundreds of thousands of variants to produce a training set. CADD and AlphaMissense are left +# missing, exactly as they are when the pipeline runs without plugin data. +# +# Only 2-star-and-above records are kept: "criteria provided, multiple submitters, no conflicts" +# or better. Labels come from CLNSIG (see docs/data.md on evaluating this honestly). +set -eu + +IMAGE=${BCFTOOLS_IMAGE:-quay.io/biocontainers/bcftools:1.20--h8b25389_0} +CLINVAR=${CLINVAR:-https://ftp.ncbi.nlm.nih.gov/pub/clinvar/vcf_GRCh38/clinvar.vcf.gz} +REGION=${REGION:-} # empty means the whole genome +OUT=${OUT:-data/clinvar-training.vep.tsv} + +mkdir -p "$(dirname "$OUT")" +docker run --rm -v "$PWD/$(dirname "$OUT"):/out" "$IMAGE" bash -eu -c " + echo '==> streaming ClinVar ${REGION:-(whole genome)}' >&2 + bcftools query ${REGION:+-r '$REGION'} \ + -f '%CHROM\t%POS\t%REF\t%ALT\t%INFO/MC\t%INFO/GENEINFO\t%INFO/AF_EXAC\t%INFO/CLNSIG\t%INFO/CLNREVSTAT\n' \ + '$CLINVAR' \ + | awk -F'\t' -v OFS='\t' ' + BEGIN { + split(\"transcript_ablation splice_acceptor_variant splice_donor_variant stop_gained frameshift_variant stop_lost start_lost transcript_amplification\", h, \" \"); + for (i in h) impact[h[i]] = \"HIGH\"; + split(\"inframe_insertion inframe_deletion missense_variant protein_altering_variant\", m, \" \"); + for (i in m) impact[m[i]] = \"MODERATE\"; + split(\"splice_region_variant synonymous_variant start_retained_variant stop_retained_variant\", l, \" \"); + for (i in l) impact[l[i]] = \"LOW\"; + print \"#Uploaded_variation\", \"Location\", \"Allele\", \"Consequence\", \"IMPACT\", \"SYMBOL\", \"gnomADe_AF\", \"CLIN_SIG\", \"CADD_PHRED\", \"am_pathogenicity\"; + } + \$9 !~ /multiple_submitters|expert_panel|practice_guideline/ { next } # 2 stars and up + \$5 == \".\" || \$6 == \".\" { next } + { + split(\$5, mc, \",\"); split(mc[1], so, \"|\"); csq = so[2]; + if (csq == \"\") next; + split(\$6, gi, \"|\"); split(gi[1], g, \":\"); gene = g[1]; + af = (\$7 == \".\" ? \"-\" : \$7); + imp = (csq in impact ? impact[csq] : \"MODIFIER\"); + print \$1 \"_\" \$2 \"_\" \$3 \"_\" \$4, \$1 \":\" \$2, \$4, csq, imp, gene, af, \$8, \"-\", \"-\"; + }' > /out/$(basename "$OUT") +" +echo "wrote $OUT: $(($(wc -l < "$OUT") - 1)) labelled variants" +awk -F'\t' 'NR>1 {print $8}' "$OUT" | sort | uniq -c | sort -rn | head -6