diff --git a/docs/api/models/pyhealth.models.MedLink.rst b/docs/api/models/pyhealth.models.MedLink.rst index d80f351c8..15913c2aa 100644 --- a/docs/api/models/pyhealth.models.MedLink.rst +++ b/docs/api/models/pyhealth.models.MedLink.rst @@ -3,6 +3,14 @@ The complete MedLink model. +.. note:: + + ``pyhealth.models.medlink.utils.collate_fn`` (used by + ``get_train_dataloader``) drops the ``s_n`` (hard negative) field for an + entire batch if any sample in it lacks one, rather than producing a + partially-present, misaligned list -- ``MedLink.forward`` consumes + ``s_n`` as a whole-batch field (``corpus = s_p + s_n``), so a + per-sample-optional value would corrupt that concatenation. .. autoclass:: pyhealth.models.MedLink :members: diff --git a/examples/patient_linkage_mimic3_medlink.py b/examples/patient_linkage_mimic3_medlink.py index 237d23f14..8a373acc4 100644 --- a/examples/patient_linkage_mimic3_medlink.py +++ b/examples/patient_linkage_mimic3_medlink.py @@ -17,10 +17,16 @@ """ IMPORTANT: This implementation differs from the original paper in order to make it work with the PyHealth framework. Specifically, we do not use the -pre-trained GloVe embeddings. And we only monitor the loss on the validation +pre-trained GloVe embeddings. And we only monitor the loss on the validation set instead of the ranking metrics. As a result, the performance of this model is different from the original paper. To reproduce the results in the paper, please use the official GitHub repo: https://github.com/zzachw/MedLink. + +Note: get_train_dataloader emits one training sample per (query, positive) +pair, so a query with multiple positives can end up in the same batch as a +query with none (or vice versa). collate_fn handles this by dropping the +hard-negative field (s_n) for the whole batch if it isn't present on every +sample, rather than producing a misaligned batch. """ USE_BM25_HARDNEGS = False diff --git a/pyhealth/models/medlink/utils.py b/pyhealth/models/medlink/utils.py index 6b3263655..235b9f239 100644 --- a/pyhealth/models/medlink/utils.py +++ b/pyhealth/models/medlink/utils.py @@ -143,11 +143,36 @@ def get_bm25_hard_negatives(bm25_model, corpus, queries, qrels): def collate_fn(samples): - outputs = {k: [] for k in samples[0].keys()} - for sample in samples: - for k, v in sample.items(): - outputs[k].append(v) - return outputs + """Batches a list of per-sample dicts into a dict of lists. + + Note: + Initializing output keys from only ``samples[0]`` breaks when + samples have heterogeneous keys -- e.g. ``get_train_dataloader``'s + ``s_n`` key, present only for samples where a hard negative was + mined. Depending on sample order this either raised a KeyError (the + first sample lacked a key a later one had) or silently produced a + shorter, misaligned list for that key (the first sample had it, a + later one didn't). ``MedLink.forward`` consumes ``s_n`` as a + whole-batch field (``corpus = s_p + s_n``), so a partially-present + ``s_n`` would corrupt that concatenation rather than just misalign + cleanly -- if any sample in the batch is missing it, drop it for + the whole batch instead (equivalent to training that batch without + hard negatives, a mode the model already supports via s_n=None). + + Examples: + >>> samples = [ + ... {"query_id": "q1", "id_p": "p1", "s_q": "Q1", "s_p": "P1", "s_n": "N1"}, + ... {"query_id": "q2", "id_p": "p2", "s_q": "Q2", "s_p": "P2"}, + ... ] + >>> collate_fn(samples) # s_n dropped for the whole batch: not every sample has one + {'query_id': ['q1', 'q2'], 'id_p': ['p1', 'p2'], 's_q': ['Q1', 'Q2'], 's_p': ['P1', 'P2']} + """ + # dict.fromkeys preserves first-seen order (deterministic, unlike a + # set) while still deduplicating across samples. + keys = dict.fromkeys(k for sample in samples for k in sample) + if "s_n" in keys and not all("s_n" in sample for sample in samples): + del keys["s_n"] + return {k: [sample[k] for sample in samples] for k in keys} def get_train_dataloader( diff --git a/tests/core/test_medlink.py b/tests/core/test_medlink.py index 9a530a11c..e4e2910e4 100644 --- a/tests/core/test_medlink.py +++ b/tests/core/test_medlink.py @@ -129,5 +129,75 @@ def test_feature_key_inference(self): self.assertEqual(model.feature_key, "conditions") +class TestMedLinkCollateFn(unittest.TestCase): + """Regression tests for collate_fn's handling of samples with + heterogeneous keys -- specifically get_train_dataloader's "s_n" key, + present only for samples where a hard negative was mined. + + The original implementation initialized output keys from only + samples[0], so depending on sample order it either raised a KeyError + (samples[0] lacked a key a later sample had) or silently produced a + shorter, misaligned list for that key (samples[0] had it, a later + sample didn't) -- both reproduced below. + """ + + def _common_samples(self, first_has_s_n, second_has_s_n): + s1 = {"query_id": "q1", "id_p": "p1", "s_q": "Q1", "s_p": "P1"} + s2 = {"query_id": "q2", "id_p": "p2", "s_q": "Q2", "s_p": "P2"} + if first_has_s_n: + s1["s_n"] = "N1" + if second_has_s_n: + s2["s_n"] = "N2" + return [s1, s2] + + def test_s_n_present_first_absent_second_does_not_crash(self): + from pyhealth.models.medlink.utils import collate_fn + + samples = self._common_samples(first_has_s_n=True, second_has_s_n=False) + out = collate_fn(samples) # would previously misalign s_n to length 1 + self.assertNotIn("s_n", out) + self.assertEqual(out["id_p"], ["p1", "p2"]) + self.assertEqual(out["s_q"], ["Q1", "Q2"]) + + def test_s_n_absent_first_present_second_does_not_crash(self): + from pyhealth.models.medlink.utils import collate_fn + + samples = self._common_samples(first_has_s_n=False, second_has_s_n=True) + out = collate_fn(samples) # would previously raise KeyError('s_n') + self.assertNotIn("s_n", out) + self.assertEqual(out["id_p"], ["p1", "p2"]) + + def test_s_n_present_in_all_samples_preserved_and_aligned(self): + from pyhealth.models.medlink.utils import collate_fn + + samples = self._common_samples(first_has_s_n=True, second_has_s_n=True) + out = collate_fn(samples) + self.assertEqual(out["s_n"], ["N1", "N2"]) + self.assertEqual(len(out["s_n"]), len(out["id_p"])) + + def test_s_n_absent_in_all_samples_unaffected(self): + from pyhealth.models.medlink.utils import collate_fn + + samples = self._common_samples(first_has_s_n=False, second_has_s_n=False) + out = collate_fn(samples) + self.assertNotIn("s_n", out) + self.assertEqual(out["id_p"], ["p1", "p2"]) + + def test_all_output_lists_same_length_as_batch(self): + """General invariant: every key in a collated batch must have + exactly one entry per input sample, regardless of which samples + contributed which keys.""" + from pyhealth.models.medlink.utils import collate_fn + + for first, second in [(True, False), (False, True), (True, True), (False, False)]: + samples = self._common_samples(first, second) + out = collate_fn(samples) + for key, values in out.items(): + self.assertEqual( + len(values), len(samples), + f"key {key!r} has {len(values)} entries, expected {len(samples)}", + ) + + if __name__ == "__main__": unittest.main()