diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c8dfc9f..f2cc269 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,19 +34,27 @@ jobs: # `logger=csv`. One extra entry turns both on: `test_every_rendered_logger_composes` is then the # job that would catch `logger=wandb` raising MissingConfigException again. loggers: [false] + data_backend: [mtd] include: - profile: supervised loggers: true + data_backend: mtd + # The custom-featurization lane: its venv deliberately omits meds-torch-data, so the lane + # itself proves nothing imports it (including `meds-model commands` — the import-guard test). + - profile: supervised + loggers: false + data_backend: custom_featurization steps: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v6 with: enable-cache: true - - name: Render the ${{ matrix.profile }} DAG (loggers=${{ matrix.loggers }}) + - name: Render the ${{ matrix.profile }} DAG (loggers=${{ matrix.loggers }}, backend=${{ matrix.data_backend }}) run: | uvx copier copy --vcs-ref=HEAD --defaults --trust \ --data model_slug=demo_model --data model_name="Demo Model" \ --data profile=${{ matrix.profile }} \ + --data data_backend=${{ matrix.data_backend }} \ --data use_wandb=${{ matrix.loggers }} --data use_mlflow=${{ matrix.loggers }} \ . /tmp/demo_model - name: Install torch (CPU) + run the generated repo's tests diff --git a/CLAUDE.md b/CLAUDE.md index 0bcc29d..89d1609 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -71,10 +71,14 @@ cd /tmp/demo && uv venv \ && uv run pytest -m "not slow" -q -rs ``` -Expect ~37 passed and 2 skipped, and both skips are the correct result — everything that does not need a -model runs, including real MTD tensorization and task materialization. One skip is `skip_if_stub`; the -other is `test_unsupported_command_fails_clearly`, whose parameter set is empty for `probe` because that -DAG registers all five commands. +Expect ~75 passed and 1–3 skipped (supervised; `probe` also skips `test_unsupported_command_fails_clearly`, +whose parameter set is empty because that DAG registers all five commands), and every skip is the correct +result — everything that does not need a model runs, including real MTD tensorization, a real predicates +featurization (both workspaces are built, and the equivalence guard compares their label partitioning), +and task materialization. One skip is `skip_if_stub`; with `'data_backend': 'custom_featurization'` in +the `data` dict the repo installs **without meds-torch-data**, the two MTD-only tests skip via +`importorskip`, and the CLI smoke test doubles as the regression test for the import guard in +`lightning/__init__.py`. Add `'use_wandb': True, 'use_mlflow': True` to that `data` dict to render the optional logger configs — otherwise `logger=csv` is the only thing the generated suite ever composes. @@ -224,6 +228,16 @@ class exists that no profile registers — which is how `MaterializedPredictComm `commands.py.jinja` builds an `entries` list in Jinja and derives its import block from it, so a profile can never import a class it does not register (an F401 in the generated repo). +`data_backend` is a second, orthogonal copier axis (`mtd` | `custom_featurization`) — a representation +choice, not a DAG choice. **[`docs/design-featurization.md`](docs/design-featurization.md) is its spec.** +It gates only dependencies (meds-torch-data), rendered configs (`datamodule/*.yaml`, +`preprocess_data.yaml`'s `featurization` default), stubs (`datamodule.py`) and the `model.yaml` +predicates wiring; `src/meds_model_base/` ships both representations unconditionally and tolerates +either dependency set (import guards in `lightning/__init__.py` and `lightning/modules.py` — do not +"clean up" those try/excepts; a custom_featurization repo has no meds_torchdata to import). The model's +`predicates.yaml` is user-owned and is the ONE predicates file: generated tests, MEDS-DEV runs and +production all read it. + `_templates_suffix: .jinja` means **only** `.jinja` files are rendered; everything else under `template/` is copied byte-for-byte (path segments like `{{ model_slug }}` are still substituted). A file whose rendered *name* is empty is skipped — that is how `predict.py` is made conditional. diff --git a/copier.yml b/copier.yml index cde7a73..020c2b7 100644 --- a/copier.yml +++ b/copier.yml @@ -22,9 +22,11 @@ _skip_if_exists: - "src/{{ model_slug }}/model.py" - "src/{{ model_slug }}/commands.py" - "src/{{ model_slug }}/predict.py" + - "src/{{ model_slug }}/datamodule.py" - "src/{{ model_slug }}/configs/model/**" - "src/{{ model_slug }}/configs/paths/**" - "src/{{ model_slug }}/configs/profile/**" + - "predicates.yaml" _exclude: - "*.pyc" @@ -113,6 +115,18 @@ profile: "custom — choose each command yourself": "custom" default: "supervised" +# The data backend decides how `preprocess_data` materializes patient data — a representation choice, +# orthogonal to the profile (which is a DAG-shape choice). It gates only dependencies, configs and +# stubs; `src/meds_model_base/` ships both paths unconditionally, so flipping the answer later and +# running `copier update` backfills what the new backend needs. +data_backend: + type: str + help: "How preprocess_data materializes patient data" + choices: + "mtd — MEDS-transforms + meds-torch-data tensorization (the standard sequence-model path)": "mtd" + "custom_featurization — 0/1 predicate columns on MEDS parquet; you write the datamodule": "custom_featurization" + default: "mtd" + # Command toggles, only *asked* for the custom profile. `preprocess_data` is always supported: every # chain starts by materializing patient data. Tasks are not a stage — commands take external_labels_dir. implements_pretrain: diff --git a/docs/CONFIG-AUDIT.md b/docs/CONFIG-AUDIT.md index 7ac81a3..38e09ff 100644 --- a/docs/CONFIG-AUDIT.md +++ b/docs/CONFIG-AUDIT.md @@ -50,7 +50,8 @@ Three keys look inert and are not: - **`callbacks`** — `instantiate_group` branches on the node being empty, not on its value. Benign, but note that dropping `model_checkpoint` silently changes which weights get published (see `_persist_checkpoint`, §2). -- **`trainer.deterministic` / `trainer.benchmark`** — the exception to "nothing here can be quietly +- **~~`trainer.deterministic` / `trainer.benchmark`~~ Fixed; both now ship in `trainer/default.yaml` + alongside an explicit `precision`.** The exception to "nothing here can be quietly wrong", and the reason `trainer.*` is qualified above. Neither appears in `configs/trainer/default.yaml`, so Lightning runs with `deterministic=False` and cuDNN autotunes its algorithms by timing. Nothing raises; the run simply is not reproducible, which makes every number it @@ -96,9 +97,19 @@ These are the ones that matter. "Branch" column gives the exact site. The complete fix is to record a config hash in the work directory and refuse to resume when it differs; flipping the default buys most of that safety for one line, and the hash can follow if resume is ever used in anger. -2. **Determinism settings** (`trainer.deterministic` / `trainer.benchmark`, §1, plus the hardcoded matmul - precision below). Equally silent, and it degrades every other result: two runs at the same seed can - disagree, so no measurement in the repository is reproducible evidence of anything. +2. **~~Determinism settings~~ (`trainer.deterministic` / `trainer.benchmark`, plus the hardcoded matmul + precision below). Fixed.** Equally silent, and it degraded every other result: two runs at the same + seed could disagree, so no measurement in the repository was reproducible evidence of anything. + + `configs/trainer/default.yaml` now ships `deterministic: warn`, `benchmark: false` and an explicit + `precision: 32-true`, and `configure_cublas_workspace()` sets `CUBLAS_WORKSPACE_CONFIG` from the + dispatcher so `trainer.deterministic=true` is a usable strict mode on CUDA instead of an error at the + first matmul. `precision` is listed rather than left implicit because configs are struct mode: without + the key, `trainer.precision=bf16-mixed` needs Hydra's `+` and nobody guesses that. + + What is now claimed, in the generated README and in the porting procedure, is *same seed + same config + + same environment → same metric* — not "reproducible". The claim can only be **shipped** here, not + **proven**: proving it needs two real training runs, which needs a model (#7). 3. **`pipeline` / `pipeline_overrides`.** Known broken until 2026-08-01; see §3. 4. **`attach_labels=false`.** Produces a file `meds_evaluation` rejects. Loud, but nothing tells a user the key exists or why they would want it. @@ -116,13 +127,20 @@ Not config-driven, but in the same blind spot and worth inspecting together: manifest, never asserted. - `load_pretrained_weights`'s **zero-match `RuntimeError`** (`train.py:280`) — the load-bearing guard of the `finetune` profile: it is what stops a renamed encoder from becoming a silent from-scratch run. -- **`torch.set_float32_matmul_precision("medium")`** (`train.py:60`, and again at `train.py:219`) — set - unconditionally, with no config key and no mention in any document. On Ampere and later this runs fp32 - matmuls in bfloat16. It is stable run-to-run, so it does not break same-machine reproducibility, but it - silently diverges from any source implementation that ran at full precision — and because it lives in - the vendored contract, a user can only change it by editing a file `copier update` overwrites. For the - porting procedure in `template/docs/PORTING-A-MODEL.md.jinja`, whose deliverable is a ledger of - deviations that affect results, this is a deviation the template imposes invisibly. +- **~~`torch.set_float32_matmul_precision("medium")`~~ Removed.** Set unconditionally at `train.py:60` and + again at `train.py:219`, with no config key and no mention in any document. On Ampere and later this ran + fp32 matmuls in bfloat16 — stable run-to-run, so it did not break same-machine reproducibility, but it + silently diverged from any source implementation that ran at full precision, and living in the vendored + contract meant a user could only change it by editing a file `copier update` overwrites. For a porting + procedure whose deliverable is a ledger of deviations that affect results, that was a deviation the + template imposed invisibly. + + Deleting the call was the whole fix: torch's own default is `highest` with `allow_tf32` off, so the + contract now simply does not take a position, and the environment matches the one every source + implementation ran in. No config key was needed — promoting it to one would have meant carrying it in + four command roots (training *and* inference; `zero_shot_direct` and `packaged` never train at all) and + keeping them agreed. `trainer.precision` is the sanctioned speed knob, and a port that specifically + needs TF32 sets it in `model.py`, which is the user's file and already requires a ledger row. - **`CoverageError`** (`predict.py:175`) — reachable in normal use; nothing exercises it. - **`InferenceKind.scores` / `MaterializedPredictCommand`** — reachable only through `zero_shot_materialized`, which `skip_if_stub`s. diff --git a/docs/design-featurization.md b/docs/design-featurization.md new file mode 100644 index 0000000..242e030 --- /dev/null +++ b/docs/design-featurization.md @@ -0,0 +1,557 @@ +# Predicate Featurization and the `data_backend` Option + +**Status: §§1–9 implemented (2026-08-04) — steps 1–5 of §11; TECO adoption (step 6) pending.** +Companion to [design-interface.md](design-interface.md), which remains the authoritative spec for the +command graph, artifact layout and manifests. This document specifies an *addition*: a second data +backend for `preprocess_data` and the changes that make the rest of the contract survive it. + +Drafted 2026-08-04 from the TECO port discussion (Florent Pollet / Claude). + +--- + +## 1. Motivation + +The template's only data path today is MEDS-transforms (optional) → meds-torch-data tensorization +("MTD"). Every downstream command consumes the tensorized layout through the MTD datamodule, and the +feature axis a model sees is the **code vocabulary** of whatever dataset it is handed. + +That is the right default for foundation-style sequence models, and the wrong shape for the large class +of published clinical models built on **named clinical variables**. The concrete case is TECO +(`models/teco`): the paper uses 11 named variables on a 15-minute grid; the MEDS port substitutes the +whole code vocabulary because a MEDS model cannot resolve "heart rate" against an arbitrary dataset. + +The MEDS ecosystem already has the resolution layer: per-dataset ACES predicates +(`MEDS-DEV/src/MEDS_DEV/datasets//predicates.yaml`) map abstract names (`creatinine`, `sodium`) +to dataset codes, and task configs bind to them through `???` placeholders. But predicates are only ever +consumed for **task extraction**. This proposal reuses them for **featurization**: `preprocess_data` +stamps one 0/1 column per predicate onto the MEDS data, and the model decides everything else — +whether to read `numeric_value` on active rows, how to aggregate, what temporal structure to impose. + +What this deliberately is *not* (v1 non-goals): + +* **No value channels, units, or transformations.** Presence only. The original rows and columns are + untouched, so a model can read `numeric_value` wherever a predicate column is 1. A future value/ + expression layer (candidate: [dftly](https://github.com/mmcdermott/dftly)) composes on top without + changing this contract. +* **No temporal aggregation.** Grids, windows, carry-forward, imputation are model-side. +* **No MEDS-DEV changes.** See §9. +* **No `es-aces` dependency.** See §5. + +--- + +## 2. The `data_backend` copier option + +New `copier.yml` question: + +```yaml +data_backend: + type: str + help: "How preprocess_data materializes patient data" + choices: + "mtd — MEDS-transforms + meds-torch-data tensorization (the default)": "mtd" + "custom_featurization — predicate columns on MEDS parquet; you write the datamodule": "custom_featurization" + default: "mtd" +``` + +A choice rather than a boolean (`include_mtd`) so a third backend can land later without a second +vocabulary, and so the answer maps one-to-one onto the runtime `featurization` config value and the +manifest's `representation` field. + +### What the option gates (render-time, following the `use_wandb` pattern) + +| Rendered surface | `mtd` | `custom_featurization` | +|---|---|---| +| `meds-torch-data` in `pyproject.toml` / `requirements.txt` | present | **absent** | +| `configs/datamodule/patients.yaml`, `task.yaml` (MTD targets) | rendered | not rendered | +| `configs/datamodule/` custom stub + `src//datamodule.py` stub | not rendered | rendered, in `_skip_if_exists` | +| `configs/preprocess_data.yaml` default | `featurization: mtd` | `featurization: predicates` | +| `predicates.yaml` starter (user-owned, `_skip_if_exists`) | rendered (the equivalence guard uses it) | rendered and wired into the test fixtures | +| `model.yaml` | as today | first command of each chain adds `external_predicates_file={predicates_path}` (§9) | +| Rendered `CLAUDE.md` / `PORTING-A-MODEL.md` backend sections | MTD text | custom text | + +The dependency removal is the real payoff: MEDS-DEV builds each model's venv from `requirements.txt`, +so a featurized model's benchmark environment does not install torch-adjacent MTD machinery it never +imports. + +### What the option must NOT do + +**No Jinja conditionals inside `src/meds_model_base/`.** The vendored contract stays one unconditional +codebase containing both paths, copied byte-for-byte as today. Reasons, in order of weight: + +1. `copier update` 3-way merges the contract into existing repos; a render-conditional contract adds a + second variant axis to every merge. +2. The render/test matrix doubles (7 profiles × 2 backends) for every structural guarantee in + `tests/test_render.py`, and Jinja-in-Python is the exact artefact class that produced the + docstring-collapse bug recorded in CLAUDE.md. +3. There is no runtime win. Import-weight discipline already defers the heavy imports; with the import + guard of §6.2, unused MTD branches cost nothing when the package is absent — the same price + `ProbeTrainCommand` costs a supervised model today. + +Invariant: **one `meds_model_base`, tolerant of either dependency set; `data_backend` decides what is +installed and which configs and stubs render.** A repo that flips its answer later runs `copier update` +and the configs backfill, exactly like `use_wandb`. + +--- + +## 3. `preprocess_data`: the `predicates` branch + +New config keys on `preprocess_data` (both backends parse them; the rendered default differs): + +```yaml +featurization: mtd # mtd | predicates +external_predicates_file: null +``` + +* `featurization=mtd` — byte-identical to today's behavior. `external_predicates_file` is ignored. +* `featurization=predicates` with `external_predicates_file=null` — **error**, before any work. A + silent passthrough would publish a valid-looking artifact with zero feature columns. +* `featurization=predicates` with a file — the branch below. + +`external_meds_dir` is **read-only**, as every `external_` input is: the featurizer reads input shards +and writes augmented copies into the `write_artifact` staging directory under `output_data_dir` (the +temp sibling of `patients/` that is renamed into place atomically — the same mechanism the MTD branch +uses, and where the optional pipeline's intermediate already lives). Nothing is ever written next to, +or into, the source dataset. + +Branch behavior, replacing `_run_mtd` + `_validate_tensorized` after the (unchanged, still optional) +MEDS-transforms `pipeline:` stage: + +1. Parse `external_predicates_file` with the reader of §4. +2. For each shard `data//.parquet` of the (possibly pipeline-transformed) input: append + the predicate columns, write the same relative path into the staging directory. Pure augmentation — + no rows dropped, no columns modified, no resharding. (Row filtering is a MEDS-transforms pipeline + job that already exists; it does not get a second home here.) +3. Copy `metadata/codes.parquet` and `metadata/dataset.json` through. **Do not copy + `metadata/subject_splits.parquet`.** Split membership travels as the shard layout and nothing else — + the same invariant the MTD branch has, asserted by `test_workspace_is_published_with_manifests`: a + copied splits file describes the *input* to preprocessing, and a filtering pipeline makes those two + different things. +4. Write `features.json` at the artifact root — the authoritative, **ordered** definition of this + artifact's feature space: + + ```json + { + "version": 1, + "features": [ + {"name": "icu_admission", "column": "predicate//icu_admission"}, + {"name": "creatinine", "column": "predicate//creatinine"} + ] + } + ``` + + Consumers: the datamodule (`vocab_size = len(features)`; columns selected **in this order**), + `_validate_featurized` (every listed column present in every shard), and the tests. The file exists + instead of "scan a shard for `predicate//` columns" for one load-bearing reason: **feature order is + part of what a trained checkpoint means.** Feature index *i* at training time must be the same + predicate at predict time; parquet column order is incidental, `features.json` is the contract. It + also keeps consumers decoupled from the column-naming convention and makes the artifact + self-describing without opening parquet. It is data-plane metadata, which is why it is a file the + datamodule reads rather than a manifest field: the manifest is provenance and gating + (`read_manifest`), and holds only `n_features` and the digest (§ below). +5. Validate (`_validate_featurized`, the sibling of `_validate_tensorized`): every declared predicate + column present in every shard, shard layout split-sharded, `features.json` consistent. Per-predicate + match counts are logged, and a predicate that matched **nothing** across all shards is a prominent + **warning, not an error**: on real data it signals a binding mismatch worth reading, but it is also + the expected outcome when real-vocabulary predicates run over synthetic test data (§8), and the + machinery being exercised there does not depend on columns being non-zero. +6. Publish through the same `write_artifact` context manager, same artifact name (`patients`), same + `ArtifactType.data` — so `predict`'s `read_manifest` check passes unchanged. Manifest `extras` gain: + +```yaml +representation: predicates # "mtd" on the other branch — new field on BOTH branches +featurization: + predicates_file: + predicates_digest: + n_features: + skipped: [] # predicate names read but not featurized (§4 skip policy) — the feature-space + # provenance beyond the run's log; the §8 validation test asserts this would be + # empty for the repo's own predicates.yaml + match_counts: {} # events matched per predicate across all shards — makes degenerate (all-zero) + # features visible in the artifact itself, not only in the run's warning +``` + +`representation` is the dispatch key for §6.1 and the fail-fast guard replacing `_validate_tensorized` +downstream: a command whose datamodule expects the other flavor fails at the manifest, not minutes in. + +`predicates_digest` exists because the manifest otherwise records only a *path* to a file that lives +outside the artifact and can change after the build: edit `predicates.yaml`, rerun without +`do_overwrite`, and the kept artifact silently describes predicates that no longer exist. The digest +makes "which predicates produced this artifact" answerable and two experiment runs distinguishable — +the same reason `train.py` already digests `external_labels_dir` into its manifests. The full column +list deliberately does **not** appear in the manifest: `features.json` is the single machine-readable +definition of the feature space (datamodules read it, `vocab_size` derives from it, validation checks +shards against it), and a manifest copy would be a second source of truth waiting to drift. + +### Resulting artifact layout + +```text +patients/ + manifest.yaml # representation: predicates + features.json + data/train/0.parquet # original columns + predicate// Int8 columns + data/tuning/0.parquet + data/held_out/0.parquet + metadata/codes.parquet + metadata/dataset.json +``` + +Still a valid MEDS dataset (the MEDS `DataSchema` is open to extra columns), so any MEDS tool — +including a later MTD run — can consume it. + +--- + +## 4. The predicates reader + +A small, hand-rolled parser in `meds_model_base` (new module, e.g. `featurize.py`). Input is a YAML +file whose `predicates:` mapping uses ACES syntax; the reader is **pure code matching plus `or()`**: + +| Form | Example | Semantics | +|---|---|---| +| exact code | `code: LAB//50912//mg/dL` | `code == literal` | +| regex | `code: { regex: "^ICU_ADMISSION//.*" }` | `str.contains` — anchors live in the pattern, matching ACES's own semantics | +| any-of | `code: { any: [HR, PULSE] }` | membership | +| derived `or` | `expr: or(creatinine_1, creatinine_2)` | column-wise max of already-computed plain columns | + +`or()` is in scope from v1 because multi-code concepts are the norm, not the exception — `creatinine` +is two codes and `sodium` three in the MIMIC-IV predicates file; without it the concept layer fragments +back into per-source-code columns, which is the problem this feature exists to solve. + +**Value bounds (`value_min`/`value_max` + inclusivity) are deliberately NOT parsed in v1**, although +ACES counts them as plain. A bounded predicate (`abnormally_high_creatinine`) bakes a clinical +threshold into featurization, and thresholds are exactly the semantics this design declares model-side: +the model reads `numeric_value` on active rows and learns thresholds itself. Dropping bounds also keeps +the reader free of numeric edge cases (null values, inclusivity). They return with the future value- +channel layer (§10), where thresholds belong. Consequence: bounded entries fall under the +unsupported-entry policy below — skipped by default with a warning and a manifest record — so a raw +dataset file remains usable, while a **curated featurization predicates file** stays the recommended +input and the required form for the repo's own file: that file is the model's declaration of its +concepts (TECO's 11 variables), not a dump of everything the dataset can express. +Concretely it is the repo's own `predicates.yaml` (§2): rendered as a user-owned starter with working +example predicates over the synthetic test vocabulary, then replaced by the author with their real +concepts. **The tests always read this file** (§8) — there is no separate test-only predicates file, +so the file the model runs with in production is the file the tests exercise. + +The unsupported-entry policy, stated precisely because it matters: the reader always parses the +**entire** file first, then applies one of two deterministic behaviors. **Default (skip)**: entries +using an unsupported form — value bounds, `and()`, sequential/temporally-scoped derived predicates, +`???` placeholders, any key the reader does not recognize — are skipped with a **warning listing them +by name**, **cascading** (an `expr` referencing a skipped predicate is itself skipped and logged), and +the skipped names are recorded in the manifest (`featurization.skipped`, §3) so the feature-space +provenance outlives the log. With `featurization_strict: true`: any unsupported entry is instead a +**hard error listing every offending predicate by name**, before any shard is touched. In **both** +modes, ending with zero featurizable predicates is a hard error — an empty feature space is never a +valid artifact. There is no third mode — entries are never dropped without having been read and named. + +Skip is a safe default only because the loss is not silent — it is named in the log and in the +manifest — and because strictness is enforced where the author's intent actually lives: the rendered +predicates-file validation test (§8) asserts the repo's **own** `predicates.yaml` parses with **zero +skips**. Unsupported forms in your own declaration are a bug in your file; the same forms in a foreign +file at runtime (a raw dataset `predicates.yaml` in a MEDS-DEV run, §9) degrade gracefully, with the +presence signal flowing through the plain sibling predicates those files already contain. + +**Considered and rejected: reading bounded predicates but ignoring their bounds** (treating +`abnormally_high_creatinine_1: {code: ..., value_min: 1.3}` as its code matcher). Two reasons. First, +the predicate *name* carries the threshold semantics: a `predicate//abnormally_high_creatinine` column +that actually means "any creatinine measurement" is a lie under a trustworthy name, propagated through +`features.json` to everyone who interprets the model — a feature lost visibly (skip) is strictly +better than a feature kept invisibly wrong. Second, in presence-only space it buys nothing: bounded +predicates share codes with their unbounded siblings (MIMIC-IV's `creatinine_1` vs +`abnormally_high_creatinine_1` are the same code), so ignoring bounds manufactures byte-identical +duplicate columns under different names, while skip mode already keeps the presence signal via the +plain siblings. The one losing case — a file defining only the bounded form of a concept — is fixed by +one line of curation (add the plain predicate; the skip log should say exactly that), and when value +bounds return with the value-channel layer they will apply *as stated*, with no version boundary where +a column changes meaning under an unchanged name. + +Column naming: `predicate//`, dtype `Int8`, dense 0/1 (not null/1 — dense columns keep model-side +logic trivial). The reader validates that no generated column collides with an existing column and that +no two predicates share a name. + +**Why not depend on `es-aces`:** the template's standing rule is that it never parses ACES *task* +definitions (they need dataset predicates it does not have). That reasoning does not forbid parsing a +predicates file handed in explicitly — but the dependency does drag in its own `meds` pin (MEDS-DEV +pins `meds==0.3.3` against the template's `meds~=0.4`), and the supported subset above is ~50 lines of +polars. If the subset ever grows toward real ACES semantics, revisit; do not reimplement ACES here. + +--- + +## 5. What breaks without MTD — the four seams + +Established by tracing every consumer of the `patients` artifact. The contract layer — manifests, +`write_artifact`/`read_manifest`, `read_labels`/`split_labels`, coverage checking, `PredictionSchema`, +`resolve_workspace` — is representation-agnostic and needs **no changes**. Four seams are MTD-shaped: + +1. **`tasks.tokenized_cohort()` — the one hard failure.** Reads + `patients/tokenization/schemas//*.parquet`; raises "not a tensorized cohort" otherwise. + Called from `materialize_labels`, i.e. from **`supervised_train`, `infer` and `predict`** — every + task-conditioned command dies before touching the model. +2. **The datamodule.** All commands call `build_datamodule(cfg)`; the shipped datamodule configs + target MTD's `Datamodule`; commands patch `cfg.datamodule.config.task_labels_dir` at runtime. +3. **`vocab_size`.** Both training commands' `build_module` pass + `vocab_size=datamodule.config.vocab_size` — an MTD config attribute. +4. **Prediction key alignment.** `_runtime.run_predict_step` takes `(subject_id, prediction_time)` + from `dataset.schema_df`, in loader order, via the `SPLIT_ATTRS` attribute names. + +--- + +## 6. Contract changes (all in `meds_model_base`, both backends, unconditional) + +### 6.1 `cohort_subjects()` — manifest-driven split membership + +Replace the direct `tokenized_cohort()` call inside `materialize_labels` with a dispatcher: + +```python +def cohort_subjects(patients_dir: Path) -> dict[str, set[int]]: + match read_manifest(patients_dir).extras["representation"]: # sketch + case "mtd": + return tokenized_cohort(patients_dir) # today's reader + case "predicates": + return featurized_cohort(patients_dir) # subject_id off data//*.parquet +``` + +`featurized_cohort` reads subject ids from the data shards themselves — the same +shard-path-is-the-authority philosophy, different files. The documented guarantee of `split_labels` +(labels partitioned by the cohort the model will actually see, so `CoverageError` cannot fire two +training runs late) carries over verbatim. A manifest without `representation` (artifact predating this +change) is treated as `mtd`. + +### 6.2 Import guard in `lightning.register_structured_configs` + +The dispatcher calls it before composing **every** command's config, and it imports `meds_torchdata` +unconditionally — with MTD uninstalled, every command dies at import, including `--help`. Change: if +`meds_torchdata` is not importable, skip registration and return. The `MEDSTorchDataConfig` config-store +group is referenced only by the MTD datamodule yamls, which a custom-backend repo does not render; if a +user composes one anyway, Hydra's missing-group error names the config, which is the right message. +`build_datamodule` itself is `instantiate(cfg.datamodule)` and needs nothing. + +### 6.3 `vocab_size` stays, by protocol + +The protocol (§7) requires `config.vocab_size`; a featurized datamodule sets it to the feature count +read from `features.json`. The name is slightly wrong for features, but renaming a contract kwarg +ripples into every generated model's `build_module` signature for zero behavior — not now. A model +wanting a different signature already has the designed escape hatch: subclass the command in user-owned +`commands.py` and override `build_module`. + +### 6.4 Predict alignment stays, by protocol + +`run_predict_step`, `stack_outputs`, `_check_coverage` and the whole `_PredictRunMixin` flow are +untouched provided the dataset exposes `schema_df` (§7). Two error strings that say "tensorized cohort" +get softened to name both layouts. + +--- + +## 7. The datamodule protocol + +The duck-typed surface the commands already rely on, promoted to a named, documented `Protocol` in +`meds_model_base` so a custom datamodule knows exactly what to implement. A conforming datamodule: + +* is a `lightning.pytorch.LightningDataModule`; +* has a `config` object with: + * `task_labels_dir` — **settable at runtime**; consumes the `{split}.parquet` layout that + `materialize_labels` writes (this layout is the interface; `boolean_value` may be absent — that is + how inference-without-ground-truth reaches the batch); + * `vocab_size` — feature-space size (`len(features.json)` for the predicates backend); +* accepts `batch_size` and `num_workers` (the conformance harness and every command pass them); +* exposes the `SPLIT_ATTRS` surface: `train_dataset`/`train_dataloader`, `val_dataset`/`val_dataloader` + (tuning), `test_dataset`/`test_dataloader` (held_out); +* each dataset exposes `schema_df`: a polars frame with `subject_id`, `prediction_time` rows **in + loader iteration order** (predict loaders must not shuffle) — this is what lets `predict` align model + outputs to timepoints without models hand-aligning anything. + +The template ships the protocol and a rendered stub implementing none of it (`datamodule.py`, in +`_skip_if_exists`) — consistent with the no-model rule: the aggregation strategy over predicate columns +*is* modeling, so it lives in user-owned code. + +--- + +## 8. Tests + +Guiding principle: both backends' contract code ships in every repo, so **contract tests run +everywhere**; only dependency-bound tests condition on the environment, via +`pytest.importorskip("meds_torchdata")` — the dependency analogue of `skip_if_stub`: green, but honest. + +### Tier 1 — template repo `tests/` (structural, no torch, seconds) + +* Parametrize the render tests over the full **7 profiles × 2 backends** (cheap here). +* Per-backend assertions: dependency present/absent in rendered `pyproject.toml`/`requirements.txt`; + correct datamodule configs/stubs rendered; correct `featurization` default; fixtures wired. +* `test_rendered_repo_is_clean_as_written` runs on both — the new Jinja conditionals are exactly the + artefact class it exists to catch. + +### Tier 2 — generated-repo conformance suite + +The predicates path needs **no extra dependencies**, so it is testable in every repo regardless of the +copier answer; the MTD path is testable only where MTD is installed. + +* **One predicates file, the model's own.** The tests read the repo's `predicates.yaml` — the same + file production runs use — never a test-only copy; a parallel test fixture would mean the model's + real predicates are the one thing the suite never exercises. The rendered *starter* content of that + file doubles as the working example: every reader form over the stable `meds_testing_helpers` + synthetic vocab — exact (`HR`), regex (`^ADMISSION//.*`), any-of (`{ any: [HR, TEMP] }`), + `or(hr, temp)` — so a fresh repo passes out of the box, and an author who replaces the content keeps + the same tests running against *their* concepts (all-zero columns over synthetic data are a logged + warning, §3, and the machinery assertions do not depend on match counts). The session `data_dir` + fixture passes `external_predicates_file=/predicates.yaml` via a new `extra_args` passthrough + on `build_workspace` (the only harness change). +* **Reader unit tests**: each grammar form against literal YAML snippets inline in the test (grammar is + the template's, not the model's — literals are correct here); the skip cascade and manifest + recording; the strict-mode errors; the always-error cases (colliding names, zero remaining + predicates, missing file with `featurization=predicates`). +* **Predicates-file validation test**: parse the repo's `predicates.yaml` and fail if **any predicate + would be skipped**, listing them. This is where strictness lives under the skip-by-default policy + (§4): unsupported forms in the model's own declaration are a bug in the file even though runtime + tolerates them in foreign files. Cheap, always-run, and the first thing to break — directly — when + an author's edit introduces an unsupported form, rather than indirectly through a warning nobody + reads. +* **Featurizer tests**: 0/1 values against hand-computed rows; original columns untouched; shard layout + and `metadata/` handling per §3 (including the `subject_splits.parquet` **absence**); `features.json` + and manifest fields. +* **`test_smoke_pipeline`, test by test**: + * `test_workspace_is_published_with_manifests` — swap `tokenized_cohort` for `cohort_subjects`; + assertions otherwise verbatim on both backends. + * `test_inference_never_sees_ground_truth` — split: the `materialize_labels` half (does + `include_labels=False` omit `boolean_value` from the written parquet) is backend-agnostic and + load-bearing — `predict` enforces the property at the command level; it runs everywhere. The + `MEDSPytorchDataset` half is an MTD-behavior check; `importorskip`. The equivalent guarantee for a + custom datamodule becomes a protocol obligation plus a `skip_if_stub`-style test that activates + when `datamodule.py` is implemented. + * `test_conflicting_sources_are_rejected` — backend-agnostic, unchanged. + * `test_end_to_end` — **unchanged**. It skips on the stub in a fresh repo on either backend (the + template's documented "nothing trains" gap is backend-independent). Every argument `run_chain` + passes is backend-neutral, so when a model and datamodule exist the same chain runs without MTD. +* **The final end-to-end using features** (what `test_end_to_end` executes once a model exists): + 1. `preprocess_data external_meds_dir= featurization=predicates + external_predicates_file=/predicates.yaml output_data_dir=ws` → + `ws/patients` with `predicate//*` columns and `representation: predicates`; + 2. `supervised_train input_data_dir=ws external_labels_dir= ...` — the custom + datamodule reads featurized shards + materialized `{split}.parquet`, `vocab_size` = count from + `features.json`, 2 CPU epochs; + 3. `predict input_supervised_model_dir=... splits=[held_out] ...` — keys via `schema_df`, coverage + `n_written == n_expected`; + 4. assertions verbatim from today: `PredictionSchema.align`, manifests, coverage recorded. + + The stronger variant is `test_property` (designed-signal learnability + negative control); its + featurized form has one constraint worth stating: the planted signal must be **predicate-visible** — + the model only sees the feature columns its datamodule selects, so the label-carrying code must be + matched by a declared predicate, or the model is structurally blind to the label and the learnability + assertion fails with nothing actually broken. Consistent with the one-file principle, the featurized + variant does **not** ship a signal-specific predicates fixture. Instead the dependency inverts: + `build_signal_dataset` gains a from-predicates mode that generates the synthetic dataset **out of the + model's own `predicates.yaml`** — it selects the predicates with literal codes (exact and any-of; a + regex cannot be reverse-instantiated into a code), designates one as the signal predicate and the + rest as distractors, emits their codes with the signal's presence carrying the label and the + distractors label-independent, and keeps the existing anti-leak construction (random insertion + position, label-independent sequence length, null `numeric_value`). The feature space the model is + tested on is therefore its *production* feature space — same names, same order, same + `features.json`. Distractors matter for the same reason as before: with a single feature the column + trivially equals the answer. Preconditions, stated rather than hidden: the file must contain **at + least two literal-code predicates** (one signal + one distractor); otherwise the test skips with a + message telling the author to add one or provide example codes (a per-predicate example-code sidecar + for regex predicates is possible future work, not v1). The rendered starter content satisfies the + precondition out of the box. The negative control is unaffected (shuffled labels defeat every + feature equally). Where it first runs: TECO is the real-world proof, but the featurized backend is also the + cheapest possible **reference implementation** for the template's agreed `examples/` follow-up + (predicate counts → linear head; no MTD, no heavy datamodule) — landing that would make this e2e run + in template CI itself, closing the "nothing trains" gap and proving the custom path in one move. +* **Equivalence guard** (MTD repos only, both workspaces built in-session): run the same labels through + `materialize_labels` against the MTD workspace and the predicates workspace; assert **identical split + partitions**. This is the test that stops the two cohort readers from drifting — the failure mode + that would otherwise surface as `CoverageError` two training runs late. +* **`test_cli_smoke` as the import-guard test, for free**: in a custom-rendered repo, + `meds_torchdata` is genuinely not installed, so `meds-model commands` / `--help` succeeding *is* the + regression test for §6.2. No mocking; the environment is the test. + +### Tier 3 — CI (`rendered-smoke`) + +Structural tier already covers all 14 combos; executing them all (~1 min/profile) is waste. Execution +matrix: the existing full-profile sweep on `mtd`, plus **`supervised` on `custom_featurization`** +(TECO's profile). The custom lane's venv deliberately omits meds-torch-data — the lane itself proves +nothing imports it. Update the documented rendered-suite expectations ("~37 passed, 2 skipped") so the +new skips stay legible under `-rs`. + +--- + +## 9. MEDS-DEV integration (via MEDS-DEV PR #325) + +[MEDS-DEV#325](https://github.com/Medical-Event-Data-Standard/MEDS-DEV/pull/325) adds the missing +plumbing: an optional `predicates_path` argument to `meds-dev-model`, exposed to model commands as the +`{predicates_path}` template variable. Properties that shape the integration: + +* **Pure explicit passthrough.** The caller of `meds-dev-model` supplies the path; there is no + automatic fallback to the registered dataset's `predicates.yaml` (unlike `meds-dev-task`). Auto- + binding is a possible MEDS-DEV follow-up, noted in the PR. +* **Backwards compatible**: models that do not reference the variable are unaffected. A model that + references it without the argument fails with a `KeyError` at command-formatting time — blunt but + early. + +### Template wiring (gated on #325 merging) + +* **`model.yaml`** under `custom_featurization`: the first command of each chain renders with + `external_predicates_file={predicates_path}` appended to `preprocess_data`. This makes + `predicates_path` *required* for such a model in MEDS-DEV — the honest contract, since + `featurization=predicates` without a file is an error by design (§3). +* **`meds-model-add-to-meds-dev`**: when registering a `custom_featurization` model, copy the repo's + `predicates.yaml` into the model's MEDS-DEV directory alongside `model.yaml` and `requirements.txt`, + so benchmark users have a reference file to pass: + `meds-dev-model ... predicates_path=/models//predicates.yaml`. +* **`test_meds_dev_e2e`**: the rendered e2e adds `predicates_path=/predicates.yaml` to its + `meds-dev-model mode=full dataset_type=full` invocation — the one-file principle extended to the + outermost test. Expectation to state plainly: over the MEDS-DEV demo dataset's vocabulary the + model's predicates may match nothing, so features can be all-zero (a §3 warning) and predictions + near-constant — the e2e asserts *plumbing* (venv build, placeholder fill, chain execution, + `PredictionSchema`, coverage), not learnability, and all of that holds on all-zero features. + Until #325 merges, this lane runs only against a checkout of the PR branch (`MEDS_DEV_DIR` + override); it must not be wired to clone `main`. + +### What predicates file does a MEDS-DEV caller pass? + +The contract: a file whose parseable predicates define the model's feature space *for that dataset*. +The model's shipped `predicates.yaml` is the reference (correct for the dataset it was curated +against); a caller running another dataset supplies a binding file for it. Passing a raw MEDS-DEV +dataset `predicates.yaml` **works by default**: its value-bounded task predicates (e.g. MIMIC-IV's +`abnormally_high_*`) are skipped under §4's policy — warned, cascaded, recorded in the manifest — +while the presence signal flows through the plain sibling predicates those files already define. +Callers who want a hard failure on any unsupported entry pass `featurization_strict=true`. Note the consequence of a trained model's feature space being +`features.json` (§3): predict-time featurization must use the same bindings the model was trained +with, and the digest in the manifest is what makes a mismatch detectable. + +--- + +## 10. Future work (recorded, not scoped) + +* **Value channels**: per-predicate value extraction/transformation (units, parsing, clipping) — + dftly is the candidate expression language (row-wise, YAML-native, same author, no temporal + semantics; the temporal grammar would be ours). Value-bounded predicates (dropped from v1, §4) + return here, where thresholds belong. +* **Reference implementation in `examples/`** (predicate counts → linear head over the + `custom_featurization` backend): closes the template's "nothing trains" gap and gives the features + e2e a home in template CI (§8). +* **Temporal spec**: grid/window/aggregation/imputation declarations — at that point this stops being + a template feature and becomes an ecosystem spec ("what ACES is to cohorts"), with MEDS-Tab's + windowed-aggregation machinery and TECO's `IntervalGrid` as the two existing consumers to unify. +* **MEDS-DEV auto-binding**: `predicates_path` defaulting to the registered dataset's + `predicates.yaml` when `dataset_name` is set (the follow-up MEDS-DEV#325 explicitly leaves open), + plus per-dataset featurization compatibility metadata (which concepts a dataset can bind, validated + at registry time — requires a `Metadata` schema change, which currently accepts exactly + `description`/`contacts`/`links` and raises on anything else at import). +* **Benchmark fairness**: concept-mapped vs vocab-generic models are different lanes; MEDS-DEV results + should eventually record which lane a result came from. + +--- + +## 11. Implementation order + +1. **Seam refactor** — `cohort_subjects` dispatch (+ `representation: mtd` written by the existing + branch), import guard, error-string softening, protocol module. Template stays green; MTD behavior + byte-identical. +2. **Featurizer** — reader + shard rewriter + `features.json` + manifest fields + `_validate_featurized`; + unit and featurizer tests (these run in MTD repos too). +3. **`data_backend` copier option** — copier.yml question, dependency/config/stub gating, per-backend + conftest rendering, tier-1 test parametrization. +4. **CI lane** — `supervised` × `custom_featurization` rendered-smoke job; docs expectation updates. +5. **MEDS-DEV wiring** (gated on MEDS-DEV#325 merging) — `model.yaml` renders + `external_predicates_file={predicates_path}`, `meds-model-add-to-meds-dev` ships `predicates.yaml`, + `test_meds_dev_e2e` passes `predicates_path` (§9). +6. **TECO adoption** — datamodule implementing the protocol over predicate columns; the first + end-to-end run of the custom path, and the point where `IntervalGrid`'s feature axis becomes the + predicate set instead of the code vocabulary. diff --git a/template/CLAUDE.md.jinja b/template/CLAUDE.md.jinja index 6ee7ec1..0c9e049 100644 --- a/template/CLAUDE.md.jinja +++ b/template/CLAUDE.md.jinja @@ -22,6 +22,26 @@ This DAG's chain: `README.md` documents the commands, their arguments and the artifact layout — read it for usage. This file covers what is easy to get wrong from the inside. +## Data backend + +Generated with `data_backend={{ data_backend }}`. +{% if data_backend == 'custom_featurization' %} +`preprocess_data` runs presence featurization: 0/1 `predicate//` columns stamped onto +split-sharded MEDS parquet, plus `features.json` (the artifact's feature space — order matters, a +trained checkpoint depends on it). The predicates come from `predicates.yaml` at the repo root — the +model's own concept declaration; the tests, MEDS-DEV runs (`predicates_path=`) and production all read +that ONE file, so edit it, don't fork it. **The datamodule over the featurized artifact is yours**: +implement `src/{{ model_slug }}/datamodule.py` against the surface in +`src/meds_model_base/lightning/protocol.py` (settable `config.task_labels_dir`, `config.vocab_size`, +the split dataset/dataloader attributes, and `schema_df` in loader order). meds-torch-data is +deliberately not installed here; the two tests that need it skip, honestly. +{% else %} +`preprocess_data` tensorizes via meds-torch-data and the shipped datamodule configs consume that +artifact. The presence-featurization path (`featurization=predicates` + `predicates.yaml`) also ships +and is exercised by the tests; a repo that should featurize instead can flip the answer and run +`copier update`. +{% endif %} + ## Porting an existing model? Read the procedure first If the task is to reimplement a published model (a package, a paper repo) here, read @@ -41,6 +61,8 @@ test fixture is the failure mode that document exists to prevent** — fix the f | Path | Owner | On `copier update` | |---|---|---| | `src/{{ model_slug }}/model.py` | you | never touched (`_skip_if_exists`) | +{% if data_backend == 'custom_featurization' %}| `src/{{ model_slug }}/datamodule.py` | you | never touched | +{% endif %}| `predicates.yaml` | you | never touched | {% if profile in ['zero_shot_direct', 'packaged'] %}| `src/{{ model_slug }}/predict.py` | you | never touched | {% endif %}| `src/{{ model_slug }}/commands.py` | you | never touched | | `src/{{ model_slug }}/configs/{model,paths,profile}/` | you | never touched | diff --git a/template/README.md.jinja b/template/README.md.jinja index bc565f8..81ba042 100644 --- a/template/README.md.jinja +++ b/template/README.md.jinja @@ -114,6 +114,28 @@ A logger group member is a mapping of *name → logger config*: `build_trainer` `_target_` child of `cfg.logger`, so one file can enable several loggers at once, and a file with a bare top-level `_target_` enables none of them without failing. +### Reproducibility + +`pretrain` and `supervised_train` seed with `seed_everything(seed, workers=True)` (`seed: 0`; `null` skips +seeding entirely), `num_workers` defaults to `0`, and the Trainer ships `deterministic: warn` and +`benchmark: false`. What that buys is worth stating exactly, because it is less than "reproducible": + +> **same seed + same config + same environment → same metric.** + +Changing `num_workers` or `batch_size` alters RNG consumption order and legitimately changes results, and +reproducibility across different hardware or library versions is not reachable at all. + +```bash +meds-model supervised_train trainer.deterministic=true # raise instead of warn on a non-deterministic op +meds-model supervised_train trainer.precision=bf16-mixed # trade precision for speed +``` + +`CUBLAS_WORKSPACE_CONFIG=:4096:8` is set for you — without it `deterministic=true` raises on CUDA at the +first matmul rather than at startup. Float32 matmul precision is left at torch's default (`highest`); this +repository does not quietly run fp32 matmuls in bfloat16 to buy speed. If you need that, `precision` above +is the better lever, and a model port that requires TF32 specifically should set it in `model.py` and +record it. + ### Two rules worth knowing **Sources are alternatives, not layers.** `supervised_train` and `predict` each accept several optional diff --git a/template/docs/PORTING-A-MODEL.md.jinja b/template/docs/PORTING-A-MODEL.md.jinja index 5130fbd..4886f16 100644 --- a/template/docs/PORTING-A-MODEL.md.jinja +++ b/template/docs/PORTING-A-MODEL.md.jinja @@ -73,11 +73,14 @@ Config paths below are relative to `src/{{ model_slug }}/configs/`. | source element | where it goes in this repository | |---|---| | preprocessing / tokenisation pipeline | `preprocess_data.yaml` → `pipeline:` (a MEDS-transforms YAML), or a documented pre-pass | +{% if data_backend == 'custom_featurization' %}| named clinical variables / feature list | `predicates.yaml` — one predicate per source variable, bound to the target dataset's codes; each becomes a `predicate//` column | +| feature construction (grids, aggregation, imputation, value use) | `src/{{ model_slug }}/datamodule.py` — over the featurized artifact, against `meds_model_base/lightning/protocol.py` | +{% endif %}| model architecture | `src/{{ model_slug }}/model.py` — wrap the source package as a dependency, do not vendor a copy | | model architecture | `src/{{ model_slug }}/model.py` — wrap the source package as a dependency, do not vendor a copy | | pretraining objective | `Model.compute_loss` when `not batch.has_labels` — delegate to the source's own loss where possible | | task objective | `Model.compute_loss` when `batch.has_labels` | | optimiser / schedule / LR | `optimizer/*.yaml`, and override `configure_optimizers` if the source groups parameters | -| seeding / determinism / precision | `seed` in `pretrain.yaml` and `supervised_train.yaml`; `deterministic` and `benchmark` in `trainer/`; matmul precision (hardcoded in the vendored contract today — see [MEDS_model_template#6](https://github.com/florian6973/MEDS_model_template/issues/6)) | +| seeding / determinism / precision | `seed` in `pretrain.yaml` and `supervised_train.yaml`; `deterministic`, `benchmark` and `precision` in `trainer/`. Float32 matmul precision is torch's own default (`highest`) — the contract does not set it, so if the source ran TF32, call `torch.set_float32_matmul_precision` in `model.py` and give it a ledger row | | generation / inference | `Model.predict_step`, `infer_step`, or `predict.py` depending on profile | | size presets | `model/default.yaml` + `DEFAULT_*` in `model.py` | | logged metrics | the metrics dict returned from `compute_loss` | diff --git a/template/model.yaml.jinja b/template/model.yaml.jinja index 8d61aa8..8ee1911 100644 --- a/template/model.yaml.jinja +++ b/template/model.yaml.jinja @@ -20,6 +20,14 @@ ## extra key here raises TypeError and takes down the whole `meds-dev-*` CLI, not just this model. ## In particular `requirements.txt` must NOT be declared: the loader discovers it from the model ## directory on disk. +{% set preprocess_extra = ' external_predicates_file={predicates_path}' if data_backend == 'custom_featurization' else '' -%} +{% if data_backend == 'custom_featurization' -%} +## +## This model featurizes via predicates, so `{predicates_path}` is REQUIRED: pass +## `predicates_path=` to `meds-dev-model` (MEDS-DEV >= PR#325; the model's own predicates.yaml — +## shipped into this model's MEDS-DEV directory by `meds-model-add-to-meds-dev` — is the reference). +## Without it, command formatting fails with a KeyError naming predicates_path. +{% endif -%} metadata: description: "{{ model_description }}" contacts: @@ -29,13 +37,13 @@ metadata: commands: {% if profile == 'supervised' %} supervised: train: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model supervised_train input_data_dir={output_dir}/data external_labels_dir={labels_dir} output_supervised_model_dir={output_dir}/model predict: |- meds-model predict input_supervised_model_dir={model_initialization_dir}/model external_labels_dir={labels_dir} splits=[{split}] output_predictions_dir={output_dir} do_overwrite=true {% elif profile == 'finetune' %} unsupervised: train: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model pretrain input_data_dir={output_dir}/data output_pretrained_model_dir={output_dir}/pretrained supervised: train: |- @@ -44,7 +52,7 @@ commands: meds-model predict input_supervised_model_dir={model_initialization_dir}/model external_labels_dir={labels_dir} splits=[{split}] output_predictions_dir={output_dir} do_overwrite=true {% elif profile == 'probe' %} unsupervised: train: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model pretrain input_data_dir={output_dir}/data output_pretrained_model_dir={output_dir}/pretrained supervised: train: |- @@ -54,7 +62,7 @@ commands: meds-model predict input_supervised_model_dir={model_initialization_dir}/model external_labels_dir={labels_dir} splits=[{split}] output_predictions_dir={output_dir} do_overwrite=true {% elif profile == 'zero_shot_materialized' %} unsupervised: train: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model pretrain input_data_dir={output_dir}/data output_pretrained_model_dir={output_dir}/pretrained supervised: train: null @@ -65,11 +73,11 @@ commands: supervised: train: null predict: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model predict input_data_dir={output_dir}/data external_labels_dir={labels_dir} splits=[{split}] output_predictions_dir={output_dir} do_overwrite=true {% else %} unsupervised: train: |- - meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data do_overwrite=true + meds-model preprocess_data external_meds_dir={dataset_dir} output_data_dir={output_dir}/data{{ preprocess_extra }} do_overwrite=true meds-model pretrain input_data_dir={output_dir}/data output_pretrained_model_dir={output_dir}/pretrained supervised: train: null diff --git a/template/predicates.yaml.jinja b/template/predicates.yaml.jinja new file mode 100644 index 0000000..7fc2fe5 --- /dev/null +++ b/template/predicates.yaml.jinja @@ -0,0 +1,32 @@ +# Featurization predicates for {{ model_name }} — THE MODEL'S declaration of its concepts. +# +# `preprocess_data featurization=predicates external_predicates_file=` stamps one 0/1 +# `predicate//` column per predicate onto the MEDS data. This file is user-owned (never +# overwritten by `copier update`) and is the ONE predicates file: the tests read it, MEDS-DEV runs pass +# it (predicates_path), and production runs use it — replace the starter content below with your real +# concepts, bound to your dataset's codes. +# +# Supported forms (ACES predicate syntax, presence subset — see src/meds_model_base/featurize.py): +# exact: code: SOME//CODE +# regex: code: { regex: "^PREFIX//.*" } # str.contains — put anchors in the pattern +# any-of: code: { any: [CODE_A, CODE_B] } +# derived: expr: or(name_a, name_b) +# Anything else (value bounds, and(), ...) is skipped with a warning; `pytest` fails if THIS file needs +# skipping, because your own declaration should never contain forms featurization cannot use. +# +# The starter predicates below bind the synthetic vocabulary of the meds-testing-helpers dataset the +# test suite runs on, so a fresh repository passes out of the box. At least two predicates with literal +# codes (exact / any-of) must remain for the designed-signal property test to run. +predicates: + heart_rate: + code: HR + temperature: + code: TEMP + admission: + code: { regex: "^ADMISSION//.*" } + discharge: + code: DISCHARGE + any_vital: + code: { any: [HR, TEMP] } + hr_or_temp: + expr: or(heart_rate, temperature) diff --git a/template/pyproject.toml.jinja b/template/pyproject.toml.jinja index fbde7b4..709e620 100644 --- a/template/pyproject.toml.jinja +++ b/template/pyproject.toml.jinja @@ -13,8 +13,8 @@ authors = [{ name = "{{ author_name }}", email = "{{ author_email }}" }] dependencies = [ # Core MEDS ecosystem (the vendored `meds_model_base` package builds on these). "meds~=0.4", - "meds-torch-data[lightning]~=0.9.0", - "MEDS-transforms>=0.6.7", +{% if data_backend == 'mtd' %} "meds-torch-data[lightning]~=0.9.0", +{% endif %} "MEDS-transforms>=0.6.7", "meds-evaluation~=0.0.6", "hydra-core>=1.3,<2", "lightning>=2.5,<3", diff --git a/template/src/meds_model_base/commands/infer.py b/template/src/meds_model_base/commands/infer.py index a55b535..328d428 100644 --- a/template/src/meds_model_base/commands/infer.py +++ b/template/src/meds_model_base/commands/infer.py @@ -117,7 +117,7 @@ def infer(self, cfg: DictConfig, module: pl_light.LightningModule, index: pl.Dat if not frames: raise RuntimeError( "The datamodule produced no rows for any requested split. Check that external_labels_dir " - "and the tensorized cohort refer to the same subjects." + "and the patients artifact refer to the same subjects." ) return pl.concat(frames, how="vertical_relaxed").unique(subset=KEYS, maintain_order=True) diff --git a/template/src/meds_model_base/commands/preprocess_data.py b/template/src/meds_model_base/commands/preprocess_data.py index 966b328..462be92 100644 --- a/template/src/meds_model_base/commands/preprocess_data.py +++ b/template/src/meds_model_base/commands/preprocess_data.py @@ -1,8 +1,15 @@ """``preprocess_data`` — external MEDS → this model's patient representation. Optionally runs a MEDS-transforms pipeline (``cfg.pipeline``, for model-specific enrichment such as -time-derived tokens or value binning), then meds-torch-data's ``MTD_preprocess`` to normalize, tokenize and -tensorize the dataset into the layout the training datamodule consumes. +time-derived tokens or value binning), then one of two representations (``cfg.featurization``, recorded +in the manifest as ``representation``): + +- ``mtd`` (default): meds-torch-data's ``MTD_preprocess`` normalizes, tokenizes and tensorizes the + dataset into the layout the MTD training datamodule consumes. +- ``predicates``: presence featurization (:mod:`meds_model_base.featurize`) — the data keeps its MEDS + layout, augmented with one 0/1 ``predicate//`` column per predicate in + ``cfg.external_predicates_file``, plus a ``features.json`` defining the feature space. Needs no + meds-torch-data; the model brings its own datamodule (see ``lightning/protocol.py``). Both stages are shelled out (they are Hydra applications with their own console scripts) and streamed live, so long runs show progress. The result is published atomically as ``/patients``: either the @@ -51,6 +58,27 @@ def run(self, cfg: DictConfig) -> Path: data_dir = Path(cfg.output_data_dir) patients_dir = data_dir / PATIENTS_SUBDIR pipeline = cfg.get("pipeline") + featurization = str(cfg.get("featurization") or "mtd") + if featurization not in ("mtd", "predicates"): + raise ValueError(f"featurization must be 'mtd' or 'predicates', got {featurization!r}.") + + parsed = predicates_file = digest = None + if featurization == "predicates": + from ..featurize import predicates_digest, read_predicates_file + + predicates_file = cfg.get("external_predicates_file") + if not predicates_file: + # An error, not a passthrough: silently publishing an artifact with zero feature + # columns would look exactly like a finished one. + raise ValueError( + "featurization=predicates requires external_predicates_file. Pass the model's " + "predicates file (e.g. external_predicates_file=predicates.yaml)." + ) + # Parsed before any heavy work: a bad predicates file fails in seconds, not after a pipeline. + parsed = read_predicates_file( + predicates_file, strict=bool(cfg.get("featurization_strict", False)) + ) + digest = predicates_digest(predicates_file) # Before the artifact is staged and long before a pipeline runs: this is the one precondition # whose violation is otherwise reported by MTD, several minutes in and in its own vocabulary. @@ -65,7 +93,7 @@ def run(self, cfg: DictConfig) -> Path: config=cfg, do_overwrite=bool(cfg.get("do_overwrite", False)), ) as (staging, extras): - mtd_input = external_meds_dir + source_input = external_meds_dir if pipeline: intermediate = staging.parent / f"{staging.name}.transforms" self._run_meds_transforms(pipeline, external_meds_dir, intermediate, cfg) @@ -73,14 +101,32 @@ def run(self, cfg: DictConfig) -> Path: # deliberately reshard. Checking anyway costs a directory listing and turns "MTD found no # schema files" back into a statement about the pipeline that actually caused it. _require_split_sharded(intermediate, "the pipeline output") - mtd_input = intermediate - - self._run_mtd(mtd_input, staging) - _validate_tensorized(staging) + source_input = intermediate extras["source"] = {"external_meds_dir": str(external_meds_dir)} - extras["tensorization"] = {"pipeline": str(pipeline) if pipeline else None} - extras.update(_describe_cohort(external_meds_dir, staging)) + if featurization == "predicates": + from ..featurize import featurize_dataset + from ..tasks import featurized_cohort + + counts = featurize_dataset(source_input, staging, parsed) + _validate_featurized(staging, parsed) + extras["representation"] = "predicates" + extras["featurization"] = { + "pipeline": str(pipeline) if pipeline else None, + "predicates_file": str(predicates_file), + "predicates_digest": digest, + "n_features": len(parsed.order), + "skipped": sorted(parsed.skipped), + "match_counts": counts, + } + cohort = featurized_cohort(staging) + else: + self._run_mtd(source_input, staging) + _validate_tensorized(staging) + extras["representation"] = "mtd" + extras["tensorization"] = {"pipeline": str(pipeline) if pipeline else None} + cohort = None + extras.update(_describe_cohort(external_meds_dir, staging, cohort=cohort)) logger.info("Patient data ready at %s.", patients_dir) return patients_dir @@ -189,7 +235,37 @@ def _validate_tensorized(output_dir: Path) -> None: ) -def _describe_cohort(meds_dir: Path, tensorized: Path) -> dict: +def _validate_featurized(output_dir: Path, parsed) -> None: + """The featurized twin of :func:`_validate_tensorized`: re-verify the artifact before publishing. + + Every declared predicate column must be present in every shard, and ``features.json`` must agree + with what was parsed — cheap re-reads of what :func:`~meds_model_base.featurize.featurize_dataset` + just wrote, standing between a partial write and a published artifact. + """ + import json + + import polars as pl + + from ..featurize import FEATURES_FILENAME + + features_fp = output_dir / FEATURES_FILENAME + if not features_fp.is_file(): + raise FileNotFoundError(f"Expected {features_fp}; featurization did not complete.") + declared = [f["column"] for f in json.loads(features_fp.read_text())["features"]] + if declared != parsed.columns: + raise RuntimeError( + f"{features_fp} disagrees with the parsed predicates: {declared} != {parsed.columns}." + ) + shards = sorted((output_dir / "data").rglob("*.parquet")) + if not shards: + raise FileNotFoundError(f"No data shards under {output_dir}/data; featurization wrote nothing.") + for fp in shards: + missing = set(declared) - set(pl.scan_parquet(fp).collect_schema().names()) + if missing: + raise RuntimeError(f"Shard {fp} is missing predicate column(s): {sorted(missing)}.") + + +def _describe_cohort(meds_dir: Path, tensorized: Path, cohort: dict | None = None) -> dict: """Cohort statistics for the manifest, counted from the tensorized output. Counting the *artifact* rather than ``external_meds_dir`` is what makes a filtering ``pipeline`` @@ -206,7 +282,8 @@ def _describe_cohort(meds_dir: Path, tensorized: Path) -> dict: from ..tasks import tokenized_cohort described: dict = {} - cohort = tokenized_cohort(tensorized) + if cohort is None: + cohort = tokenized_cohort(tensorized) if cohort: described["splits"] = {split: len(subjects) for split, subjects in sorted(cohort.items())} diff --git a/template/src/meds_model_base/commands/train.py b/template/src/meds_model_base/commands/train.py index 2ddebc0..5c98921 100644 --- a/template/src/meds_model_base/commands/train.py +++ b/template/src/meds_model_base/commands/train.py @@ -46,7 +46,6 @@ class _TrainFlow: output_key: ClassVar[str] def run(self, cfg: DictConfig) -> Path: - import torch from lightning.pytorch import seed_everything from ..lightning import build_datamodule, build_trainer @@ -57,7 +56,6 @@ def run(self, cfg: DictConfig) -> Path: if cfg.get("seed") is not None: seed_everything(cfg.seed, workers=True) - torch.set_float32_matmul_precision("medium") label_extras = self.prepare_labels(cfg, work_dir) datamodule = build_datamodule(cfg) @@ -196,7 +194,6 @@ def build_module(self, cfg, datamodule, source): # pragma: no cover - unused; s raise NotImplementedError("ProbeTrainCommand builds its module inside run().") def run(self, cfg: DictConfig) -> Path: - import torch from hydra.utils import instantiate from lightning.pytorch import seed_everything @@ -216,7 +213,6 @@ def run(self, cfg: DictConfig) -> Path: if cfg.get("seed") is not None: seed_everything(cfg.seed, workers=True) - torch.set_float32_matmul_precision("medium") labels_dir, label_summary = materialize_labels( cfg.external_labels_dir, Path(cfg.input_data_dir) / PATIENTS_SUBDIR, work_dir / "labels" diff --git a/template/src/meds_model_base/dispatch.py b/template/src/meds_model_base/dispatch.py index 9ec9487..671fd09 100644 --- a/template/src/meds_model_base/dispatch.py +++ b/template/src/meds_model_base/dispatch.py @@ -167,11 +167,18 @@ def _run_with_hydra(command: MEDSModelCommand, config_dir: str) -> None: paths (``meds-model commands`` / top-level ``--help``) stay cheap. The command is invoked via ``__call__``, not ``run``, so argument arbitration cannot be skipped. + + ``configure_cublas_workspace`` runs first, and this is the reason it lives here rather than in + ``train.py``: it must precede any CUDA work, and every command reaches torch through this function. + It sets an environment variable and imports nothing, so the torch-free introspection paths are + unaffected — they never get here. """ import hydra from .lightning import register_structured_configs + from .utils import configure_cublas_workspace + configure_cublas_workspace() register_structured_configs() hydra.main(version_base="1.3", config_path=config_dir, config_name=command.config_name)(command)() diff --git a/template/src/meds_model_base/featurize.py b/template/src/meds_model_base/featurize.py new file mode 100644 index 0000000..0881a06 --- /dev/null +++ b/template/src/meds_model_base/featurize.py @@ -0,0 +1,402 @@ +"""Presence featurization: stamp one 0/1 column per ACES predicate onto MEDS data. + +This is the ``featurization=predicates`` branch of ``preprocess_data`` (see +``docs/design-featurization.md`` in MEDS_model_template). A predicates YAML file — the same per-dataset +binding syntax MEDS-DEV ships for task extraction — is parsed down to its **presence subset**, and each +predicate becomes a dense ``predicate//`` Int8 column: 1 where the event matches, 0 elsewhere. +Everything else about the data is untouched: no rows dropped, no columns modified, same shard layout. +The model decides downstream whether to read ``numeric_value`` on active rows and how to aggregate. + +Supported predicate forms (deliberately *pure code matching* plus ``or()``): + +- exact: ``code: LAB//50912//mg/dL`` +- regex: ``code: { regex: "^ICU_ADMISSION//.*" }`` (``str.contains`` — anchors live in the pattern, + matching ACES semantics) +- any-of: ``code: { any: [HR, PULSE] }`` +- derived: ``expr: or(creatinine_1, creatinine_2)`` (column-wise max of its inputs) + +Everything else — value bounds, ``and()``, ``???`` placeholders, unknown keys — is *unsupported*: +skipped with a warning by default (cascading through ``or()`` expressions that reference a skipped +predicate), or a hard error under ``strict=True``. Skips are never silent: they are logged, returned to +the caller, and recorded in the artifact manifest. Value bounds in particular are a deliberate +omission, not a gap: a threshold is model-side semantics under presence featurization (the model reads +``numeric_value`` on active rows), and a bounded predicate whose bounds were quietly dropped would be a +lie under a trustworthy name. + +The template does not depend on ``es-aces`` for this: the supported subset is a few polars expressions, +and the dependency would drag in its own ``meds`` pin. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import re +from dataclasses import dataclass, field +from pathlib import Path + +import polars as pl +import yaml + +from .schemas import code_metadata_filepath, dataset_metadata_filepath + +logger = logging.getLogger(__name__) + +#: Every generated feature column is ``predicate//``; the prefix is what keeps the columns out of +#: the way of MEDS's own (and any pipeline-added) columns, and what tests key on. +PREDICATE_COLUMN_PREFIX = "predicate//" + +#: The ordered feature-space definition written at the artifact root. Feature order is part of what a +#: trained checkpoint means (feature *i* at training time must be the same predicate at predict time), +#: so consumers read this file rather than scanning parquet column order. +FEATURES_FILENAME = "features.json" + +_OR_EXPR = re.compile(r"^or\(\s*([A-Za-z0-9_]+(?:\s*,\s*[A-Za-z0-9_]+)+)\s*\)$") + + +class PredicatesError(ValueError): + """A predicates file (or the featurization it implies) is unusable, with the reason spelled out.""" + + +@dataclass +class ParsedPredicates: + """The outcome of parsing a predicates file down to the supported presence subset. + + ``exprs`` maps predicate name → boolean polars expression **over the base MEDS columns** for plain + predicates, and → list of input names for derived ``or()`` predicates (resolved at featurize time, + after the plain columns exist). ``order`` is the canonical feature order (file order, plains and + deriveds interleaved as declared); ``skipped`` maps skipped predicate name → reason. + """ + + plain: dict[str, pl.Expr] = field(default_factory=dict) + derived: dict[str, list[str]] = field(default_factory=dict) + order: list[str] = field(default_factory=list) + skipped: dict[str, str] = field(default_factory=dict) + + @property + def names(self) -> list[str]: + return list(self.order) + + def column(self, name: str) -> str: + return f"{PREDICATE_COLUMN_PREFIX}{name}" + + @property + def columns(self) -> list[str]: + return [self.column(n) for n in self.order] + + +def _parse_plain(name: str, spec: dict) -> pl.Expr | str: + """One plain predicate → a boolean expression, or a reason string if unsupported.""" + extra = sorted(set(spec) - {"code"}) + if extra: + return f"unsupported key(s) {', '.join(extra)} (presence featurization is code matching only)" + code = spec["code"] + match code: + case str() if code: + return pl.col("code") == code + case {"regex": str() as pattern} if pattern and len(code) == 1: + return pl.col("code").str.contains(pattern) + case {"any": list() as options} if ( + options and len(code) == 1 and all(isinstance(o, str) for o in options) + ): + return pl.col("code").is_in(options) + case _: + return f"unsupported code form {code!r} (expected a string, {{regex: ...}} or {{any: [...]}})" + + +def parse_predicates(raw: dict, *, strict: bool = False, source: str = "predicates") -> ParsedPredicates: + """Parse a loaded predicates mapping down to the supported presence subset. + + Args: + raw: the ``predicates:`` mapping of a predicates YAML file (name → spec). + strict: if True, any unsupported entry is a :class:`PredicatesError` listing every offender; + if False (default), unsupported entries are skipped with a warning, cascading through + ``or()`` expressions that reference them. + source: where the mapping came from, for error messages. + + Raises: + PredicatesError: on a non-mapping input, in strict mode on any unsupported entry, and always + when nothing featurizable remains — an empty feature space is never a valid outcome. + + Examples: + >>> parsed = parse_predicates({ + ... "hr": {"code": "HR"}, + ... "adm": {"code": {"regex": "^ADMISSION//.*"}}, + ... "vitals": {"code": {"any": ["HR", "TEMP"]}}, + ... "hr_or_adm": {"expr": "or(hr, adm)"}, + ... }) + >>> parsed.names + ['hr', 'adm', 'vitals', 'hr_or_adm'] + >>> parsed.skipped + {} + + Unsupported forms (here: value bounds, and()) skip with a cascade — ``high_or_low`` references a + skipped predicate, so it is skipped too: + + >>> parsed = parse_predicates({ + ... "hr": {"code": "HR"}, + ... "high_hr": {"code": "HR", "value_min": 110}, + ... "low_hr": {"code": "HR", "value_max": 40}, + ... "high_or_low": {"expr": "or(high_hr, low_hr)"}, + ... "both": {"expr": "and(hr, high_hr)"}, + ... }) + >>> parsed.names + ['hr'] + >>> sorted(parsed.skipped) + ['both', 'high_hr', 'high_or_low', 'low_hr'] + + The same file under ``strict=True`` is an error listing every offender: + + >>> parse_predicates({"hr": {"code": "HR"}, "high_hr": {"code": "HR", "value_min": 110}}, + ... strict=True) + Traceback (most recent call last): + ... + meds_model_base.featurize.PredicatesError: predicates has unsupported predicate(s) ...high_hr... + + Nothing featurizable is always an error: + + >>> parse_predicates({"high_hr": {"code": "HR", "value_min": 110}}) + Traceback (most recent call last): + ... + meds_model_base.featurize.PredicatesError: No featurizable predicate remains ... + """ + if not isinstance(raw, dict) or not raw: + raise PredicatesError(f"{source} must be a non-empty mapping of predicate name -> spec.") + + parsed = ParsedPredicates() + pending_derived: dict[str, list[str]] = {} + + for name, spec in raw.items(): + if not isinstance(spec, dict): + parsed.skipped[name] = f"spec is {type(spec).__name__}, expected a mapping" + elif "code" in spec: + result = _parse_plain(name, spec) + if isinstance(result, str): + parsed.skipped[name] = result + else: + parsed.plain[name] = result + elif "expr" in spec: + extra = sorted(set(spec) - {"expr"}) + m = _OR_EXPR.match(str(spec["expr"]).strip()) if not extra else None + if extra: + parsed.skipped[name] = f"unsupported key(s) {', '.join(extra)} on a derived predicate" + elif m is None: + parsed.skipped[name] = f"unsupported expr {spec['expr']!r} (only or(a, b, ...) is supported)" + else: + pending_derived[name] = [p.strip() for p in m.group(1).split(",")] + else: + parsed.skipped[name] = "neither a plain (code:) nor a derived (expr:) predicate" + + # Resolve derived predicates against what survived, iterating so or() over or() works; anything + # referencing a skipped or unknown name cascades into the skip list with the reference named. + progressed = True + while pending_derived and progressed: + progressed = False + for name, inputs in list(pending_derived.items()): + if all(i in parsed.plain or i in parsed.derived for i in inputs): + parsed.derived[name] = inputs + del pending_derived[name] + progressed = True + for name, inputs in pending_derived.items(): + missing = [i for i in inputs if i not in parsed.plain and i not in parsed.derived] + parsed.skipped[name] = "references skipped or unknown predicate(s) " + ", ".join(missing) + + parsed.order = [n for n in raw if n in parsed.plain or n in parsed.derived] + + if parsed.skipped: + listing = "; ".join(f"{n}: {why}" for n, why in parsed.skipped.items()) + if strict: + raise PredicatesError( + f"{source} has unsupported predicate(s) under featurization_strict: {listing}" + ) + logger.warning( + "Skipping %d unsupported predicate(s) from %s (featurization_strict=false): %s", + len(parsed.skipped), + source, + listing, + ) + + if not parsed.order: + raise PredicatesError( + f"No featurizable predicate remains in {source}: every entry was skipped " + f"({'; '.join(f'{n}: {why}' for n, why in parsed.skipped.items())}). " + "An empty feature space is never a valid artifact." + ) + return parsed + + +def read_predicates_file(path: Path | str, *, strict: bool = False) -> ParsedPredicates: + """Read and parse a predicates YAML file (a mapping with a ``predicates:`` key, as ACES files are).""" + path = Path(path) + if not path.is_file(): + raise PredicatesError(f"external_predicates_file {path} does not exist.") + try: + loaded = yaml.safe_load(path.read_text()) + except yaml.YAMLError as e: + raise PredicatesError(f"Could not parse {path}: {e}") from e + if not isinstance(loaded, dict) or "predicates" not in loaded: + raise PredicatesError( + f"{path} has no top-level 'predicates:' mapping; it does not look like a predicates file." + ) + return parse_predicates(loaded["predicates"], strict=strict, source=str(path)) + + +def predicates_digest(path: Path | str) -> str: + """sha256 of the predicates file bytes — the manifest's record of *which* predicates built an artifact. + + The manifest otherwise records only a path to a file that lives outside the artifact and can change + after the build; the digest is what makes two runs distinguishable and a drifted file detectable. + """ + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def literal_code_predicates(parsed: ParsedPredicates, raw: dict) -> dict[str, list[str]]: + """The featurized predicates whose codes can be *emitted* (exact / any-of), name → codes. + + A regex cannot be reverse-instantiated into a code, so this is the subset the synthetic + signal-dataset builder can plant events for (see ``testing.synthetic``). + + Examples: + >>> raw = {"hr": {"code": "HR"}, "vitals": {"code": {"any": ["HR", "TEMP"]}}, + ... "adm": {"code": {"regex": "^ADM.*"}}} + >>> literal_code_predicates(parse_predicates(raw), raw) + {'hr': ['HR'], 'vitals': ['HR', 'TEMP']} + """ + out: dict[str, list[str]] = {} + for name in parsed.order: + if name not in parsed.plain: + continue + code = raw[name]["code"] + if isinstance(code, str): + out[name] = [code] + elif isinstance(code, dict) and "any" in code: + out[name] = list(code["any"]) + return out + + +def featurize_frame(df: pl.DataFrame, parsed: ParsedPredicates) -> pl.DataFrame: + """Append the predicate columns to one MEDS frame. Pure augmentation — existing columns untouched. + + Examples: + >>> df = pl.DataFrame({"subject_id": [1, 1, 2], "code": ["HR", "DISCHARGE", "TEMP"], + ... "numeric_value": [99.0, None, 37.0]}) + >>> parsed = parse_predicates({"hr": {"code": "HR"}, "temp": {"code": "TEMP"}, + ... "vital": {"expr": "or(hr, temp)"}}) + >>> out = featurize_frame(df, parsed) + >>> out.columns + ['subject_id', 'code', 'numeric_value', 'predicate//hr', 'predicate//temp', 'predicate//vital'] + >>> out["predicate//hr"].to_list(), out["predicate//vital"].to_list() + ([1, 0, 0], [1, 0, 1]) + >>> out["predicate//hr"].dtype + Int8 + """ + collisions = sorted(set(parsed.columns) & set(df.columns)) + if collisions: + raise PredicatesError( + f"Input already carries predicate column(s) {', '.join(collisions)} — refusing to overwrite " + "(was this data already featurized?)." + ) + df = df.with_columns( + [expr.cast(pl.Int8).fill_null(0).alias(parsed.column(n)) for n, expr in parsed.plain.items()] + ) + for name in parsed.order: # file order; or() over or() resolves because inputs precede by iteration + if name in parsed.derived: + df = df.with_columns( + pl.max_horizontal([pl.col(parsed.column(i)) for i in parsed.derived[name]]) + .cast(pl.Int8) + .alias(parsed.column(name)) + ) + return df + + +def featurize_dataset(input_dir: Path, staging: Path, parsed: ParsedPredicates) -> dict[str, int]: + """Featurize every ``data//…`` shard of ``input_dir`` into the same layout under ``staging``. + + Copies ``metadata/codes.parquet`` and ``metadata/dataset.json`` through. Deliberately does **not** + copy ``metadata/subject_splits.parquet``: split membership travels as the shard layout and nothing + else — a copied splits table describes the *input* to preprocessing, and a filtering pipeline makes + those two different things (the same invariant the MTD branch keeps). + + Returns per-predicate match counts across all shards (events where the column is 1), which the + caller logs and folds into the manifest. + """ + import shutil + + data_in = Path(input_dir) / "data" + counts: dict[str, int] = dict.fromkeys(parsed.order, 0) + n_shards = 0 + for fp in sorted(data_in.rglob("*.parquet")): + rel = fp.relative_to(data_in) + if any(part.startswith(".") for part in rel.parts): + continue + out = featurize_frame(pl.read_parquet(fp), parsed) + dest = staging / "data" / rel + dest.parent.mkdir(parents=True, exist_ok=True) + out.write_parquet(dest) + n_shards += 1 + for name in parsed.order: + counts[name] += int(out[parsed.column(name)].sum()) + + if not n_shards: + raise PredicatesError(f"No data shards found under {data_in}; nothing to featurize.") + + for rel in (code_metadata_filepath, dataset_metadata_filepath): + src = Path(input_dir) / rel + if src.is_file(): + dest = staging / rel + dest.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(src, dest) + + (staging / FEATURES_FILENAME).write_text( + json.dumps( + { + "version": 1, + "features": [{"name": n, "column": parsed.column(n)} for n in parsed.order], + }, + indent=2, + ) + + "\n" + ) + + unmatched = sorted(n for n, c in counts.items() if c == 0) + if unmatched: + # A warning, not an error: on real data it signals a binding mismatch worth reading, but it is + # also the expected outcome when real-vocabulary predicates run over synthetic test data. + logger.warning( + "%d of %d predicate(s) matched no event at all: %s. If this is real data, check the " + "bindings; over a synthetic or foreign-vocabulary dataset this is expected.", + len(unmatched), + len(counts), + ", ".join(unmatched), + ) + return counts + + +def load_features(patients_dir: Path | str) -> list[dict[str, str]]: + """The ordered feature definitions of a featurized patients artifact (``features.json``). + + This — not parquet column order, not the predicates YAML — is what a datamodule reads to learn the + feature space: ``vocab_size = len(load_features(dir))``, columns selected in this order. + """ + fp = Path(patients_dir) / FEATURES_FILENAME + if not fp.is_file(): + raise PredicatesError( + f"No {FEATURES_FILENAME} in {patients_dir}; it is not a featurized patients artifact." + ) + return json.loads(fp.read_text())["features"] + + +__all__ = [ + "FEATURES_FILENAME", + "PREDICATE_COLUMN_PREFIX", + "ParsedPredicates", + "PredicatesError", + "featurize_dataset", + "featurize_frame", + "literal_code_predicates", + "load_features", + "parse_predicates", + "predicates_digest", + "read_predicates_file", +] diff --git a/template/src/meds_model_base/lightning/__init__.py b/template/src/meds_model_base/lightning/__init__.py index 6765c50..4620fa1 100644 --- a/template/src/meds_model_base/lightning/__init__.py +++ b/template/src/meds_model_base/lightning/__init__.py @@ -28,7 +28,15 @@ def register_structured_configs(group: str = "datamodule/config") -> None: global _STRUCTURED_CONFIGS_REGISTERED if _STRUCTURED_CONFIGS_REGISTERED: return - from meds_torchdata import MEDSTorchDataConfig + try: + from meds_torchdata import MEDSTorchDataConfig + except ImportError: + # meds-torch-data is optional: a repo generated with data_backend=custom_featurization does not + # install it. The config-store group registered here is referenced only by the MTD datamodule + # configs, which such a repo does not render — composing one anyway fails with Hydra's + # missing-group error, which names the config that needs the package. + _STRUCTURED_CONFIGS_REGISTERED = True + return MEDSTorchDataConfig.add_to_config_store(group) _STRUCTURED_CONFIGS_REGISTERED = True diff --git a/template/src/meds_model_base/lightning/modules.py b/template/src/meds_model_base/lightning/modules.py index 2a3d6bc..219a061 100644 --- a/template/src/meds_model_base/lightning/modules.py +++ b/template/src/meds_model_base/lightning/modules.py @@ -22,14 +22,23 @@ import lightning.pytorch as L import torch -from meds_torchdata import MEDSTorchBatch from torch import Tensor, nn if TYPE_CHECKING: # pragma: no cover - typing only from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler -PAD_INDEX = MEDSTorchBatch.PAD_INDEX +try: + from meds_torchdata import MEDSTorchBatch + + PAD_INDEX = MEDSTorchBatch.PAD_INDEX +except ImportError: + # meds-torch-data is optional (a data_backend=custom_featurization repo does not install it). This + # module must still *import* — the generated model stub subclasses BaseLightningModule — but the + # MTD-batch helpers (padding_mask, CodeEmbedder) are then unusable, which is correct: there is no + # MEDSTorchBatch to hand them. PAD_INDEX keeps MTD's value; annotations are strings (PEP 563). + MEDSTorchBatch = None + PAD_INDEX = 0 def padding_mask(batch: MEDSTorchBatch) -> Tensor: diff --git a/template/src/meds_model_base/lightning/protocol.py b/template/src/meds_model_base/lightning/protocol.py new file mode 100644 index 0000000..66736b6 --- /dev/null +++ b/template/src/meds_model_base/lightning/protocol.py @@ -0,0 +1,78 @@ +"""The datamodule surface the commands rely on, promoted from duck typing to a named contract. + +Every command constructs the datamodule through ``build_datamodule(cfg)`` (``hydra.utils.instantiate`` +of ``cfg.datamodule``) and then interacts with it **only** through the surface below. The shipped MTD +configs satisfy it via meds-torch-data; a repo generated with ``data_backend=custom_featurization`` +implements it in its own ``datamodule.py``. Conformance is structural — nothing checks +``isinstance``; these classes exist so the contract is written down in one place instead of scattered +across the commands that consume it. + +Beyond this surface, a conforming datamodule must be a ``lightning.pytorch.LightningDataModule`` (it is +handed to ``trainer.fit``) and must accept ``batch_size`` and ``num_workers`` constructor arguments — +every command config and the conformance harness pass them. + +Deliberately torch-free: importing this module must stay cheap (see the import-weight discipline note +in ``dispatch.py``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: # pragma: no cover - typing only + import polars as pl + + +class DataModuleConfig(Protocol): + """The ``config`` object commands read from and write to. + + ``task_labels_dir`` is **set at runtime**: ``supervised_train`` / ``infer`` / ``predict`` + materialize ``external_labels_dir`` into a work directory and point the datamodule at it via + ``cfg.datamodule.config.task_labels_dir = str(labels_dir)`` *before* instantiation. The layout at + that path is ``{split}.parquet`` files written by :func:`meds_model_base.tasks.materialize_labels`; + ``boolean_value`` may be absent (inference), which is how "prediction never reads ground truth" + reaches the batch a model sees. + + ``vocab_size`` is the feature-space size handed to ``build_module`` as the model's ``vocab_size`` + kwarg. For the MTD backend it is the code-vocabulary size; for a featurized backend it is the + feature count of the patients artifact (``len(features.json)``, + :func:`meds_model_base.featurize.load_features`). + """ + + task_labels_dir: str | None + vocab_size: int + + +class SplitDataset(Protocol): + """One split's dataset, as ``predict``/``infer`` consume it. + + ``schema_df`` is the alignment contract: a polars frame with ``subject_id`` and ``prediction_time`` + rows in **loader iteration order** (prediction loaders must not shuffle). ``run_predict_step`` + zips it against the model's per-batch outputs, which is what lets models never hand-align their + predictions to timepoints. + """ + + schema_df: pl.DataFrame + + def __len__(self) -> int: ... # pragma: no cover - protocol + + +class MEDSModelDataModule(Protocol): + """The datamodule attribute surface, keyed by MEDS split name via ``SPLIT_ATTRS``: + + ========== ================= ==================== + split dataset attribute dataloader attribute + ========== ================= ==================== + train ``train_dataset`` ``train_dataloader`` + tuning ``val_dataset`` ``val_dataloader`` + held_out ``test_dataset`` ``test_dataloader`` + ========== ================= ==================== + """ + + config: DataModuleConfig + train_dataset: SplitDataset + val_dataset: SplitDataset + test_dataset: SplitDataset + + +__all__ = ["DataModuleConfig", "MEDSModelDataModule", "SplitDataset"] diff --git a/template/src/meds_model_base/meds_dev.py b/template/src/meds_model_base/meds_dev.py index 5194dbf..7cd8f8c 100644 --- a/template/src/meds_model_base/meds_dev.py +++ b/template/src/meds_model_base/meds_dev.py @@ -44,6 +44,8 @@ SPEC_FILE = "model.yaml" REQUIREMENTS_FILE = "requirements.txt" README_FILE = "README.md" +#: Shipped alongside when the spec references {predicates_path}: the reference featurization bindings. +PREDICATES_FILE = "predicates.yaml" _PROBE = ( "import importlib.util as u; s = u.find_spec('MEDS_DEV'); " @@ -325,6 +327,11 @@ def main(argv: list[str] | None = None) -> int: payload = {SPEC_FILE: spec_text, REQUIREMENTS_FILE: requirements} if args.force or not (dest / README_FILE).is_file(): payload[README_FILE] = readme_stub(name, spec_description(spec_text), source) + # A model whose commands reference {predicates_path} needs a reference predicates file benchmark + # users can pass; shipping the repo's own predicates.yaml alongside model.yaml is what makes + # `meds-dev-model ... predicates_path=/models//predicates.yaml` possible. + if (repo / PREDICATES_FILE).is_file() and "{predicates_path}" in spec_text: + payload[PREDICATES_FILE] = (repo / PREDICATES_FILE).read_text() if args.dry_run: print(f"\nwould write {len(payload)} files under {dest}:") diff --git a/template/src/meds_model_base/tasks.py b/template/src/meds_model_base/tasks.py index 55ccc97..9eb309e 100644 --- a/template/src/meds_model_base/tasks.py +++ b/template/src/meds_model_base/tasks.py @@ -22,6 +22,7 @@ import polars as pl +from .manifest import read_manifest from .schemas import LabelSchema logger = logging.getLogger(__name__) @@ -144,6 +145,50 @@ def tokenized_cohort(patients_dir: Path | str) -> dict[str, set[int]]: return cohort +def featurized_cohort(patients_dir: Path | str) -> dict[str, set[int]]: + """Subjects present in a featurized (MEDS-layout) patients artifact, keyed by split. + + Read from ``data//*.parquet`` — in a featurized artifact the shard layout is the only record + of split membership, exactly as ``tokenization/schemas/`` is for a tensorized one (the artifact + deliberately carries no ``subject_splits.parquet``; see ``preprocess_data``). A shard with no leading + split component belongs to no split and is skipped rather than guessed at. + """ + data_dir = Path(patients_dir) / "data" + if not data_dir.is_dir(): + raise TaskMaterializationError( + f"No data/ shards under {patients_dir}; it is not a featurized patients artifact." + ) + cohort: dict[str, set[int]] = {} + for fp in sorted(data_dir.rglob("*.parquet")): + shard = fp.relative_to(data_dir) + if len(shard.parts) < 2: + continue + subjects = pl.read_parquet(fp, columns=["subject_id"])["subject_id"].to_list() + cohort.setdefault(shard.parts[0], set()).update(subjects) + return cohort + + +def cohort_subjects(patients_dir: Path | str) -> dict[str, set[int]]: + """Split membership of a patients artifact, dispatched on its manifest's ``representation``. + + The manifest is the authority on how the artifact stores split membership: ``mtd`` artifacts carry it + in ``tokenization/schemas/`` (read by :func:`tokenized_cohort`), ``predicates`` artifacts in the + ``data//`` shard layout (read by :func:`featurized_cohort`). A manifest without the field + predates the field and can only be an MTD artifact, so it is treated as one. + """ + representation = read_manifest(patients_dir).get("representation", "mtd") + match representation: + case "mtd": + return tokenized_cohort(patients_dir) + case "predicates": + return featurized_cohort(patients_dir) + case other: + raise TaskMaterializationError( + f"{patients_dir} declares representation {other!r}, which this version does not know how " + "to read splits from. Known representations: mtd, predicates." + ) + + def split_labels(labels: pl.DataFrame, cohort: dict[str, set[int]]) -> dict[str, pl.DataFrame]: """Partition labels by the split each subject was actually tokenized into. @@ -173,13 +218,13 @@ def split_labels(labels: pl.DataFrame, cohort: dict[str, set[int]]) -> dict[str, if not out: raise TaskMaterializationError( - "No label row matched any subject in the tensorized cohort. Either these labels are for a " + "No label row matched any subject in the patients cohort. Either these labels are for a " "different dataset, or the preprocess_data pipeline filtered the whole cohort away." ) if kept < labels.height: dropped = labels.filter(~pl.col("subject_id").is_in(sorted(set().union(*cohort.values())))) logger.warning( - "Dropping %d label row(s) for %d subject(s) absent from the tensorized cohort: either " + "Dropping %d label row(s) for %d subject(s) absent from the patients cohort: either " "filtered out by the preprocess_data pipeline (its manifest records how many), or not part " "of this dataset.", labels.height - kept, @@ -211,7 +256,7 @@ def materialize_labels( dest = Path(dest) dest.mkdir(parents=True, exist_ok=True) - by_split = split_labels(read_labels(external_labels_dir), tokenized_cohort(patients_dir)) + by_split = split_labels(read_labels(external_labels_dir), cohort_subjects(patients_dir)) for split, df in by_split.items(): if not include_labels: diff --git a/template/src/meds_model_base/testing/__init__.py b/template/src/meds_model_base/testing/__init__.py index 8e19331..e263b59 100644 --- a/template/src/meds_model_base/testing/__init__.py +++ b/template/src/meds_model_base/testing/__init__.py @@ -6,6 +6,8 @@ vacuously. - :func:`build_signal_dataset` — a classifier signal: a marker code deterministically predicts the label. +- :func:`signal_dataset_from_predicates` — the same, with codes drawn from the model's own predicates + file, so a featurized model is tested on its production feature space. - :func:`build_pattern_dataset` — a generative signal: a fixed repeating code pattern. - :func:`binary_auroc`, :func:`assert_learns_signal` — learnability assertions (+ negative control). - :func:`skip_if_stub` — skip a conformance test while ``model.py`` is still the generated stub. @@ -15,7 +17,12 @@ from .harness import build_workspace, run_chain, run_cli, supported_sources from .property import assert_learns_signal, binary_auroc from .stub import is_stub, skip_if_stub -from .synthetic import SIGNAL_CODE, build_pattern_dataset, build_signal_dataset +from .synthetic import ( + SIGNAL_CODE, + build_pattern_dataset, + build_signal_dataset, + signal_dataset_from_predicates, +) __all__ = [ "SIGNAL_CODE", @@ -29,4 +36,5 @@ "skip_if_stub", "build_pattern_dataset", "build_signal_dataset", + "signal_dataset_from_predicates", ] diff --git a/template/src/meds_model_base/testing/harness.py b/template/src/meds_model_base/testing/harness.py index e755d08..4485d11 100644 --- a/template/src/meds_model_base/testing/harness.py +++ b/template/src/meds_model_base/testing/harness.py @@ -49,12 +49,15 @@ def supported_sources(commands: Mapping[CommandName, type[MEDSModelCommand]], na return cls.supported_sources if cls.supported_sources is not None else frozenset(cls.sources) -def build_workspace(meds_root: Path, workspace: Path) -> Path: +def build_workspace(meds_root: Path, workspace: Path, extra_args: list[str] | None = None) -> Path: """Run ``preprocess_data``; return the shared ``data_dir``. No model is involved, so this half of the contract is exercisable from the moment a repository is generated — before its ``model.py`` stub is implemented. Labels are not preprocessed: each command materializes them from ``external_labels_dir`` itself. + + ``extra_args`` are appended verbatim — how a fixture selects the predicates backend + (``featurization=predicates external_predicates_file=...``) without this harness knowing about it. """ run_cli( [ @@ -62,6 +65,7 @@ def build_workspace(meds_root: Path, workspace: Path) -> Path: f"external_meds_dir={meds_root}", f"output_data_dir={workspace}", "do_overwrite=true", + *(extra_args or []), ] ) return workspace diff --git a/template/src/meds_model_base/testing/synthetic.py b/template/src/meds_model_base/testing/synthetic.py index 39997bd..3301cca 100644 --- a/template/src/meds_model_base/testing/synthetic.py +++ b/template/src/meds_model_base/testing/synthetic.py @@ -82,10 +82,12 @@ def build_signal_dataset( signal_rate: float = 0.5, seed: int = 0, shuffle_labels: bool = False, + signal_code: str = SIGNAL_CODE, + background: list[str] | None = None, ) -> Path: - """Write a classifier-signal MEDS dataset: the presence of :data:`SIGNAL_CODE` determines the label. + """Write a classifier-signal MEDS dataset: the presence of ``signal_code`` determines the label. - Each subject gets a short run of background codes; with probability ``signal_rate`` the ``SIGNAL_CODE`` + Each subject gets a short run of background codes; with probability ``signal_rate`` the ``signal_code`` is inserted. The boolean label is exactly whether the subject has the signal (so a model that reads the sequence can reach AUROC ≈ 1.0). ``shuffle_labels=True`` breaks the signal↔label link (negative control: nothing to learn, AUROC ≈ 0.5). @@ -93,11 +95,16 @@ def build_signal_dataset( **Only the signal code is informative.** Sequence length and the signal's position within the sequence are both label-independent by construction; see the comment in the loop for why that matters. + ``signal_code`` / ``background`` default to this module's fixed vocabulary; a featurized repo passes + codes drawn from its own predicates file instead (:func:`signal_dataset_from_predicates`), so the + model is tested on its *production* feature space. + Returns ``root`` (a MEDS dataset with a ``signal_task`` task-labels directory). """ rng = random.Random(seed) + background = list(background) if background else list(_BACKGROUND) counts = {train_split: n_train, tuning_split: n_tuning, held_out_split: n_held_out} - codes = [*_BACKGROUND, SIGNAL_CODE] + codes = [*background, signal_code] per_split: dict[str, list[dict]] = {} task_labels: dict[str, pl.DataFrame] = {} @@ -114,13 +121,13 @@ def build_signal_dataset( # over: the code would always sit at position 0, a positive subject would always have exactly # one more event than a negative one — a length-only predictor measured AUROC 0.59 — and since # prediction_time is derived from the event count, it would differ by class too. A model could - # then score well here without ever reading SIGNAL_CODE, which is the one thing this dataset - # exists to check. - events = [rng.choice(_BACKGROUND) for _ in range(rng.randint(4, 9))] + # then score well here without ever reading the signal code, which is the one thing this + # dataset exists to check. + events = [rng.choice(background) for _ in range(rng.randint(4, 9))] if has_signal: - events.insert(rng.randrange(len(events) + 1), SIGNAL_CODE) + events.insert(rng.randrange(len(events) + 1), signal_code) else: - events.append(rng.choice(_BACKGROUND)) + events.append(rng.choice(background)) for i, code in enumerate(events): rows.append( { @@ -149,6 +156,49 @@ def build_signal_dataset( return _write_meds(root, per_split, codes, task_labels) +def signal_dataset_from_predicates( + root: Path, + predicates_file: Path | str, + *, + seed: int = 0, + shuffle_labels: bool = False, +) -> Path: + """A signal dataset whose codes come from **the model's own predicates file** (the one-file principle). + + The dependency inverts: instead of a predicates fixture matched to a fixed synthetic vocabulary, the + dataset adapts to the predicates. The predicates with literal codes (exact / any-of — a regex cannot + be reverse-instantiated into a code) are collected; the first becomes the signal predicate (its first + code planted with the label), the rest are distractors (their codes emitted label-independently). The + feature space the model is then tested on is its production one: same names, same order, same + ``features.json``. + + Distractors are why at least **two** literal-code predicates are required: with a single feature the + column trivially equals the answer and the test stops checking that the model weights the right + feature. Callers should ``pytest.skip`` on the ValueError this raises. + + Raises: + ValueError: if fewer than two predicates with literal codes remain after parsing. + """ + import yaml + + from ..featurize import literal_code_predicates, parse_predicates + + raw = yaml.safe_load(Path(predicates_file).read_text())["predicates"] + literal = literal_code_predicates(parse_predicates(raw), raw) + if len(literal) < 2: + raise ValueError( + f"The designed-signal test needs at least two predicates with literal codes (exact or " + f"any-of) in {predicates_file}; found {len(literal)} ({', '.join(literal) or 'none'}). " + "Add one, or provide example codes for your regex predicates." + ) + names = list(literal) + signal_code = literal[names[0]][0] + background = sorted({c for name in names[1:] for c in literal[name]} - {signal_code}) + return build_signal_dataset( + root, seed=seed, shuffle_labels=shuffle_labels, signal_code=signal_code, background=background + ) + + def build_pattern_dataset( root: Path, *, n_train: int = 200, n_tuning: int = 40, n_held_out: int = 40, seed: int = 0 ) -> Path: diff --git a/template/src/meds_model_base/utils.py b/template/src/meds_model_base/utils.py index e4f0eb5..416a138 100644 --- a/template/src/meds_model_base/utils.py +++ b/template/src/meds_model_base/utils.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import os import shutil from pathlib import Path @@ -19,10 +20,41 @@ #: Name of the checkpoint copied into a published model artifact. BEST_CKPT_FILENAME = "checkpoint" +#: cuBLAS workspace layout required for deterministic matmuls on CUDA. See :func:`configure_cublas_workspace`. +CUBLAS_WORKSPACE_CONFIG = ":4096:8" + #: Suffix for the scratch directory holding in-progress training state. WORK_DIR_SUFFIX = ".work" +def configure_cublas_workspace() -> str: + """Set ``CUBLAS_WORKSPACE_CONFIG`` so strict determinism is *available* on CUDA. Returns the value. + + ``trainer.deterministic=true`` makes Lightning call ``torch.use_deterministic_algorithms(True)``, which + raises on CUDA unless this variable is set — and it raises at the first cuBLAS matmul, minutes into a + run rather than at startup. Flipping the flag alone is therefore not a usable strict mode, which is why + this is set unconditionally rather than only when the flag is on. + + Placement matters, and not in the obvious way. cuBLAS reads this when it creates its workspace, at the + first cuBLAS call — *not* at ``import torch`` — so this only has to run before any CUDA work, which is + why the dispatcher can call it before composing the config. Do not "fix" this by moving it next to a + torch import: by then the dispatcher has already resolved the command class, which imports torch. + + ``setdefault``, so a value the user exported still wins. + + Examples: + >>> import os + >>> _ = os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) + >>> configure_cublas_workspace() + ':4096:8' + >>> os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + >>> configure_cublas_workspace() # an explicitly exported value wins + ':16:8' + >>> _ = os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) + """ + return os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", CUBLAS_WORKSPACE_CONFIG) + + def work_dir_for(output_dir: Path | str) -> Path: """Resolve the scratch directory for a training run. diff --git a/template/src/{{ model_slug }}/configs/datamodule/patients.yaml b/template/src/{{ model_slug }}/configs/datamodule/patients.yaml deleted file mode 100644 index 035d977..0000000 --- a/template/src/{{ model_slug }}/configs/datamodule/patients.yaml +++ /dev/null @@ -1,13 +0,0 @@ -defaults: - - config: MEDSTorchDataConfig - - _self_ - -# Task-free view of the cohort, used by pretrain: random windows over each subject's whole timeline. -_target_: meds_torchdata.extensions.lightning_datamodule.Datamodule -config: - tensorized_cohort_dir: ${input_data_dir}/patients - max_seq_len: ${max_seq_len} - seq_sampling_strategy: RANDOM - static_inclusion_mode: OMIT -batch_size: ${batch_size} -num_workers: ${num_workers} diff --git a/template/src/{{ model_slug }}/configs/datamodule/patients.yaml.jinja b/template/src/{{ model_slug }}/configs/datamodule/patients.yaml.jinja new file mode 100644 index 0000000..5817a93 --- /dev/null +++ b/template/src/{{ model_slug }}/configs/datamodule/patients.yaml.jinja @@ -0,0 +1,25 @@ +{% if data_backend == 'mtd' -%} +defaults: + - config: MEDSTorchDataConfig + - _self_ + +# Task-free view of the cohort, used by pretrain: random windows over each subject's whole timeline. +_target_: meds_torchdata.extensions.lightning_datamodule.Datamodule +config: + tensorized_cohort_dir: ${input_data_dir}/patients + max_seq_len: ${max_seq_len} + seq_sampling_strategy: RANDOM + static_inclusion_mode: OMIT +batch_size: ${batch_size} +num_workers: ${num_workers} +{%- else -%} +# Task-free view of the featurized cohort, used by pretrain. The datamodule class is YOURS to implement +# (src/{{ model_slug }}/datamodule.py) against the surface in meds_model_base/lightning/protocol.py. +_target_: {{ model_slug }}.datamodule.FeaturizedDataModule +config: + _target_: {{ model_slug }}.datamodule.FeaturizedDataModuleConfig + patients_dir: ${input_data_dir}/patients + task_labels_dir: null +batch_size: ${batch_size} +num_workers: ${num_workers} +{%- endif %} diff --git a/template/src/{{ model_slug }}/configs/datamodule/task.yaml b/template/src/{{ model_slug }}/configs/datamodule/task.yaml deleted file mode 100644 index 226a48d..0000000 --- a/template/src/{{ model_slug }}/configs/datamodule/task.yaml +++ /dev/null @@ -1,19 +0,0 @@ -defaults: - - config: MEDSTorchDataConfig - - _self_ - -# Task-conditioned view of the cohort: one sample per (subject_id, prediction_time), with the sequence -# ending at that timepoint (TO_END). Used by supervised_train, infer and predict. -# -# `task_labels_dir` is deliberately null here and set at runtime: the command materializes -# external_labels_dir into its own work directory first, and that path is derived from the output -# directory rather than expressible in config. -_target_: meds_torchdata.extensions.lightning_datamodule.Datamodule -config: - tensorized_cohort_dir: ${input_data_dir}/patients - task_labels_dir: null - max_seq_len: ${max_seq_len} - seq_sampling_strategy: TO_END - static_inclusion_mode: OMIT -batch_size: ${batch_size} -num_workers: ${num_workers} diff --git a/template/src/{{ model_slug }}/configs/datamodule/task.yaml.jinja b/template/src/{{ model_slug }}/configs/datamodule/task.yaml.jinja new file mode 100644 index 0000000..f5bddea --- /dev/null +++ b/template/src/{{ model_slug }}/configs/datamodule/task.yaml.jinja @@ -0,0 +1,36 @@ +{% if data_backend == 'mtd' -%} +defaults: + - config: MEDSTorchDataConfig + - _self_ + +# Task-conditioned view of the cohort: one sample per (subject_id, prediction_time), with the sequence +# ending at that timepoint (TO_END). Used by supervised_train, infer and predict. +# +# `task_labels_dir` is deliberately null here and set at runtime: the command materializes +# external_labels_dir into its own work directory first, and that path is derived from the output +# directory rather than expressible in config. +_target_: meds_torchdata.extensions.lightning_datamodule.Datamodule +config: + tensorized_cohort_dir: ${input_data_dir}/patients + task_labels_dir: null + max_seq_len: ${max_seq_len} + seq_sampling_strategy: TO_END + static_inclusion_mode: OMIT +batch_size: ${batch_size} +num_workers: ${num_workers} +{%- else -%} +# Task-conditioned view of the featurized cohort: one sample per (subject_id, prediction_time). Used by +# supervised_train, infer and predict. The datamodule class is YOURS to implement (src/{{ model_slug }}/ +# datamodule.py) against the surface in meds_model_base/lightning/protocol.py. +# +# `task_labels_dir` is deliberately null here and set at runtime: the command materializes +# external_labels_dir into its own work directory first, and that path is derived from the output +# directory rather than expressible in config. +_target_: {{ model_slug }}.datamodule.FeaturizedDataModule +config: + _target_: {{ model_slug }}.datamodule.FeaturizedDataModuleConfig + patients_dir: ${input_data_dir}/patients + task_labels_dir: null +batch_size: ${batch_size} +num_workers: ${num_workers} +{%- endif %} diff --git a/template/src/{{ model_slug }}/configs/preprocess_data.yaml b/template/src/{{ model_slug }}/configs/preprocess_data.yaml deleted file mode 100644 index fada270..0000000 --- a/template/src/{{ model_slug }}/configs/preprocess_data.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# @package _global_ -# preprocess_data: external MEDS -> this model's patient representation. -# Creates /patients, the immutable base of the shared workspace. -defaults: - - paths: default - - _self_ - - profile: default - -external_meds_dir: ??? -output_data_dir: ${data_dir} - -#: Optional MEDS-transforms pipeline YAML applied before tensorization, plus extra overrides for it. -pipeline: null -pipeline_overrides: [] - -#: Artifacts are immutable; set true to replace an existing patients/ directory. -do_overwrite: false - -hydra: - run: {dir: "${output_data_dir}/.logs"} - job: {chdir: false} diff --git a/template/src/{{ model_slug }}/configs/preprocess_data.yaml.jinja b/template/src/{{ model_slug }}/configs/preprocess_data.yaml.jinja new file mode 100644 index 0000000..c888116 --- /dev/null +++ b/template/src/{{ model_slug }}/configs/preprocess_data.yaml.jinja @@ -0,0 +1,34 @@ +# @package _global_ +# preprocess_data: external MEDS -> this model's patient representation. +# Creates /patients, the immutable base of the shared workspace. +defaults: + - paths: default + - _self_ + - profile: default + +external_meds_dir: ??? +output_data_dir: ${data_dir} + +#: Optional MEDS-transforms pipeline YAML applied before the representation step, plus extra overrides +#: for it. +pipeline: null +pipeline_overrides: [] + +#: Which representation to materialize: `mtd` (meds-torch-data tensorization) or `predicates` +#: (presence featurization: 0/1 predicate columns on MEDS parquet; see meds_model_base/featurize.py). +featurization: {{ 'predicates' if data_backend == 'custom_featurization' else 'mtd' }} + +#: The predicates YAML binding concept names to dataset codes (ACES syntax, presence subset). Required +#: when featurization=predicates; the repo's own predicates.yaml is the reference file. +external_predicates_file: null + +#: Unsupported predicate forms (value bounds, and(), ...) are skipped with a warning by default and +#: recorded in the manifest; set true to make any unsupported entry a hard error instead. +featurization_strict: false + +#: Artifacts are immutable; set true to replace an existing patients/ directory. +do_overwrite: false + +hydra: + run: {dir: "${output_data_dir}/.logs"} + job: {chdir: false} diff --git a/template/src/{{ model_slug }}/configs/trainer/default.yaml b/template/src/{{ model_slug }}/configs/trainer/default.yaml index af25f53..06e2112 100644 --- a/template/src/{{ model_slug }}/configs/trainer/default.yaml +++ b/template/src/{{ model_slug }}/configs/trainer/default.yaml @@ -5,3 +5,17 @@ devices: 1 log_every_n_steps: 1 gradient_clip_val: 1.0 enable_progress_bar: false + +#: Determinism. `warn` asks for deterministic kernels and warns where none exists, rather than killing the +#: run; `benchmark: false` stops cuDNN picking algorithms by timing. With `seed`, these give the claim the +#: porting procedure asks you to make: same seed + same config + same environment -> same metric. For a +#: hard guarantee use `trainer.deterministic=true`, which raises instead of warning. +deterministic: warn +benchmark: false + +#: Full fp32, which is torch's own default — this repository does not quietly trade precision for speed. +#: Listed explicitly so `trainer.precision=bf16-mixed` composes without Hydra's `+` (configs are struct +#: mode) and so the tradeoff is visible rather than implicit. That is the speed knob to reach for; if the +#: model you are porting ran TF32 matmuls, call `torch.set_float32_matmul_precision` in your `model.py` +#: and record it in the implementation report. +precision: 32-true diff --git a/template/src/{{ model_slug }}/model.py.jinja b/template/src/{{ model_slug }}/model.py.jinja index e780931..693881c 100644 --- a/template/src/{{ model_slug }}/model.py.jinja +++ b/template/src/{{ model_slug }}/model.py.jinja @@ -1,3 +1,4 @@ +{% set batch_type = 'MEDSTorchBatch' if data_backend == 'mtd' else 'Batch' -%} """Your model. **This file is a stub — nothing here is implemented.** The template generates the *command DAG*, not the model: the ``{{ profile }}`` profile decides which @@ -18,17 +19,23 @@ The DAG calls exactly these hooks, and nothing else: Once you implement them, **delete the ``is_stub`` marker**. The conformance tests in ``tests/`` skip while it is set and start running when it is gone; they are the contract your model has to satisfy. -``meds_model_base.lightning.modules`` provides the pieces that must know the MEDS batch format — +{% if data_backend == 'mtd' %}``meds_model_base.lightning.modules`` provides the pieces that must know the MEDS batch format — ``CodeEmbedder``, ``padding_mask``, ``masked_mean``. Everything else is yours. -""" +{% else %}The batches these hooks receive are whatever your datamodule (``datamodule.py``) collates from the +featurized patients artifact — the ``Batch`` alias below is yours to replace along with it. +{% endif %}""" {% if profile != 'packaged' %} from collections.abc import Callable {% if profile == 'probe' %}import lightning.pytorch as L {% endif %}from meds_model_base.lightning.modules import BaseLightningModule -from meds_torchdata import MEDSTorchBatch -from torch import Tensor -{% endif %}{% if profile == 'packaged' %} +{% if data_backend == 'mtd' %}from meds_torchdata import MEDSTorchBatch +{% endif %}from torch import Tensor +{% if data_backend == 'custom_featurization' %} +#: Your datamodule's batch type. `object` only marks the stub: replace it (a TypedDict, a dataclass, a +#: plain Tensor — whatever datamodule.py yields) when you implement the model. +Batch = object +{% endif %}{% endif %}{% if profile == 'packaged' %} class Model: """This profile ships trained weights with the repository, so there is no module to train. @@ -64,7 +71,7 @@ class Model(BaseLightningModule): self.save_hyperparameters(ignore=["optimizer", "scheduler"]) self.vocab_size = vocab_size - def compute_loss(self, batch: MEDSTorchBatch) -> tuple[Tensor, dict[str, Tensor]]: + def compute_loss(self, batch: {{ batch_type }}) -> tuple[Tensor, dict[str, Tensor]]: """Return ``(loss, metrics)`` for a training batch.{% if profile in ['finetune', 'probe'] %} This DAG trains twice. ``batch.has_labels`` distinguishes them: pretraining sees no labels, while @@ -72,7 +79,7 @@ class Model(BaseLightningModule): {% endif %}""" raise NotImplementedError("Implement compute_loss in src/{{ model_slug }}/model.py.") {% if profile == 'probe' %} - def encode(self, batch: MEDSTorchBatch) -> Tensor: + def encode(self, batch: {{ batch_type }}) -> Tensor: """Return a ``[B, D]`` representation per prediction timepoint. ``infer`` materializes this as the ``embedding`` column; ``Probe`` is then trained on it. Nothing @@ -80,7 +87,7 @@ class Model(BaseLightningModule): """ raise NotImplementedError("Implement encode in src/{{ model_slug }}/model.py.") {% elif profile == 'zero_shot_materialized' %} - def infer_step(self, batch: MEDSTorchBatch) -> dict[str, Tensor]: + def infer_step(self, batch: {{ batch_type }}) -> dict[str, Tensor]: """Return ``{"predicted_boolean_probability": [B]}`` — the task scores ``infer`` materializes. ``predict`` reads these back from the artifact instead of re-running the model, so the scores stay @@ -88,7 +95,7 @@ class Model(BaseLightningModule): """ raise NotImplementedError("Implement infer_step in src/{{ model_slug }}/model.py.") {% else %} - def predict_step(self, batch: MEDSTorchBatch, batch_idx: int = 0) -> dict[str, Tensor]: + def predict_step(self, batch: {{ batch_type }}, batch_idx: int = 0) -> dict[str, Tensor]: """Return ``{"predicted_boolean_probability": [B]}`` aligned to the dataloader's sample order. Optionally add ``predicted_boolean_value``. ``predict`` validates the result against diff --git a/template/src/{{ model_slug }}/{% if data_backend == 'custom_featurization' %}datamodule.py{% endif %}.jinja b/template/src/{{ model_slug }}/{% if data_backend == 'custom_featurization' %}datamodule.py{% endif %}.jinja new file mode 100644 index 0000000..bd654fa --- /dev/null +++ b/template/src/{{ model_slug }}/{% if data_backend == 'custom_featurization' %}datamodule.py{% endif %}.jinja @@ -0,0 +1,80 @@ +"""Your datamodule over the featurized patients artifact. THE STUB IS DELIBERATE — implement it here. + +``preprocess_data`` (with ``featurization=predicates``) publishes ``/patients`` as MEDS +parquet augmented with one 0/1 ``predicate//`` column per predicate, plus ``features.json`` +defining the feature space and its order. How those event-level columns become model inputs — +aggregation, windowing, imputation, whether to read ``numeric_value`` on active rows — is modeling, so +it lives here, not in the template. + +The commands interact with this class only through the surface documented in +``meds_model_base/lightning/protocol.py``. In short, you must provide: + +- ``config.task_labels_dir`` (set at runtime by the commands): read ``{split}.parquet`` label files + from it — one sample per ``(subject_id, prediction_time)``, using events at or before that time. + ``boolean_value`` may be absent (inference); the timepoints must still yield samples. +- ``config.vocab_size``: already implemented below from ``features.json``. +- ``train_dataset``/``val_dataset``/``test_dataset`` (+ the matching ``*_dataloader`` methods), + for the train / tuning / held_out splits respectively. +- each dataset's ``schema_df``: a polars frame of ``subject_id``, ``prediction_time`` in loader + iteration order (prediction loaders must not shuffle) — this is how ``predict`` aligns model outputs + to timepoints. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import lightning.pytorch as pl +from meds_model_base.featurize import load_features + + +@dataclass +class FeaturizedDataModuleConfig: + """What the shipped datamodule configs instantiate; extend freely (it is yours).""" + + patients_dir: str + task_labels_dir: str | None = None + + @property + def vocab_size(self) -> int: + """Feature-space size, from the artifact's ``features.json`` — never from parquet column order.""" + return len(load_features(self.patients_dir)) + + +class FeaturizedDataModule(pl.LightningDataModule): + """A deliberate stub: every hook raises until you implement it. See the module docstring.""" + + is_stub = True + + def __init__(self, config: FeaturizedDataModuleConfig, batch_size: int, num_workers: int = 0): + super().__init__() + self.config = config + self.batch_size = batch_size + self.num_workers = num_workers + + def _unimplemented(self, what: str): + raise NotImplementedError( + f"{type(self).__name__}.{what} is the generated stub. Implement your datamodule in " + f"src/{{ model_slug }}/datamodule.py (see the module docstring for the contract)." + ) + + @property + def train_dataset(self): + self._unimplemented("train_dataset") + + @property + def val_dataset(self): + self._unimplemented("val_dataset") + + @property + def test_dataset(self): + self._unimplemented("test_dataset") + + def train_dataloader(self): + self._unimplemented("train_dataloader") + + def val_dataloader(self): + self._unimplemented("val_dataloader") + + def test_dataloader(self): + self._unimplemented("test_dataloader") diff --git a/template/tests/conftest.py.jinja b/template/tests/conftest.py.jinja index 6c5d374..98eec68 100644 --- a/template/tests/conftest.py.jinja +++ b/template/tests/conftest.py.jinja @@ -16,6 +16,13 @@ from meds_model_base.testing import build_workspace, run_chain, run_cli from {{ model_slug }}.commands import COMMANDS +#: The model's own predicates file — the ONE predicates file (tests, MEDS-DEV runs and production all +#: read the same declaration; there is deliberately no test-only fixture copy). +PREDICATES_FILE = Path(__file__).resolve().parent.parent / "predicates.yaml" + +#: The overrides that select the predicates backend for a `preprocess_data` run. +FEATURIZE_ARGS = ["featurization=predicates", f"external_predicates_file={PREDICATES_FILE}"] + @pytest.fixture def cli(): @@ -48,8 +55,29 @@ def labels_dir(meds_root) -> Path: return tasks[0] +{% if data_backend == 'custom_featurization' %}@pytest.fixture(scope="session") +def data_dir(meds_root, tmp_path_factory) -> Path: + """A shared workspace with patients/ (predicates representation), built once for the session.""" + workspace = tmp_path_factory.mktemp("workspace") / "data" + return build_workspace(meds_root, workspace, extra_args=FEATURIZE_ARGS) + + @pytest.fixture(scope="session") +def featurized_data_dir(data_dir) -> Path: + """In this repo the default workspace already IS the featurized one.""" + return data_dir +{% else %}@pytest.fixture(scope="session") def data_dir(meds_root, tmp_path_factory) -> Path: """A shared workspace with patients/, built once for the session. Labels are not preprocessed.""" workspace = tmp_path_factory.mktemp("workspace") / "data" return build_workspace(meds_root, workspace) + + +@pytest.fixture(scope="session") +def featurized_data_dir(meds_root, tmp_path_factory) -> Path: + """A predicates-representation workspace. Cheap and dependency-free, so it is built in MTD repos + too: the featurization contract tests run everywhere, and the equivalence guard can compare the two + representations' label partitioning in the same session.""" + workspace = tmp_path_factory.mktemp("featurized_workspace") / "data" + return build_workspace(meds_root, workspace, extra_args=FEATURIZE_ARGS) +{%- endif %} diff --git a/template/tests/test_featurization.py b/template/tests/test_featurization.py new file mode 100644 index 0000000..96d3ed7 --- /dev/null +++ b/template/tests/test_featurization.py @@ -0,0 +1,206 @@ +"""Presence featurization: the reader's grammar, the artifact contract, and the model's own predicates. + +Three layers, all dependency-free (no meds-torch-data, no model), so they run in every generated repo +whatever its ``data_backend``: + +- **Grammar** tests use inline YAML/frames: the grammar belongs to the template, so literals are correct + here — everywhere else the suite reads the repository's own ``predicates.yaml``, never a test copy. +- **The file** is validated strictly: runtime skips unsupported forms in *foreign* files, but your own + declaration containing forms featurization cannot use is a bug in the file, surfaced here directly + rather than through a warning nobody reads. +- **The artifact** tests assert the published contract (columns, ``features.json``, manifest, + split-layout-as-authority) against whatever predicates the file declares — they hardcode none of its + names, so they keep passing when you replace the starter content with your real concepts. + +The equivalence guard at the bottom is the anti-drift test between the two representations' cohort +readers; it needs an MTD workspace, so it self-skips in a ``custom_featurization`` repo. +""" + +import json +from pathlib import Path + +import polars as pl +import pytest +import yaml +from meds_model_base.featurize import ( + FEATURES_FILENAME, + PREDICATE_COLUMN_PREFIX, + PredicatesError, + featurize_frame, + load_features, + parse_predicates, + read_predicates_file, +) +from meds_model_base.manifest import read_manifest +from meds_model_base.tasks import cohort_subjects, featurized_cohort, materialize_labels + +PREDICATES_FILE = Path(__file__).resolve().parent.parent / "predicates.yaml" + + +# --- Grammar (inline literals: the grammar is the template's, not the model's) ------------------------- + + +def _frame(): + return pl.DataFrame( + { + "subject_id": [1, 1, 2, 2], + "code": ["HR", "ADMISSION//CARDIAC", "TEMP", "DISCHARGE"], + "numeric_value": [99.0, None, 37.0, None], + } + ) + + +def test_every_supported_form_matches_what_it_says(tmp_path): + parsed = parse_predicates( + { + "hr": {"code": "HR"}, + "adm": {"code": {"regex": "^ADMISSION//.*"}}, + "vitals": {"code": {"any": ["HR", "TEMP"]}}, + "hr_or_adm": {"expr": "or(hr, adm)"}, + } + ) + out = featurize_frame(_frame(), parsed) + assert out["predicate//hr"].to_list() == [1, 0, 0, 0] + assert out["predicate//adm"].to_list() == [0, 1, 0, 0] + assert out["predicate//vitals"].to_list() == [1, 0, 1, 0] + assert out["predicate//hr_or_adm"].to_list() == [1, 1, 0, 0] + assert all(out[c].dtype == pl.Int8 for c in parsed.columns) + + +def test_original_columns_are_untouched(): + df = _frame() + out = featurize_frame(df, parse_predicates({"hr": {"code": "HR"}})) + assert out.select(df.columns).equals(df), "featurization is pure augmentation" + + +def test_unsupported_forms_skip_with_a_cascade(): + """Value bounds are model-side semantics; the or() referencing them must fall with them.""" + parsed = parse_predicates( + { + "hr": {"code": "HR"}, + "high_hr": {"code": "HR", "value_min": 110}, + "high_or_hr": {"expr": "or(high_hr, hr)"}, + } + ) + assert parsed.names == ["hr"] + assert set(parsed.skipped) == {"high_hr", "high_or_hr"} + + +def test_strict_mode_errors_listing_the_offenders(): + with pytest.raises(PredicatesError, match="high_hr"): + parse_predicates({"hr": {"code": "HR"}, "high_hr": {"code": "HR", "value_min": 110}}, strict=True) + + +def test_an_empty_feature_space_is_always_an_error(): + """Skipping everything must not degrade into publishing a featureless artifact.""" + with pytest.raises(PredicatesError, match="No featurizable predicate remains"): + parse_predicates({"high_hr": {"code": "HR", "value_min": 110}}) + + +def test_refeaturizing_featurized_data_is_refused(): + parsed = parse_predicates({"hr": {"code": "HR"}}) + with pytest.raises(PredicatesError, match="already carries predicate column"): + featurize_frame(featurize_frame(_frame(), parsed), parsed) + + +def test_a_missing_predicates_file_names_itself(tmp_path): + with pytest.raises(PredicatesError, match="does not exist"): + read_predicates_file(tmp_path / "nope.yaml") + + +# --- The model's own predicates file ------------------------------------------------------------------- + + +def test_the_repos_predicates_file_parses_with_zero_skips(): + """Strictness lives here, not at runtime: skip-with-warning is for *foreign* files (a raw dataset + predicates.yaml in a MEDS-DEV run); your own declaration should never contain forms featurization + cannot use — that is a model silently narrower than its author believes.""" + parsed = read_predicates_file(PREDICATES_FILE) + assert parsed.skipped == {}, ( + f"predicates.yaml contains unsupported predicate(s): {parsed.skipped}. Rewrite them in the " + "supported presence subset (exact / regex / any-of / or)." + ) + assert parsed.names, "predicates.yaml must declare at least one predicate" + + +# --- The published artifact (generic over whatever the file declares) ---------------------------------- + + +@pytest.fixture(scope="session") +def patients(featurized_data_dir): + return featurized_data_dir / "patients" + + +def test_manifest_declares_the_representation_and_the_provenance(patients): + manifest = read_manifest(patients, require_type="data") + assert manifest["representation"] == "predicates" + feat = manifest["featurization"] + assert feat["predicates_digest"], "the digest is what makes a drifted predicates file detectable" + assert feat["n_features"] == len(load_features(patients)) + assert feat["skipped"] == [], "the shipped file parses clean, so nothing may be skipped" + + +def test_features_json_is_the_feature_space(patients): + """Order and spelling come from features.json — never from parquet column order or the YAML.""" + features = load_features(patients) + declared = set(yaml.safe_load(PREDICATES_FILE.read_text())["predicates"]) + for f in features: + assert f["name"] in declared + assert f["column"] == f"{PREDICATE_COLUMN_PREFIX}{f['name']}" + + columns = [f["column"] for f in features] + for shard in sorted((patients / "data").rglob("*.parquet")): + df = pl.read_parquet(shard) + for c in columns: + assert df[c].dtype == pl.Int8, f"{c} in {shard.name}" + assert set(df[c].unique().to_list()) <= {0, 1} + raw = json.loads((patients / FEATURES_FILENAME).read_text()) + assert raw["version"] == 1 + + +def test_match_counts_in_the_manifest_are_the_column_sums(patients): + counts = read_manifest(patients)["featurization"]["match_counts"] + shards = sorted((patients / "data").rglob("*.parquet")) + for f in load_features(patients): + total = sum(int(pl.read_parquet(s)[f["column"]].sum()) for s in shards) + assert counts[f["name"]] == total + + +def test_split_membership_travels_as_shard_layout_only(patients): + """The featurized artifact keeps the MTD branch's invariant: no subject_splits.parquet copy.""" + assert not (patients / "metadata" / "subject_splits.parquet").exists() + cohort = featurized_cohort(patients) + assert set(cohort) and all(cohort.values()) + assert cohort == cohort_subjects(patients), "the manifest dispatch must land on the shard reader" + assert read_manifest(patients)["splits"] == {s: len(m) for s, m in sorted(cohort.items())} + + +def test_labels_materialize_against_the_featurized_cohort(patients, labels_dir, tmp_path): + dest, summary = materialize_labels(labels_dir, patients, tmp_path / "labels") + assert summary, "at least one split must receive labels" + for split in summary: + assert (dest / f"{split}.parquet").is_file() + + +# --- Equivalence guard: the two cohort readers must never drift ---------------------------------------- + + +def test_both_representations_partition_labels_identically( + data_dir, featurized_data_dir, labels_dir, tmp_path +): + """The anti-drift test. materialize_labels resolves splits from tokenization/schemas/ on an MTD + artifact and from the data/ shard layout on a featurized one; if those ever disagree, the failure + would otherwise surface as a CoverageError two training runs later.""" + pytest.importorskip("meds_torchdata") + if read_manifest(data_dir / "patients").get("representation", "mtd") != "mtd": + pytest.skip("no MTD workspace in this repo (data_backend=custom_featurization)") + + _, mtd_summary = materialize_labels(labels_dir, data_dir / "patients", tmp_path / "mtd") + _, feat_summary = materialize_labels( + labels_dir, featurized_data_dir / "patients", tmp_path / "featurized" + ) + mtd = {s: pl.read_parquet(tmp_path / "mtd" / f"{s}.parquet") for s in mtd_summary} + feat = {s: pl.read_parquet(tmp_path / "featurized" / f"{s}.parquet") for s in feat_summary} + assert set(mtd) == set(feat) + for split in mtd: + assert sorted(mtd[split]["subject_id"].to_list()) == sorted(feat[split]["subject_id"].to_list()) diff --git a/template/tests/test_label_cohort.py b/template/tests/test_label_cohort.py index c6ac73f..4f8ddf4 100644 --- a/template/tests/test_label_cohort.py +++ b/template/tests/test_label_cohort.py @@ -1,7 +1,8 @@ -"""Labels are partitioned by the split each subject was actually tokenized into. +"""Labels are partitioned by the split each subject actually landed in the patients artifact under. Needs no model, so unlike most of this suite it runs from the moment a repository is generated: the -``data_dir`` fixture is a real MTD tensorization of the meds-testing-helpers dataset. +``data_dir`` fixture is a real ``preprocess_data`` run over the meds-testing-helpers dataset, in +whichever representation this repo's ``data_backend`` selected. The bug being pinned is quiet by construction. A ``preprocess_data`` pipeline that filters subjects tokenizes them away while their labels survive; ``pretrain`` and ``supervised_train`` read training data @@ -14,12 +15,13 @@ import polars as pl import pytest +from meds_model_base.manifest import read_manifest from meds_model_base.tasks import ( TaskMaterializationError, + cohort_subjects, materialize_labels, read_labels, split_labels, - tokenized_cohort, ) #: Far outside any fixture's id space, so it can only ever be "absent from the cohort". @@ -33,7 +35,7 @@ def patients_dir(data_dir): @pytest.fixture def cohort(patients_dir): - return tokenized_cohort(patients_dir) + return cohort_subjects(patients_dir) def _labels(subject_ids): @@ -46,10 +48,13 @@ def _labels(subject_ids): ) -def test_cohort_is_read_from_the_schema_dirs(patients_dir, cohort): - """The cohort comes from where meds-torch-data reads splits, not from ``subject_splits.parquet``.""" - schema_dir = patients_dir / "tokenization" / "schemas" - expected = {p.parts[0] for p in (f.relative_to(schema_dir) for f in schema_dir.rglob("*.parquet"))} +def test_cohort_is_read_from_the_representations_own_layout(patients_dir, cohort): + """The cohort comes from where the representation itself stores splits — the tokenization schema + directories for an MTD artifact, the ``data//`` shard layout for a featurized one — and never + from ``subject_splits.parquet``.""" + representation = read_manifest(patients_dir).get("representation", "mtd") + layout = patients_dir / ("tokenization/schemas" if representation == "mtd" else "data") + expected = {f.relative_to(layout).parts[0] for f in layout.rglob("*.parquet")} assert set(cohort) == expected assert all(cohort.values()), "every split should contribute at least one subject" @@ -75,7 +80,7 @@ def test_absent_subjects_are_dropped_with_a_warning(cohort, caplog): out = split_labels(_labels([*present, ABSENT_SUBJECT]), cohort) assert sorted(out[split]["subject_id"].to_list()) == present - assert "Dropping 1 label row(s) for 1 subject(s) absent from the tensorized cohort" in caplog.text + assert "Dropping 1 label row(s) for 1 subject(s) absent from the patients cohort" in caplog.text def test_a_complete_cohort_is_passed_through_silently(cohort, caplog): diff --git a/template/tests/test_meds_dev_e2e.py.jinja b/template/tests/test_meds_dev_e2e.py.jinja index 15eae1d..a8d5ab0 100644 --- a/template/tests/test_meds_dev_e2e.py.jinja +++ b/template/tests/test_meds_dev_e2e.py.jinja @@ -183,6 +183,23 @@ def test_meds_dev_runs_this_model_end_to_end(meds_dev, meds_root, labels_dir, tm f"package, so this means the model directory did not land there. Registered: {listed.stdout!r}" ) +{% if data_backend == 'custom_featurization' %} # This model's commands reference {predicates_path}, which needs MEDS-DEV's passthrough (PR#325). + # Over the demo vocabulary the model's predicates may match nothing — that is fine: this test + # asserts plumbing (venv, placeholder fill, chain, schema, coverage), not learnability. + run_model_cfg = meds_dev / "src" / "MEDS_DEV" / "configs" / "_run_model.yaml" + if "predicates_path" not in run_model_cfg.read_text(): + pytest.skip( + "this MEDS-DEV checkout has no predicates_path passthrough (MEDS-DEV PR#325); point " + "MEDS_DEV_DIR at a checkout that includes it" + ) + shipped_predicates = meds_dev / "src" / "MEDS_DEV" / "models" / MODEL_NAME / "predicates.yaml" + assert shipped_predicates.is_file(), ( + "meds-model-add-to-meds-dev did not ship predicates.yaml alongside model.yaml; a benchmark " + "user would have no reference file to pass as predicates_path" + ) + predicates_args = [f"predicates_path={shipped_predicates}"] +{% else %} predicates_args: list[str] = [] +{% endif %} out = tmp_path / "run" _run( [ @@ -196,6 +213,7 @@ def test_meds_dev_runs_this_model_end_to_end(meds_dev, meds_root, labels_dir, tm "mode=full", "dataset_type=full", "do_overwrite=true", + *predicates_args, ], timeout=RUN_TIMEOUT, ) diff --git a/template/tests/test_property.py.jinja b/template/tests/test_property.py.jinja index 812fb6b..0cc724c 100644 --- a/template/tests/test_property.py.jinja +++ b/template/tests/test_property.py.jinja @@ -11,8 +11,36 @@ a supervised model, a fine-tune, a probe and a zero-shot model alike. Deselected in the fast CI job (`-m "not slow"`); run it before you claim the model works. """ +{% if data_backend == 'custom_featurization' %}from pathlib import Path + import polars as pl import pytest +from meds_model_base.testing import ( + assert_learns_signal, + binary_auroc, + build_workspace, + signal_dataset_from_predicates, + skip_if_stub, +) + +from {{ model_slug }}.model import Model + +pytestmark = pytest.mark.slow + +#: The model's own predicates file (the one-file principle): the signal dataset is generated FROM it, +#: so the model is tested on its production feature space — same names, same order, same features.json. +PREDICATES_FILE = Path(__file__).resolve().parent.parent / "predicates.yaml" + +_WORKSPACE_ARGS = ["featurization=predicates", f"external_predicates_file={PREDICATES_FILE}"] + + +def _signal_dataset(root, *, seed: int, shuffle_labels: bool): + try: + return signal_dataset_from_predicates(root, PREDICATES_FILE, seed=seed, shuffle_labels=shuffle_labels) + except ValueError as e: # fewer than two literal-code predicates — a curation gap, not a model bug + pytest.skip(str(e)) +{% else %}import polars as pl +import pytest from meds_model_base.testing import ( assert_learns_signal, binary_auroc, @@ -25,12 +53,18 @@ from {{ model_slug }}.model import Model pytestmark = pytest.mark.slow +_WORKSPACE_ARGS: list[str] = [] + + +def _signal_dataset(root, *, seed: int, shuffle_labels: bool): + return build_signal_dataset(root, seed=seed, shuffle_labels=shuffle_labels) +{% endif %} def _predictions_for(chain, tmp_path, name: str, *, shuffle_labels: bool, seed: int): """Build a synthetic dataset, run the DAG over it, and return (predictions, held-out labels).""" - raw = build_signal_dataset(tmp_path / f"raw_{name}", seed=seed, shuffle_labels=shuffle_labels) + raw = _signal_dataset(tmp_path / f"raw_{name}", seed=seed, shuffle_labels=shuffle_labels) labels_dir = raw / "task_labels" / "signal_task" - data_dir = build_workspace(raw, tmp_path / f"data_{name}") + data_dir = build_workspace(raw, tmp_path / f"data_{name}", extra_args=_WORKSPACE_ARGS) artifacts = chain(data_dir, labels_dir, tmp_path / f"out_{name}", epochs=25, batch_size=16) assert "predictions" in artifacts, "this DAG does not register `predict`, so it cannot be evaluated" return ( diff --git a/template/tests/test_smoke_pipeline.py.jinja b/template/tests/test_smoke_pipeline.py.jinja index d36fd55..044519e 100644 --- a/template/tests/test_smoke_pipeline.py.jinja +++ b/template/tests/test_smoke_pipeline.py.jinja @@ -14,7 +14,7 @@ import pytest import yaml from meds_evaluation.schema import PredictionSchema from meds_model_base.commands import CommandName -from meds_model_base.tasks import tokenized_cohort +from meds_model_base.tasks import cohort_subjects from meds_model_base.testing import skip_if_stub from {{ model_slug }}.commands import COMMANDS @@ -29,31 +29,28 @@ def test_workspace_is_published_with_manifests(data_dir): # Split membership must travel with the artifact, since every command that materializes labels needs # it and nothing downstream may depend on the raw MEDS directory still existing. It travels as the - # shard layout, which is where meds-torch-data reads it from, rather than as a copy of the source - # dataset's subject_splits.parquet -- a copy describes the input to preprocessing, and a pipeline - # that filters subjects makes those two different things. - assert set(tokenized_cohort(data_dir / "patients")) + # representation's own layout (tokenization schemas for mtd, data/ shards for predicates) — never as + # a copy of the source dataset's subject_splits.parquet: a copy describes the input to preprocessing, + # and a pipeline that filters subjects makes those two different things. + assert set(cohort_subjects(data_dir / "patients")) assert not (data_dir / "patients" / "metadata" / "subject_splits.parquet").exists() # Counted from the artifact, so a filtering pipeline shows up here rather than being invisible. assert patients["splits"] == { - split: len(subjects) for split, subjects in sorted(tokenized_cohort(data_dir / "patients").items()) + split: len(subjects) for split, subjects in sorted(cohort_subjects(data_dir / "patients").items()) } def test_inference_never_sees_ground_truth(data_dir, labels_dir, tmp_path): """At inference, the labels are not merely dropped from the output — they never reach the model. - `predict` and `infer` materialize the index without `boolean_value`. meds-torch-data distinguishes a - task *index* from task *labels*, so the batch handed to `predict_step` has no `boolean_value` at all, - and a model cannot read the answer it is about to be scored against even by accident. - + `predict` and `infer` materialize the index without `boolean_value`; whatever datamodule consumes + the materialized layout can only hand the model what is in it. This is the command-level, load- + bearing half of the property, so it is backend-agnostic and needs no model and no meds-torch-data. Training is the opposite case and is checked here too, so the difference is deliberate rather than - incidental. Needs no model. + incidental. """ from meds_model_base.tasks import materialize_labels - from meds_torchdata import MEDSPytorchDataset, MEDSTorchDataConfig - from meds_torchdata.types import SubsequenceSamplingStrategy patients = data_dir / "patients" @@ -64,6 +61,28 @@ def test_inference_never_sees_ground_truth(data_dir, labels_dir, tmp_path): written = pl.read_parquet(dest / "held_out.parquet") assert ("boolean_value" in written.columns) is expect_labels + +def test_mtd_distinguishes_a_task_index_from_task_labels(data_dir, labels_dir, tmp_path): + """The MTD half of the property above: without `boolean_value` on disk, the batch a model receives + has no `boolean_value` at all — it cannot read the answer it is about to be scored against even by + accident. A custom datamodule owes the same guarantee (see meds_model_base/lightning/protocol.py); + that can only be tested against its implementation, so this test is MTD-specific and skips cleanly + where meds-torch-data is not installed. + """ + pytest.importorskip("meds_torchdata") + from meds_model_base.manifest import read_manifest + from meds_model_base.tasks import materialize_labels + from meds_torchdata import MEDSPytorchDataset, MEDSTorchDataConfig + from meds_torchdata.types import SubsequenceSamplingStrategy + + patients = data_dir / "patients" + if read_manifest(patients).get("representation", "mtd") != "mtd": + pytest.skip("the workspace is not an MTD tensorization") + + for include_labels, expect_labels in [(False, False), (True, True)]: + dest = tmp_path / f"labels_{include_labels}" + materialize_labels(labels_dir, patients, dest, include_labels=include_labels) + dataset = MEDSPytorchDataset( MEDSTorchDataConfig( tensorized_cohort_dir=str(patients), diff --git a/tests/test_render.py b/tests/test_render.py index 374f4bb..f60ac72 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -87,11 +87,12 @@ def _registry(dst: Path, slug: str) -> dict[str, str]: return dict(re.findall(r"CommandName\.(\w+): (\w+),", text)) +@pytest.mark.parametrize("data_backend", ["mtd", "custom_featurization"]) @pytest.mark.parametrize("profile,commands", sorted(PROFILE_COMMANDS.items())) @pytest.mark.render -def test_render_profile(tmp_path, profile, commands): +def test_render_profile(tmp_path, profile, commands, data_backend): dst = tmp_path / profile - slug = _render(dst, profile) + slug = _render(dst, profile, data_backend=data_backend) for rel in [ "pyproject.toml", @@ -103,6 +104,8 @@ def test_render_profile(tmp_path, profile, commands): "docs/PORTING-A-MODEL.md", "model.yaml", ".copier-answers.yml", + # The model's one predicates file: rendered on BOTH backends (MTD repos featurize in tests too). + "predicates.yaml", f"src/{slug}/__main__.py", f"src/{slug}/model.py", f"src/{slug}/commands.py", @@ -392,6 +395,49 @@ def test_rendered_repo_is_clean_as_written(tmp_path, profile): ) +@pytest.mark.render +def test_the_contract_takes_no_position_on_matmul_precision(): + """The vendored contract must not call ``set_float32_matmul_precision``. + + It used to, unconditionally, with ``"medium"`` — running fp32 matmuls in bfloat16 on Ampere and later. + That is stable run-to-run, so it never broke same-machine reproducibility; what it broke was + comparison against the source a port is reproducing, silently, from a file `copier update` overwrites. + + Torch's own default is ``highest``, so *not calling it* is the fix. The pull is to re-add it for + speed — that belongs in ``model.py``, which the user owns and which the implementation report already + requires a ledger row for. ``trainer.precision`` is the sanctioned knob. + """ + payload = REPO / "template/src/meds_model_base" + offenders = [ + str(fp.relative_to(REPO)) + for fp in sorted(payload.rglob("*.py")) + if "set_float32_matmul_precision" in fp.read_text() + ] + assert not offenders, f"the contract must not set matmul precision: {offenders}" + + +#: Keys `configs/trainer/default.yaml` must carry. Each is listed rather than left to Lightning's default +#: because configs are struct mode: an absent key needs Hydra's `+` to override, which nobody guesses. +TRAINER_REPRODUCIBILITY_KEYS = {"deterministic": "warn", "benchmark": False, "precision": "32-true"} + + +@pytest.mark.parametrize("profile", sorted(PROFILE_COMMANDS)) +@pytest.mark.render +def test_trainer_config_ships_the_determinism_keys(tmp_path, profile): + """Seeding alone does not give the same numbers twice on GPU. + + ``seed_everything(seed, workers=True)`` is called and `num_workers` defaults to 0, but without + ``deterministic`` cuDNN autotunes by timing and picks non-deterministic kernels — so two runs at one + seed can disagree, which is exactly the check `docs/PORTING-A-MODEL.md` Step 6 requires a port to pass. + """ + dst = tmp_path / f"determinism_{profile}" + slug = _render(dst, profile) + cfg = yaml.safe_load((dst / f"src/{slug}/configs/trainer/default.yaml").read_text()) + + for key, expected in TRAINER_REPRODUCIBILITY_KEYS.items(): + assert cfg.get(key) == expected, f"trainer/default.yaml: {key} should be {expected!r}" + + @pytest.mark.parametrize("profile", sorted(PROFILE_COMMANDS)) @pytest.mark.render def test_rendered_configs_parse(tmp_path, profile): @@ -479,6 +525,45 @@ def test_optional_loggers_render_with_their_extra(tmp_path, enabled): assert not _bad_eof(dst), f"files not ending in exactly one newline: {_bad_eof(dst)}" +@pytest.mark.parametrize("data_backend", ["mtd", "custom_featurization"]) +@pytest.mark.render +def test_data_backend_gates_dependencies_configs_and_stubs(tmp_path, data_backend): + """``data_backend`` is implemented the ``use_wandb`` way: it gates the meds-torch-data dependency, + which datamodule configs render, the datamodule stub, the preprocess default and the model.yaml + predicates wiring — while ``src/meds_model_base/`` ships both paths unconditionally (Jinja-stripping + the vendored contract would double every merge and render matrix for zero runtime win).""" + custom = data_backend == "custom_featurization" + dst = tmp_path / data_backend + slug = _render(dst, "supervised", data_backend=data_backend) + + pyproject = (dst / "pyproject.toml").read_text() + assert ("meds-torch-data" in pyproject) is not custom, "the MTD dependency must follow the backend" + + task_cfg = (dst / f"src/{slug}/configs/datamodule/task.yaml").read_text() + patients_cfg = (dst / f"src/{slug}/configs/datamodule/patients.yaml").read_text() + for cfg in (task_cfg, patients_cfg): + assert ("meds_torchdata" in cfg) is not custom + assert (f"{slug}.datamodule" in cfg) is custom + + assert (dst / f"src/{slug}/datamodule.py").exists() is custom, "the stub renders only when it is yours" + + preprocess = (dst / f"src/{slug}/configs/preprocess_data.yaml").read_text() + expected_default = "featurization: predicates" if custom else "featurization: mtd" + assert expected_default in preprocess + + model_yaml = (dst / "model.yaml").read_text() + assert ("{predicates_path}" in model_yaml) is custom, ( + "a featurized model must bind the dataset predicates via MEDS-DEV's {predicates_path}; " + "an MTD model must not require an argument it does not use" + ) + + # Both backends ship the model's one predicates file and the vendored contract's featurizer: the + # featurization tests run in MTD repos too, and flipping the answer later needs no contract change. + assert (dst / "predicates.yaml").is_file() + assert (dst / "src/meds_model_base/featurize.py").is_file() + assert not _bad_eof(dst), f"files not ending in exactly one newline: {_bad_eof(dst)}" + + @pytest.mark.render def test_logger_group_members_nest_under_their_own_name(tmp_path): """Each logger config must be a mapping of *name → logger*, not a bare ``_target_``.