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

[AxelNoun]

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 simply the
only one whose import breaks, which is why it is the only one with an issue.
The others fail later, and quietly, 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.

Affected modules

Group A: reachable runtime failures

Module	1.x symbol	Breaks at	Why
medcode/.../kg_emb	SampleBaseDataset	import	SampleKGDataset subclassed a name removed without a stub, and its constructor passed a list into the 2.0 streaming SampleDataset. Addressed in #1202.
datasets/mimicextract.py	BaseEHRDataset, Visit, parallel_apply	runtime	The stub constructor only warns, so parse_basic_info fails before any of the four parallel_apply calls is reached. Visit is a stub too, and 2.0 Patient has no add_visit.
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.
tasks/drug_recommendation.py	Visit, patient[i]	runtime	for i in range(len(patient)) followed by visit: Visit = patient[i]; 2.0 Patient defines neither __len__ nor __getitem__.
metrics/fairness_utils/utils.py	BaseEHRDataset, dataset.patients	when called	Reads dataset.patients[patient_id]; the 2.0 Patient surface differs.

Group B: stale references, no runtime impact

Module	1.x symbol	Nature	Why
models/gan.py	BaseSignalDataset	dead import	Imported, never used.
models/vae.py	BaseSignalDataset, SampleSignalDataset	annotation and example	BaseSignalDataset annotates the dataset parameter; the __main__ example builds a SampleSignalDataset, whose stub does not store samples.
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.

Two API-level observations

Not module bugs, but they shape how a migration would be written.

1. get_dataloader is streaming-only. It calls dataset.set_shuffle(shuffle),
   a litdata.StreamingDataset method, so it cannot accept a plain
   torch.utils.data.Dataset even though it is presented as a general helper.
   Any module migrating away from a streaming 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 sharp
   edge rather than a current bug.

Suggested triage

Each module needs one of three decisions, and I don't think that call 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 arguably the worst of the three, since the stubs
neither work nor fail clearly. Making the four stub __init__ methods raise
NotImplementedError with a migration pointer, instead of warning and returning
a half-built object, is a small change that would convert every Group A entry
from a confusing downstream error into an actionable one, and would surface the
real scope immediately.

One inconsistency worth settling at the same time: four 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 you pick, 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 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.
Happy to paste individual tracebacks if useful.

Related: #952, #1202.

[AxelNoun]

Follow-up PR: #1202.