PyHealth 2.0: several modules still target the 1.x dataset API

Start by reading the affected-module table and the linked issues #952 and #1202; the issue does not name specific test files or a chosen migration policy. The first decision is whether each affected module should migrate, be removed, or fail loudly, then scope a focused change rather than treating this as one large PR. Done means the selected module no longer silently constructs or uses the incompatible 1.x API, with behavior verified by an appropriate test.

Summary

While preparing a fix for #952 I audited the modules that still reference the
PyHealth 1.x dataset API. kg_emb is not an isolated case. It is the only
one whose import breaks, which is why it is the only one with an issue.
The others fail later, and quietly, mostly because the 1.x base classes now
resolve to deprecation stubs whose __init__ only emits a warning:

class BaseEHRDataset:
    """This class is deprecated and should not be used."""

    def __init__(self, *args, **kwargs):
        import warnings
        warnings.warn(
            "The BaseEHRDataset class is deprecated and will be removed in a future version.",
            DeprecationWarning,
        )

A subclass of such a stub constructs successfully, sets none of its attributes,
and fails somewhere downstream, often several calls later, with an error that
does not point at the cause. All of the modules below are still exported from
their packages, so users can import and instantiate them today.

This is a tracker, not a request for one large PR. I am opening it so the scope
is visible in one place; I am not proposing to fix all of it myself.

Last checked on 2026-10-06 against master at 5568155 (17 commits past
v2.0.2). The only later commit on master, e94ca1d (#1273), does not change
any of the code cited below. What changed since the first version is listed at
the end.

Affected modules

Group A: reachable runtime failures

Module	1.x symbol	Breaks at	Why
medcode/.../kg_emb	SampleBaseDataset	import	SampleKGDataset subclasses SampleBaseDataset, a name removed without a stub. Renaming it to the 2.0 SampleDataset only moves the failure to construction: that class is a streaming dataset that expects a directory, and the constructor passes it a list of samples. Addressed in #1202.
datasets/mimicextract.py	BaseEHRDataset, Visit, parallel_apply, 1.x Event signature	runtime	The stub constructor only warns and never runs the 1.x parse_tables pipeline, so nothing is loaded. Called directly, parse_basic_info fails at the first of the four parallel_apply calls (l. 199), since the stub does not initialise pandarallel (no longer a declared dependency); without PyTables, also undeclared, it fails earlier, at read_hdf. Visit is a stub too, and 2.0 Patient has no add_visit. The three Event(...) calls (l. 267, 377, 487) pass 1.x keyword arguments and no event_type, which 2.0 Event.__init__(event_type, timestamp=None, **kwargs) rejects with TypeError (also noted in #1247).
datasets/shhs.py	BaseSignalDataset	runtime	No local __init__, so the stub never sets root.
datasets/isruc.py	BaseSignalDataset	runtime	Same as SHHS.
datasets/cardiology.py	BaseSignalDataset	runtime (partial)	Reassigns self.root itself after the no-op super().__init__, so construction survives; the 1.x set_task pipeline it then relies on is gone. Its task functions in tasks/cardiology_detect.py (cardiology_is*_fn) are 1.x-style; the review of #1241 (28 Sep) says they should move to the object-oriented task API.
tasks/drug_recommendation.py	Visit, patient[i]	runtime	The three legacy functions (drug_recommendation_mimic3_fn, drug_recommendation_mimic4_fn, drug_recommendation_omop_fn) loop for i in range(len(patient)) and read visit: Visit = patient[i]; 2.0 Patient defines neither __len__ nor __getitem__. Since #1203 the API docs state that they cannot run through BaseDataset.set_task() and point to the task classes; with DrugRecommendationOMOP added there, each function now has a class replacement.
models/vae.py	BaseSignalDataset, SampleSignalDataset, 1.x BaseModel signature	construction	VAE.__init__ passes feature_keys, label_key and mode to BaseModel.__init__, which now accepts only dataset, so VAE(...) always raises TypeError. BaseSignalDataset annotates the dataset parameter; SampleSignalDataset is imported in __main__ but not used. #1210 proposed deleting the module; on 31 Aug lehendo asked to make it 2.0-compatible instead, and #1210 was closed.

Group B: stale references, no runtime impact

Module	1.x symbol	Nature	Why
models/gan.py	BaseSignalDataset	dead import	Imported, never used.
models/safedrug.py	SampleEHRDataset	annotation	Runtime goes through BaseModel, so only the hint is stale.
models/contrawr.py	SampleSignalDataset	docstring	A 1.x >>> example that would fail if doctests were run.

Resolved since the first version

Module	Resolved by
metrics/fairness_utils/utils.py	#1200 (in v2.0.2): sensitive_attributes_from_patient_ids now takes a BaseDataset and reads demographics through get_patient() and get_events(), with a doctest.

Two API-level observations

These are not module bugs, but they shape how a migration would be written.

1. get_dataloader assumes the SampleDataset hierarchy. It calls
   dataset.set_shuffle(shuffle), a litdata.StreamingDataset method that
   InMemorySampleDataset also implements, so it cannot accept a plain
   torch.utils.data.Dataset even though its docstring presents it as a general
   helper. Any module migrating to a plain Dataset has to bypass it.
2. collate_fn_dict_with_padding dispatches on tuple length. Any value that
   is a 2-tuple is treated as a temporal (time, values) pair, so a task whose
   samples legitimately contain 2-tuples would be silently mis-collated. KG
   triples are 3-tuples and pass through unaffected, so this is a latent hazard
   rather than a current bug.

Suggested triage

Each module needs one of three decisions, and I do not think that decision is mine:

• migrate to the 2.0 BaseDataset / SampleDataset pipeline;
• remove, with a clear error pointing at the 2.0 replacement;
• keep as documented-deprecated, but fail loudly.

The current situation is, in my view, the worst of the three, since the stubs
neither work nor fail clearly. Making the five stub __init__ methods (the four
in pyhealth/datasets/__init__.py and Visit in pyhealth/data/__init__.py)
raise NotImplementedError with a migration pointer, instead of warning and
returning a half-built object, is a small change that would convert the
stub-backed Group A entries (mimicextract, shhs, isruc, cardiology) from
a confusing downstream error into an actionable one, and would surface the real
scope immediately.

One inconsistency could be settled at the same time: five 1.x names were retired
behind warning stubs, but SampleBaseDataset was removed outright with no stub
at all. That asymmetry is why kg_emb was the only module to fail at import,
and therefore the only one that got reported. Whichever policy is chosen, applying
it uniformly would make the remaining breakage visible instead of silent.

How this was verified

Read-only audit against master at 0a75f99: import-chain tracing plus
construction attempts on the concrete classes. Group A entries, as first listed,
were reproduced; Group B entries are static findings from usage-site greps,
excluding the definitions in pyhealth/datasets/__init__.py and
pyhealth/data/__init__.py. Individual tracebacks are available on request.

Re-check on 2026-10-06 against master at 5568155, on a local install of that
commit (Windows, Python 3.12, NumPy 2.2.6, PyTables added for the HDF5 read):
the kg_emb import still raises ImportError (SampleBaseDataset); with that
name aliased to SampleDataset in memory (plus a placeholder for pandarallel,
which kg_emb imports but 2.0 does not declare), SampleKGDataset(samples=[...])
raises ValueError from litdata (a list where a directory is expected);
MIMICExtractDataset.parse_basic_info() on a minimal synthetic
all_hourly_data.h5 raises AttributeError at its first parallel_apply call;
Event(...) called with the keyword arguments used at mimicextract.py l. 267,
377 and 487 raises TypeError (missing event_type); SHHSDataset and
ISRUCDataset construct, then process_EEG_data() raises AttributeError;
CardiologyDataset constructs, keeps root, and has no set_task; the three
legacy drug recommendation functions raise
TypeError: object of type 'Patient' has no len() on a 2.0 Patient;
VAE(...) raises TypeError from BaseModel.__init__; the fairness_utils
doctest passes.

Changes since the first version

• medcode/.../kg_emb: reworded. The first version said the constructor passed a list into SampleDataset; that happens only once SampleBaseDataset is renamed to SampleDataset.
• metrics/fairness_utils: fixed by #1200, moved to "Resolved".
• tasks/drug_recommendation: reworded after #1203 (legacy note in the API docs, DrugRecommendationOMOP).
• datasets/mimicextract: added the three Event(...) calls without event_type. The first version said parse_basic_info fails before any parallel_apply call; with PyTables installed, it fails at the first one.
• datasets/cardiology: linked to the review of #1241.
• models/vae: moved from Group B to Group A, since construction fails. The first version also said that the __main__ example builds a SampleSignalDataset; it only imports it.
• Stubs: five, not four (Visit was not counted). The triage paragraph now names the Group A entries that the change would affect.
• API observation 1: InMemorySampleDataset also implements set_shuffle. The "general helper" remark now refers to the docstring, since the dataset parameter is annotated litdata.StreamingDataset.
• Adjacent item from my 29 Aug comment below: np.in1d is called at kg_base.py lines 155 and 171 on master (the comment first said 167, the line on the #1202 branch at the time). It works under the pinned numpy~=2.2.0; checked present in NumPy 2.2.0, 2.3.0 and 2.3.5, and absent (AttributeError) in 2.4.0.

Related: #952, #1202, #1200, #1203, #1210, #1241, #1247.