MedLink Bounty

[Rian354]

PR for MedLink bounty

Tests:
To run the MedLink unit tests, from the project root run:

pytest tests/core/test_medlink.py (locally, 3 passed & 1 warning)

Model Implementation:

• Implemented the MedLink retrieval model on top of the current BaseModel / dataset API.
• Added unit tests with small synthetic data for MedLink.
• Added a Jupyter notebook that trains and evaluates MedLink on the MIMIC-III demo dataset.

Additions to "pyhealth/models/medlink/model.py":

• BaseModel-compatible "MedLink" class that takes a task-generated dataset (e.g., "SampleDataset" from "set_task") and "feature_keys".
• Vocabulary construction from the underlying task dataset using "dataset.get_all_tokens(...)" for queries and documents.
• Query and corpus encoders ("encode_queries", "encode_corpus") that produce sparse multi-hot representations.
• BM25-style scoring in "compute_scores", compatible with the IR-format data produced by the MedLink utilities.
• Combined retrieval and prediction loss in forward / get_loss, returning a scalar loss for training.

Other changes:

• Extended SampleDataset w/ get_all_tokens(key: str) to collect unique tokens across samples, used by MedLink for vocabulary building.
• Implemented BM25 and IR helpers in the pyhealth.models.medlink package:
  • BM25Okapi
  • convert_to_ir_format, tvt_split
  • generate_candidates, filter_by_candidates
  • get_bm25_hard_negatives, get_train_dataloader, get_eval_dataloader
• Exported MedLink via pyhealth.models.init, so users can do: from pyhealth.models import MedLink

Added examples/medlink_mimic3.ipynb, a runnable notebook that:

Loads the MIMIC-III demo dataset via MIMIC3Dataset.

Defines a patient linkage task to generate query–candidate pairs.

Uses the MedLink helpers to build IR-format data and PyTorch dataloaders.

Trains and evaluates MedLink and reports ranking metrics.

Locally ran:

examples/medlink_mimic3.ipynb runs end-to-end on the MIMIC-III demo dataset.

The notebook includes a note on how to run the MedLink unit tests from project root.

Files to review:

pyhealth/datasets/sample_dataset.py – SampleDataset.get_all_tokens helper for vocabulary construction.

pyhealth/models/medlink/model.py – core MedLink model implementation.

pyhealth/models/medlink/bm25.py – BM25Okapi implementation used in the retrieval pipeline.

pyhealth/models/medlink/utils.py – IR-format conversion, TVT split, candidate generation, dataloaders.

pyhealth/models/init.py – export of MedLink.

tests/core/test_medlink.py – synthetic unit tests for MedLink (forward pass, encoders, score shapes).

examples/medlink_mimic3.ipynb – Jupyter notebook for training and evaluating MedLink on the MIMIC-III demo dataset.

[jhnwu3]

I'll probably add more comments as I have more time to dig deeper into this, but nice first attempt at actually a pretty hard bounty.

[jhnwu3]

Some quick thoughts that:

• Can we move the medlink task into the pyhealth.tasks module too? I actually think it'd be really helpful also to further have detailed documentation surrounding the query/document identifiers. It'd be good to link it up with the original paper's task of mapping records to a master known patient record.

[jhnwu3]

It would also be nice to have it in the docs/ as that'll actually be a pretty nice to have for anyone working on record linkage problems.

[jhnwu3]

Can we also try to see if we can't build new processors here to pass to the MedLink model.

Actually, I think the sequence processors should have built-in vocabularies here. But, it would be nice to update the EmbeddingModel to better support things like initialized Glove vectors or just use randomly initailized embeddings for now. This way medlink can be better integrated with the rest of PyHealth, and I think it'd be a nice lesson in replicating the original implementation. (A lot of the techniques are pretty relevant to clinical predictive modeling, so I think it's a good learning exercise).

Example of a PR working with the processors instead of the previous old PyHealth tokenizer approach here: https://github.com/sunlabuiuc/PyHealth/pull/610/files

Glove vectors from the original implementation: https://github.com/zzachw/MedLink here.

[Rian354]

Summary

Updated with content that completes the Medlink bounty via adding the processor native implementation as requested (integrated w/ Pyhealth 2.x processors), unit tests with synthetic data, and an end to end MIMIC-III demo. It also has optional pretrained embedding initialization (glove style vectors) to better match the original Medlink implementation.

Changes

1) MedLink model implementation (processor native)

• Refactored pyhealth/models/medlink/model.py to use SequenceProcessor vocab over the legacy tokenizer approach.
• Builds a shared vocabulary across query and document processors (conditions and d_conditions) and remaps indices so embeddings are consistent between both.
• Contains randomly initialized embeddings by default, and optional pretrained embedding initialization (glove style) when given.

2) Pretrained embeddings support

• Updated EmbeddingModel to support loading pretrained embedding vectors from a text file and optionally freezing them.
• This allows for reproducing the original Medlink workflow (optional) while keeping default behavior as just random initialization.

3) Unit tests

• Added/updated tests/core/test_medlink.py with synthetic data to validate:
  • forward pass returns a scalar loss
  • query/corpus encoding shapes are consistent w/ vocab
  • score matrix shape matches (num_queries, num_docs)
• Tests are lightweight and do not require MIMIC.

4) End-to-end example notebook

• Added examples/medlink_mimic3.ipynb showing:
  • loading MIMIC-III demo data
  • building patient linkage samples
  • training Medlink for a few epochs
  • evaluating with ranking metrics

Bounty requirement coverage from feedback

• Implement MedLink model (using PyHealth 2.x patterns; processor-driven)
• Add a unit test with pseudo data that runs a forward pass (tests/core/test_medlink.py)
• Add an example notebook demonstrating training on MIMIC-III (examples/medlink_mimic3.ipynb)

How to run tests

From repo root:

pip install -e .
pytest -q tests/core/test_medlink.py

[jhnwu3]

Very close! Thanks for the hardwork!

[jhnwu3]

Can we make sure our documentation (docstrings) clearly documents what the exact inputs and outputs are for the model?

[Rian354]

The tests seem to initialize sampledataset w/ samples parameter, but its expecting a schema file. So I added a helper to build the dataset w/ in memory samples. Should pass tests now