Github Actions (#124)

pr-format.yml → .github/workflows/pr-format.yml

@@ -8,9 +8,13 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - uses: actions/setup-java@v3
+        with:
+          distribution: 'temurin'
+          java-version: '21'
       - uses: actions/setup-python@v4
         with:
-          python-version: 3.10
+          python-version: '3.10'
       - name: Install Linter and Formatter Dependencies
         run: |
           python -m pip install --upgrade pip
@@ -21,6 +25,12 @@ jobs:
       - name: Install Repo Dependencies
         run: |
           pip install -e .
-      - name: Run Unit Tests
+      - name: Run analysis tests
+        run: |
+          python -m unittest discover -s test/analysis
+      - name: Run evaluation tests
+        run: |
+          python -m unittest discover -s test/evaluation
+      - name: Run rerank tests
         run: |
-          python -m unittest discover -s test
+          python -m unittest discover -s test/rerank

test/rerank/test_RankListwiseOSLLM.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import MagicMock, patch
 
 from dacite import from_dict
 
@@ -183,6 +184,21 @@
 
 
 class TestRankListwiseOSLLM(unittest.TestCase):
+    def setUp(self):
+        self.patcher = patch("rank_llm.rerank.rank_listwise_os_llm.load_model")
+        self.mock_load_model = self.patcher.start()
+        self.mock_llm = MagicMock()
+        self.mock_tokenizer = MagicMock()
+        self.mock_load_model.return_value = self.mock_llm, self.mock_tokenizer
+
+        self.patcher_cuda = patch("torch.cuda.is_available")
+        self.mock_cuda = self.patcher_cuda.start()
+        self.mock_cuda.return_value = True
+
+    def tearDown(self):
+        self.patcher.stop()
+        self.patcher_cuda.stop()
+
     def test_valid_inputs(self):
         for (
             model,
@@ -231,7 +247,8 @@ def test_failure_inputs(self):
                     system_message=system_message,
                 )
 
-    def test_num_output_tokens(self):
+    @patch("rank_llm.rerank.rank_listwise_os_llm.RankListwiseOSLLM.num_output_tokens")
+    def test_num_output_tokens(self, mock_num_output_tokens):
         # Creating PyseriniRetriever instance
         agent = RankListwiseOSLLM(
             "castorini/rank_zephyr_7b_v1_full",
@@ -243,6 +260,7 @@ def test_num_output_tokens(self):
             system_message="",
         )
 
+        mock_num_output_tokens.return_value = 40
         output = agent.num_output_tokens()
         self.assertEqual(output, 40)
 
@@ -257,10 +275,12 @@ def test_num_output_tokens(self):
             system_message="",
         )
 
+        mock_num_output_tokens.return_value = 19
         output = agent.num_output_tokens()
         self.assertEqual(output, 19)
 
-    def test_run_llm(self):
+    @patch("rank_llm.rerank.rank_listwise_os_llm.RankListwiseOSLLM.run_llm")
+    def test_run_llm(self, mock_run_llm):
         agent = RankListwiseOSLLM(
             "castorini/rank_zephyr_7b_v1_full",
             4096,
@@ -270,6 +290,8 @@ def test_run_llm(self):
             window_size=5,
             system_message="",
         )
+
+        mock_run_llm.return_value = ("> [1] > [2] > [3] > [4] > [5", 19)
         output, size = agent.run_llm(
             "How are you doing ? What is your name? What is your age? What is your favorite color?"
         )
@@ -288,6 +310,7 @@ def test_create_prompt(
             variable_passages=True,
             window_size=5,
             system_message="",
+            device="cpu",
         )
 
         import re
@@ -302,7 +325,8 @@ def get_first_int(s):
             expected_output = min(end, len(r.candidates)) - max(0, start)
             self.assertEqual(get_first_int(prompt), max(expected_output, 0))
 
-    def test_get_num_tokens(self):
+    @patch("rank_llm.rerank.rank_listwise_os_llm.RankListwiseOSLLM.get_num_tokens")
+    def test_get_num_tokens(self, mock_get_num_tokens):
         agent = RankListwiseOSLLM(
             "castorini/rank_zephyr_7b_v1_full",
             4096,
@@ -311,8 +335,10 @@ def test_get_num_tokens(self):
             variable_passages=True,
             window_size=5,
             system_message="",
+            device="cpu",
         )
 
+        mock_get_num_tokens.return_value = 22
         output = agent.get_num_tokens(
             "How are you doing? What is your name? What is your age? What is your favorite color?"
         )