rank_results function abstraction for testing

config.ini

@@ -28,4 +28,4 @@ Dropout = 0.2
 Quantize = False
 
 # Number of results to return on each query
-NumberOfResults = 10
+NumberOfResults = 9

main.py

@@ -8,7 +8,6 @@
 import time
 import argparse
 import torch
-import random
 from configparser import ConfigParser
 from ipv8.community import Community, CommunitySettings
 from ipv8.configuration import ConfigBuilder, Strategy, WalkerDefinition, default_bootstrap_defs
@@ -94,7 +93,6 @@ def started(self) -> None:
 
             inferred_ranking = list(self.ltr.query(query).keys())
             print(fmt(f'nDCG: {round(ndcg(ranked_result_ids, inferred_ranking), 3)}', 'yellow'))
-            print(fmt(f'Random nDCG: {round(ndcg(random.sample(ranked_result_ids, len(ranked_result_ids)), inferred_ranking), 3)}', 'yellow'))
 
         async def app() -> None:
             threading.Thread(target=self.input_thread, daemon=True).start()

p2p_ol2r/ltr.py

@@ -45,6 +45,18 @@ def __init__(self, cfg: Config):
         self.embeddings = Embeddings({ 'path': 'allenai/specter' })
         self.embeddings.load('data/embeddings_index.tar.gz')
 
+    def _get_results(self, query: str) -> list[str]:
+            """
+            Returns an unordered list of relevant search results for the given query.
+
+            Args:
+                query: The search query.
+
+            Returns:
+                A list of document IDs
+            """
+            return [x for x, _ in self.embeddings.search(query, self.cfg.number_of_results)]
+
     def _get_result_pairs(self, query: str) -> list[str]:
         """
         Retrieve top-k results from semantic search and generate all possible combination pairs.
@@ -89,40 +101,50 @@ def gen_train_data(self, query: str, results: list[str], selected_res: int = Non
     def result_ids_to_titles(self, results: list[str]) -> list[str]:
         return [self.metadata[x] for x in results]
 
-    def query(self, query: str) -> dict[str, str]:
+    def rank_results(self, query: np.ndarray, result_ids: list[str]) -> list[str]:
         """
         Determine ranking of results based on pairwise comparisons on the model.
 
+        Args:
+            query: The query to rank the results against.
+            result_ids: The list of doc IDs to rank.
+
+        Returns:
+            The list of result IDs sorted by their relevance to the query.
+        """
+        result_pairs = list(combinations(result_ids, 2))
+        result_scores = {}
+
+        # aggregate inferred probabilities for each result pair
+        for (d1_id, d2_id) in result_pairs:
+            d1 = self.embeddings_map[d1_id]
+            d2 = self.embeddings_map[d2_id]
+            _, v = self.infer(ModelInput(query, d1, d2))
+            prob_1_over_2 = v[0] if type(v) == tuple else v
+            prob_2_over_1 = v[1] if type(v) == tuple else 1-v
+
+            result_scores[d1_id] = result_scores.get(d1_id, 0) + prob_1_over_2
+            result_scores[d2_id] = result_scores.get(d2_id, 0) + prob_2_over_1
+
+        # order result scores
+        result_scores = dict(sorted(result_scores.items(), key=itemgetter(1), reverse=True))
+
+        return [res_id for res_id, _ in result_scores.items()]
+
+    def query(self, query: str) -> dict[str, str]:
+        """
+        Returns ranked results for the given query.
+
         Args:
             query: query string
         
         Returns:
             Dict of result IDs to titles ordered by their ranking
         """
         query_vector = self.embed(query)
-        results_combs = self._get_result_pairs(query)
-        results_scores = {}
-
-        self.model.eval()
-        with torch.no_grad():
-            for result_pair in results_combs:
-                vec1 = self.embeddings_map[result_pair[0]]
-                vec2 = self.embeddings_map[result_pair[1]]
-
-                _, v = self.infer(ModelInput(query_vector, vec1, vec2))
-                prob_1_over_2 = v[0] if type(v) == tuple else v
-                prob_2_over_1 = v[1] if type(v) == tuple else 1-v
-
-
-                k = result_pair[0]
-                results_scores[k] = results_scores.get(k, 0) + prob_1_over_2
-                k = result_pair[1]
-                results_scores[k] = results_scores.get(k, 0) + prob_2_over_1
-
-        # aggregating results by summing the probabilities of being superior
-        results_scores = dict(sorted(results_scores.items(), key=itemgetter(1), reverse=True))
-        ranked_results = {res_id: self.metadata[res_id] for res_id, _ in results_scores.items()}
-        return ranked_results
+        docs = self._get_results(query)
+        ranked_results = self.rank_results(query_vector, docs)
+        return {id: self.metadata[id] for id in ranked_results}
 
     def on_result_selected(self, query: str, results: list[str], selected_res: int):
         """

p2p_ol2r/model.py

@@ -115,6 +115,10 @@ def _train_step(self, model_input: ModelInput, label: bool) -> float:
     def train(self, true_train_data: list[ModelInput], epochs: int = 1):
         """
         Trains the model on the given training data and its inverse (if i>j is true then j>i must be false).
+
+        Args:
+            true_train_data: The training data to train on to be classified as true.
+            epochs: The number of epochs to train each training data item on.
         """
         self.model.train()
         shuffle(true_train_data)
@@ -134,9 +138,9 @@ def infer(self, model_input: ModelInput) -> (bool, float|tuple):
         Infer the relative relevance of two documents given a query from the model.
 
         Args:
-            query (np.ndarray): The query vector.
-            sup_doc (np.ndarray): The supposedly superior document vector.
-            inf_doc (np.ndarray): The supposedly inferior document vector.
+            query: The query vector.
+            sup_doc: The supposedly superior document vector.
+            inf_doc: The supposedly inferior document vector.
 
         Returns:
             bool: True if the superior document is more relevant than the inferior document, False otherwise.

tests/test_ltr.py

@@ -2,15 +2,41 @@
 from p2p_ol2r.ltr import *
 from tests import cfg
 
-class TestGenTrainData(unittest.TestCase):
+class TestLTR(unittest.TestCase):
 
     ltr = LTR(cfg)
 
-    def test_lengths(self):
+    def test_gen_train_data(self):
         query = 'molecular tumor'
         results = [x for x, _ in self.ltr.embeddings.search(query, cfg.number_of_results)]
         train_data = self.ltr.gen_train_data(query, results, 0)
         self.assertEqual(len(train_data), cfg.number_of_results - 1)
+
+    def test_rank_results(self):
+        q = np.random.rand(768)
+        k = cfg.number_of_results
+        doc_ids = [f'id{i}abcd' for i in range(k)]
+        self.ltr.embeddings_map = {id: np.random.rand(768) for id in doc_ids}
+
+        # We mock `infer` to respond positively for each query.
+        # We expect the results to be ordered as they are in `doc_ids`.
+        for r in [
+            (True, (1.0, 0.0)), (True, (0.51, 0.49)),
+            (True, 1.0), (True, 0.51)
+        ]:
+            with unittest.mock.patch('p2p_ol2r.model.LTRModel.infer', new=lambda _, __: r):
+                ordered_docs = self.ltr.rank_results(q, doc_ids)
+                self.assertListEqual(ordered_docs, doc_ids)
+
+        # We mock `infer` to respond negatively for each query.
+        # We expect the results to be ordered reversed to what they are in `doc_ids`.
+        for r in [
+            (False, (0.0, 1.0)), (False, (0.49, 0.51)),
+            (False, 0.0), (False, 0.49)
+        ]:
+            with unittest.mock.patch('p2p_ol2r.model.LTRModel.infer', new=lambda _, __: r):
+                ordered_docs = self.ltr.rank_results(q, doc_ids)
+                self.assertListEqual(ordered_docs, list(reversed(doc_ids)))
 
 if __name__ == "__main__":
     unittest.main()

tests/test_model.py

@@ -51,6 +51,8 @@ def test_recent_bias(self):
         self.assertFalse(res)
 
     def test_how_much_it_takes(self):
+        # We test 100 inputs of docs[0] > docs[1] vs. 90 inputs of docs[1] > docs[0], 
+        # expecting docs[0] still to be ranked higher than docs[1].
         ltr_model, _, q, docs = setUp()
         train_data = [ModelInput(q, docs[0], docs[1])] * 100
         train_data.extend([ModelInput(q, docs[1], docs[0])] * 90)
@@ -61,7 +63,7 @@ def test_how_much_it_takes(self):
         self.assertFalse(res)
 
 
-    @unittest.skip("the ultimate test - doesn't work yet 🥲")
+    #@unittest.skip("the ultimate test - doesn't work yet 🥲")
     def test_full_ranking(self):
         ltr_model, k, q, docs = setUp()
 
@@ -70,33 +72,17 @@ def test_full_ranking(self):
         for i in range(k-1):
             # docs[i] to be above all others
             i_over_all = [ModelInput(q, docs[i], docs[j]) for j in range(k) if i != j]
-            epochs = max(0, k*10 - i*10)
+            epochs = max(0, k*10 - i*10) # with k=9 and epochs+100, this test will fail
             train_data.extend(i_over_all * epochs)
 
         with silence(): ltr_model.train(train_data)
 
-        # expected: docs[0] > docs[1] > ... > docs[k-1]
         for i in range(k-1):
-            print(f"docs[{i}] > docs[{i+1}]")
-            res, values = ltr_model.infer(ModelInput(q, docs[i], docs[i+1]))
-            print(values)
-            self.assertTrue(res)
-
-            #res, values = ltr_model.infer(ModelInput(q, docs[i+1], docs[i]))
-            #print(values)
-            #self.assertFalse(res)
-
-        # for i in range(k-1):
-        #     for j in range(i+1, k):
-        #         # self.assertGreater(
-        #         #     ltr_model.model(torch.from_numpy(mki(q, docs[j], docs[i]))).item(),
-        #         #     0.5
-        #         # )
-        #         self.assertLess(
-        #             ltr_model.model(torch.from_numpy(mki(q, docs[j], docs[i]))).item(),
-        #             0.5
-        #         )
-        #         break
+            for j in range(i+1, k):
+                res, _ = ltr_model.infer(ModelInput(q, docs[i], docs[j]))
+                self.assertTrue(res)
+                res, _ = ltr_model.infer(ModelInput(q, docs[j], docs[i]))
+                self.assertFalse(res)
 
 if __name__ == "__main__":
     unittest.main()