fix: handle all test_range types in SizePartitioner.get_cost

opencompass/partitioners/size.py

@@ -157,7 +157,7 @@ def split_dataset(self, dataset_cfg: ConfigDict) -> List[ConfigDict]:
             cfg = copy.deepcopy(dataset_cfg)
             cfg['abbr'] = abbr + f'_{part}'
             test_range = cfg['reader_cfg'].get('test_range', '')
-            cfg['reader_cfg']['test_range'] = f'{test_range}[{i}:{i+step}]'
+            cfg['reader_cfg']['test_range'] = f'{test_range}[{i}: {i + step}]'
             split_configs.append(cfg)
         return split_configs
 
@@ -185,6 +185,39 @@ def get_factor(self, dataset: ConfigDict) -> int:
 
         return factor
 
+    def _get_actual_size(self, test_range: Optional[Union[int, float, str]],
+                         total_size: int) -> int:
+        """Compute the actual number of test samples after applying test_range.
+
+        This mirrors the slicing logic in
+        :func:`opencompass.openicl.icl_dataset_reader.load_partial_dataset`
+        so that :meth:`get_cost` returns an accurate estimate without
+        loading the dataset into memory.
+
+        Args:
+            test_range: The ``test_range`` value from the dataset reader
+                config.  May be ``None``, an empty string, an ``int``, a
+                ``float``, or a slice string such as ``"[:100]"``.
+            total_size (int): The full size of the dataset.
+
+        Returns:
+            int: The number of samples that will actually be evaluated.
+        """
+        if not test_range:
+            return total_size
+
+        if isinstance(test_range, (int, float)):
+            if test_range <= 0 or test_range >= total_size:
+                return total_size
+            if isinstance(test_range, float) and test_range < 1:
+                return int(test_range * total_size)
+            return int(test_range)
+
+        if isinstance(test_range, str):
+            return len(eval(f'range(total_size){test_range}'))
+
+        return total_size
+
     def get_cost(self,
                  dataset: ConfigDict,
                  get_raw_factors: bool = False) -> Union[int, Tuple[int, int]]:
@@ -205,8 +238,8 @@ def get_cost(self,
         factor = self.get_factor(dataset)
 
         if dataset_abbr in self.dataset_size:
-            actual_size = eval('len(range(self.dataset_size[dataset_abbr])'
-                               f'{test_range})')
+            actual_size = self._get_actual_size(
+                test_range, self.dataset_size[dataset_abbr])
             if get_raw_factors:
                 return actual_size, factor
             return factor * actual_size
@@ -220,8 +253,8 @@ def get_cost(self,
                       indent=4,
                       ensure_ascii=False)
 
-        actual_size = eval('len(range(self.dataset_size[dataset_abbr])'
-                           f'{test_range})')
+        actual_size = self._get_actual_size(test_range,
+                                            self.dataset_size[dataset_abbr])
         if get_raw_factors:
             return actual_size, factor
         return factor * actual_size

tests/partitioners/test_size.py

@@ -0,0 +1,82 @@
+"""Unit tests for SizePartitioner."""
+
+import tempfile
+import unittest
+
+from mmengine.config import ConfigDict
+
+from opencompass.partitioners.size import SizePartitioner
+
+
+class TestSizePartitioner(unittest.TestCase):
+    """Test cases for SizePartitioner."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.partitioner = SizePartitioner(
+            out_dir=self.temp_dir.name,
+            dataset_size_path=f'{self.temp_dir.name}/dataset_size.json')
+        self.partitioner._dataset_size = {'demo': 1000}
+
+    def tearDown(self):
+        """Clean up test fixtures."""
+        self.temp_dir.cleanup()
+
+    def _make_dataset_cfg(self, test_range):
+        return ConfigDict({
+            'abbr': 'demo',
+            'reader_cfg': {
+                'test_range': test_range
+            },
+            'infer_cfg': {
+                'prompt_template': {
+                    'template': '{question}'
+                }
+            }
+        })
+
+    def test_get_actual_size_handles_supported_test_range_types(self):
+        """Test actual size calculation for supported test_range types."""
+        cases = [
+            (None, 1000),
+            ('', 1000),
+            (0, 1000),
+            (-1, 1000),
+            (1000, 1000),
+            (100, 100),
+            (0.5, 500),
+            ('[:100]', 100),
+            ('[100:200]', 100),
+            ('[::2]', 500),
+            ('[:100][10:20]', 10),
+        ]
+
+        for test_range, expected_size in cases:
+            with self.subTest(test_range=test_range):
+                self.assertEqual(
+                    self.partitioner._get_actual_size(test_range, 1000),
+                    expected_size)
+
+    def test_get_cost_handles_non_string_test_range(self):
+        """Test get_cost does not fail on non-string test_range values."""
+        cases = [
+            (None, (1000, 20), 20000),
+            (100, (100, 20), 2000),
+            (0.5, (500, 20), 10000),
+            (0, (1000, 20), 20000),
+        ]
+
+        for test_range, raw_factors, expected_cost in cases:
+            dataset_cfg = self._make_dataset_cfg(test_range)
+            with self.subTest(test_range=test_range):
+                self.assertEqual(
+                    self.partitioner.get_cost(dataset_cfg,
+                                              get_raw_factors=True),
+                    raw_factors)
+                self.assertEqual(self.partitioner.get_cost(dataset_cfg),
+                                 expected_cost)
+
+
+if __name__ == '__main__':
+    unittest.main()