Fix medlink collate_fn: robust to heterogeneous per-sample keys

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:

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

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(

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()