Fix kg_emb broken import: SampleBaseDataset -> SampleDataset, add pan…

[userjuma]

Closes #952

Fixes the broken import in kg_emb caused by SampleBaseDataset being renamed to SampleDataset in PyHealth 2.0. Updated across all 7 affected files (2 in kg_emb/datasets, 5 in kg_emb/models).

Also adds pandarallel (>=1.6.5) to pyproject.toml. It is used in kg_emb/datasets/base_kg_dataset.py but was previously missing from declared dependencies, which caused a ModuleNotFoundError immediately after the dataset rename was resolved.

Added tests/core/test_kg_emb.py covering:
1.Clean import of the kg_emb module
2. Importability of all five kg_emb model classes
3.A regression check ensuring TransE's dataset parameter is explicitly typed as SampleDataset

Manual verification:
Reproduced the original crash via python -c "import pyhealth.medcode.pretrained_embeddings" and confirmed it now resolves cleanly end-to-end.

[AxelNoun]

Thanks for this. The ImportError diagnosis is right, and the failure chain in
your description, SampleBaseDataset first and then pandarallel, matches what
I reproduced exactly.

I think one step is missing though. SampleDataset in 2.0 isn't the successor
to SampleBaseDataset: it's a litdata.StreamingDataset that expects a
directory containing schema.pkl. SampleKGDataset.__init__ still passes a
list, so:

ValueError: dir_path must be either a string, Path, or Dir, got: <class 'list'>

The module imports, but set_task() can't return anything usable. The three
tests here don't catch it because none of them instantiates the class.

Two smaller notes:

• pandarallel: initialize() is indeed called in umls.py, but the import in
  base_kg_dataset.py, the one cited in the description, is dead, and there is
  no parallel_apply anywhere in kg_emb. It is also not declared in
  pyproject.toml on master, which is what produced the ModuleNotFoundError
  in the first place.
• The test comparing TransE.__init__'s annotation to SampleDataset locks in
  an incorrect hint: KGEBaseModel reads dataset.entity_num, which only
  exists on SampleKGDataset.

Worth saying that #952 asked for the rename across eight files and your PR does
exactly that, so this is really about how the issue was framed.

I've opened #1202 going after the root cause, and credited you in the
description. Would be glad to have your review on it.

[userjuma]

Thanks for the thorough review and for taking the time to reproduce the issue directly.

We reached the same root cause independently: SampleKGDataset was calling SampleDataset.__init__ using the deprecated in-memory list interface, whereas SampleDataset is now backed by litdata and expects paths generated via SampleBuilder.save(). A recent commit addresses this by bypassing super().__init__(), assigning the required base attributes directly, and implementing a dedicated __len__ method.

The accompanying feedback has also been addressed:

1. Corrected dataset type annotations across KGEBaseModel, TransE, RotatE, DistMult, ComplEx, and split() to explicitly require SampleKGDataset, which supplies the necessary entity_num and relation_num attributes.
2. Expanded the test suite beyond basic import and annotation checks to instantiate SampleKGDataset, execute split(), and initialize all five model classes end-to-end.
   3.Clarified the dependency requirement for pandarallel; while the import in base_kg_dataset.py was unused, the dependency remains declared in pyproject.toml to support the active call inside umls.py during initialize().

Since #1202 addresses the same underlying issue, this PR can be closed in favor of yours if that branch is further along, or we can move forward with whichever implementation is cleaner.

[AxelNoun]

Thanks for jumping on the constructor fix and tests so fast, that definitely closed the gap I flagged.

I think I'll take you up on your offer (if you're still cool with it). It really just came down to scope: #1202 completely drops the SampleDataset inheritance. Since kg_emb doesn't actually use the streaming contract, this decouples it so it won't break when that layer changes. I put the full context in the PR description so we don't have to rehash it here.

Since this all started from your initial diagnosis, would you mind giving #1202 a quick review when you have a sec? I'd love to get your eyes on the redesign before we ping a maintainer. Please be brutal if anything looks off, I'd way rather catch it now than after merge!

[AxelNoun]

Update, and I need to walk something back: please don't close this in favor of #1202.

After I'd already asked you to close, I checked the target architecture with John Wu, and he wants SampleKGDataset to stay inside the SampleDataset hierarchy the opposite of what #1202 currently does. That means #1202 needs a rework, and in the meantime the rename here is needed no matter which hierarchy we land on. #1192 is scoped exactly to #952 as filed, so I think it's the right thing to merge now rather than wait on me.

I'd like to offer you the fixes from #1202 that don't depend on the hierarchy question split()'s reproducibility/validation fixes, the entity2id=None constructor crash, stat() returning None, and the phantom SampleKGDataset import in the __main__ blocks. Whichever you'd prefer: a PR against this branch, or a follow-up on master once you've merged.

Worth noting: I branched these off master and four of my tests fail on len(dataset) your dedicated __len__ is what makes them runnable at all. So they land on top of yours either way.

[userjuma]

Good call checking with John before committing to the redesig definitely better to confirm the direction upfront than rework things later.

Glad #1192 can land as filed. Good to know the __len__ fix is pulling double duty for your tests, too.

As for the other four bugs you spotted, let’s handle those in a dedicated follow-up PR against master once this one merges, rather than continuing to pile onto this branch. Between the initial rename, the docstring examples, the init fix, and the type hints, #1192 has already grown past its original scope. Keeping your four fixes as their own focused PR ensures they get reviewed on their own merits instead of getting buried inside an unrelated diff.

Ping me when its up

[AxelNoun]

It's very good, congratulations!

I just have two tiny, non-blocking nits whenever you have a second:

• pyproject.toml: The linear-attention-transformer line picked up an extra 2-space indent. (The new pandarallel line matches the existing style, so it's just that one line).
• Trailing newlines: Both examples/medcode.py and tests/core/test_kg_emb.py are missing a trailing newline at the end of the file.

@jhnwu3 It seems to me that everything is good!