Fix MedLink BM25 hard-negative mining

pyhealth/models/medlink/utils.py

@@ -127,18 +127,29 @@ def get_bm25_hard_negatives(bm25_model, corpus, queries, qrels):
 
     Returns:
         qrels_w_neg: Updated qrels dictionary containing both positives (1) and negatives (-1).
+
+    Examples:
+        >>> # bm25_model.get_scores(query) -> {doc_id: score}
+        >>> corpus = {"d0": ["fever", "cough"], "d1": ["fever", "rash"]}
+        >>> queries = {"q0": ["fever", "cough"]}
+        >>> qrels = {"q0": {"d0": 1}}  # d0 is q0's positive match
+        >>> qrels_w_neg = get_bm25_hard_negatives(bm25_model, corpus, queries, qrels)
+        >>> qrels_w_neg["q0"]  # positive kept; top query-ranked non-positive labeled -1
+        {'d0': 1, 'd1': -1}
     """
     qrels_w_neg = {}
     for q_id, q in tqdm.tqdm(queries.items()):
         d_ids = [d_id for d_id in qrels[q_id] if qrels[q_id][d_id] > 0]
-        ds = [corpus[d_id] for d_id in d_ids]
-        for d_id, d in zip(d_ids, ds):
-            scores = bm25_model.get_scores(d)
-            for (ned_d_id, neg_s) in sorted(scores.items(), key=lambda x: x[1],
-                                            reverse=True):
-                if ned_d_id != d_id:
-                    qrels_w_neg[q_id] = {d_id: 1, ned_d_id: -1}
-                    break
+
+        qrels_w_neg[q_id] = {d_id: 1 for d_id in d_ids}
+
+        scores = bm25_model.get_scores(q)
+        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+        for neg_d_id, neg_s in ranked:
+            if neg_d_id not in d_ids: # exclude every positive, not just d_id
+                qrels_w_neg[q_id][neg_d_id] = -1
+                break
+
     return qrels_w_neg
 
 
@@ -167,39 +178,37 @@ def get_train_dataloader(
 
     Returns:
         DataLoader returning batches of dicts.
+
+    Examples:
+        >>> corpus = {"p1": "positive one", "p2": "positive two", "n": "negative"}
+        >>> queries = {"q1": "query"}
+        >>> qrels = {"q1": {"p1": 1, "p2": 1, "n": -1}}
+        >>> loader = get_train_dataloader(corpus, queries, qrels, batch_size=2, shuffle=False)
+        Loaded 2 training pairs.
+        >>> next(iter(loader))["id_p"]
+        ['p1', 'p2']
     """
 
     query_ids = list(queries.keys())
     train_samples = []
     for query_id in query_ids:
         s_q = queries[query_id]
-        id_p, s_p, s_n = None, None, None
-        assert len(qrels[query_id]) <= 2
+        positive_ids, s_n = [], None
         for corpus_id, score in qrels[query_id].items():
             if score == 1:
-                id_p = corpus_id
-                s_p = corpus[corpus_id]
+                positive_ids.append(corpus_id)
             if score == -1:
                 s_n = corpus[corpus_id]
-        if s_n is not None:
-            train_samples.append(
-                {
-                    "query_id": query_id,
-                    "id_p": id_p,
-                    "s_q": s_q,
-                    "s_p": s_p,
-                    "s_n": s_n,
-                }
-            )
-        else:
-            train_samples.append(
-                {
-                    "query_id": query_id,
-                    "id_p": id_p,
-                    "s_q": s_q,
-                    "s_p": s_p,
-                }
-            )
+        for id_p in positive_ids:
+            sample = {
+                "query_id": query_id,
+                "id_p": id_p,
+                "s_q": s_q,
+                "s_p": corpus[id_p],
+            }
+            if s_n is not None:
+                sample["s_n"] = s_n
+            train_samples.append(sample)
     print("Loaded {} training pairs.".format(len(train_samples)))
     train_dataloader = DataLoader(
         train_samples, shuffle=shuffle, batch_size=batch_size, collate_fn=collate_fn

tests/core/test_medlink.py

@@ -128,6 +128,77 @@ def test_feature_key_inference(self):
         )
         self.assertEqual(model.feature_key, "conditions")
 
+    def test_hard_negatives_score_query_not_positive_doc(self):
+        """Regression: hard negatives must be mined by scoring the QUERY, not
+        the positive document. The old code called get_scores(d) on the
+        positive doc, so "hard negatives" were docs similar to the answer.
+        """
+        from pyhealth.models.medlink.utils import get_bm25_hard_negatives
+
+        class FakeBM25:
+            def get_scores(self, text):
+                if text == "QUERY":  # query ranks pos top, then neg_q
+                    return {"pos": 10.0, "neg_q": 9.0, "neg_d": 1.0}
+                return {"pos": 10.0, "neg_d": 9.0, "neg_q": 1.0}  # doc-scoring picks neg_d
+
+        corpus = {"pos": "POS", "neg_q": "NQ", "neg_d": "ND"}
+        queries = {"q1": "QUERY"}
+        qrels = {"q1": {"pos": 1}}
+        out = get_bm25_hard_negatives(FakeBM25(), corpus, queries, qrels)
+        # scoring the query picks neg_q; scoring the positive doc would pick neg_d
+        self.assertEqual(out["q1"], {"pos": 1, "neg_q": -1})
+
+    def test_hard_negatives_exclude_all_positives(self):
+        """Regression: with multiple positives, no positive may be chosen as a
+        negative. Scoring the query ranks the positives on top, so excluding
+        only the current positive would pick another positive as a false negative.
+        """
+        from pyhealth.models.medlink.utils import get_bm25_hard_negatives
+
+        class FakeBM25:
+            def get_scores(self, text):
+                return {"pos1": 10.0, "pos2": 9.0, "neg": 8.0}
+
+        corpus = {"pos1": "P1", "pos2": "P2", "neg": "N"}
+        queries = {"q1": "QUERY"}
+        qrels = {"q1": {"pos1": 1, "pos2": 1}}
+        out = get_bm25_hard_negatives(FakeBM25(), corpus, queries, qrels)
+        neg_ids = [d for d, lbl in out["q1"].items() if lbl == -1]
+        self.assertNotIn("pos1", neg_ids)
+        self.assertNotIn("pos2", neg_ids)
+        self.assertEqual(neg_ids, ["neg"])
+
+    def test_hard_negatives_preserves_all_positives(self):
+        """Regression: with multiple positives, all positives must remain in the
+        output when a hard negative is added.
+        """
+        from pyhealth.models.medlink.utils import (
+            get_bm25_hard_negatives,
+            get_train_dataloader,
+        )
+
+        class FakeBM25:
+            def get_scores(self, text):
+                return {"pos1": 10.0, "pos2": 9.0, "neg": 8.0}
+
+        corpus = {"pos1": "P1", "pos2": "P2", "neg": "N"}
+        queries = {"q1": "QUERY"}
+        qrels = {"q1": {"pos1": 1, "pos2": 1}}
+
+        out = get_bm25_hard_negatives(FakeBM25(), corpus, queries, qrels)
+
+        self.assertEqual(
+            out["q1"],
+            {"pos1": 1, "pos2": 1, "neg": -1},
+        )
+
+        dataloader = get_train_dataloader(
+            corpus, queries, out, batch_size=2, shuffle=False
+        )
+        batch = next(iter(dataloader))
+        self.assertEqual(batch["id_p"], ["pos1", "pos2"])
+        self.assertEqual(batch["s_n"], ["N", "N"])
+
 
 if __name__ == "__main__":
     unittest.main()