[Move DISCO queue to core]:

- Move load_anchor_points to DISCOQueue

Author: arubique

maseval/benchmark/mmlu/mmlu.py

@@ -24,19 +24,9 @@
 """
 
 import json
-import pickle
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
-# numpy is optional - only needed for anchor points processing
-try:
-    import numpy as np
-
-    HAS_NUMPY = True
-except ImportError:
-    np = None  # type: ignore[assignment]
-    HAS_NUMPY = False
-
 from maseval import (
     AgentAdapter,
     DISCOQueue,
@@ -561,29 +551,6 @@ def get_model_adapter(self, model_id: str, **kwargs: Any) -> ModelAdapter:
 # =============================================================================
 
 
-def load_pickle(path: Union[str, Path]) -> Any:
-    """Load a pickle file."""
-    with open(path, "rb") as f:
-        return pickle.load(f)
-
-
-def load_anchor_points(path: Union[str, Path]) -> List[int]:
-    """Load anchor points from a .json or .pkl file. Returns a list of doc_ids."""
-    path = Path(path)
-    if not path.exists():
-        raise FileNotFoundError(f"Anchor points file not found: {path}")
-    if path.suffix.lower() == ".json":
-        with open(path) as f:
-            anchor_points = json.load(f)
-    else:
-        anchor_points = load_pickle(path)
-    if HAS_NUMPY and isinstance(anchor_points, np.ndarray):
-        anchor_points = anchor_points.tolist()
-    elif not HAS_NUMPY and hasattr(anchor_points, "tolist"):
-        anchor_points = anchor_points.tolist()
-    return list(anchor_points)
-
-
 def load_tasks(
     data_path: Union[str, Path],
     anchor_points_path: Optional[Union[str, Path]] = None,
@@ -601,8 +568,6 @@ def load_tasks(
     Returns:
         TaskQueue containing MMLU tasks.
 
-    Raises:
-        ImportError: If anchor_points_path is provided but numpy is not installed.
     """
     data_path = Path(data_path)
 
@@ -642,14 +607,8 @@ def load_tasks(
         )
         tasks.append(task)
 
-    # Load anchor points if provided
-    anchor_points = None
     if anchor_points_path is not None:
-        anchor_points = load_anchor_points(anchor_points_path)
-
-    # Create appropriate queue
-    if anchor_points is not None:
-        return DISCOQueue(tasks, anchor_points)
+        return DISCOQueue(tasks, anchor_points_path=anchor_points_path)
     else:
         return SequentialTaskQueue(tasks)
 

maseval/core/task.py

@@ -5,6 +5,7 @@
 from collections.abc import Sequence
 from typing import Iterable, List, Union, Iterator, Optional
 import json
+import pickle
 from pathlib import Path
 from enum import Enum
 
@@ -339,25 +340,75 @@ class DISCOQueue(InformativeSubsetQueue):
     Example:
         ```python
         queue = DISCOQueue(tasks, anchor_points=[0, 5, 12])
+        # or load from file:
+        queue = DISCOQueue(tasks, anchor_points_path="anchor_points.pkl")
 
         for task in queue:
-            result = execute(task)  # Only 3 tasks
+            result = execute(task)  # Only anchor-point tasks
         ```
     """
 
-    def __init__(self, tasks: Iterable[Task], anchor_points: Optional[List[int]] = None) -> None:
+    def __init__(
+        self,
+        tasks: Iterable[Task],
+        anchor_points: Optional[List[int]] = None,
+        anchor_points_path: Optional[Union[str, Path]] = None,
+    ) -> None:
         """Initialize DISCO task queue.
 
+        Anchor points can be supplied directly via ``anchor_points`` or loaded
+        from a file via ``anchor_points_path``.  Providing both is an error.
+
         Args:
             tasks: Full list of tasks (ordered by index).
             anchor_points: Diversity-selected indices into ``tasks``.
-                Typically loaded from a DISCO anchor-points file or
-                downloaded from a HuggingFace DISCO model repo.
-                If ``None``, evaluates all tasks in order.
+                Typically downloaded from a HuggingFace DISCO model repo.
+                If ``None`` and ``anchor_points_path`` is also ``None``,
+                evaluates all tasks in order.
+            anchor_points_path: Path to a ``.json`` or ``.pkl`` file
+                containing anchor-point indices.  Mutually exclusive with
+                ``anchor_points``.
         """
+        if anchor_points is not None and anchor_points_path is not None:
+            raise ValueError("Provide either anchor_points or anchor_points_path, not both.")
+
+        if anchor_points_path is not None:
+            anchor_points = self.load_anchor_points(anchor_points_path)
+
         self._anchor_points: Optional[List[int]] = anchor_points
         super().__init__(tasks, indices=anchor_points)
 
+    @staticmethod
+    def load_anchor_points(path: Union[str, Path]) -> List[int]:
+        """Load anchor points from a ``.json`` or ``.pkl`` file.
+
+        Args:
+            path: Path to anchor points file. JSON files should contain a
+                list of integer indices. Pickle files may contain a list or
+                a numpy array.
+
+        Returns:
+            List of integer anchor-point indices.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+        """
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"Anchor points file not found: {path}")
+
+        if path.suffix.lower() == ".json":
+            with open(path) as f:
+                anchor_points = json.load(f)
+        else:
+            with open(path, "rb") as f:
+                anchor_points = pickle.load(f)
+
+        if hasattr(anchor_points, "tolist"):
+            anchor_points = anchor_points.tolist()
+
+        return list(anchor_points)
+
 
 class PriorityTaskQueue(BaseTaskQueue):
     """Execute tasks ordered by priority.

tests/test_core/test_queue.py

@@ -20,6 +20,16 @@
 )
 
 
+class _FakeArray:
+    """Pickle-serializable array-like for testing .tolist() conversion."""
+
+    def tolist(self):
+        return [1, 2, 3]
+
+    def __iter__(self):
+        return iter([1, 2, 3])
+
+
 # ==================== Fixtures ====================
 
 
@@ -318,6 +328,86 @@ def test_len_matches_anchor_count(self):
         assert len(queue) == 3
 
 
+@pytest.mark.core
+class TestDISCOQueueLoadAnchorPoints:
+    """Tests for DISCOQueue.load_anchor_points static method."""
+
+    def test_load_from_json(self, tmp_path):
+        """Should load anchor points from a JSON file."""
+        import json
+
+        path = tmp_path / "anchors.json"
+        path.write_text(json.dumps([0, 5, 12, 99]))
+
+        result = DISCOQueue.load_anchor_points(path)
+
+        assert result == [0, 5, 12, 99]
+
+    def test_load_from_pickle(self, tmp_path):
+        """Should load anchor points from a pickle file."""
+        import pickle
+
+        path = tmp_path / "anchors.pkl"
+        with open(path, "wb") as f:
+            pickle.dump([2, 7, 15], f)
+
+        result = DISCOQueue.load_anchor_points(path)
+
+        assert result == [2, 7, 15]
+
+    def test_load_converts_tolist(self, tmp_path):
+        """Should call .tolist() on array-like objects (e.g. numpy arrays)."""
+        import pickle
+
+        path = tmp_path / "anchors.pkl"
+        with open(path, "wb") as f:
+            pickle.dump(_FakeArray(), f)
+
+        result = DISCOQueue.load_anchor_points(path)
+
+        assert result == [1, 2, 3]
+
+    def test_file_not_found(self, tmp_path):
+        """Should raise FileNotFoundError for missing files."""
+        with pytest.raises(FileNotFoundError, match="not found"):
+            DISCOQueue.load_anchor_points(tmp_path / "nonexistent.json")
+
+    def test_accepts_string_path(self, tmp_path):
+        """Should accept a string path, not just Path objects."""
+        import json
+
+        path = tmp_path / "anchors.json"
+        path.write_text(json.dumps([10, 20]))
+
+        result = DISCOQueue.load_anchor_points(str(path))
+
+        assert result == [10, 20]
+
+    def test_init_with_anchor_points_path(self, tmp_path):
+        """DISCOQueue should load anchor points from file when anchor_points_path is given."""
+        import json
+
+        tasks = [Task(query=f"Q{i}") for i in range(10)]
+        path = tmp_path / "anchors.json"
+        path.write_text(json.dumps([2, 5, 8]))
+
+        queue = DISCOQueue(tasks, anchor_points_path=path)
+
+        assert len(queue) == 3
+        assert queue._anchor_points == [2, 5, 8]
+
+    def test_init_rejects_both_anchor_args(self, tmp_path):
+        """DISCOQueue should raise ValueError when both anchor_points and anchor_points_path are given."""
+        import json
+
+        tasks = [Task(query=f"Q{i}") for i in range(5)]
+        path = tmp_path / "anchors.json"
+        path.write_text(json.dumps([0, 1]))
+
+        with pytest.raises(ValueError, match="not both"):
+            DISCOQueue(tasks, anchor_points=[0, 1], anchor_points_path=path)
+
+
 # ==================== PriorityTaskQueue Tests ====================