cotorra 26.6.0__tar.gz

cotorra-26.6.0/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Michael C. Burkhart, et al.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cotorra-26.6.0/PKG-INFO

@@ -0,0 +1,579 @@
+Metadata-Version: 2.4
+Name: cotorra
+Version: 26.6.0
+Summary: configurable trainer for generative event models
+Author-email: "Michael C. Burkhart" <burkh4rt@uchicago.edu>, Luke Solo <lsolo@uchicago.edu>, Inhyeok Lee <ihlee@uchicago.edu>, S'Khaja Charles <skhaja@uchicago.edu>, "Brett K. Beaulieu-Jones" <beaulieujones@uchicago.edu>
+License-Expression: MIT
+Project-URL: Repository, https://github.com/bbj-lab/cotorra
+Keywords: model training,generative event models,configurable training
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: accelerate
+Requires-Dist: datasets
+Requires-Dist: huggingface-hub
+Requires-Dist: lightgbm
+Requires-Dist: numpy
+Requires-Dist: omegaconf
+Requires-Dist: opacus
+Requires-Dist: optuna
+Requires-Dist: polars
+Requires-Dist: rich
+Requires-Dist: scikit-learn
+Requires-Dist: torch
+Requires-Dist: torchaudio
+Requires-Dist: torchvision
+Requires-Dist: tqdm
+Requires-Dist: transformers
+Requires-Dist: typer
+Requires-Dist: xgboost
+Requires-Dist: wandb
+Provides-Extra: all
+Requires-Dist: cotorra[dev,gen]; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pip; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+<p align="center">
+<img src="https://raw.githubusercontent.com/burkh4rt/cotorra/master/img/cotorra.png" width="400" style="display: block;
+margin: 0 auto; -webkit-mask-image: radial-gradient(
+    ellipse at center,
+    rgba(0,0,0,1) 50%,
+    rgba(0,0,0,0) 100%
+  );
+  mask-image: radial-gradient(
+    ellipse at center,
+    rgba(0,0,0,1) 50%,
+    rgba(0,0,0,0) 100%
+  );"/>
+</p>
+
+# Cotorra: a configurable trainer
+
+[![DOI](https://raw.githubusercontent.com/burkh4rt/cotorra/master/img/1193885071.svg)](https://doi.org/10.5281/zenodo.20414127)
+[![SWH](https://archive.softwareheritage.org/badge/origin/https://github.com/bbj-lab/cotorra/)](https://archive.softwareheritage.org/browse/origin/?origin_url=https://github.com/bbj-lab/cotorra)
+
+> 🦜 the wild parakeet of Chicago's south side
+
+## About
+
+This repo provides a configurable trainer for generative event models on
+tokenized timelines. _Cotorra_ is a Spanish term for a small-to-medium sized
+parrot, particularly the Monk parakeet. Monk parakeets were introduced to the
+south side of Chicago, where they have flourished. [^1] It benefits from previous
+experience training foundation models on tokenized electronic health records.
+[^2] [^3] [^4] [^5]
+
+## Installation
+
+You can download and install this package as follows:
+
+```sh
+git clone git@github.com:bbj-lab/cotorra.git
+cd cotorra
+python -m venv .venv
+. .venv/bin/activate
+pip install -e ".[gen]" \
+  --index-url https://download.pytorch.org/whl/cu128 \
+  --extra-index-url https://pypi.org/simple
+```
+
+## Context
+
+Suppose you have a dataset of tokenized timelines `tokens_times.parquet` as a
+parquet table with columns:
+
+- `subject_id`
+- `tokens` — the integer token sequence for the subject's timeline.
+- `times` — a parallel list of timestamps, one per token, indicating when each
+  event occurred.
+
+The table will look something like this:
+
+```
+┌────────────────────┬─────────────────┬─────────────────────────────────┐
+│ subject_id         ┆ tokens          ┆ times                           │
+│ ---                ┆ ---             ┆ ---                             │
+│ str                ┆ list[u32]       ┆ list[datetime[μs]]              │
+╞════════════════════╪═════════════════╪═════════════════════════════════╡
+│ 20002103           ┆ [20, 350, … 21] ┆ [2116-05-08 02:45:00, 2116-05-… │
+│ 20008372           ┆ [20, 350, … 21] ┆ [2110-10-30 13:03:00, 2110-10-… │
+│ …                  ┆ …               ┆ …                               │
+│ 29994865           ┆ [20, 364, … 21] ┆ [2111-01-28 21:49:00, 2111-01-… │
+└────────────────────┴─────────────────┴─────────────────────────────────┘
+```
+
+You also have a `tokenizer.yaml`, a plain yaml file that contains information
+about the configuration, learned vocabulary, and bins. This file is sufficient to
+reconstitute the tokenizer object. We only need this file to contain a lookup
+table:
+
+```yaml
+lookup:
+  UNK: 0
+  ADMN//direct: 1
+  ADMN//ed: 2
+  ADMN//elective: 3
+  AGE//age_Q0: 4
+  ...
+```
+
+Finally, we need `subject_splits.parquet` which is a table listing out all
+subject_id's and their corresponding split assignment (with splits: `train`,
+`tuning`, and `held_out`). This table may include additional demographic
+information provided as pass-through-columns to
+[cocoa](https://github.com/bbj-lab/cocoa).
+
+```
+┌────────────┬──────────┐
+│ subject_id ┆ split    │
+│ ---        ┆ ---      │
+│ str        ┆ str      │
+╞════════════╪══════════╡
+│ 21081215   ┆ train    │
+│ 20302177   ┆ train    │
+│ …          ┆ …        │
+│ 28150003   ┆ held_out │
+│ 22151813   ┆ held_out │
+└────────────┴──────────┘
+```
+
+For extraction and scoring workflows, we also need split-specific inference
+tables in the same `processed_data_home` directory:
+
+- `train_for_inference.parquet`
+- `tuning_for_inference.parquet`
+- `held_out_for_inference.parquet`
+
+These tables are expected to include at least:
+
+- `tokens_past` (the model context used for extraction/scoring)
+- `s_elapsed_past` (if using `time_based_rope`)
+- token-specific label columns such as `<TOKEN>_past` and `<TOKEN>_future` used
+  by generative and representation-based scoring.
+
+The `cocoa winnow` command provides these.
+
+<!-- prettier-ignore-start -->
+> [!TIP]
+> For getting your data to this point, check out our configurable
+> collator / tokenizer: [☕️ cocoa](https://github.com/bbj-lab/cocoa)
+<!-- prettier-ignore-end -->
+
+Given these things, we want to train a model to predict the next token in a
+subject's timeline given their complete history or context up to this point. This
+package is designed to do that in a configurable way.
+
+## Configuration
+
+This library can be extensively customized through yaml configuration files. Each
+command has its own default config under `src/cotorra/config/`, which you can
+override by passing a config file via the appropriate CLI flag. Any value can
+also be overridden programmatically via `**kwargs` which are merged on top of the
+YAML config via OmegaConf.
+
+### Training configuration ([example](src/cotorra/config/training.yaml))
+
+Used by `cotorra train` and `cotorra tune`.
+
+- **model**:
+  - **model_name**: Name or path of the HuggingFace model (e.g.,
+    `meta-llama/Llama-3.2-1B`).
+  - **model_args**: Model architecture parameters passed directly to
+    HuggingFace's
+    [`AutoConfig`](https://huggingface.co/docs/transformers/en/model_doc/auto).
+
+  _Note: The bundled config defines reusable model presets under
+  `model_presets`._
+
+- **max_seq_len**: Maximum sequence length for model input.
+- **n_epochs**: Number of epochs (handled in the dataloader, not the trainer).
+- **run_name**: Name for the current run (referenced by `wandb` and
+  `training_args`).
+- **tokens_of_interest**: List of special tokens to upweight during training
+  (referenced by loss config). Supports patterns specified with fnmatch.
+- **wandb**:
+  - **project**: Weights & Biases project name for experiment tracking.
+  - **run_name**: Name for the current run.
+- **custom_loss**: Boolean flag to enable custom loss functions (default:
+  `false`).
+- **quantile_token_loss** _(optional)_: Upweights loss on quantile boundary
+  tokens.
+  - **qt_weight**: Weight multiplier for quantile tokens.
+- **label_weighted_loss** _(optional)_: Upweights loss on specific tokens of
+  clinical interest.
+  - **tokens_of_interest**: List of token labels to upweight. Supports patterns
+    specified with fnmatch.
+  - **toi_weight**: Weight multiplier applied to those tokens.
+- **time_based_rope** _(optional)_: Enables time-aware rotary position
+  embeddings.
+  - **sec_per_pos_id**: Number of seconds represented by one position id
+    increment.
+- **training_args**: Arguments passed to HuggingFace's
+  [`TrainingArguments`](https://huggingface.co/docs/transformers/en/main_classes/trainer#transformers.TrainingArguments).
+- **tuning_args**: Arguments passed to HuggingFace's
+  [`hyperparameter_search`](https://huggingface.co/docs/transformers/hpo_train?backends=Optuna)
+  when `cotorra tune` is called.
+
+#### Model presets
+
+We offer the following presents
+
+| designator     | base model                | # params w/ ~1340-token vocab |
+| -------------- | ------------------------- | ----------------------------- |
+| `llama_32`     | `meta-llama/Llama-3.2-1B` | ~76.9M                        |
+| `llama_32_mid` | `meta-llama/Llama-3.2-1B` | ~8.2M                         |
+| `qwen_3`       | `Qwen/Qwen3-1.7B-Base`    | ~74.1M                        |
+| `qwen_3_mid`   | `Qwen/Qwen3-1.7B-Base`    | ~8.4M                         |
+| `gemma_3`      | `google/gemma-3-1b-pt`    | ~75.7M                        |
+| `gemma_3_mid`  | `google/gemma-3-1b-pt`    | ~7.8M                         |
+
+Use the `model` key to select one of these presets and then override any
+individual `model_args` entries as needed.
+
+<!-- prettier-ignore-start -->
+> [!TIP]
+> Training supports the `--resume-from-checkpoint` (`-r`) flag. When set,
+> `cotorra train` will attempt to resume from the latest HuggingFace checkpoint
+> saved under `--output-home`. If no checkpoint is found (or resumption fails),
+> it automatically falls back to training from scratch — so the flag is safe to
+> pass unconditionally in scripts. Use `save_steps` in `training_args` in the
+> [training.yaml](src/cotorra/config/training.yaml) file to control the frequency
+> of checkpointing.
+<!-- prettier-ignore-end -->
+
+#### Differential privacy
+
+We wrap [opacus](https://opacus.ai) to support training with differential privacy
+(see `train-private` below). The following relevant parameters can be modified in
+the configuration:
+
+```
+privacy_parameters:
+  noise_multiplier: !!float 1.0
+  max_grad_norm: !!float 1.0
+```
+
+### Extraction configuration ([example](src/cotorra/config/extraction.yaml))
+
+Used by `cotorra extract`.
+
+- **max_seq_len**: Maximum sequence length.
+- **time_based_rope** _(optional)_: Enables time-aware position ids during
+  extraction (must match the setting used at training time).
+  - **sec_per_pos_id**: Number of seconds represented by one position id
+    increment.
+- **extract**:
+  - **max_len**: Maximum input length (tokens) during extraction.
+  - **batch_size**: Batch size for inference.
+  - **shard_size** _(optional)_: Number of samples per output parquet shard. Omit
+    to write a single file per split.
+
+### Scoring configuration ([example](src/cotorra/config/scoring.yaml))
+
+Used by `cotorra generative-score` and `cotorra rep-based-score`.
+
+- **run_name**: Name for the current run, used to label output files.
+- **tokens_of_interest**: List of token-based outcomes of interest. Supports
+  patterns specified with fnmatch. (Referenced by target tokens.)
+- **score**:
+  - **max_len**: Maximum input length (tokens) during scoring.
+  - **n_samp**: Number of Monte Carlo samples per input per trajectory type.
+  - **target_tokens**: Token-based outcomes of interest to score. Supports
+    patterns specified with fnmatch.
+  - **end_tokens**: Tokens that naturally terminate a generated sequence (e.g.
+    `EOS`).
+  - **suppressed_tokens**: Tokens to suppress via logit bias during generation
+    (e.g. `PAD`).
+  - **trunc_id**: Token id forced after the time horizon is exceeded.
+  - **max_time**: Maximum time horizon in minutes.
+  - **batch_size**: Batch size for inference.
+
+## Usage
+
+We provide a CLI:
+
+```
+ Usage: cotorra [OPTIONS] COMMAND [ARGS]...
+
+ Configurable training for generative event models (vXX.X.X)
+
+╭─ Options ───────────────────────────────────────────────────────────────────╮
+│ --install-completion            Install completion for the current shell.   │
+│ --show-completion               Show completion for the current shell, to   │
+│                                 copy it or customize the installation.      │
+│ --help                -h        Show this message and exit.                 │
+╰─────────────────────────────────────────────────────────────────────────────╯
+╭─ Commands ──────────────────────────────────────────────────────────────────╮
+│ train             Train a model on tokenized data. For tokenization,        │
+│                   consult the cocoa package.                                │
+│ train-private     Train a model with differential privacy on tokenized      │
+│                   data.                                                     │
+│ tune              Run hyperparameter tuning while training a model.         │
+│ extract           Extract representations from a trained model.             │
+│ generative-score  Generate SCORE/REACH metrics from a trained model and     │
+│                   save them to parquet.                                     │
+│ rep-based-score   Generate rep-based scores for the token-based outcomes of │
+│                   interest.                                                 │
+│                   Note: this requires that features have already been       │
+│                   extracted and saved                                       │
+╰─────────────────────────────────────────────────────────────────────────────╯
+```
+
+with commands:
+
+- `cotorra train`
+
+  ```
+  Usage: cotorra train [OPTIONS]
+
+  Train a model on tokenized data. For tokenization, consult the cocoa package.
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --training-config         -t      PATH  Training configuration file      │
+  │                                            (overrides default)              │
+  │ *  --processed-data-home     -p      TEXT  Processed data directory         │
+  │                                            (overrides config)               │
+  │                                            [required]                       │
+  │ *  --output-home             -o      TEXT  Output directory for trained     │
+  │                                            models                           │
+  │                                            [required]                       │
+  │    --resume-from-checkpoint  -r            Try to resume training from the  │
+  │                                            latest checkpoint in             │
+  │                                            --output-home.                   │
+  │    --verbose                 -v            Verbose logging                  │
+  │    --help                    -h            Show this message and exit.      │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+- `cotorra tune`
+
+  ```
+  Usage: cotorra tune [OPTIONS]
+
+  Run hyperparameter tuning while training a model.
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --training-config      -t      PATH  Training configuration file         │
+  │                                         (overrides default)                 │
+  │ *  --processed-data-home  -p      TEXT  Processed data directory (overrides │
+  │                                         config)                             │
+  │                                         [required]                          │
+  │ *  --output-home          -o      TEXT  Output directory for trained models │
+  │                                         [required]                          │
+  │    --verbose              -v            Verbose logging         │
+  │    --help                 -h            Show this message and exit.         │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+- `cotorra generative-score`
+
+  ```
+  Usage: cotorra generative-score [OPTIONS]
+
+  Generate SCORE/REACH metrics from a trained model and save them to parquet.
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --scoring-config       -s      PATH  Scoring configuration file          │
+  │                                         (overrides default)                 │
+  │ *  --processed-data-home  -p      TEXT  Processed data directory [required] │
+  │ *  --model-home           -m      TEXT  Directory of the trained model to   │
+  │                                         score with                          │
+  │                                         [required]                          │
+  │    --output-home          -o      TEXT  Output directory for scores,        │
+  │                                         defaults to processed-data-home     │
+  │    --verbose              -v            Verbose logging         │
+  │    --help                 -h            Show this message and exit.         │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+- `cotorra extract`
+
+  ```
+  Usage: cotorra extract [OPTIONS]
+
+  Extract representations from a trained model.
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --extraction-config    -e      PATH  Extraction configuration file       │
+  │                                         (overrides default)                 │
+  │ *  --processed-data-home  -p      TEXT  Processed data directory [required] │
+  │ *  --model-home           -m      TEXT  Directory of the trained model to   │
+  │                                         extract from                        │
+  │                                         [required]                          │
+  │    --output-home          -o      TEXT  Output directory for extracted      │
+  │                                         features, defaults to               │
+  │                                         processed-data-home                 │
+  │    --all-times            -a            Extract features for all time steps │
+  │                                         (instead of just the final one)?    │
+  │    --help                 -h            Show this message and exit.         │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+- `cotorra rep-based-score` (note: you need to run `extract` first)
+
+  ```
+  Usage: cotorra rep-based-score [OPTIONS]
+
+  Generate rep-based scores for the token-based outcomes of interest. Note:
+  this requires that features have already been extracted and saved
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --scoring-config      -s      PATH                 Scoring configuration │
+  │                                                       file (overrides       │
+  │                                                       default)              │
+  │ *  --processed-data-ho…  -p      TEXT                 Processed data        │
+  │                                                       directory             │
+  │                                                       [required]            │
+  │ *  --model-home          -m      TEXT                 Directory of the      │
+  │                                                       trained model to      │
+  │                                                       score with            │
+  │                                                       [required]            │
+  │    --output-home         -o      TEXT                 Output directory for  │
+  │                                                       scores, defaults to   │
+  │                                                       processed-data-home   │
+  │                                                       [default: None]       │
+  │    --estimator           -e      [k-NN|lightGBM|logi  Estimator to use for  │
+  │                                  stic|logistic-z|log  rep-based scoring     │
+  │                                  istic-CV|logistic-C  [default: lightGBM]   │
+  │                                  V-z|XGBoost]                               │
+  │    --verbose             -v                           Verbose logging       │
+  │    --help                -h                           Show this message and │
+  │                                                       exit.                 │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+- `cotorra train-private`
+
+  ```
+  Usage: cotorra train-private [OPTIONS]
+
+  Train a model with differential privacy on tokenized data.
+
+  ╭─ Options ───────────────────────────────────────────────────────────────────╮
+  │    --training-config      -t      PATH   Training configuration file        │
+  │                                          (overrides default)                │
+  │ *  --processed-data-home  -p      TEXT   Processed data directory           │
+  │                                          (overrides config)                 │
+  │                                          [required]                         │
+  │ *  --output-home          -o      TEXT   Output directory for trained       │
+  │                                          models                             │
+  │                                          [required]                         │
+  │    --noise-multiplier     -n      FLOAT  Noise multiplier (overrides        │
+  │                                          configuration)                     │
+  │    --max-grad-norm        -m      FLOAT  Max grad norm (overrides           │
+  │                                          configuration)                     │
+  │    --verbose              -v             Verbose logging                    │
+  │    --help                 -h             Show this message and exit.        │
+  ╰─────────────────────────────────────────────────────────────────────────────╯
+  ```
+
+[^1]:
+    L. Gersony, "The Quiet Victory of Chicago’s Monk Parakeets," _The Chicago
+    Maroon_, 23 January 2022,
+    https://chicagomaroon.com/28830/grey-city/quiet-protest-chicagos-monk-parakeets/
+
+[^2]:
+    M. Burkhart, B. Ramadan, Z. Liao, K. Chhikara, J. Rojas, W. Parker, & B.
+    Beaulieu-Jones, Foundation models for electronic health records:
+    representation dynamics and transferability,
+    [arXiv:2504.10422](https://doi.org/10.48550/arXiv.2504.10422)
+
+[^3]:
+    M. Burkhart, B. Ramadan, L. Solo, W. Parker, & B. Beaulieu-Jones,
+    [Quantifying surprise in clinical care: Detecting highly informative events in electronic health records with foundation models](https://doi.org/10.1142/9789819824755_0013),
+    Pacific Symposium on Biocomputing 31 (2026), 173–188
+
+[^4]:
+    L. Solo, M. McDermott, W. Parker, B. Ramadan, M. Burkhart, & B.
+    Beaulieu-Jones, Efficient generative prediction for EHR foundation models:
+    the SCOPE and REACH estimators,
+    [arXiv:2602.03730](https://doi.org/10.48550/arXiv.2602.03730)
+
+[^5]:
+    I. Lee, L. Solo, M. Burkhart, B. Ramadan, W. Parker, & B. Beaulieu-Jones,
+    Representation before training: a fixed-budget benchmark for generative
+    medical event models,
+    [arXiv:2604.16775](https://doi.org/10.48550/arXiv.2604.16775)
+
+<!--
+
+Run in tmux:
+```
+tmux new -s co || tmux a -t co
+```
+
+Format:
+```sh
+ruff format .
+ruff check . --fix
+```
+
+Send to bbj-lab1:
+```
+rsync -avht \
+ --delete \
+ --exclude "processed" \
+ --exclude "data-raw" \
+ --exclude "output" \
+ --exclude "wandb" \
+ --exclude ".venv" \
+ --exclude ".idea" \
+ ~/Documents/chicago/cotorra \
+ bbj-lab1:~
+```
+
+```
+for d in data-raw processed; do ln -s /mnt/bbj-lab/users/burkh4rt/$d $d; done
+ds='mimic-icu'
+cotorra train \
+		--training-config src/cotorra/config/training.yaml \
+		--processed-data-home processed/mimic-icu \
+		--output-home output/test \
+    --resume-from-checkpoint \
+    --verbose
+```
+
+Send to randi:
+```
+for d in data-raw processed; do
+	ln -s /gpfs/data/bbj-lab/users/burkh4rt/$d $d
+done
+```
+```
+rsync -avh \
+ --exclude "output" \
+ --exclude "processed" \
+ --exclude "data-raw" \
+ --exclude "logs" \
+ --exclude "wandb" \
+ --exclude ".venv/" \
+ --exclude ".idea/" \
+ ~/Documents/chicago/cotorra \
+ randi:/gpfs/data/bbj-lab/users/burkh4rt
+```
+
+```
+srun -p gpuq \
+ --gres=gpu:1 \
+ --time=8:00:00 \
+ --job-name=adhoc \
+ --pty bash -i
+. .venv/bin/activate
+```
+
+Send to pypi:
+```
+rm -rf dist
+python3 -m pip install --upgrade build
+python3 -m build
+python3 -m pip install --upgrade twine
+python3 -m twine upload --repository pypi dist/*
+```
+
+-->