lazy import datasets

pyproject.toml

@@ -30,6 +30,7 @@ classifiers = [
 dependencies = [
     "datasets>=3.0.0",
     "jinja2>=3.1.6",
+    "numpy",
     "math-verify>=0.8.0",
     "mcp>=1.14.1",
     "nest-asyncio>=1.6.0", # for jupyter notebooks

verifiers/envs/env_group.py

@@ -1,14 +1,15 @@
+from __future__ import annotations
+
 import time
 from typing import TYPE_CHECKING, AsyncContextManager, Mapping, final
 
-from datasets import Dataset, concatenate_datasets
 from openai import AsyncOpenAI
 
 import verifiers as vf
 from verifiers.types import RolloutInput, SamplingArgs
 
 if TYPE_CHECKING:
-    pass
+    from datasets import Dataset
 
 
 class EnvGroupRubric(vf.Rubric):
@@ -142,6 +143,8 @@ def __init__(
                       If not provided, uses "env_0", "env_1", etc.
             **kwargs: Additional arguments passed to parent Environment
         """
+        from datasets import concatenate_datasets
+
         if not envs:
             raise ValueError("EnvGroup requires at least one environment")
 

verifiers/envs/environment.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import atexit
 import functools
@@ -24,8 +26,10 @@
     final,
 )
 
-from datasets import Dataset
 from openai import AsyncOpenAI, BadRequestError, OpenAI
+
+if TYPE_CHECKING:
+    from datasets import Dataset
 from openai.types.chat import ChatCompletion
 from openai.types.chat.chat_completion import Choice
 from openai.types.completion_choice import CompletionChoice
@@ -907,6 +911,8 @@ async def generate(
         """
         Generate rollouts for a set of inputs.
         """
+        from datasets import Dataset
+
         if isinstance(inputs, Dataset):
             inputs_list = inputs.to_list()
         elif isinstance(inputs, list):

verifiers/types.py

@@ -1,16 +1,21 @@
+from __future__ import annotations
+
 import sys
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     Awaitable,
     Callable,
     Literal,
+    TypeAlias,
 )
 
-from datasets import Dataset
-
 from verifiers.errors import Error
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
 if sys.version_info < (3, 12):
     from typing_extensions import TypedDict
 else:
@@ -50,7 +55,7 @@
 IndividualRewardFunc = Callable[..., float | Awaitable[float]]
 GroupRewardFunc = Callable[..., list[float] | Awaitable[list[float]]]
 RewardFunc = IndividualRewardFunc | GroupRewardFunc
-DatasetBuilder = Callable[[], Dataset]
+DatasetBuilder: TypeAlias = "Callable[[], Dataset]"
 
 
 class TrajectoryStepTokens(TypedDict):

verifiers/utils/data_utils.py

@@ -1,12 +1,15 @@
 # NOTE: Helper functions for example datasets. Not intended for core functionality.
 
-import random
-from typing import Any, Callable, cast
+from __future__ import annotations
 
-from datasets import Dataset, concatenate_datasets, load_dataset
+import random
+from typing import TYPE_CHECKING, Any, Callable, cast
 
 from verifiers.types import ChatMessage
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
 ### PROMPTS ###
 
 THINK_BOXED_SYSTEM_PROMPT = "Think step-by-step inside <think>...</think> tags. \
@@ -259,6 +262,8 @@ def preprocess_prime_code(x: dict[str, Any]) -> dict[str, Any]:
 def load_example_dataset(
     name: str = "gsm8k", split: str | None = None, n: int | None = None, seed: int = 0
 ) -> Dataset:
+    from datasets import Dataset, concatenate_datasets, load_dataset
+
     if name == "aime2024":
         if split is None:
             split = "train"

verifiers/utils/display_utils.py

@@ -16,7 +16,6 @@
 from collections import deque
 from typing import Any
 
-from datasets import disable_progress_bar, enable_progress_bar
 from rich.console import Console
 from rich.live import Live
 from rich.panel import Panel
@@ -126,6 +125,8 @@ def _make_log_panel(self) -> Panel:
     def start(self) -> None:
         """Start the live display."""
         # Suppress datasets progress bars (e.g. from .map())
+        from datasets import disable_progress_bar
+
         disable_progress_bar()
 
         # Suppress console output from existing handlers but capture logs for display
@@ -171,6 +172,8 @@ def stop(self) -> None:
             self._live = None
 
         # Restore datasets progress bar
+        from datasets import enable_progress_bar
+
         enable_progress_bar()
 
         # Remove our log handler and restore original handler levels

verifiers/utils/eval_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import importlib.util
 import json
@@ -6,18 +8,19 @@
 from collections import Counter, defaultdict
 from contextlib import contextmanager
 from pathlib import Path
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 try:
     import tomllib  # type: ignore[import-not-found]
 except ImportError:
     import tomli as tomllib  # type: ignore[import-not-found]
 
 import numpy as np
-from datasets import Dataset, disable_progress_bar, enable_progress_bar
-from datasets.utils import logging as ds_logging
 
 import verifiers as vf
+
+if TYPE_CHECKING:
+    from datasets import Dataset
 from verifiers.types import (
     Endpoints,
     EvalConfig,
@@ -536,6 +539,8 @@ def get_hf_hub_dataset_name(results: GenerateOutputs) -> str:
 
 
 def make_dataset(results: GenerateOutputs, **kwargs) -> Dataset:
+    from datasets import Dataset
+
     clean_prompts = [messages_to_printable(p) for p in results["prompt"]]
     clean_prompts = [sanitize_tool_calls(p) for p in clean_prompts]
     clean_completions = [messages_to_printable(c) for c in results["completion"]]
@@ -578,6 +583,9 @@ def make_dataset(results: GenerateOutputs, **kwargs) -> Dataset:
 
 @contextmanager
 def quiet_datasets():
+    from datasets import disable_progress_bar, enable_progress_bar
+    from datasets.utils import logging as ds_logging
+
     prev_level = ds_logging.get_verbosity()
     ds_logging.set_verbosity(ds_logging.WARNING)
     disable_progress_bar()