eval-framework 0.2.7__py3-none-any.whl

eval_framework/utils/file_ops.py

@@ -0,0 +1,245 @@
+import atexit
+import logging
+import os
+import shutil
+import signal
+import tempfile
+import time
+from pathlib import Path
+from types import FrameType
+from typing import Any
+
+import wandb
+
+logger = logging.getLogger(__name__)
+
+
+class WandbFs:
+    REGISTRY_MODEL_ROOT = "wandb-registry-model"
+    """
+    WandbFs provides an interface to interact with Weights & Biases artifacts.
+
+    WandB provides a unified API to access artifacts with artifact.download().
+
+    Several issues with the standard WandB artifact handling motivated the creation of this class:
+
+    1. Artifacts may not always be in a HuggingFace-compatible format and may have extra directories.
+    This class includes methods to find HuggingFace checkpoints in downloaded artifacts.
+
+    2. Custom download paths with clean up upon failure: Rather than downloading to a
+    WandB-managed cache directory, this class allows users to specify a custom download path
+    (which can be a persistent directory). If no path is provided, a temporary directory is
+    created and cleaned up automatically.
+
+    3. Cleanup is handled in two ways, and exit handles are registered accordingly:
+        - when not used as a context manager atexit handlers ensure cleanup on normal script termination
+        - when used as a context manager, cleanup is handled in __exit__ directly
+
+    Args:
+        user_supplied_download_path: Optional path to download artifacts to. If not given, WANDB_ARTIFACT_DIR is used.
+        This acts as a cache, so if the artifact is already present, it will not be re-downloaded.
+
+    example usage:
+    >> with WandbFs("./my-download-path/") as wandb_fs:
+    ..    artifact = wandb_fs.get_artifact("my-artifact", version="v1")
+    ..    download_path = wandb_fs.download_artifact(artifact)
+    ..    file_root = wandb_fs.find_hf_checkpoint_root_from_path_list()
+
+    >> wandb_fs = WandbFs("./my-download-path/")
+    .. wandb_fs.setup_cleanup_handlers()
+    .. artifact = wandb_fs.get_artifact("my-artifact", version="v1")
+    .. download_path = wandb_fs.download_artifact(artifact)
+    .. file_root = wandb_fs.find_hf_checkpoint_root_from_path_list()
+    .. wandb_fs.restore_cleanup_handlers()
+    .. some_function_that_uses_the_artifact(file_root)
+    """
+
+    def __init__(self, download_path: str | Path | None = None):
+        self.api = wandb.Api()
+        if download_path is None and os.getenv("WANDB_ARTIFACT_DIR") is not None:
+            download_path = os.getenv("WANDB_ARTIFACT_DIR")
+        self.user_supplied_download_path = Path(download_path) if download_path is not None else None
+        self._temp_dir: tempfile.TemporaryDirectory | None = None
+        self.download_path: Path | None = None
+        self._setup_s3_client()
+
+    def setup_cleanup_handlers(self) -> None:
+        """
+        because wandbfs deals with downloading files, we will need to
+        make sure that at exit and at failure, the directory does not persist
+
+        these are present to ensure cleanup on normal script termination if not a context manager.
+        """
+        # we only want to register the appropriate cleanup function
+        if self.user_supplied_download_path:
+            atexit.register(self._cleanup_user_dir)
+        else:
+            atexit.register(self._cleanup_temp_dir)
+
+    def restore_cleanup_handlers(self) -> None:
+        """
+        unregister the cleanup handlers
+        """
+        if self.user_supplied_download_path:
+            atexit.unregister(self._cleanup_user_dir)
+        else:
+            atexit.unregister(self._cleanup_temp_dir)
+
+    def _clean_on_signal(self, signum: int, frame: FrameType | None) -> None:
+        # we need to re-raise the signal to terminate gracefully with __exit__
+        # if we call cleanup directly, then the first time we try to rmtree
+        # we get an OSError
+        self.__exit__(None, None, None)
+        signal.signal(signum, signal.SIG_DFL)
+        os.kill(os.getpid(), signum)
+
+    def _setup_s3_client(self) -> None:
+        required_env_vars = ["AWS_ENDPOINT_URL", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"]
+        for var in required_env_vars:
+            if var not in os.environ:
+                raise ValueError(f"Missing required environment variable: {var}")
+        endpoint = os.environ["AWS_ENDPOINT_URL"]
+        if not endpoint.startswith(("http://", "https://")):
+            os.environ["AWS_ENDPOINT_URL"] = f"https://{endpoint}"
+
+    def get_artifact(self, artifact_id: str, version: str = "latest") -> wandb.Artifact:
+        name = f"{artifact_id}:{version}"
+        if "/" not in artifact_id:
+            name = f"{self.REGISTRY_MODEL_ROOT}/{name}"
+
+        # Prefer checkpoints refering to local file system (speed & immediate availability)
+        local_name = name if name.endswith("-local") else f"{name}-local"
+        if self.api.artifact_exists(local_name):
+            artifact = self.api.artifact(local_name)
+            availabilities = []
+            for f in artifact.files():
+                local_path = Path(f.path_uri.removeprefix("file://"))
+                availabilities.append(local_path.exists() and local_path.stat().st_size == f.size)
+            if all(availabilities):
+                if local_name != name:
+                    logger.info(f"Local artifact '{local_name}' available, using it instead of '{version}'.")
+                return artifact
+
+        # Otherwise fall back to the non-local version (requires download)
+        final_name = name.removesuffix("-local")
+        if final_name != name:
+            logger.info(f"Local artifact '{name}' NOT available, using '{final_name}' instead.")
+
+        if self.api.artifact_exists(local_name):
+            # Wait for a non-local artifact to pop-up, which is expected when a local one exists
+            end_time = int(os.getenv("WANDB_ARTIFACT_WAIT_TIMEOUT_SEC", "3600")) + time.time()
+            while end_time > time.time():
+                if self.api.artifact_exists(final_name):
+                    return self.api.artifact(final_name)
+                logger.info(f"Artifact '{final_name}' not found, retrying in 30 seconds...")
+                time.sleep(30)
+            raise RuntimeError(f"Timed out waiting for {final_name}, perhaps increase WANDB_ARTIFACT_WAIT_TIMEOUT_SEC.")
+        else:
+            # Return artifact or crash if not found
+            return self.api.artifact(final_name)
+
+    def download_artifact(
+        self,
+        artifact: wandb.Artifact,
+    ) -> Path:
+        """
+        download_artifact downloads the specified artifact to either a user-specified
+        directory or a temporary directory. If the user-specified directory already
+        contains the artifact, it will not be re-downloaded.
+
+        Args:
+            artifact: The WandB artifact object to download.
+        Returns:
+            Path: The path to the downloaded artifact.
+        """
+        if artifact.qualified_name.endswith("-local"):
+            self.download_path = Path(
+                os.path.commonpath([Path(f.path_uri.removeprefix("file://")) for f in artifact.files()])
+            )
+            return self.download_path if self.download_path.is_dir() else self.download_path.parent
+
+        # create the base path for either a temp or user dir
+        if self.user_supplied_download_path is None:
+            self._temp_dir = tempfile.TemporaryDirectory()
+            base_path = Path(self._temp_dir.name)
+        else:
+            base_path = self.user_supplied_download_path
+
+        artifact_subdir = "/".join(artifact.name.split(":"))
+        self.download_path = base_path / artifact_subdir
+        if self.user_supplied_download_path and self.download_path.exists():
+            return self.download_path
+
+        logger.info(f"Downloading artifact to {self.download_path}")
+        # Since the cache lives inside the docker container, it is unused in future
+        # runs. Skipping the cache also avoids file duplication and extra copying.
+        self._artifact_downloaded = False
+        skip_cache = os.getenv("WANDB_CACHE_SKIP", "true").lower() in ["true", "1", "yes"]
+        artifact_path = artifact.download(root=str(self.download_path), skip_cache=skip_cache)
+        self._artifact_downloaded = True
+
+        return Path(artifact_path)
+
+    def find_hf_checkpoint_root_from_path_list(self) -> Path | None:
+        """Find HuggingFace checkpoint root from a list of file paths.
+
+        Args:
+            file_paths: List of file paths (can be S3 URIs or local paths)
+
+        Returns:
+            str | None: Path to the HuggingFace checkpoint root folder, or None if not found
+        """
+
+        # self.download_path can be None if download_artifact was never called
+        if self.download_path and self.download_path.exists():
+            checkpoint_roots = [x for x in Path(self.download_path).glob("**/config.json")]
+            if checkpoint_roots:
+                assert len(checkpoint_roots) == 1, "Multiple checkpoints found"
+                logger.info(f"Checkpoint found in {checkpoint_roots[0].parent}.")
+                return checkpoint_roots[0].parent
+
+        return None
+
+    def __enter__(self) -> "WandbFs":
+        self._original_sigterm = signal.signal(signal.SIGTERM, self._clean_on_signal)
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """
+        exit the context manager, cleaning up the temporary directory if it was used
+        or the user directory if it was specified and failed to download.
+        """
+        if self.user_supplied_download_path:
+            self._cleanup_user_dir()
+        else:
+            self._cleanup_temp_dir()
+
+        signal.signal(signal.SIGTERM, self._original_sigterm)
+
+    def _cleanup_user_dir(self) -> None:
+        """
+        _cleanup_user_dir will remove the contents of the user specified
+        download path if there was an attempt to download the artifact and it failed.
+        """
+        if (
+            # check to make sure this flag was set; if not, there was no attempt to download.
+            not getattr(self, "_artifact_downloaded", True)
+            and self.user_supplied_download_path
+            and self.download_path
+            and self.download_path.exists()
+        ):
+            # remove the contents of the download path.
+            logger.info(f"Cleaning up user-specified download path...{self.download_path}")
+            # ignore errors because the directory is not empty
+            shutil.rmtree(self.download_path, ignore_errors=True)
+
+    def _cleanup_temp_dir(self) -> None:
+        if hasattr(self, "_temp_dir") and self._temp_dir is not None:
+            try:
+                self._temp_dir.cleanup()
+            except (OSError, FileNotFoundError):
+                # Directory might already be cleaned up or removed
+                pass
+            finally:
+                self._temp_dir = None
+                self.download_path = None

eval_framework/utils/generate_task_docs.py

@@ -0,0 +1,244 @@
+import argparse
+import inspect
+import os
+import re
+from pathlib import Path
+
+import tqdm
+
+from eval_framework.tasks.registry import get_task, registered_task_names
+from eval_framework.tasks.task_loader import load_extra_tasks
+from template_formatting.formatter import BaseFormatter, ConcatFormatter, Llama3Formatter
+
+DEFAULT_OUTPUT_DOCS_DIRECTORY = Path("docs/tasks")
+
+EXCLUDED_TASKS: list[str] = []
+
+# Base URL for the main repository to ensure links work even in external/companion repos
+REPO_URL = "https://github.com/Aleph-Alpha-Research/eval-framework/blob/main"
+
+
+def parse_args(cli_args: list[str] | None = None) -> argparse.Namespace:
+    """Parse command line arguments for the script."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--add-prompt-examples",
+        action="store_true",
+        default=False,
+        required=False,
+        help="If set, examples prompts for each of the formatters will be added in the generated docs.",
+    )
+    parser.add_argument(
+        "--exclude-tasks",
+        nargs="*",
+        type=str,
+        default=[],
+        required=False,
+        help="List of task names to exclude from documentation generation.",
+    )
+    parser.add_argument(
+        "--extra-task-modules",
+        nargs="*",
+        type=str,
+        default=[],
+        required=False,
+        help="List of files and folders containing additional task definitions.",
+    )
+    parser.add_argument(
+        "--formatter",
+        nargs="*",
+        type=str,
+        required=False,
+        default=["ConcatFormatter", "Llama3Formatter"],
+        help="Specify which formatter to use for formatting the task samples. "
+        "If not explicitly specified, default formatters will be used.",
+    )
+    parser.add_argument(
+        "--only-tasks",
+        nargs="*",
+        type=str,
+        default=[],
+        required=False,
+        help="List of task names to generate documentation for. If empty, all tasks will be processed.",
+    )
+    return parser.parse_args(args=cli_args)
+
+
+def generate_docs_for_task(
+    output_docs_directory: Path, task_name: str, formatters: list[BaseFormatter], add_prompt_examples: bool
+) -> None:
+    """Generate documentation for a specific task."""
+    task_class = get_task(task_name)
+
+    try:
+        num_fewshot = 1
+        task = task_class(num_fewshot=num_fewshot)
+    except Exception:
+        try:
+            num_fewshot = 0
+            task = task_class(num_fewshot=num_fewshot)
+        except Exception as e:
+            print(f"Failed to instantiate task {task_name}: {e}")
+            return
+
+    with open(f"{output_docs_directory}/{task_name}.md", "w") as f:
+        f.write(f"# {task_name}\n\n")
+        http_path = f"https://huggingface.co/datasets/{task.DATASET_PATH}" if task.DATASET_PATH else None
+
+        f.write("````\n")  # fence with 4 thicks because some prompts have code blocks with 3 thicks
+        f.write(f"NAME = {task_name}".strip() + "\n")
+        if hasattr(task, "DATASET_PATH"):
+            f.write(f"DATASET_PATH = {task.DATASET_PATH}".strip() + "\n")
+        if hasattr(task, "SAMPLE_SPLIT"):
+            f.write(f"SAMPLE_SPLIT = {task.SAMPLE_SPLIT}".strip() + "\n")
+        if hasattr(task, "FEWSHOT_SPLIT"):
+            f.write(f"FEWSHOT_SPLIT = {task.FEWSHOT_SPLIT}".strip() + "\n")
+        if hasattr(task, "RESPONSE_TYPE"):
+            f.write(f"RESPONSE_TYPE = {task.RESPONSE_TYPE.name}".strip() + "\n")
+        if hasattr(task, "METRICS"):
+            metrics_list = [f"{m.__name__}" for m in task.METRICS]
+            f.write(f"METRICS = [{', '.join(metrics_list)}]".strip() + "\n")
+        if hasattr(task, "SUBJECTS"):
+            f.write(f"SUBJECTS = {repr(task.SUBJECTS)}".strip() + "\n")
+        if hasattr(task, "LANGUAGE"):
+            f.write(f"LANGUAGE = {repr(task.LANGUAGE)}".strip() + "\n")
+        f.write("````\n\n")
+
+        f.write(f"- Module: `{task_class.__module__}`\n\n")
+
+        try:
+            raw_file_path = inspect.getfile(task_class)
+            # Find the package root 'eval_framework' in the path
+            match = re.search(r"eval_framework.*", raw_file_path)
+            if match:
+                # Reconstruct relative path assuming standard 'src' structure
+                task_file = f"src/{match.group(0)}"
+                # Provide a local relative link (for VS Code) and an absolute link (for GitHub/Web)
+                f.write(f"- File: [{task_file}](../../{task_file}) | [View on GitHub]({REPO_URL}/{task_file})\n\n")
+            else:
+                # Fallback for tasks defined outside the main package (e.g., custom local tasks)
+                f.write(f"- File: `{raw_file_path}`\n\n")
+        except Exception:
+            f.write("- File: `Dynamic or Built-in`\n\n")
+
+        if http_path:
+            f.write(f"- Link to dataset: [{http_path}]({http_path})\n\n")
+
+        if not add_prompt_examples:
+            f.write(
+                f"More detailed documentation, with prompt examples and ground truth completions, can be generated "
+                f"with `uv run -m eval_framework.utils.generate_task_docs --add-prompt-examples "
+                f'--only-tasks "{task_name}"`.\n'
+            )
+
+        else:
+            s = next(iter(task.iterate_samples(1)))
+            for split in task.dataset:
+                f.write(f"- `{split}` has {len(task.dataset[split])} samples\n\n")
+
+            for formatter in formatters:
+                f.write(f"## Example prompt with {formatter.__class__.__name__} ({num_fewshot}-shot)\n\n")
+                formatted_sample = formatter.format(s.messages, output_mode="string")
+                f.write("````\n")
+                f.write(f'"{formatted_sample}"')
+                f.write("\n````\n\n")
+
+            f.write("## Possible completions:\n\n")
+            f.write("````\n")
+            if s.possible_completions:
+                for item in (
+                    s.possible_completions if isinstance(s.possible_completions, list) else [s.possible_completions]
+                ):
+                    f.write(f'- "{item}"\n')
+            else:
+                f.write("None\n")
+            f.write("````\n\n")
+
+            f.write("## Ground truth:\n\n")
+            f.write("````\n")
+            if s.ground_truth:
+                for item in s.ground_truth if isinstance(s.ground_truth, list) else [s.ground_truth]:
+                    f.write(f'- "{item}"\n')
+            else:
+                f.write("None\n")
+            f.write("````\n")
+
+
+def generate_readme_list(output_docs_directory: Path, total_tasks: int) -> None:
+    """Generate a README file listing all tasks with total count."""
+
+    with open(f"{output_docs_directory}/README.md", "w") as f:
+        f.write(
+            "# Task documentation\n\n"
+            "This directory contains the generated documentation for all benchmark tasks available in the package.\n\n"
+            f"**Total number of tasks: {total_tasks}**\n\n"
+            "The documentation can be generated or updated with "
+            "`uv run -m eval_framework.utils.generate_task_docs`.\n\n"
+            "NOTE: This is an automatically generated file. Any manual modifications will not be preserved when "
+            "the file is updated.\n\n"
+        )
+
+        f.write("## List of tasks\n\n")
+        # sort files alphabetically and ignore README.md
+        for file in sorted(os.listdir(output_docs_directory)):
+            if file.endswith(".md") and file != "README.md":
+                task_name = file[:-3]
+                f.write(f"- [{task_name}]({task_name}.md)\n")
+
+
+def generate_all_docs(args: argparse.Namespace, output_docs_directory: Path) -> None:
+    # Load extra tasks if specified
+    if args.extra_task_modules:
+        print(f"Loading extra tasks from: {args.extra_task_modules}")
+        load_extra_tasks(args.extra_task_modules)
+
+    # List the tasks to process
+    filtered_tasks = []
+    for task_name in registered_task_names():
+        if args.only_tasks and task_name not in args.only_tasks:
+            continue
+        if task_name in args.exclude_tasks or task_name in EXCLUDED_TASKS:
+            continue
+        filtered_tasks.append(task_name)
+    filtered_tasks.sort()
+
+    print(f"Found {len(filtered_tasks)} tasks to process: {', '.join([task_name for task_name in filtered_tasks])}")
+
+    # List the formatters to use
+    supported_formatters = {f.__class__.__name__: f for f in [ConcatFormatter(), Llama3Formatter()]}
+    formatters = []
+    for f in args.formatter:
+        if f in supported_formatters:
+            formatters.append(supported_formatters[f])
+        else:
+            raise ValueError(f"Unsupported formatter: {f}")
+
+    # Create the output directory if it does not exist
+    os.makedirs(output_docs_directory, exist_ok=True)
+
+    for task_name in tqdm.tqdm(filtered_tasks, desc="Generating documentation for tasks"):
+        try:
+            generate_docs_for_task(
+                output_docs_directory=output_docs_directory,
+                task_name=task_name,
+                formatters=formatters,
+                add_prompt_examples=args.add_prompt_examples,
+            )
+
+        except Exception as e:
+            print("---")
+            print(f"failed generating documentation for task {task_name}: {e}")
+            file_path = f"{output_docs_directory}/{task_name}.md"
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            print("---")
+
+    # Pass the total number of processed tasks to the README generator
+    generate_readme_list(output_docs_directory=output_docs_directory, total_tasks=len(filtered_tasks))
+
+
+if __name__ == "__main__":
+    print("Generating task documentation...")
+    args = parse_args()
+    generate_all_docs(args, output_docs_directory=DEFAULT_OUTPUT_DOCS_DIRECTORY)

eval_framework/utils/helpers.py

@@ -0,0 +1,32 @@
+import numpy as np
+
+
+def count_bytes(text: str, /, *, encoding: str = "utf-8") -> int:
+    """Count the number of bytes in a string."""
+    return len(bytes(text.encode(encoding)))
+
+
+def cosine_similarity(embedding_a: list[float], embedding_b: list[float]) -> float:
+    """Computes the cosine similarity between two embeddings."""
+    return pairwise_cosine_similarity([embedding_a], [embedding_b])[0][0]
+
+
+def pairwise_cosine_similarity(embeddings_a: list[list[float]], embeddings_b: list[list[float]]) -> list[list[float]]:
+    """
+    Computes the pairwise cosine similarity matrix between two lists of embeddings.
+    Output[i][j] is the cosine similarity between embeddings_a[i] and embeddings_b[j].
+    """
+    # M, D = len(embeddings_a), len(embeddings_a[0])
+    # N, D = len(embeddings_b), len(embeddings_b[0])
+    A = np.array(embeddings_a)  # shape (M, D)
+    B = np.array(embeddings_b)  # shape (N, D)
+
+    norms_a = np.linalg.norm(A, axis=1, keepdims=True)
+    norms_b = np.linalg.norm(B, axis=1, keepdims=True)
+
+    epsilon = 1e-9
+    A_normalized = A / (norms_a + epsilon)
+    B_normalized = B / (norms_b + epsilon)
+
+    similarity_matrix = A_normalized @ B_normalized.T  # (M, D) @ (D, N) -> (M, N)
+    return similarity_matrix.tolist()

eval_framework/utils/logging.py

@@ -0,0 +1,62 @@
+import logging
+import sys
+from pathlib import Path
+
+VERBOSITY_MAP = {
+    0: logging.CRITICAL,
+    1: logging.INFO,
+    2: logging.DEBUG,
+}
+
+
+def setup_logging(
+    output_dir: Path | None = None, log_level: int = 1, log_filename: str = "evaluation.log"
+) -> logging.Logger:
+    """
+    Set up centralized logging configuration for the entire framework.
+
+    Args:
+        output_dir: Directory to save log files. If None, logs only to console.
+        log_level: Logging level (default: INFO)
+        log_filename: Name of the log file
+
+    Returns:
+        Configured root logger
+    """
+    # Map verbosity integer to logging level
+    mapped_log_level = VERBOSITY_MAP.get(log_level, logging.INFO)
+
+    # Basic configuration
+    logging.basicConfig(level=mapped_log_level)
+
+    # Create formatter
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+
+    # Get root logger and clear any existing handlers
+    root_logger = logging.getLogger()
+    root_logger.handlers.clear()
+    root_logger.setLevel(mapped_log_level)
+
+    # Console handler (always present)
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(mapped_log_level)
+    console_handler.setFormatter(formatter)
+    root_logger.addHandler(console_handler)
+
+    # File handler (if output directory provided)
+    if output_dir:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        log_file = output_dir / log_filename
+
+        file_handler = logging.FileHandler(log_file, mode="w")
+        file_handler.setLevel(mapped_log_level)
+        file_handler.setFormatter(formatter)
+        root_logger.addHandler(file_handler)
+
+        root_logger.info(f"Logging configured. File: {log_file}")
+    else:
+        root_logger.info("Logging configured (console only)")
+
+    root_logger.info(f"Output directory for logs: {output_dir if output_dir else 'None'}")
+
+    return root_logger

eval_framework/utils/packaging.py

@@ -0,0 +1,52 @@
+import importlib
+import importlib.metadata
+from collections.abc import Sequence
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+
+def validate_package_extras(extras: str | Sequence[str], /, *, package: str = "eval_framework") -> Sequence[str]:
+    """Validate that the specified extras are valid for the given package."""
+    if isinstance(extras, str):
+        extras = [extras]
+
+    metadata = importlib.metadata.metadata(package)
+    package_extras = set(metadata.get_all("Provides-Extra") or [])
+    for extra in extras:
+        if extra not in package_extras:
+            raise ValueError(f"Invalid extra: {extra}. Options are {package_extras}")
+
+    return extras
+
+
+def extra_requires(extra: str, /, *, package: str = "eval_framework") -> list[str]:
+    """Return a list of requirements for the specified extra."""
+    validate_package_extras(extra, package=package)
+    dist = importlib.metadata.distribution(package)
+    requires = dist.requires or []
+    extra_str = f"extra == '{extra}'"
+    return [r.split(";")[0].strip() for r in requires if r.endswith(extra_str)]
+
+
+def _dependency_satisfied(dep: str, /) -> bool:
+    """Return True if the dependency string is satisfied.
+
+    Args:
+        A dependency string: for example "torch~=2.0".
+    """
+    try:
+        dist = importlib.metadata.distribution(Requirement(dep).name)
+        installed_version = Version(dist.version)
+        req = Requirement(dep)
+        return installed_version in req.specifier
+    except (importlib.metadata.PackageNotFoundError, Exception):
+        return False
+
+
+def is_extra_installed(extra: str, package: str = "eval_framework") -> bool:
+    """Return `True` if all dependencies for a given extra are installed."""
+    for req in extra_requires(extra, package=package):
+        if not _dependency_satisfied(req):
+            return False
+    return True

eval_framework/utils/tqdm_handler.py

@@ -0,0 +1,14 @@
+import logging
+
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+def safe_tqdm_write(msg: str, level: int = logging.INFO) -> None:
+    if logger.isEnabledFor(level):
+        tqdm.write(msg)
+
+
+def get_disable_bar_flag() -> bool:
+    return logger.getEffectiveLevel() >= logging.WARNING