EveryQuery

Given a MEDS dataset, 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.

EveryQuery is built on the MEDS ecosystem leveraging meds-torch-data for tensorization and MEDS-transforms for preprocessing.

Install

As a dependency:

pip install EveryQuery

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 + EQ_sample_task_tracking_pairs (TaskQuerySchema parquets: scattered for PT, dense for eval, pos/neg pairs for in-training AUROC)
├── 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)
└── 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.

Research-only, paper-specific code (ID/OOD code sampling, ablations, the results notebook, figure code, the ETHOS comparison) lives in the separate EveryQueryExperiments repo, which depends on EveryQuery as an installed library. The split is tracked in #186.

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.

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 the cross-model leaderboard, which now lives in the EveryQueryExperiments repo.

Pipeline

Current (on dev)

flowchart TD
    meds[MEDS cohort] --> process[EQ_process_data]
    process --> intermediate[("MEDS event shards<br/>($TOKENIZED_EVENTS_DIR)")]
    process --> cohort[("tensorized cohort<br/>($TENSORIZED_COHORT_DIR)")]

    intermediate --> train_tasks[EQ_generate_training_tasks<br/><i>scattered, random tasks</i>]
    intermediate --> eval_tasks[EQ_generate_evaluation_tasks<br/><i>dense grid: codes × durations</i>]

    train_tasks -- TaskQuerySchema parquets --> train[EQ_train]
    cohort -- tensorized cohort --> train
    train --> ckpt[/best_model.ckpt/]

    ckpt --> predict[EQ_predict]
    eval_tasks -- TaskQuerySchema parquets --> predict

    predict -- PredictionSchema parquet --> evaluate[EQ_evaluate]
    evaluate --> metrics[("per-(query, duration_days)<br/>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="$DATA_DIR" \
	intermediate_dir="$TOKENIZED_EVENTS_DIR" \
	output_dir="$TENSORIZED_COHORT_DIR"

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

2a. Generate pre-training task labels

EQ_generate_training_tasks \
	split=train \
	num_queries=4000000 \
	num_contexts_per_query=1 \
	max_workers=1 \
	data_dir="$TOKENIZED_EVENTS_DIR" \
	out_dir="$TRAINING_TASKS_DIR" \
	query_codes="$TENSORIZED_COHORT_DIR"

data_dir is the MEDS dataset root (event shards read from {data_dir}/data/{split}/*.parquet) and out_dir is the final-dataset root. Both are required Hydra args (no .env fallback — see #235); pass them as shell-expanded vars after source env.sh.

One command runs the whole 5-stage sampler in a single process (Stages 0–3 inline, then Stage 4 labels shards in parallel). The dataset lands at $TRAINING_TASKS_DIR/{split}/{shard}.parquet, with intermediates in the sibling *_artifacts dir (see generate_tasks/README.md). Columns conform to TaskQuerySchema — subject_id, prediction_time, query, duration_days, boolean_value — where boolean_value is three-valued: True (query code occurs in (prediction_time, prediction_time + duration_days]), False (window fully observed, no occurrence), or null (censored — window extends past the subject's last recorded time).

max_workers sets how many shards are labeled in parallel, so raising it raises peak RAM. If Stage 4 OOMs, set max_workers=1.

Note: The total number of training samples generated will be num_queries * num_contexts_per_query

query_codes= is required for training. Set it to a metadata root (query_codes=$TENSORIZED_COHORT_DIR) to sample from {dir}/metadata/codes.parquet, or 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 \
	'query_codes=[HR, TEMP]' \
	'durations=[1, 7, 30, 90, 365]' \
	data_dir="$TOKENIZED_EVENTS_DIR" \
	out_dir=$EVAL_TASKS_DIR

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

The endpoint writes one parquet per (split, input_shard) worker, so to label a whole split use Hydra multirun (-m) to sweep input_shard — there's no auto-discovery, so enumerate the range explicitly. Count the shards for your split first (N = this number):

ls "$TOKENIZED_EVENTS_DIR"/data/held_out/*.parquet | wc -l

Then sweep input_shard=range(0,N):

# Sequential (basic launcher) — one shard after another in a single process:
EQ_generate_evaluation_tasks -m \
	input_shard=range(0,16) \
	split=held_out \
	prediction_times_per_subject=5 \
	'query_codes=[HR, TEMP]' \
	'durations=[1, 7, 30, 90, 365]'

# Parallel on SLURM (submitit launcher — already a dependency):
EQ_generate_evaluation_tasks -m \
	hydra/launcher=submitit_slurm \
	input_shard=range(0,16) \
	split=held_out …

A comma list (input_shard=0,1,2) works too; range(0,16) is just shorthand for 0..15. The prediction-time sampler is deterministic in (seed, input_shard, split), so a swept run and the equivalent per-shard runs produce identical output.

As with training, data_dir / out_dir are required Hydra args (pass them as shell-expanded vars). query_codes is also required — it is the evaluation query universe.

query_codes= accepts an inline list (as above), a metadata root / codes.parquet path (query_codes=$TENSORIZED_COHORT_DIR reads {dir}/metadata/codes.parquet), 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 \
	datamodule.config.tensorized_cohort_dir="$TENSORIZED_COHORT_DIR" \
	datamodule.config.task_labels_dir="$TRAINING_TASKS_DIR" \
	output_dir="$TRAINING_OUTPUT_DIR"

• output_dir is a required Hydra arg that is a base path you supply with output_dir=, e.g. =$TRAINING_OUTPUT_DIR. Hydra appends <YYYY-MM-DD>/<HH-MM-SS> for per-run uniqueness.
• If you want to override more parameters for training refer to src/every_query/train/configs/config.yaml

Optional: in-training macro-AUROC tracking

EQ_train can log a cheap, macro-averaged per-task AUROC estimate every validation pass (tuning/occurs_auroc_macro_sampled, plus ..._n_tasks for how many tasks contributed) via TaskAurocTrackingCallback. Instead of scoring the whole tuning split, it scores a fixed, offline-sampled parquet of exactly one positive + one negative row per (query, duration_days) task. Because a task's AUROC equals P(score(pos) > score(neg)) for a random pos/neg pair, the win/tie/loss on that one pair is an unbiased (high-variance) estimate; macro-averaging across tasks gives macro AUROC at O(n_tasks) forward examples. It is tracking-only — it does not affect tuning/loss-driven checkpointing or early stopping.

First sample the tracking pairs once from the dense tuning-split eval labels (output of EQ_generate_evaluation_tasks):

EQ_sample_task_tracking_pairs \
	eval_labels_dir="$EVAL_TASKS_DIR/eval" \
	out_dir="$TASK_TRACKING_DIR" \
	split=tuning
# pairs land at $TASK_TRACKING_DIR/tuning/0.parquet

Then enable the callback via its Hydra config group and point task_labels_dir at $TASK_TRACKING_DIR:

EQ_train +callbacks=task_auroc \
	++trainer.callbacks.task_auroc_tracking.config.task_labels_dir="$TASK_TRACKING_DIR"

+callbacks=task_auroc composes configs/callbacks/task_auroc.yaml into the callbacks list. It is off by default — config.yaml ships task_auroc_tracking: null, and values_as_list drops None entries, so a plain EQ_train (e.g. when EveryQuery is used as a dependency) runs without it and without needing any tracking pairs. Notes:

• The split is locked to tuning on purpose — tracking mirrors tuning/loss checkpointing, so it is not configurable. Sample the pairs with split=tuning.
• Out-of-vocab query codes (not in the model vocab) are skipped at scoring; the callback warns at setup and you'll see a smaller ..._n_tasks. An empty tracking set logs a warning and simply never logs the metric (training is unaffected).
• Under DDP it scores on rank 0 only (the tracking set is identical on every rank), and the parquet is read once at setup — re-sampling mid-run requires a restart.
• To disable it entirely, just omit +callbacks=task_auroc (the default). The mandatory task_labels_dir: ??? only errors when the group is actually enabled.

4. Predict

EQ_predict \
	model_run_dir="$TRAINING_OUTPUT_DIR/YYYY-MM-DD/HH-MM-SS" \
	tasks_dir="$EVAL_TASKS_DIR/eval/held_out" \
	output_parquet="$TRAINING_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="$TRAINING_OUTPUT_DIR/predictions.parquet" \
	metrics_parquet="$TRAINING_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.

Paths & environment

Path roots are plain Hydra args, not env vars read by the package (the .env/load_dotenv layer was removed in #235). The shell owns the vars: source env.sh (copied from env.example.sh) exports them, and you pass them into each CLI as shell-expanded key=$VAR overrides — source-ing one file is all that's needed to move machines (SLURM scripts source the same file). EQ_train validates only the values it actually resolves — validate_training_config() in train.py checks the resolved cohort/task dirs exist and that WANDB_ENTITY is set only when a wandb logger is enabled.

The genuine env read that remains in the train config: WANDB_ENTITY (read natively by wandb, backs ${oc.env:WANDB_ENTITY,null}). output_dir is now a required arg with no env fallback (pass output_dir=$TRAINING_OUTPUT_DIR if you keep that var in env.sh). The preprocessing subprocess bridge (RAW_MEDS_DIR, MTD_INPUT_DIR, MIN_SUBJECTS_PER_CODE, MIN_EVENTS_PER_SUBJECT) and the optional aces_to_eq pipeline (ACES_SHARDS_DIR) also use ${oc.env:...} — see those submodules.

env.example.sh is the reference — copy to env.sh, edit, and source it.

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)
├── sampler/                        (unit: per-stage sampler tests — stage0-4, pure helpers, orchestration)
├── test_sampler_dataset_integration.py  (integration: sampler output is drop-in for EveryQueryPytorchDataset)
├── 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)
└── 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

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.