PyHealth 2.0: modules still targeting the 1.x API

[AxelNoun]

Several PyHealth modules still target the 1.x in-memory dataset API. In 2.0 those base classes are empty deprecation stubs (BaseEHRDataset, BaseSignalDataset, Visit, SampleBaseDataset removed without a stub). The modules import cleanly and fail later, except kg_emb, whose import already breaks.

This issue tracks the leftovers. A follow-up PR will restore kg_emb as a standalone torch.utils.data.Dataset (replacing #1192 / related to #952). Other rows are out of that PR on purpose.

Inventory

Module	1.x symbol	Breaks at	Notes
medcode/.../kg_emb	SampleBaseDataset	import	SampleKGDataset subclasses a removed name; constructor still passes a list into 2.0 SampleDataset (streaming + schema.pkl)
datasets/mimicextract.py	BaseEHRDataset, Visit, parallel_apply	runtime	Constructor only warns; parse_basic_info never reaches parallel_apply; Visit is a stub; Patient 2.0 has no add_visit
datasets/shhs.py	BaseSignalDataset	runtime	No local __init__; stub never sets root
datasets/isruc.py	BaseSignalDataset	runtime	Same as SHHS
datasets/cardiology.py	BaseSignalDataset	runtime (partial)	Sets self.root after a no-op super(); 1.x set_task pipeline is gone
models/vae.py	BaseSignalDataset, SampleSignalDataset	__main__ example	Stub does not store samples
models/gan.py	BaseSignalDataset	unused import	Dead 1.x import
models/safedrug.py	SampleEHRDataset	annotation only	Runtime still goes through BaseModel
models/contrawr.py	SampleSignalDataset	docstring	1.x doctest example
metrics/fairness_utils/utils.py	BaseEHRDataset, dataset.patients	if called	Patient 2.0 surface differs
tasks/drug_recommendation.py	Visit, patient[i]	runtime	2.0 Patient has no visit __len__ / __getitem__

Also observed (not a 1.x class, but 2.0 API drift)

• get_dataloader now requires litdata.StreamingDataset.set_shuffle(). It is still documented as a general helper. Map-style datasets (including a restored kg_emb) should use torch.utils.data.DataLoader + collate_fn_dict_with_padding directly.
• collate_fn_dict_with_padding treats every 2-tuple as (time, values). KG triples are 3-tuples, so they currently fall through to the list branch. Latent fragility if a 2-tuple feature is added later.

Suggested order

1. Restore kg_emb without inheriting SampleDataset (map-style Dataset + Protocol).
2. Drop unused pandarallel (initialize() only in kg_emb/datasets/umls.py; no parallel_apply in kg_emb; mimicextract cannot reach parallel_apply).
3. Separate PRs for mimicextract / signal datasets / drug_recommendation.

[AxelNoun]

Superseded by #1201, which carries the same inventory in a cleaner form. Closing the discussion here to keep it in one place.