EveryQuery

Important

This fork (conditional-queries branch) adds a conditional query-sequence model — a bidirectional patient encoder + block-autoregressive decoder that answers an ordered list of queries, each conditioned on the patient state and the teacher-forced answers of all earlier queries. Binary observed-occurrence labels; censoring is expressed by querying the real TIMELINE//END code rather than via a separate head. See CONDITIONAL_QUERIES.md for the design, the new EQ_generate_query_sequences / EQ_predict_sequences / EQ_evaluate_sequences CLIs, the macro per-task evaluation methodology, and results.

A framework for training and evaluating foundation models over structured EHR data, built on the MEDS ecosystem — meds-torch-data for tensorization, MEDS-transforms for preprocessing, PyTorch Lightning for training.

Given a tensorized MEDS cohort, EveryQuery trains a ModernBERT-style encoder to answer "query" prediction tasks of the form: given a subject's history up to time t, will code c occur within d days? The same trained model is then evaluated against arbitrary (code, duration) combinations.

Note

The Phase-1 + Phase-2 refactor from #54 has landed: the full preprocess → generate_training_tasks / generate_evaluation_tasks → train → predict → evaluate pipeline uses the cross-stage TaskQuerySchema throughout. EQ_evaluate now resolves to the single-stage evaluator that consumes PredictionSchema parquets; the legacy four-stage evaluator has been deleted (recover from git history if needed). See Roadmap for the remaining #83 cleanup.

Install

For development (recommended):

git clone git@github.com:payalchandak/EveryQuery.git
cd EveryQuery
uv sync --group dev
cp .env.example .env # then edit paths for your machine

As a dependency:

# not yet on PyPI — installable from git for now:
pip install "git+https://github.com/payalchandak/EveryQuery.git@main"

Repository layout

Every production module lives under a submodule that reflects its role:

src/every_query/
├── preprocessing/      → EQ_process_data        (raw MEDS → tensorized cohort)
├── generate_tasks/     → EQ_generate_training_tasks + EQ_generate_evaluation_tasks (TaskQuerySchema parquets: scattered for PT, dense for eval)
├── train/              → EQ_train               (train the model)
├── predict/            → EQ_predict             (inference; consumes TaskQuerySchema, emits PredictionSchema)
│   └── external_tasks/                         (ACES + composite aggregation — currently `python -m` only;
│                                                  [#62](https://github.com/payalchandak/EveryQuery/issues/62) tracks promoting to console scripts, draft PR [#95](https://github.com/payalchandak/EveryQuery/pull/95))
├── evaluate/           → EQ_evaluate           (metrics on a PredictionSchema parquet)
├── model/              (shared: nn.Module + LightningModule)
├── data/               (shared: PyTorch Dataset + Batch types + TaskQuerySchema)
├── paper_experiments/  (research-only: ID/OOD splits, ablations, figure code)
│   └── sample_codes/   (query-code sampling for paper experiments; dataset-agnostic on #97)
└── utils/              (helpers: seeds, code slugs, env-var validation, model_loader)

Every submodule has its own README.md explaining what belongs there, its pipeline position, and the tracking issues for remaining work.

Console scripts

pip install exposes the CLIs below, all Hydra-configurable. Run any with --help or --cfg job to inspect the resolved config. The Tests column summarises the coverage that lands with each CLI on dev today — unit tests (fast, tests/test_<name>_logic.py or tests/test_<module>.py), CLI smoke tests (tests/test_cli_smoke.py, --help-exits-0), and end-to-end subprocess tests that run the real script against a fixture cohort.

Script	Stage	Purpose	Tests
`EQ_process_data`	preprocessing	Orchestrate MEDS-transforms + `meds-torch-data` tensorization	smoke; E2E via `test_process_data.py` + `test_e2e_foundation.py`
`EQ_generate_training_tasks`	PT task labels	Sample `N` tasks × `M` contexts (scattered `(query, duration_days)`), label via single-pass asof	smoke; unit `test_sample_tasks.py`; E2E `test_generate_tasks.py`
`EQ_generate_evaluation_tasks`	eval task labels	Sample `K` prediction times per subject, cross-join with `(codes × durations)` grid for dense evaluation shape	smoke; E2E `test_generate_evaluation_tasks_cli.py`
`EQ_train`	training	Train the ModernBERT encoder on the labeled tasks	smoke; unit `test_training.py`; E2E `test_train_cli.py` + `test_train.py`; signal test `tests/training_validity/` (slow)
`EQ_predict`	inference	Consume a `TaskQuerySchema` parquet dir + checkpoint, emit a `PredictionSchema` parquet (`censor_prob`, `occurs_prob`)	smoke; E2E `test_predict_cli.py` (row-order preserved); exercised by `tests/training_validity/` (slow)
`EQ_evaluate`	metrics	Consume a `PredictionSchema` parquet, write per-`(query, duration_days)` metrics (`occurs_auroc`, `censor_auroc`, etc.)	smoke; E2E `test_evaluate_cli.py`; exercised by `tests/training_validity/` (slow)

The legacy four-stage evaluator (every_query.evaluate.eval, with gen_index_times, gen_task, select_model siblings) has been deleted; recover from git history if needed. #83 tracks any paper_experiments/leaderboard/ relocation for cross-model comparison.

Pipeline

Current (on `dev`)

    MEDS cohort  ──►  EQ_process_data  ──►  tensorized cohort ($FINAL_DATA_DIR)
                                                          │
                                                          ├─────────────────────────────┐
                                                          ▼                             ▼
                                            EQ_generate_training_tasks       EQ_generate_evaluation_tasks
                                            (scattered, random tasks)        (dense grid: codes × durations)
                                                          │                             │
                                                          │ TaskQuerySchema parquets    │
                                                          ▼                             │
                                                     EQ_train ──► best_model.ckpt       │
                                                                           │            │
                                                                           ▼            │
                                                                      EQ_predict ◄──────┘
                                                                           │
                                                                           │ PredictionSchema parquet
                                                                           ▼
                                                                     EQ_evaluate
                                                                           │
                                                                           ▼
                                                               per-(query, duration_days) metrics parquet

Both task-generation endpoints emit TaskQuerySchema-conformant parquets. Training uses the scattered shape (one random (query, duration_days) per row); evaluation uses the dense shape (every held-out (subject, time) × every (query × duration) the user wants metrics for) so EQ_predict + EQ_evaluate cover a full grid without having to run inference twice.

1. Preprocess

EQ_process_data \
	input_dir="$RAW" \
	intermediate_dir="$INTERMEDIATE" \
	output_dir="$FINAL_DATA_DIR"

Produces a tensorized MEDS cohort under $FINAL_DATA_DIR. $INTERMEDIATE is a staging directory for the MEDS-transforms stages; $PROCESSED holds cross-shard metadata ($PROCESSED/metadata/codes.parquet is the query-code universe the sampler draws from).

2a. Generate pre-training task labels

EQ_generate_training_tasks \
	split=train \
	input_shard=0 \
	task_shard=0 \
	n_tasks=1024 \
	contexts_per_task=1

Sweep across shards with python -m every_query.generate_tasks.sample_tasks -m input_shard=0,1,2,… task_shard=range(0,K). Each worker writes labeled task parquets under $TASK_DIR/{split}/*.parquet idempotently. Output columns conform to TaskQuerySchema — subject_id, prediction_time, query, duration_days, boolean_value — where boolean_value is a nullable three-valued label (null = censored, True = event occurred in [prediction_time, prediction_time + duration_days), False = observed-but-no-event).

query_codes= is optional for training. Leave it unset/null to sample query codes from $PROCESSED/metadata/codes.parquet, or set it to an inline list / YAML path to restrict which codes can be sampled as queries. YAML files may be a flat list or a mapping with a codes: key. This does not remove codes from patient histories.

EQ_generate_training_tasks query_codes=/path/to/train_query_codes.yaml …

# train_query_codes.yaml
codes:
  - HR
  - TEMP

2b. Generate evaluation task labels

EQ_generate_evaluation_tasks \
	split=held_out \
	input_shard=0 \
	prediction_times_per_subject=5 \
	'codes=[HR, TEMP]' \
	'durations=[1, 7, 30, 90, 365]'

Samples K prediction times per subject, cross-joins with the full (codes × durations) grid, labels via the same primitive as training. Output lands under $TASK_DIR/eval/{split}/*.parquet (separate eval/ subdir so it doesn't collide with the training-task output).

codes= accepts an inline list (as above), a metadata root / codes.parquet path, or — for reproducible pre-sampled code universes kept out of git — a path to a YAML file. The YAML is either a bare list or a mapping with a codes: key:

# sampled_codes.yaml
codes:
  - HR
  - TEMP
  - ICD//A01

EQ_generate_evaluation_tasks codes=/path/to/sampled_codes.yaml …

3. Train

EQ_train \
	output_dir="$OUTPUT_DIR/outputs/\${run_id:}" \
	datamodule.config.task_labels_dir="$TASK_DIR" \
	datamodule.config.tensorized_cohort_dir="$FINAL_DATA_DIR"

EQ_train reads the long-format labels written by EQ_generate_training_tasks directly — the inline collation step that lived in train.py was removed in #76.

Seeding: cfg.seed (default 140799) is passed through lightning.seed_everything before model + datamodule instantiation (fix landed in #124), so model weight initialization is byte-reproducible across Python versions and platforms for a given seed.

4. Predict

EQ_predict \
	model_run_dir="$OUTPUT_DIR/outputs/YYYY-MM-DD/HH-MM-SS" \
	tasks_dir="$TASK_DIR/eval/held_out" \
	output_parquet="$OUTPUT_DIR/predictions.parquet" \
	split=held_out

Reads every *.parquet under tasks_dir (TaskQuerySchema-conformant), runs the checkpoint's predict_step over the chosen split, writes a single PredictionSchema parquet with censor_prob + occurs_prob per input row. See predict/README.md for details.

5. Evaluate

EQ_evaluate \
	predictions_parquet="$OUTPUT_DIR/predictions.parquet" \
	metrics_parquet="$OUTPUT_DIR/metrics.parquet"

Per-(query, duration_days) metrics from the predictions parquet — n_rows, n_occurs_labeled, n_positive, prevalence, occurs_auroc (on non-censored rows), censor_auroc. See evaluate/README.md.

Configuration

All CLIs are @hydra.main entry points; every config knob is overridable on the command line with key=value or +new_key=value. The config directory is resolved via importlib.resources.files("every_query"), so package-shipped YAMLs work identically whether you run from a source checkout or a pip installed wheel.

Environment variables

ensure_env() (in utils/_env.py) requires these be set before EQ_train and the eval CLIs. Scope of this gate was tightened in #127 — PROCESSED and INTERMEDIATE were dropped because no Hydra config interpolates them (they were only read by a dotenv fallback in the sampler, which already tolerates missing env vars when CLI config values are supplied).

Var	Purpose
`PROJECT_DIR`	Repo root (for relative output paths in a few configs)
`OUTPUT_DIR`	Where training run dirs land
`TASK_DIR`	Where task parquets read / write
`FINAL_DATA_DIR`	Tensorized cohort (output of `EQ_process_data`)
`WANDB_ENTITY`	W&B entity for training telemetry

.env.example is the reference — copy to .env and edit. Both Python (via python-dotenv) and the SLURM wrappers under scripts/ source it. Further phases of #117 will migrate the remaining gated vars to ${oc.env:VAR,???} / ${oc.env:VAR,default} form (Hydra-native required or optional-with-fallback) and eventually retire ensure_env() entirely.

Development

uv sync --group dev
uv run pytest                         # full suite, excluding slow tests (~2 min)
uv run pytest -m 'slow or not slow'   # full suite incl. slow training-validity test (~8-10 min extra)
uv run pytest tests/test_cli_smoke.py # CLI smoke tests only
uv run pre-commit run --all-files     # lint, format, codespell

CI runs the full pytest -m "slow or not slow" (both slow-marked and unmarked tests) on Python 3.11 and 3.12, plus ruff check and ruff format --check on every PR; coverage is uploaded to Codecov. Full CI session: ~10-11 min typical.

Test layout

tests/
├── test_cli_smoke.py               (every EQ_* CLI; --help exits 0)
├── test_process_data.py            (E2E: EQ_process_data output shape + metadata)
├── test_generate_tasks.py          (E2E: EQ_generate_training_tasks ground-truth label recompute + reproducibility)
├── test_generate_evaluation_tasks_cli.py  (E2E: EQ_generate_evaluation_tasks dense-grid shape + determinism)
├── test_sample_tasks.py            (unit: sampler primitives, determinism, edge cases)
├── test_train_cli.py               (E2E: EQ_train CLI, resume flow, overwrite flag)
├── test_train.py                   (E2E: resume-actually-loads-ckpt two-stage differential)
├── test_training.py                (unit: single training step, checkpoint roundtrip, demo-mode checks)
├── test_predict_cli.py             (E2E: EQ_predict against a trained checkpoint + row-order preservation)
├── test_evaluate_cli.py            (E2E: EQ_evaluate on a synthetic PredictionSchema parquet)
├── test_e2e_foundation.py          (E2E: full preprocess → generate_training_tasks → train pipeline chains)
├── test_dataset_logic.py           (unit: EveryQueryPytorchDataset + EveryQueryBatch)
├── test_lightning_logic.py         (unit: LightningModule loss wiring, mask semantics)
├── test_model_logic.py             (unit: model heads, censored/occurs loss flip sensitivity)
├── test_run_id.py                  (unit: run_id resolver determinism)
└── training_validity/              (E2E @pytest.mark.slow: model actually learns; runs the full EQ_predict → EQ_evaluate chain; see its README)
    ├── __init__.py
    ├── conftest.py
    ├── README.md
    └── test_training_validity.py

Roadmap

Overall refactor umbrella: #54 — target architecture is preprocess → generate_tasks → train → predict → evaluate with a shared cross-stage task-query schema.

Phase 2 status

Sub-phase	Issue	State
2.1: TaskQuerySchema design	#80	✅ merged via #96 (also closes #122)
2.2: EQ_predict	#81	✅ merged via #99
2.3: eval-suite inventory	#82	Decisions captured on the issue + reflected in #100's scope; no code change needed
2.4: EQ_evaluate consolidation	#83	✅ new `evaluate.py` is the `EQ_evaluate` entry point (rewired in this PR); `every_query.evaluate.eval` + siblings deleted — recover from git history if needed
2.5: EQ_generate_evaluation_tasks	(part of this PR)	✅ new dense-grid task-generator endpoint to feed `EQ_predict`; training-task endpoint renamed to `EQ_generate_training_tasks` for clarity

E2E testing status (#104)

Subprocess test	Issue	State
`test_process_data.py`	(pre-104)	✅ merged
`test_generate_tasks.py`	#107	✅ merged via #112 (training-task shape)
`test_generate_evaluation_tasks_cli.py`	(part of this PR)	✅ dense-grid + determinism coverage for the new eval-task endpoint
`test_train.py`	#108	✅ merged via #113
`test_predict_cli.py`	(part of #99)	✅ merged via #99 (row-order preservation covered)
`test_evaluate_cli.py`	#109	✅ merged via #100; rewired to the `EQ_evaluate` console script in this PR
training-validity (model learns)	#118	✅ merged via #119; runs the full `EQ_predict` → `EQ_evaluate` chain as subprocesses (slow)

Hygiene / follow-ups

Issue	Description
#62	Promote `aces_to_eq` / `process_composite` to entry points — draft PR #95
#64	Drop gitignored `{train,eval}_codes` defaults (design pick pending)
#85	Rewrite `sample_codes/` dataset-agnostic — draft PR #97
#117	Env-var audit — phase 1 merged via #127; phases 2-4 pending
#125	Adopt hypothesis-based property tests for the sampler
#129	Rename `PredictionSchema.occurs_prob` → `label_prob` post-NeurIPS once non-occurrence task types land
#59	Docs: final rewrite after the refactor settles

Model / architecture research (non-blocking)

#101 / #102 — RoPE for time-deltas
#103 — Evaluate alternatives to ModernBERT as the encoder backbone

Acknowledgements

EveryQuery sits on top of MEDS, meds-torch-data, MEDS-transforms, and MEDS_EIC_AR (architectural reference). It uses Hydra for configuration, PyTorch Lightning for training, and W&B for telemetry.

License

MIT — see LICENSE.

Name		Name	Last commit message	Last commit date
.github		.github
analysis		analysis
configs/hydra/launcher		configs/hydra/launcher
reports		reports
scripts		scripts
src/every_query		src/every_query
tests		tests
.env.example		.env.example
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
COHORT_INFERENCE_NOTES.md		COHORT_INFERENCE_NOTES.md
CONDITIONAL_QUERIES.md		CONDITIONAL_QUERIES.md
EveryQuery_Conditional_Report.pdf		EveryQuery_Conditional_Report.pdf
EveryQuery_Conditional_Report_v2.pdf		EveryQuery_Conditional_Report_v2.pdf
HANDOFF.md		HANDOFF.md
LICENSE		LICENSE
README.md		README.md
conftest.py		conftest.py
pyproject.toml		pyproject.toml
uv.lock		uv.lock

Navigation Menu

Repository files navigation

EveryQuery

Install

Repository layout

Console scripts

Pipeline

Current (on `dev`)

1. Preprocess

2a. Generate pre-training task labels

2b. Generate evaluation task labels

3. Train

4. Predict

5. Evaluate

Configuration

Environment variables

Development

Test layout

Roadmap

Phase 2 status

E2E testing status (#104)

Hygiene / follow-ups

Model / architecture research (non-blocking)

Acknowledgements

License

About

Releases

Packages

Contributors

Languages

Latest commit

History

Folders and files

Repository files navigation

EveryQuery

Install

Repository layout

Console scripts

Pipeline

Current (on dev)

1. Preprocess

2a. Generate pre-training task labels

2b. Generate evaluation task labels

3. Train

4. Predict

5. Evaluate

Configuration

Environment variables

Development

Test layout

Roadmap

Phase 2 status

E2E testing status (#104)

Hygiene / follow-ups

Model / architecture research (non-blocking)

Acknowledgements

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages

Current (on `dev`)