[halo-pr-prep] Overview: HALO PyHealth Integration Tasks

[jalengg]

HALO PyHealth Integration - Task Overview

This issue tracks the integration of the HALO synthetic data generator with PyHealth standards. The work is broken down into 9 sequential tasks.

Goal

Integrate HALO to follow PyHealth conventions:

• Use `BaseModel` for the model class
• Use task functions instead of custom datasets
• Use `NestedSequenceProcessor` for data handling
• Follow PyHealth example patterns
• Maintain HALO's custom training loop

Architecture

Hybrid Integration Approach:

• ✅ HALO inherits from `BaseModel` for compatibility
• ✅ Uses PyHealth's task/processor infrastructure for data loading
• ✅ Preserves custom training loop (doesn't use Trainer)
• ✅ Keeps proven HALO transformer implementation unchanged

Tasks

Phase 1: Merge with Master

• [halo-pr-prep] Task 1: Complete upstream/master merge #8 - Task 1: Complete upstream/master merge ⚠️ CRITICAL

Phase 2: Create Task Function

• [halo-pr-prep] Task 2: Create halo_generation task function #9 - Task 2: Create halo_generation task function

Phase 3: Refactor HALO Model

• [halo-pr-prep] Task 3: Refactor HALO model to inherit BaseModel #10 - Task 3: Refactor HALO model to inherit BaseModel

Phase 4: Update Examples

• [halo-pr-prep] Task 4: Update training example script #11 - Task 4: Update training example script
• [halo-pr-prep] Task 5: Update generation example script #12 - Task 5: Update generation example script

Phase 5: Cleanup

• [halo-pr-prep] Task 6: Remove old HALO_MIMIC3Dataset class #13 - Task 6: Remove old HALO_MIMIC3Dataset class
• [halo-pr-prep] Task 7: Update docstrings and documentation #14 - Task 7: Update docstrings and documentation

Phase 6: Testing

• [halo-pr-prep] Task 8: Create end-to-end integration test #15 - Task 8: Create end-to-end integration test

Phase 7: Final Verification

• [halo-pr-prep] Task 9: Run all tests and final verification #16 - Task 9: Run all tests and final verification

Dependencies

```
Task 1 (merge) → Task 2 (task function) → Task 3 (refactor model)
↓
Task 4 (training example)
↓
Task 5 (generation example)
↓
Task 6 (cleanup)
↓
Task 8 (integration tests) ←──────────── Task 7 (docs)
↓
Task 9 (final verification)
```

Documentation

• Implementation Plan: `docs/plans/2026-02-16-halo-pyhealth-integration.md`
• Design Document: `docs/plans/2026-02-16-halo-pyhealth-integration-design.md`
• Branch: `halo-pr-528`

Current Status

⚠️ Task 1 has CRITICAL issues - The merge is incomplete:

• Missing 9 processor files including `NestedSequenceProcessor`
• Broken imports in `pyhealth/processors/init.py`
• Must be fixed before proceeding to Task 2

Success Criteria

• HALO inherits from `BaseModel`
• Task function registered in `pyhealth/tasks/init.py`
• Uses `NestedSequenceProcessor` from master
• Examples run without errors
• Old `HALO_MIMIC3Dataset` removed
• All tests pass
• Ready for PR to master

Labels

All tasks are tagged with `halo-pr-prep` for tracking.

Related

• Original PR: HALO Implementation into PyHealth + Colab Notebook sunlabuiuc/PyHealth#812
• Upstream repo: https://github.com/sunlabuiuc/PyHealth
• Fork: https://github.com/jalengg/PyHealth

[jalengg]

HALO PyHealth Integration — Complete

All 9 tasks are done. Here is the full accounting.

---

What was built

Task 2 — pyhealth/tasks/halo_generation.py (commit a864781f)

• HaloGenerationMIMIC3(BaseTask) and HaloGenerationMIMIC4(HaloGenerationMIMIC3)
• input_schema={"visits": "nested_sequence"}, output_schema={}
• DRY design: MIMIC4 overrides only _icd_col; full __call__ logic lives in MIMIC3
• Convenience instances: halo_generation_mimic3_fn, halo_generation_mimic4_fn
• Uses PyHealth 2.0 Polars API (patient.get_events())
• Exported from pyhealth/tasks/__init__.py

Task 3 — pyhealth/models/generators/halo.py (commits fe080059, 661ec69b)

• HALO(BaseModel) — inherits nn.Module via BaseModel
• Constructor derives vocab sizes automatically from NestedSequenceProcessor (no manual HALOConfig required)
• _halo_collate_fn with pad_sequence to handle variable visit counts across patients in a batch
• forward(visits, **kwargs) → {"loss": scalar, "predictions": tensor}
• train_model(train_dataset, val_dataset=None) — custom loop, not PyHealth Trainer; saves best checkpoint
• synthesize_dataset(num_samples, random_sampling=True) → List[Dict] with string codes
• Exported from pyhealth/models/__init__.py

Task 4 — examples/halo_mimic3_training.py (commit 5fb8ce03)
Replaced 151-line argparse/pickle/HALOConfig script with 61-line PyHealth 2.0 standard pattern: MIMIC3Dataset → set_task → split_by_patient → HALO(dataset, ...) → model.train_model(...).

Task 5 — examples/generate_synthetic_mimic3_halo.py (commit 8892ddaf)
Replaced legacy script (MinimalDataset shim, raw HALOModel, pickle output) with PyHealth 2.0 pattern: same dataset reconstruction as training (to match vocab), checkpoint loading via model.halo_model.load_state_dict, JSON output, summary stats.

Task 6 — Deleted pyhealth/datasets/halo_mimic3.py (commit 7df241f7)
Removed 139 lines of dead code and its import from pyhealth/datasets/__init__.py. Zero remaining references in pyhealth/ or tests/.

Task 7 — Docstrings (commits 299d2721, 4dcff7cf)
All public method Returns: blocks in halo.py rewritten to Google/PyHealth named-return style. Missing type annotation added to dataset param in class Args:.

Task 8 — tests/integration/test_halo_end_to_end.py (commits 61424768, 1b744eb4, d5248de5)
12 tests: 8 in-memory (always run, cover every integration seam), 4 MIMIC-III (skip gracefully if data absent). See MECE rationale comment on #15.

Task 9 — Final verification
57 tests passing, 9 skipped (MIMIC-III data not on this machine), 0 failing. Branch pushed to fork/halo-pr-528.

---

Full commit log (branch-only, oldest → newest)

Commit	Description
0c9e9733	chore: merge upstream/master, resolve processor/model conflicts
40404229	fix: complete merge — add missing processor files from upstream
a864781f	feat: add halo_generation task function (HaloGenerationMIMIC3/4)
fe080059	refactor: make HALO inherit BaseModel with forward() and train_model()
6a814d1f	test: add synthesize_dataset coverage for HALO
661ec69b	fix: collate_fn for variable visit lengths and end-token position in HALO
5fb8ce03	feat: update halo training example to PyHealth 2.0 API
8892ddaf	feat: update halo generation example to PyHealth 2.0 API
7df241f7	feat: remove HALO_MIMIC3Dataset (replaced by HaloGenerationMIMIC3 task)
299d2721	docs: update HALO docstrings to Google/PyHealth style
4dcff7cf	docs: fix synthesize_dataset Returns style and dataset type annotation
61424768	test: add HALO end-to-end integration tests
1b744eb4	test: fix tearDown cleanup, env var path, and relative bootstrap paths
d5248de5	test: guard integration test against sys.modules stub contamination

---

Branch halo-pr-528 is ready for upstream PR to sunlabuiuc/PyHealth.