Add LabelProcessor: label semantics live on the label processor by Logiquo · Pull Request #1389 · sunlabuiuc/PyHealth

Logiquo · 2026-11-04T23:52:53Z

Summary

Adds LabelProcessor, a base class for output processors. Everything that depends on the kind of label is defined on the processor that defines the label, not in BaseModel, Trainer and the modules around them.

Eight subclasses:

Processor	`mode`	activation / loss	`pos_weight`
`BinaryLabelProcessor`	`binary`	sigmoid / BCE	✓
`MultiClassLabelProcessor`	`multiclass`	softmax / cross-entropy
`MultiLabelProcessor`	`multilabel`	sigmoid / BCE	✓
`RegressionLabelProcessor`	`regression`	identity / MSE
`TimeseriesBinaryLabelProcessor`	`timeseries_binary`	as binary, per label record	✓
`TimeseriesMultiClassLabelProcessor` (new)	`timeseries_multiclass`	as multiclass, per label record
`TimeseriesMultiLabelProcessor` (new)	`timeseries_multilabel`	as multilabel, per label record	✓
`TimeseriesRegressionLabelProcessor`	`timeseries_regression`	as regression, per label record

The four static processors do not subclass each other. Binary and multilabel share only a small helper for the optional pos_weight. Each time-series processor stores an instance of its static counterpart (TimeseriesLabelProcessor[B]) and applies it to every label record. They no longer subclass SequenceProcessorV2. They parse {"time", "values"} and pack batches (collate) themselves, in the same (values, time, offsets) format as before. They were added in #1337 and are not in a release, so they move to timeseries_label_processor.py without a compatibility import.

1. Why unify them

The kind of a label (binary, multiclass, ...) decided seven things, and each was decided somewhere else, by switching on a mode string:

the number of model outputs (BaseModel.get_output_size)
the activation (BaseModel.prepare_y_prob, an if/elif over six modes)
the loss (BaseModel.get_loss_function, another if/elif, plus _label_record_loss for time-series labels)
the positive-class weight (BaseModel.set_pos_weight, with hard-coded ("binary", "multilabel") checks in BaseModel and autoimprove)
how predictions are collected (Trainer.inference, a per_step branch)
which metrics apply (trainer.get_metrics_fn, an if/elif over six modes)
the mode itself (BaseModel._resolve_mode, which walked each processor class's MRO against the registry names; autoimprove's _label_mode and the interpretability metrics had their own copies)

The processor is the one place that knows what its label is. Spreading that knowledge has a cost: adding a label kind meant editing all of these places, and they could disagree. For example, adding the time-series labels needed new branches in BaseModel, Trainer and RNN, and #1358 adds another one to Trainer.inference. With LabelProcessor, a new label kind is one class, and BaseModel and Trainer do not change.

2. The methods, and why each moves to the processor

Method	Was in	Why it belongs on the processor
`mode` (class attribute)	`BaseModel._resolve_mode`, autoimprove `_label_mode`, `metrics/interpretability/base.py`	The kind of label is a fact about the processor class, not something to rediscover by walking class hierarchies against registry names. It equals the registered name, so `output_schema={"y": "binary"}` still selects `BinaryLabelProcessor`. `model.mode` is now a read-only property that reads it.
`size()`	(already on every label processor)	The width of the processed label, the shared `FeatureProcessor` convention: 1 for binary and regression, the number of classes or labels otherwise. It is now declared abstract on `LabelProcessor`, so every label processor must have it.
`output_size()`	`BaseModel.get_output_size`	The number of model outputs (logits), which models size their output layer with. `get_output_size` only forwarded to the processor's `size()`, and `output_size()` does the same for every label processor so far. It is a separate method so that a label whose logits differ from its width only overrides `output_size()`, without touching the models. Models call `self.label_processor.output_size()`.
`activation(logits)`	`BaseModel.prepare_y_prob`	Sigmoid, softmax or identity follows from the label kind alone.
`loss(logits, y_true)`	`BaseModel.get_loss_function`, `_label_record_loss`	The loss depends on the label kind and on the processed label's layout. For time-series labels, that layout is the packed `(values, time, offsets)` that the processor itself builds. The per-record shape check and the zero loss for a batch without label records are in `TimeseriesLabelProcessor.loss`.
`pos_weight` (constructor option) and `weight`	`BaseModel.set_pos_weight` / `pos_weight`, autoimprove's `pos_weight` search parameter, ProtoECGNet's `set_pos_weight` override	Weighting the positive class is part of the task, not of a model instance. `"balanced"` is a statistic of the training labels, which is exactly what `fit` computes. A `"balanced"` processor sets `learns_statistics`, so the existing warning fires when processors are fitted on all samples rather than on a training split. Set it in the schema: `("binary", {"pos_weight": "balanced"})`, or a number, or one number per label for multilabel.
`to_numpy(y_true, y_prob)` / `gather(...)`	`Trainer.inference`'s `per_step` branch	How a batch of outputs turns into metric inputs depends on the label layout: one array for per-sample labels, one array per sample for label records. The trainer no longer branches on the mode.
`metrics_fn(y_true, y_prob, metrics)` (class method)	`trainer.get_metrics_fn`	The metrics follow from the label kind. The time-series processors pool the label records of all samples, then score them with their static counterpart's metrics. `get_metrics_fn(mode)` stays, as a registry lookup, for calibration, autoimprove and existing callers.
`labels=` option (multiclass and multilabel)	`GradientBoostedTrees` raising when a class has no training samples	A fixed label set keeps an output for a class that is missing from the training split. When omitted, the vocabulary is learned as before, which drug recommendation and medical coding rely on.

BaseModel gains label_processors (the output processors that are LabelProcessors, by key) and label_processor (the single one, or a clear error). Models that score something other than their dataset's label override label_processors:

APN: a regression processor with its masked MSE.
VAE: a regression processor for the reconstruction.
TrajSurv and SurvivalMDN: {}, so the trainer reports only their loss, as before.

Removed from BaseModel: _resolve_mode, the mode setter, get_output_size, get_loss_function, prepare_y_prob, set_pos_weight, pos_weight, _label_record_loss and VALID_MODES. These are public, and set_pos_weight shipped in 2.1.0. There are no compatibility aliases. examples/ and the tutorial notebooks are left unchanged and still call set_pos_weight.

Other behaviour changes:

Assigning model.mode raises AttributeError.
An unknown label raises ValueError instead of KeyError.
pos_weight is no longer in autoimprove's training search space, since it is now a task option.
TemperatureScaling and DirichletCalibration no longer cast y_true with long if mode == "multiclass" else float. The arrays keep the dtype the processor emitted (long for multiclass, float otherwise), so the cast changed nothing.

3. Files

Call-site changes only (29 files). These only rename self.get_output_size() → self.label_processor.output_size(), self.get_loss_function()(...) → self.label_processor.loss(...) and self.prepare_y_prob(...) → self.label_processor.activation(...):

pyhealth/calib/calibration/logistic_recalibration.py, pyhealth/models/{adacare, agent, biot, cnn, concare, contrawr, deepr, ehrmamba_cehr, gnn, gradient_boosted_trees, graphcare, grasp, heartlang, jamba_ehr, llm_ehr_encoder, logistic_regression, mlp, retain, sparcnet, stagenet, stagenet_mha, tcn, tfm_tokenizer, torchvision_model, transformer, transformer_deid, transformers_model}.py, tests/core/test_gnn.py

Three more get the same renames, plus the >>> docstring example the PR rules require of a modified public class: pyhealth/models/{califorest, ehrmamba, graph_torchvision_model}.py.

Real changes:

Processors:
- pyhealth/processors/label_processor.py: LabelProcessor and the four static processors.
- pyhealth/processors/timeseries_label_processor.py: new; the four time-series processors.
- pyhealth/processors/timeseries_processor.py: the old time-series label classes are removed.
- pyhealth/processors/__init__.py: exports.
Model and trainer:
- pyhealth/models/base_model.py: label_processors, label_processor and the read-only mode; the removed methods.
- pyhealth/trainer.py: inference and evaluate delegate to the label processor; get_metrics_fn is a registry lookup.
Models with label-specific logic:
- apn.py, vae.py, trajsurv.py, survival_mdn.py: override label_processors instead of setting self.mode.
- protoecgnet.py: class_weights defaults to the task's pos_weight; the set_pos_weight override is removed.
- rnn.py: RNN, MultimodalRNN and RNN_V2 detect per-step labels with isinstance(..., TimeseriesLabelProcessor) instead of a mode tuple.
- medfuse.py, molerec.py, safedrug.py: remove their own mode resolution.
Other modules:
- pyhealth/calib/calibration/{dircal, temperature_scale}.py: the redundant dtype casts are removed.
- pyhealth/interpret/methods/base_interpreter.py, pyhealth/metrics/interpretability/base.py: read the mode from the label processor.
- pyhealth/autoimprove/{improver, evaluate, space}.py: the label mode comes from the processor; pos_weight is no longer searched.
Docs:
- docs/api/models.rst and docs/api/processors.rst: updated.
- API pages: new for LabelProcessor and the two new time-series processors; updated for TimeseriesLabelProcessor.
Tests:
- New: tests/core/test_label_processor.py.
- Rewritten for the new API: test_pos_weight, test_legacy_mode_resolution, test_model_mode.
- Updated:
  - label_record_model_utils.py
  - test_apn
  - test_autoimprove
  - test_gim and test_ig_gim: the mock datasets get a real BinaryLabelProcessor.
  - test_interp_metrics
  - test_per_step_labels
  - test_protoecgnet
  - test_timeseries_modes
  - test_trajsurv

4. Prediction types not covered yet (planned follow-ups)

These tasks and models exist in the repo but do not fit one of the eight label processors yet. They keep working as before: their models compute their own loss, and where a model has no label processor, the Trainer reports only the loss.

Prediction type	Where it is used	How it works today	Planned
Survival (time to event)	`survival_mimic4` (`{"duration": "tensor", "event": "tensor"}`), TrajSurv, SurvivalMDN	The label is two plain tensor fields. The models compute their own loss and return no label processors, so `Trainer.evaluate` reports only the loss, although `metrics/survival.py` has the c-index, IPCW Brier score and cumulative/dynamic AUC.	A `SurvivalLabelProcessor` for one `(duration, event)` label whose `metrics_fn` is `survival_metrics_fn`. The loss differs by model (Cox partial likelihood, mixture density, trajectory), so the models would keep their own loss.
Multi-output regression with missing targets (forecasting)	`physionet2012_forecasting` (`{"y": "tensor"}`, NaN where a value is not observed), APN	`RegressionLabelProcessor` has one output, so APN defines its own `_ForecastLabel` with a masked MSE.	A regression processor with a configurable output width and a masked loss that skips NaN targets. APN's override then goes away.
One label per token (sequence labeling)	`deid_ner` (`{"labels": TextProcessor}`), TransformerDeID	The label is stored as text, and the model calls `cross_entropy` with `ignore_index` itself.	Close to `timeseries_multiclass`, with token positions in place of times and padding for positions without a label. Deferred until a second task needs it.
Ranking and retrieval	`patient_linkage_mimic3` (`output_schema = {}`), MedLink, `metrics/ranking.py`	MedLink trains through `Trainer` on in-batch pairs. It is evaluated outside it: the corpus is encoded and searched, results are restricted to each query's candidates (`filter_by_candidates`), then scored with `ranking_metrics_fn`.	A `RankingLabelProcessor` (`"ranking"`): per query, `{"candidates": [...], "relevance": [...]}`, packed like the time-series labels as `(relevance [ΣC], candidate [ΣC], offsets [B + 1])`. `output_size()` is 1 (one score per candidate). The loss is softmax cross-entropy within each query's group (MedLink's in-batch retrieval loss is the case where the candidates are the batch's positives). `metrics_fn` gives NDCG, MAP, Recall and P at k per query. The packing (`collate`, splitting by offsets) moves into a helper shared with the time-series processors, rather than ranking subclassing them. Searching the full corpus stays model API (`encode_queries`, `encode_corpus`, `search`), and its results are scored with the same `metrics_fn`. Adopting it needs the linkage task to emit candidates and relevance per query.
Reconstruction (target is the input)	image VAE, MedGAN and CorGAN autoencoder pretraining, sequence VAE (#1358)	The image VAE overrides `label_processors` with a regression processor (in this PR). The others report only their loss.	No new processor: the model overrides `label_processors` with the existing processor that matches the input: `RegressionLabelProcessor` for an image, `MultiLabelProcessor(labels=<input vocabulary>)` for a multi-hot code set, `TimeseriesMultiClassLabelProcessor(labels=<vocabulary size>)` for a code sequence (position as time). The models keep their own loss (reconstruction plus KL, or adversarial). Two pieces are planned: a constructor that builds the processor from the input processor's vocabulary, so label indices equal token ids; and packed `(tokens, position, offsets)` as `y_true` for sequences, since `to_numpy` sees only `y_true` and `y_prob`.
Several labels on one model (multi-task)	none yet	`label_processors` can already hold several, but `Trainer` and `model.mode` handle only a model with one label.	Collect predictions and compute metrics per label key in the `Trainer`. This is a `Trainer` change, not a new processor.

Generation (sampling new patients: MedGAN, CorGAN, HALO, GPT-2, PromptEHR, Synthea) is out of scope because it has no per-sample label. The generators train with their own train_model loops, generate(num_samples) returns records rather than tensors, and quality is judged on the whole dataset by metrics/generative (privacy, utility and prevalence against real train and test data). Supporting it properly needs its own interface for generative models (training plus generate) and an evaluator around evaluate_synthetic_ehr, not a label processor. Their tasks keep output_schema = {}.

Testing

Focused suite. The tests of the changed modules (852 tests: processors, the models above, trainer, calibration, interpretability and autoimprove) pass locally, except two that also fail on master here:
- test_trajsurv's state-dict round trip: a CPU/CUDA device mismatch.
- test_transformers_model.TestChunking.test_windows_need_a_fast_tokenizer.
test_survival_mdn imports pytest, which is not installed in this environment, so it did not run locally.
New tests in tests/core/test_label_processor.py:
- activation and loss of every processor against the torch functions they replace
- pos_weight, including "balanced" learned in fit and its leak warning
- the labels= option, and pickles from before the new options
- time-series parsing, collate, the per-record loss, to_numpy/gather and pooled metrics
- an end-to-end Trainer run for all eight modes
Coverage. All added lines in the processors, BaseModel, Trainer and the other real changes above are covered.
Other checks. tools/check_pr_rules.py passes. Pyright reports no errors in the two processor modules, and no new errors in the other changed files.

BaseModel, Trainer, calibration, autoimprove and the interpretability modules each switched on a mode string to decide a label's output size, activation, loss, positive-class weight, how predictions are gathered and which metrics apply. LabelProcessor puts all of that on the processor that defines the label: size(), activation(), loss(), to_numpy(), gather() and metrics_fn(), with pos_weight as a task option of the BCE-based processors. The four static processors and four time-series counterparts (timeseries_multiclass and timeseries_multilabel are new) subclass it. BaseModel exposes label_processors / label_processor and a read-only mode, and drops _resolve_mode, the mode setter, get_output_size, get_loss_function, prepare_y_prob and set_pos_weight; models call the label processor directly.

It forwards to size() for every label processor so far. Models size their output layer with it, so a label whose logits differ from its width only needs to override output_size.

github-actions · 2026-11-05T00:08:16Z

Test Coverage

Total: 77.96% (25015 / 32089 lines)

Modified file	Coverage	Lines
`pyhealth/autoimprove/evaluate.py`	91.98%	149 / 162
`pyhealth/autoimprove/improver.py`	94.86%	498 / 525
`pyhealth/autoimprove/space.py`	97.45%	191 / 196
`pyhealth/calib/calibration/dircal.py`	80.46%	70 / 87
`pyhealth/calib/calibration/logistic_recalibration.py`	97.47%	77 / 79
`pyhealth/calib/calibration/temperature_scale.py`	77.19%	44 / 57
`pyhealth/interpret/methods/base_interpreter.py`	84.00%	21 / 25
`pyhealth/metrics/interpretability/base.py`	76.71%	168 / 219
`pyhealth/models/adacare.py`	88.51%	208 / 235
`pyhealth/models/agent.py`	92.72%	242 / 261
`pyhealth/models/apn.py`	97.48%	116 / 119
`pyhealth/models/base_model.py`	97.44%	38 / 39
`pyhealth/models/biot.py`	73.57%	103 / 140
`pyhealth/models/califorest.py`	89.58%	86 / 96
`pyhealth/models/cnn.py`	80.11%	141 / 176
`pyhealth/models/concare.py`	73.17%	270 / 369
`pyhealth/models/contrawr.py`	95.10%	97 / 102
`pyhealth/models/deepr.py`	88.99%	97 / 109
`pyhealth/models/ehrmamba.py`	84.43%	141 / 167
`pyhealth/models/ehrmamba_cehr.py`	95.56%	43 / 45
`pyhealth/models/gnn.py`	77.00%	241 / 313
`pyhealth/models/gradient_boosted_trees.py`	96.79%	271 / 280
`pyhealth/models/graph_torchvision_model.py`	28.47%	39 / 137
`pyhealth/models/graphcare.py`	11.76%	18 / 153
`pyhealth/models/grasp.py`	92.17%	200 / 217
`pyhealth/models/heartlang.py`	98.58%	209 / 212
`pyhealth/models/jamba_ehr.py`	74.50%	111 / 149
`pyhealth/models/llm_ehr_encoder.py`	91.16%	134 / 147
`pyhealth/models/logistic_regression.py`	71.23%	52 / 73
`pyhealth/models/medfuse.py`	75.68%	112 / 148
`pyhealth/models/micron.py`	83.33%	120 / 144
`pyhealth/models/mlp.py`	72.66%	93 / 128
`pyhealth/models/molerec.py`	66.39%	239 / 360
`pyhealth/models/protoecgnet.py`	97.87%	367 / 375
`pyhealth/models/retain.py`	87.18%	136 / 156
`pyhealth/models/rnn.py`	93.78%	347 / 370
`pyhealth/models/safedrug.py`	86.26%	270 / 313
`pyhealth/models/sparcnet.py`	95.96%	95 / 99
`pyhealth/models/stagenet.py`	96.04%	194 / 202
`pyhealth/models/stagenet_mha.py`	82.31%	228 / 277
`pyhealth/models/survival_mdn.py`	93.91%	216 / 230
`pyhealth/models/tcn.py`	95.45%	105 / 110
`pyhealth/models/tfm_tokenizer.py`	65.54%	232 / 354
`pyhealth/models/torchvision_model.py`	17.58%	32 / 182
`pyhealth/models/trajsurv.py`	95.31%	264 / 277
`pyhealth/models/transformer.py`	88.84%	191 / 215
`pyhealth/models/transformer_deid.py`	100.00%	74 / 74
`pyhealth/models/transformers_model.py`	90.00%	81 / 90
`pyhealth/models/vae.py`	96.88%	93 / 96
`pyhealth/processors/__init__.py`	94.44%	34 / 36
`pyhealth/processors/label_processor.py`	97.69%	211 / 216
`pyhealth/processors/timeseries_label_processor.py`	100.00%	123 / 123
`pyhealth/processors/timeseries_processor.py`	84.42%	65 / 77
`pyhealth/trainer.py`	71.96%	195 / 271

Workflow run · 0552d73

Logiquo added 2 commits November 4, 2026 23:51

Add LabelProcessor.output_size for the number of model outputs

0552d73

It forwards to size() for every label processor so far. Models size their output layer with it, so a label whose logits differ from its width only needs to override output_size.

Logiquo mentioned this pull request Nov 5, 2026

Move label semantics onto LabelProcessor #1363

Closed

Navigation Menu

Repository navigation

Add LabelProcessor: label semantics live on the label processor - #1389

Add LabelProcessor: label semantics live on the label processor#1389
Logiquo wants to merge 2 commits into
masterfrom
label-processor

Logiquo commented Nov 4, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Nov 5, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Logiquo commented Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

1. Why unify them

2. The methods, and why each moves to the processor

3. Files

4. Prediction types not covered yet (planned follow-ups)

Testing

Uh oh!

github-actions Bot commented Nov 5, 2026

Test Coverage

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Logiquo commented Nov 4, 2026 •

edited

Loading