Rework the imports to be more versatile

Style

Author: LysandreJik

docs/source/using-the-python-api.mdx

@@ -12,9 +12,9 @@ import lighteval
 from lighteval.logging.evaluation_tracker import EvaluationTracker
 from lighteval.models.vllm.vllm_model import VLLMModelConfig
 from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
-from lighteval.utils.imports import is_accelerate_available
+from lighteval.utils.imports import is_package_available
 
-if is_accelerate_available():
+if is_package_available("accelerate"):
     from datetime import timedelta
     from accelerate import Accelerator, InitProcessGroupKwargs
     accelerator = Accelerator(kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(seconds=3000))])

src/lighteval/logging/evaluation_tracker.py

@@ -43,13 +43,13 @@
     TaskConfigLogger,
     VersionsLogger,
 )
-from lighteval.utils.imports import NO_TENSORBOARDX_WARN_MSG, is_nanotron_available, is_tensorboardX_available
+from lighteval.utils.imports import is_package_available, not_installed_error_message
 from lighteval.utils.utils import obj_to_markdown
 
 
 logger = logging.getLogger(__name__)
 
-if is_nanotron_available():
+if is_package_available("nanotron"):
     from nanotron.config import GeneralArgs  # type: ignore
 
 try:
@@ -659,11 +659,11 @@ def recreate_metadata_card(self, repo_id: str) -> None:  # noqa: C901
     def push_to_tensorboard(  # noqa: C901
         self, results: dict[str, dict[str, float]], details: dict[str, DetailsLogger.CompiledDetail]
     ):
-        if not is_tensorboardX_available:
-            logger.warning(NO_TENSORBOARDX_WARN_MSG)
+        if not is_package_available("tensorboardX"):
+            logger.warning(not_installed_error_message("tensorboardX"))
             return
 
-        if not is_nanotron_available():
+        if not is_package_available("nanotron"):
             logger.warning("You cannot push results to tensorboard without having nanotron installed. Skipping")
             return
 

src/lighteval/logging/info_loggers.py

@@ -34,13 +34,13 @@
 from lighteval.models.model_output import ModelResponse
 from lighteval.tasks.lighteval_task import LightevalTask, LightevalTaskConfig
 from lighteval.tasks.requests import Doc
-from lighteval.utils.imports import is_nanotron_available
+from lighteval.utils.imports import is_package_available
 
 
 logger = logging.getLogger(__name__)
 
 
-if is_nanotron_available():
+if is_package_available("nanotron"):
     pass
 
 

src/lighteval/main_nanotron.py

@@ -32,11 +32,13 @@
     reasoning_tags,
     remove_reasoning_tags,
 )
+from lighteval.utils.imports import requires
 
 
 SEED = 1234
 
 
+@requires("nanotron")
 def nanotron(
     checkpoint_config_path: Annotated[
         str, Option(help="Path to the nanotron checkpoint YAML or python config file, potentially on s3.")
@@ -45,12 +47,9 @@ def nanotron(
     remove_reasoning_tags: remove_reasoning_tags.type = remove_reasoning_tags.default,
     reasoning_tags: reasoning_tags.type = reasoning_tags.default,
 ):
-    """Evaluate models using nanotron as backend."""
-    from lighteval.utils.imports import NO_NANOTRON_ERROR_MSG, is_nanotron_available
-
-    if not is_nanotron_available():
-        raise ImportError(NO_NANOTRON_ERROR_MSG)
-
+    """
+    Evaluate models using nanotron as backend.
+    """
     from nanotron.config import GeneralArgs, ModelArgs, TokenizerArgs, get_config_from_dict, get_config_from_file
 
     from lighteval.logging.evaluation_tracker import EvaluationTracker

src/lighteval/metrics/imports/data_stats_metric.py

@@ -30,7 +30,7 @@
 from typing import Literal
 
 from lighteval.metrics.imports.data_stats_utils import Fragments
-from lighteval.utils.imports import NO_SPACY_ERROR_MSG, is_spacy_available
+from lighteval.utils.imports import raise_if_package_not_available
 
 
 logger = logging.getLogger(__name__)
@@ -86,8 +86,7 @@ def __init__(
                 determines the spaCy model used for tokenization. Currently supports English,
                 German, French, and Italian.
         """
-        if not is_spacy_available():
-            raise ImportError(NO_SPACY_ERROR_MSG)
+        raise_if_package_not_available("spacy")
         import spacy
 
         self.n_gram = n_gram

src/lighteval/metrics/utils/extractive_match_utils.py

@@ -34,12 +34,12 @@
 from lighteval.tasks.requests import Doc
 from lighteval.tasks.templates.utils.formulation import ChoicePrefix, get_prefix
 from lighteval.tasks.templates.utils.translation_literals import TRANSLATION_LITERALS
-from lighteval.utils.imports import requires_latex2sympy2_extended
+from lighteval.utils.imports import requires
 from lighteval.utils.language import Language
 from lighteval.utils.timeout import timeout
 
 
-@requires_latex2sympy2_extended
+@requires("latex2sympy2_extended")
 def latex_normalization_config_default_factory():
     from latex2sympy2_extended.latex2sympy2 import NormalizationConfig
 
@@ -373,7 +373,7 @@ def get_target_type_order(target_type: ExtractionTarget) -> int:
 
 # Small cache, to catche repeated calls invalid parsing
 @lru_cache(maxsize=20)
-@requires_latex2sympy2_extended
+@requires("latex2sympy2_extended")
 def parse_latex_with_timeout(latex: str, timeout_seconds: int):
     from latex2sympy2_extended.latex2sympy2 import latex2sympy
 
@@ -428,7 +428,7 @@ def convert_to_pct(number: Number):
     return sympy.Mul(number, sympy.Rational(1, 100), evaluate=False)
 
 
-@requires_latex2sympy2_extended
+@requires("latex2sympy2_extended")
 @lru_cache(maxsize=20)
 def extract_latex(
     match: re.Match, latex_config: LatexExtractionConfig, timeout_seconds: int

src/lighteval/metrics/utils/linguistic_tokenizers.py

@@ -18,10 +18,8 @@
 from typing import Callable, Iterator
 
 from lighteval.utils.imports import (
-    NO_SPACY_TOKENIZER_ERROR_MSG,
-    NO_STANZA_TOKENIZER_ERROR_MSG,
-    can_load_spacy_tokenizer,
-    can_load_stanza_tokenizer,
+    Extras,
+    raise_if_package_not_available,
 )
 from lighteval.utils.language import Language
 
@@ -102,8 +100,8 @@ def span_tokenize(self, text: str) -> list[tuple[int, int]]:
 class SpaCyTokenizer(WordTokenizer):
     def __init__(self, spacy_language: str, config=None):
         super().__init__()
-        if not can_load_spacy_tokenizer(spacy_language):
-            raise ImportError(NO_SPACY_TOKENIZER_ERROR_MSG)
+        raise_if_package_not_available(Extras.MULTILINGUAL, language=spacy_language)
+
         self.spacy_language = spacy_language
         self.config = config
         self._tokenizer = None
@@ -140,8 +138,7 @@ def span_tokenize(self, text: str) -> list[tuple[int, int]]:
 class StanzaTokenizer(WordTokenizer):
     def __init__(self, stanza_language: str, **stanza_kwargs):
         super().__init__()
-        if not can_load_stanza_tokenizer():
-            raise ImportError(NO_STANZA_TOKENIZER_ERROR_MSG)
+        raise_if_package_not_available("stanza")
         self.stanza_language = stanza_language
         self.stanza_kwargs = stanza_kwargs
         self._tokenizer = None

src/lighteval/metrics/utils/llm_as_judge.py

@@ -34,7 +34,7 @@
 from tqdm import tqdm
 from tqdm.asyncio import tqdm_asyncio
 
-from lighteval.utils.imports import is_litellm_available, is_openai_available, is_vllm_available
+from lighteval.utils.imports import raise_if_package_not_available
 from lighteval.utils.utils import as_list
 
 
@@ -151,8 +151,7 @@ def __lazy_load_client(self):  # noqa: C901
             # Both "openai" and "tgi" backends use the OpenAI-compatible API
             # They are handled separately to allow for backend-specific validation and setup
             case "openai" | "tgi":
-                if not is_openai_available():
-                    raise RuntimeError("OpenAI backend is not available.")
+                raise_if_package_not_available("openai")
                 if self.client is None:
                     from openai import OpenAI
 
@@ -162,13 +161,11 @@ def __lazy_load_client(self):  # noqa: C901
                 return self.__call_api_parallel
 
             case "litellm":
-                if not is_litellm_available():
-                    raise RuntimeError("litellm is not available.")
+                raise_if_package_not_available("litellm")
                 return self.__call_litellm
 
             case "vllm":
-                if not is_vllm_available():
-                    raise RuntimeError("vllm is not available.")
+                raise_if_package_not_available("vllm")
                 if self.pipe is None:
                     from vllm import LLM, SamplingParams
                     from vllm.transformers_utils.tokenizer import get_tokenizer

src/lighteval/metrics/utils/math_comparison.py

@@ -51,7 +51,7 @@
 from sympy.core.function import UndefinedFunction
 from sympy.core.relational import Relational
 
-from lighteval.utils.imports import requires_latex2sympy2_extended
+from lighteval.utils.imports import requires
 from lighteval.utils.timeout import timeout
 
 
@@ -308,7 +308,7 @@ def is_equation(expr: Basic | MatrixBase) -> bool:
     return False
 
 
-@requires_latex2sympy2_extended
+@requires("latex2sympy2_extended")
 def is_assignment_relation(expr: Basic | MatrixBase) -> bool:
     from latex2sympy2_extended.latex2sympy2 import is_expr_of_only_symbols
 

src/lighteval/models/endpoints/litellm_model.py

@@ -32,12 +32,12 @@
 from lighteval.tasks.prompt_manager import PromptManager
 from lighteval.tasks.requests import Doc, SamplingMethod
 from lighteval.utils.cache_management import SampleCache, cached
-from lighteval.utils.imports import is_litellm_available
+from lighteval.utils.imports import is_package_available, requires
 
 
 logger = logging.getLogger(__name__)
 
-if is_litellm_available():
+if is_package_available("litellm"):
     import litellm
     from litellm import encode
     from litellm.caching.caching import Cache
@@ -110,6 +110,7 @@ class LiteLLMModelConfig(ModelConfig):
     concurrent_requests: int = 10
 
 
+@requires("litellm")
 class LiteLLMClient(LightevalModel):
     _DEFAULT_MAX_LENGTH: int = 4096