fix(kg_emb): migrate SampleKGDataset off the removed 1.x dataset API

[AxelNoun]

Closes #952. Supersedes #1192. Part of #1201.

What the rename does not reach

#1192 renames SampleBaseDataset to SampleDataset across 7 files. The rename
clears the ImportError, but the module stays unusable: since 2.0,
SampleDataset is a litdata.StreamingDataset whose constructor expects a
directory containing schema.pkl, while SampleKGDataset.__init__ still passes
a list of samples, an argument the rename does not touch.

>>> SampleKGDataset(samples=[{"triple": (0, 0, 1)}], entity2id={}, relation2id={})
ValueError: dir_path must be either a string, Path, or Dir, got: <class 'list'>
  litdata.streaming.resolver._resolve_dir

self.samples is then never assigned, so __getitem__ and stat() raise
AttributeError, and BaseKGDataset.set_task() cannot return anything usable.
The failure moves from import time to call time.

To be fair to the author: #952 explicitly asked for the rename across eight
files, and #1192 does exactly that. The gap is in how the issue was framed.

Root cause

SampleKGDataset inherited from SampleBaseDataset to reuse __len__ and
.samples. It required none of the parent's behavioural contract and used none
of the services the hierarchy exists to provide: a knowledge graph has no
feature schema, no processors and no patient/visit index. The inheritance was
convenience reuse, not an is-a relationship.

The 2.0 migration did not introduce this. It removed the coincidence that made
it harmless, namely the overlap between the interface actually used and the one
the parent exposed. Adapting SampleKGDataset to the streaming contract would
mean satisfying a constraint the module has no use for, so removing the
inheritance is the minimal fix rather than the drastic one.

What this PR does

SampleKGDataset becomes a standalone torch.utils.data.Dataset.

Models depend on a structural typing.Protocol (PEP 544) instead of a concrete
class. KGEBaseModel.__init__ only ever reads entity_num, relation_num and
task_spec_param, so both the old and the new annotation over-specified the
contract, and both were wrong: SampleDataset has neither attribute.
KGDatasetProtocol states the capability actually required. The annotation
becomes checkable by mypy, the model layer becomes testable with lightweight
doubles, and kg_emb stops being coupled to changes in the main pipeline. That
last point is the durable one: a rename would have held until the next
SampleDataset refactor.

Also in scope:

• split() no longer overwrites NumPy's global random state, and validates its
  input with ValueError rather than assert, which is stripped under
  python -O and should not be load-bearing for user input validation. Note:
  for a given seed=, this returns a different partition than the previous
  implementation (different RNG algorithm). There is no prior release of this
  module, so no existing caller depends on the old ordering.
• The four models' __main__ examples imported SampleKGDataset from
  pyhealth.datasets, where it has never existed. They now use
  torch.utils.data.DataLoader with collate_fn_dict_with_padding directly,
  since get_dataloader calls set_shuffle and is streaming-only.
• The pandarallel import, of a package never declared in pyproject.toml, is
  removed rather than added to the dependencies, since no parallel_apply
  exists anywhere in kg_emb. There is no lockfile entry to regenerate.
• base_kg_dataset.py and umls.py imported their siblings through the
  absolute package path, so the package only initialised correctly for one
  ordering of the imports in __init__.py, which nothing enforced. Enabling
  ruff's isort rule would have sorted base_kg_dataset first and broken it.

Tests

tests/core/test_kg_emb.py, 22 behavioural tests: construction, indexing,
vocabularies, cardinality validation, split partitioning and reproducibility,
generic collation of variable-length ground truths, BaseKGDataset.set_task()
on a synthetic graph, and scoring invariants (DistMult symmetry in head and
tail, TransE margin on an exact triple).

No test asserts on a type annotation: annotations are metadata, and under
PEP 563 such an assertion compares against a string.

Every test was validated by fault injection. Reverting the fix it protects makes
it fail:

Test	Mutation	Caught
test_construction_and_length	restore super().__init__(samples, dataset_name, task_name)	yes
test_is_a_map_style_dataset	add a fake set_shuffle	yes
test_set_task_on_a_synthetic_graph	drop **kwargs from the SampleKGDataset(...) call	yes
test_global_numpy_state_is_untouched	restore np.random.seed plus shuffle	yes
test_variable_length_ground_truth_stays_a_python_list	tensorise ground_truth_head in __getitem__	yes
...	(22/22, full table on request)	yes

These are hand-seeded mutants aimed at this PR's changes, not a systematic
mutation run, so 22/22 is not a mutation score in the usual sense. It supports a
weaker and sufficient claim: a rename-only fix cannot make this suite go green.

Scope and limits

• This decouples kg_emb from the 2.0 pipeline; it does not migrate it into
  the streaming pipeline. If you would rather see the latter, the dataset layer
  here is isolated and tested and would be the starting point.
• Correctness of the plumbing, not empirical validation: no MRR or Hits@k
  compared against the original papers. Verified on synthetic graphs only, as
  UMLS requires a licence.
• get_dataloader remains streaming-only. kg_emb no longer depends on it, so
  this is not blocking here; tracked in PyHealth 2.0: several modules still target the 1.x dataset API #1201 together with the other
  modules still targeting the 1.x API.

Credit to @userjuma for the original diagnosis in #1192. The import chain
described there is accurate and was the starting point for this work.

[AxelNoun]

Update: after discussing direction with @jhnwu3, the target is to keep
SampleKGDataset inside the SampleDataset hierarchy rather than decoupling from it,
with in-memory loading as an interim step for negative sampling. Moving this to draft
while I rework it, and suggesting #1192 merge first as the unblocking fix for #952.

Carrying over regardless of the base class: the split() fixes, the constructor
raising AttributeError on its own default entity2id=None, stat() returning
None, the phantom SampleKGDataset import in the models' __main__ blocks, and the
test suite. Not carrying over: KGDatasetProtocol and the standalone Dataset.

[AxelNoun]

Architecture & Implementation Plan Update

Just dropping a quick update here to keep a clear record of the architectural decisions we aligned on via Discord, which will guide the next commits for #1202:

1. Base Architecture (In-Memory)
As discussed, SampleKGDataset will inherit directly from InMemorySampleDataset. This keeps us correctly under the SampleDataset umbrella while cleanly sidestepping the immediate complexities of handling negative sampling in a pure streaming setup.

2. Data Shape & Serialization (The "Tensor Trick")
To handle the KG triples and the variable-length entity lists (ground_truth_head / ground_truth_tail), we are going to avoid slow JSON/Pickle serialization. I dug into the tuple_time_text_processor.py and will adapt their pattern to keep the litdata backend happy:

• I will implement a custom KGProcessor.
• fit() phase: Calculate the global max_length of the entity lists.
• process() phase: Pad the variable-length lists and convert everything directly into pure PyTorch tensors before serialization.

Next Steps:
I will be working on drafting the KGProcessor and the base class integration tonight. I'll push the new commits once the core pipeline is running so we have concrete code to review!

[AxelNoun]

CI/CD Green & Ready for Review

Everything is fully green! I've completed the implementation of the architectural plan we discussed:

• Reverted SampleKGDataset to correctly inherit from InMemorySampleDataset.
• Integrated the KGProcessor (using the "Tensor Trick" with independent max-length padding) to bypass the litdata serialization bottlenecks.
• Ensured negative sampling correctness in kg_base.py by safely reconstructing the unpadded entity lists via attention masks prior to np.in1d filtering.
• Passed all local repo linting rules (Ruff Python 3.13 typing updates and newly required doctests).

The PR is out of draft and ready for your final review whenever you have time!