eval-framework 0.2.7__py3-none-any.whl

eval_framework/result_processors/result_processor.py

@@ -0,0 +1,129 @@
+import hashlib
+import importlib.metadata
+import json
+import logging
+import os
+from datetime import datetime
+from pathlib import Path
+
+import jsonlines
+from pydantic import RootModel
+
+from eval_framework.result_processors.base import Result, ResultProcessor
+from eval_framework.shared.types import Completion, Loglikelihood
+from eval_framework.tasks.eval_config import EvalConfig
+
+MAIN = "eval_framework_results"
+logger = logging.getLogger(__name__)
+
+
+class ResultsFileProcessor(ResultProcessor):
+    def __init__(self, output_dir: Path) -> None:
+        self.output_dir = output_dir
+        os.makedirs(self.output_dir, exist_ok=True)
+
+    def save_metadata(self, metadata: dict) -> None:
+        with open(self.output_dir / "metadata.json", "w") as f:
+            json.dump(metadata, f, indent=2, sort_keys=True)
+
+    def load_metadata(self) -> dict:
+        metadata_file = self.output_dir / "metadata.json"
+        if os.path.exists(metadata_file):
+            with open(metadata_file) as f:
+                return json.load(f)
+        else:
+            logger.info("No metadata found.")
+            return {}
+
+    def save_responses(self, responses: list[Completion | Loglikelihood]) -> None:
+        responses_data = [response.model_dump(mode="json", serialize_as_any=True) for response in responses]
+        with jsonlines.open(self.output_dir / "output.jsonl", "w") as f:
+            f.write_all(responses_data)
+
+    def save_response(self, response: Completion | Loglikelihood) -> None:
+        with jsonlines.open(self.output_dir / "output.jsonl", "a") as f:
+            f.write(response.model_dump(mode="json", serialize_as_any=True))
+
+    def load_responses(self) -> list[Completion | Loglikelihood]:
+        output_file = self.output_dir / "output.jsonl"
+        broken_file = output_file.with_suffix(f".jsonl.broken.{datetime.now().strftime('%Y%m%d_%H%M%S')}")
+
+        try:
+            Item = RootModel[Loglikelihood | Completion]
+            with jsonlines.open(output_file, "r") as f:
+                responses = [Item.model_validate(x).root for x in f]
+        except FileNotFoundError:
+            logger.info("No saved completions found.")
+            responses = []
+        except (json.decoder.JSONDecodeError, jsonlines.jsonlines.InvalidLineError):
+            logger.info(f"Error decoding JSON, the file is corrupted. It will be renamed to {broken_file} and ignored")
+            output_file.rename(broken_file)
+            responses = []
+
+        ids_list = [(resp.id, resp.subject) for resp in responses]
+        if len(ids_list) != len(set(ids_list)) and "mtbench" not in str(output_file):
+            logger.info(
+                f"Error: {len(ids_list) - len(set(ids_list))} duplicate response IDs found, the file is corrupted. "
+                f"It will be renamed to {broken_file} and ignored"
+            )
+            output_file.rename(broken_file)
+            responses = []
+
+        return responses
+
+    def save_metrics_results(self, results: list[Result]) -> None:
+        result_data = [x.model_dump(mode="json") for x in results]
+        with jsonlines.open(self.output_dir / "results.jsonl", "w") as f:
+            f.write_all(result_data)
+
+    def save_metrics_result(self, result: Result) -> None:
+        with jsonlines.open(self.output_dir / "results.jsonl", "a") as f:
+            f.write(result.model_dump(mode="json"))
+
+    def save_aggregated_results(self, results: dict[str, float | None]) -> None:
+        with open(self.output_dir / "aggregated_results.json", "w") as f:
+            json.dump(results, f, indent=4, sort_keys=True)
+
+    def load_metrics_results(self) -> list[Result]:
+        results_file = self.output_dir / "results.jsonl"
+        try:
+            with jsonlines.open(results_file, "r") as f:
+                result_data = [Result.model_validate(x) for x in f]
+            return result_data
+        except FileNotFoundError:
+            logger.info("No saved metrics found.")
+            return []
+
+
+def generate_output_dir(llm_name: str, config: EvalConfig) -> Path:
+    # get the package version
+    version_str = f"v{importlib.metadata.version('eval_framework')}"
+
+    # Handle None values for num_fewshot and num_samples
+    fewshot_str = f"fewshot_{config.num_fewshot}" if config.num_fewshot is not None else "fewshot_None"
+    samples_str = f"samples_{config.num_samples}" if config.num_samples is not None else "samples_None"
+    tokens_str = f"tokens_{config.max_tokens}" if config.max_tokens is not None else ""
+
+    # Serialize key parameters for inclusion in the name
+    params_str = f"{fewshot_str}__{samples_str}"
+    if tokens_str:
+        params_str += f"__{tokens_str}"
+
+    # Serialize the full config for hashing
+    # Convert the config to a dict and sort keys to ensure consistent hashing
+    config_json = config.model_json_robust_subset_dump()
+    config_hash = hashlib.sha256(config_json.encode("utf-8")).hexdigest()[:5]  # Short hash of 5 characters
+
+    # Include the hash in the directory name
+    dir_name = f"{params_str}_{config_hash}"
+
+    # add timestamp to dir in debug mode
+    if os.environ.get("DEBUG", "FALSE").lower() == "true":
+        # Generate the timestamp
+        timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
+        dir_name += f"_{timestamp}"
+
+    # Combine all components to form the full output directory path
+    output_dir = config.output_dir / llm_name / f"{version_str}_{config.task_name}" / dir_name
+
+    return output_dir

eval_framework/result_processors/wandb_uploader.py

@@ -0,0 +1,137 @@
+"""
+Module for writing result folder to a W&B artifact
+"""
+
+import hashlib
+import logging
+import subprocess
+from collections.abc import Callable
+from pathlib import Path
+
+import wandb
+from wandb.sdk.artifacts._validators import NAME_MAXLEN
+
+from eval_framework.result_processors.base import ResultsUploader
+from eval_framework.tasks.eval_config import EvalConfig
+
+logger = logging.getLogger(__name__)
+
+ArtifactUploadFunction = Callable[[str, str, list[Path]], str | None]  # returns reference path for W&B or None
+_ARTIFACT_UPLOAD_FUNCTION: ArtifactUploadFunction | None = None
+
+
+def register_artifact_upload_function(func: ArtifactUploadFunction | None) -> None:
+    global _ARTIFACT_UPLOAD_FUNCTION
+    _ARTIFACT_UPLOAD_FUNCTION = func
+
+
+def artifact_upload_function(artifact_name: str, subpath: str, file_paths: list[Path]) -> str | None:
+    if _ARTIFACT_UPLOAD_FUNCTION is None:
+        return None
+    logger.info(f"Uploading '{artifact_name}'.")
+    reference_path = _ARTIFACT_UPLOAD_FUNCTION(artifact_name, subpath, file_paths)
+    if reference_path is None:
+        logger.warning("Failed uploading, the custom upload function returned empty destination path!")
+    else:
+        logger.info(f"Successfully uploaded to {reference_path}.")
+    return reference_path
+
+
+class WandbUploader(ResultsUploader):
+    def __init__(
+        self,
+        config: EvalConfig,
+        include_all: bool = True,
+        compress_non_json: bool = True,
+        wandb_registry: str | None = None,
+    ) -> None:
+        if not config.wandb_upload_results:
+            logger.warning("Results will not be persisted in WandB (`wandb_upload_results` not set).")
+            return
+        if config.output_dir is None:
+            raise ValueError("Output directory is not set in the configuration.")
+        if wandb.run is None or wandb.run.settings._noop:
+            logger.warning("Results will not be persisted in WandB (no WandB run active / `wandb_project` not set).")
+            return
+
+        self._include_all = include_all
+        self._compress_non_json = compress_non_json
+        self._wandb_registry = wandb_registry
+
+    def upload(self, llm_name: str, config: EvalConfig, output_dir: Path) -> bool:
+        if hasattr(self, "_wandb_registry") is False:
+            return False  # not initialized
+
+        try:
+            if self._include_all and self._compress_non_json:
+                # note: individual gz files can be easily read by `less` or `grepz`, unlike a tar of multiple files
+                subprocess.run(
+                    ["find", output_dir, "-type", "f", "!", "-name", "*.json", "-exec", "gzip", "-k", "{}", ";"],
+                    check=True,
+                )
+                file_paths = list(output_dir.glob("*.json")) + list(output_dir.glob("*.gz"))
+            elif self._include_all:
+                file_paths = list(output_dir.glob("*"))
+            else:
+                file_paths = list(output_dir.glob("*.json"))
+
+            artifact_name = self._get_artifact_name(llm_name, config)
+            alias_name = self._get_alias(output_dir)
+
+            try:
+                rel_upload_dir = str(output_dir.relative_to(config.output_dir))
+                reference_path = artifact_upload_function(artifact_name, rel_upload_dir, file_paths)
+            except Exception as e:
+                logger.error(f"Problem during artifact upload function, aborting registration: {e}.")
+                return False
+
+            artifact = wandb.Artifact(name=artifact_name, type="eval")  # note: metadata is added from run automatically
+            if reference_path:
+                artifact.add_reference(reference_path, checksum=True)
+            else:
+                logger.info("Uploading results directly to WandB.")
+                for fp in file_paths:
+                    artifact.add_file(str(fp))
+
+            # Because metadata and e.g. logs are added to the artifact, we get a new version everytime!
+            # To mitigate this, we also add an alias based on the content hash of the "pure" result files.
+            wandb.log_artifact(artifact, aliases=[alias_name])
+            if self._wandb_registry:
+                artifact.link(f"wandb-registry-{self._wandb_registry}/{artifact_name}")
+            logger.info(f"Successfully registered '{artifact_name}' in WandB")
+
+        finally:
+            for fp in file_paths:
+                if fp.suffix == ".gz" and self._compress_non_json:
+                    fp.unlink(missing_ok=True)
+
+        return True
+
+    def _get_artifact_name(self, llm_name: str, config: EvalConfig) -> str:
+        llm_name = llm_name.replace("/", "_")  # assuming this has class name and checkpoint name in it
+
+        # As in generate_output_dir() for consistency, though shorter.
+        # But we don't include the eval_framework version and timestamp here (-> handled via W&B versioning)
+        fewshot_str = f"fs{config.num_fewshot}" if config.num_fewshot is not None else ""
+        samples_str = f"s{config.num_samples}" if config.num_samples is not None else ""
+        tokens_str = f"t{config.max_tokens}" if config.max_tokens is not None else ""
+        params_str = f"{fewshot_str}{samples_str}{tokens_str}"
+
+        config_json = config.model_json_robust_subset_dump()
+        config_hash = hashlib.sha256(config_json.encode("utf-8")).hexdigest()[:5]
+
+        # Respect W&B artifact name length limit
+        eval_name = f"__{config.task_name}__{params_str}_{config_hash}"
+        max_llm_name_len = NAME_MAXLEN - len(eval_name)
+        return llm_name[:max_llm_name_len] + eval_name
+
+    def _get_alias(self, output_dir: Path) -> str:
+        digests = []
+        # These files don't contain result-irrelevant things such as timestamps or paths and are fields are ordered.
+        # This makes them good for generating a hash that identifies the actual results and not "random" metadata.
+        for filename in ["aggregated_results.json", "output.jsonl", "results.jsonl"]:
+            if (output_dir / filename).exists():
+                with open(output_dir / filename, "rb") as f:
+                    digests.append(hashlib.file_digest(f, "sha256").hexdigest())
+        hash = hashlib.sha256("".join(digests).encode("utf-8")).hexdigest()
+        return f"H-{hash[:10]}"

eval_framework/run.py

@@ -0,0 +1,369 @@
+import argparse
+import datetime
+import logging
+from pathlib import Path
+from typing import Any
+
+try:
+    from eval_framework.context.determined import DeterminedContext
+except ImportError:
+    DeterminedContext = None  # type: ignore
+
+
+from eval_framework.context.local import LocalContext
+from eval_framework.main import main
+from eval_framework.tasks.task_loader import load_extra_tasks
+from eval_framework.utils.logging import setup_logging
+
+logger = logging.getLogger(__name__)
+
+CONTEXT = {
+    "local": LocalContext,
+    "determined": DeterminedContext,
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--context",
+        type=str,
+        required=False,
+        default="local",
+        choices=["local", "determined"],
+        help="The context in which the evaluation is run.",
+    )
+    parser.add_argument(
+        "--models",
+        type=Path,
+        required=False,
+        default=Path(__file__).parent / "llm" / "models.py",
+        help="The path to the Python module file containing model classes.",
+    )
+    parser.add_argument(
+        "--extra-task-modules",
+        nargs="*",
+        default=[],
+        required=False,
+        help="List of files and folders containing additional task definitions.",
+    )
+    parser.add_argument(
+        "--llm-name",
+        type=str,
+        required=False,
+        help=(
+            "Either a full import path for a model (e.g., `eval_framework.llm.huggingface.HFLLM`) or the "
+            "name of a class derived from `eval_framework.llm.base.BaseLLM` that can be found in the "
+            "models file. The resulting model is instantiated with the arguments provided via `--llm-args`."
+        ),
+    )
+    parser.add_argument(
+        "--llm-args",
+        type=str,
+        nargs="+",
+        default=(),
+        required=False,
+        help="The arguments to pass to the LLM as key=value pairs.",
+    )
+    parser.add_argument(
+        "--num-samples", type=int, required=False, help="The number of samples per subject to evaluate."
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        required=False,
+        help="The maximum number of tokens to generate for each sample. Overwrites any task default value.",
+    )
+    parser.add_argument(
+        "--num-fewshot", type=int, required=False, default=0, help="The number of fewshot examples to use."
+    )
+    parser.add_argument("--task-name", type=str, required=False, help="The name of the task to evaluate.")
+    parser.add_argument(
+        "--randomize-judge-order",
+        action="store_true",
+        help="Randomize the order of answers presented to the LLM judge to mitigate position bias.",
+    )
+
+    # Perturbation arguments
+    parser.add_argument(
+        "--perturbation-type",
+        type=str,
+        required=False,
+        choices=[
+            "editor",
+            "permute",
+            "replace",
+            "delete",
+            "uppercase",
+        ],
+        help=(
+            "The type of perturbation to apply. Note that this may not make sense for some prompts, for example those "
+            "containing math and code."
+        ),
+    )
+    parser.add_argument(
+        "--perturbation-probability",
+        type=float,
+        required=False,
+        default=None,
+        help="The probability of applying a perturbation to each word or character (between 0.0 and 1.0).",
+    )
+    parser.add_argument(
+        "--perturbation-seed",
+        type=int,
+        required=False,
+        default=42,
+        help="Random seed controlling perturbations.",
+    )
+    parser.add_argument(
+        "--task-subjects",
+        type=str,
+        nargs="+",
+        required=False,
+        help=(
+            "The subjects of the task to evaluate. If empty, all subjects are evaluated. Subjects in the form of "
+            "tuples can be specified in a comma-delimited way, possibly using wildcard * in some dimensions of a "
+            "tuple, e.g., 'DE_DE, *' or 'FR_FR, astronomy'."
+        ),
+    )
+    parser.add_argument(
+        "--hf-revision",
+        type=str,
+        required=False,
+        default=None,
+        help="A tag name, a branch name, or commit hash for the task HF dataset.",
+    )
+    parser.add_argument(
+        "--judge-models",
+        type=Path,
+        required=False,
+        help="The path to the Python module file containing LLM judge model classes.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default="outputs",
+        required=False,
+        help="The path for the evaluation outputs.",
+    )
+    parser.add_argument(
+        "--hf-upload-repo",
+        type=str,
+        default="Aleph-Alpha/evaluation-results",
+        required=False,
+        help="Customizable path for the HuggingFace git repository where runs will be stored.",
+    )
+    parser.add_argument(
+        "--hf-upload-dir",
+        type=str,
+        default="",
+        required=False,
+        help="Folder name for the HuggingFace git repository where runs will be stored.",
+    )
+    parser.add_argument(
+        "--wandb-project",
+        type=str,
+        default=None,
+        required=False,
+        help=(
+            "The name of the Weights & Biases project to log runs to. "
+            "The environment variable WANDB_API_KEY must be set."
+        ),
+    )
+    parser.add_argument(
+        "--wandb-entity",
+        type=str,
+        default=None,
+        required=False,
+        help="The name of the Weights & Biases entity to log runs to. Defaults to the user's default entity.",
+    )
+    parser.add_argument(
+        "--wandb-run-id",
+        type=str,
+        default=None,
+        required=False,
+        help=(
+            "The ID of an existing Weights & Biases run to resume. "
+            "If not given, creates a new run. If given and exists, "
+            "will continue the run but will overwrite the Python command logged in wandb."
+        ),
+    )
+    parser.add_argument(
+        "--wandb-upload-results",
+        action=argparse.BooleanOptionalAction,
+        required=False,
+        default=True,
+        help=("Whether to upload results as an artifact to Weights & Biases (default: True). Needs `--wandb-project`."),
+    )
+    parser.add_argument(
+        "--description",
+        type=str,
+        required=False,
+        help="Description of the run. This will be added to the metadata of the run to help with bookkeeping.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        required=False,
+        help=(
+            "Size of batch of samples to send to the LLM for evaluation in parallel. "
+            "Use 1 for sequential running (default)."
+        ),
+    )
+    parser.add_argument(
+        "--save-logs",
+        action="store_true",
+        default=True,
+        required=False,
+        help="Whether to save logs to a file in the output directory (default: True).",
+    )
+
+    parser.add_argument(
+        "--judge-model-name",
+        type=str,
+        required=False,
+        help=(
+            "Either a full import path for a judge (e.g., `eval_framework.llm.huggingface.HFLLM`) or the "
+            "name of a class derived from `eval_framework.llm.base.BaseLLM` that can be found in the "
+            "models file. The resulting judge model is instantiated with the arguments provided via "
+            "`--judge-model-args`."
+        ),
+    )
+    parser.add_argument(
+        "--judge-model-args",
+        type=str,
+        required=False,
+        nargs="+",
+        default=(),
+        help="The args of the judge model used.",
+    )
+    parser.add_argument(
+        "--resource-cleanup",
+        action="store_true",
+        required=False,
+        default=False,
+        help=("Add this flag to free up GPU resources between response generation and evaluation"),
+    )
+    parser.add_argument(
+        "--delete-output-dir-after-upload",
+        action="store_true",
+        required=False,
+        default=False,
+        help=("Add this flag to remove the output directory after a successful upload to HF or WandB."),
+    )
+    parser.add_argument(
+        "-v",
+        "--verbosity",
+        type=int,
+        nargs="?",
+        default=1,
+        choices=[0, 1, 2],
+        help="Set the logging verbosity level: 0=critical, 1=info, 2=debug",
+    )
+
+    llm_args: dict[str, Any] = {}
+    args = parser.parse_args()
+
+    for arg in args.llm_args:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+
+            # Handle nested keys like "sampling_params.temperature=0.7"
+            if "." in key:
+                nested_key, sub_key = key.split(".", 1)
+                if nested_key not in llm_args:
+                    llm_args[nested_key] = {}
+                llm_args[nested_key][sub_key] = value
+            else:
+                llm_args[key] = value
+
+    args.llm_args = llm_args
+
+    judge_model_args = {}
+    for arg in args.judge_model_args:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+            judge_model_args[key] = value
+
+    args.judge_model_args = judge_model_args
+
+    # if args.extra_task_modules:
+    #     # Convert the comma-separated string into a list
+    #     args.extra_task_modules = [file_or_dir.strip() for file_or_dir in args.extra_task_modules.split(",")]
+    # else:
+    #     args.extra_task_modules = None
+
+    return args
+
+
+def run_with_kwargs(kwargs: dict) -> None:
+    # Setup logging for the output directory
+    output_dir = kwargs.get("output_dir", "results")
+    log_level = kwargs.get("verbosity", 1)
+    setup_logging(output_dir, log_level=log_level)
+
+    logger.info(kwargs)
+
+    now = datetime.datetime.now()
+    logger.info(f"starting time: {now}")
+
+    if kwargs.get("extra_task_modules"):
+        load_extra_tasks(kwargs["extra_task_modules"])
+
+    context_name = kwargs.pop("context")
+
+    context = CONTEXT[context_name](
+        llm_name=kwargs["llm_name"],
+        models_path=kwargs["models"],
+        num_samples=kwargs["num_samples"],
+        max_tokens=kwargs["max_tokens"],
+        num_fewshot=kwargs["num_fewshot"],
+        task_name=kwargs["task_name"],
+        task_subjects=kwargs["task_subjects"],
+        hf_revision=kwargs["hf_revision"],
+        output_dir=kwargs["output_dir"],
+        wandb_project=kwargs["wandb_project"],
+        wandb_entity=kwargs["wandb_entity"],
+        wandb_run_id=kwargs["wandb_run_id"],
+        wandb_upload_results=kwargs["wandb_upload_results"],
+        hf_upload_dir=kwargs["hf_upload_dir"],
+        hf_upload_repo=kwargs["hf_upload_repo"],
+        llm_args=kwargs["llm_args"],
+        judge_models_path=kwargs["judge_models"],
+        judge_model_name=kwargs["judge_model_name"],
+        judge_model_args=kwargs["judge_model_args"],
+        batch_size=kwargs["batch_size"],
+        description=kwargs["description"],
+        perturbation_type=kwargs["perturbation_type"],
+        perturbation_probability=kwargs["perturbation_probability"],
+        perturbation_seed=kwargs["perturbation_seed"],
+        randomize_judge_order=kwargs["randomize_judge_order"],
+        delete_output_dir_after_upload=kwargs["delete_output_dir_after_upload"],
+        # save_logs=kwargs["save_logs"],
+    )
+
+    with context as ctx:
+        if ctx.config is None:
+            raise ValueError(f"Context configuration is not set for '{type(ctx)}'.")
+
+        main(
+            llm=ctx.config.llm_class(**ctx.config.llm_args),
+            config=ctx.config,
+            should_preempt_callable=ctx.should_preempt,
+            trial_id=ctx.get_trial_id(),
+            resource_cleanup=kwargs.pop("resource_cleanup", False),
+            verbosity=log_level,
+        )
+
+    logger.info(f"time since start: {datetime.datetime.now() - now}")
+
+
+def run() -> None:
+    run_with_kwargs(vars(parse_args()))
+
+
+# Enable execution via `python -m eval_framework.run`. Useful for
+# debugging via `debugpy -m eval_framework.run`
+if __name__ == "__main__":
+    run()

eval_framework/run_direct.py

@@ -0,0 +1,42 @@
+import datetime
+import logging
+from pathlib import Path
+
+from eval_framework.run import run_with_kwargs
+
+logger = logging.getLogger(__name__)
+
+if __name__ == "__main__":
+    logger.info(Path("models.py"))
+    now = datetime.datetime.now()
+    logger.info("starting time:", now)
+    # insert API token here
+
+    # main block for local debugging
+    kwargs = {
+        "context": "local",
+        "models": Path("src/eval_framework/llm/models.py"),
+        "judge_models": None,
+        "judge_model_name": None,
+        "judge_model_args": {},
+        # ---
+        "llm_name": "Llama31_8B_Instruct_API",
+        "llm_args": {},
+        "num_samples": 1,
+        "max_tokens": None,
+        "num_fewshot": 4,
+        "task_name": "Math",  # complete task
+        "task_subjects": None,
+        "hf_revision": None,
+        "output_dir": Path("outputs"),
+        "hf_upload_dir": "",
+        "description": "",
+        "batch_size": 1,
+        "perturbation_type": None,
+        "perturbation_probability": None,
+        "perturbation_seed": None,
+        "save_logs": True,
+    }
+    run_with_kwargs(kwargs)
+
+    logger.info("time since start:", datetime.datetime.now() - now)