eval-framework 0.2.0__py3-none-any.whl

eval_framework/tasks/benchmarks/wmt.py

@@ -0,0 +1,160 @@
+import random
+from abc import ABC
+from typing import Any
+
+import pycountry
+import sacrebleu
+
+from eval_framework.metrics.completion.bleu import LINEWISE_BLEU
+from eval_framework.metrics.completion.chrf import LINEWISE_CHRF
+from eval_framework.metrics.completion.ter import LINEWISE_TER
+from eval_framework.tasks.base import RANDOM_SEED, BaseTask, Language, ResponseType, Sample
+
+
+class WMT(BaseTask[str], ABC):
+    """WMT dataset:"""
+
+    NAME = "WMT"
+    DATASET_PATH = ""
+    SAMPLE_SPLIT = "test"
+    FEWSHOT_SPLIT = "test"
+    RESPONSE_TYPE = ResponseType.COMPLETION
+    METRICS = [LINEWISE_BLEU, LINEWISE_CHRF, LINEWISE_TER]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["phrase"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        super().__init__(num_fewshot)
+        self.stop_sequences: list[str] = [".\n", " phrase: ", "phrase:", "phrase: ", " phrase:", "\n\n"]
+
+    def _load_dataset(self, subject: str | None) -> None:
+        src_file, ref_file, _, _, _ = sacrebleu.download_test_set(test_set=self.DATASET_PATH, langpair=subject)
+        src_data, ref_data = [[line.rstrip() for line in sacrebleu.smart_open(file)] for file in (src_file, ref_file)]
+
+        data_list = [{"source": src, "target": ref, "subject": subject} for src, ref in zip(src_data, ref_data)]
+        self.rnd = random.Random(RANDOM_SEED)
+        self.rnd.shuffle(data_list)
+        self.dataset = {"test": data_list}
+
+    def _code_to_language(self, code: str) -> str:
+        # key is alpha_2 or alpha_3 depending on the code length
+        key = f"alpha_{len(code)}"
+        language_tuple = pycountry.languages.get(**{key: code})
+        return language_tuple.name
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        language_codes = item["subject"].split("-")
+        src_lang = self._code_to_language(language_codes[0])
+
+        language_codes = item["subject"].split("-")
+        tar_lang = self._code_to_language(language_codes[1])
+        cue = f"{tar_lang} phrase:"
+
+        return f"{src_lang} phrase: {item['source']}\n{cue}"
+
+    def _get_ground_truth(self, item: dict[str, Any]) -> str | None:
+        return item["target"] if isinstance(item["target"], str) else item["target"][0]
+
+    def _get_fewshot_target_text(self, item: dict[str, Any]) -> str:
+        target = self._get_ground_truth(item)
+        assert target is not None
+        assert isinstance(target, str)
+        return f" {target}"
+
+    def post_process_generated_completion(self, completion_text: str, sample: Sample | None = None) -> str:
+        for stop_sequence in self.stop_sequences:
+            if stop_sequence in completion_text:
+                completion_text = completion_text.split(stop_sequence)[0]
+        return completion_text.strip()
+
+
+class WMT14(WMT):
+    NAME = "WMT14"
+    DATASET_PATH = "wmt14"
+    SUBJECTS = ["en-fr", "fr-en"]
+    LANGUAGE = {
+        "en-fr": (Language["ENG"], Language["FRA"]),
+        "fr-en": (Language["FRA"], Language["ENG"]),
+    }
+
+
+class WMT16(WMT):
+    NAME = "WMT16"
+    DATASET_PATH = "wmt16"
+    SUBJECTS = ["de-en", "en-de"]
+    LANGUAGE = {
+        "de-en": (Language["DEU"], Language["ENG"]),
+        "en-de": (Language["ENG"], Language["DEU"]),
+    }
+
+
+class WMT20(WMT):
+    NAME = "WMT20"
+    DATASET_PATH = "wmt20"
+    SUBJECTS = ["de-en", "de-fr", "en-de", "fr-de"]
+    LANGUAGE = {
+        "de-en": (Language["DEU"], Language["ENG"]),
+        "de-fr": (Language["DEU"], Language["FRA"]),
+        "en-de": (Language["ENG"], Language["DEU"]),
+        "fr-de": (Language["FRA"], Language["DEU"]),
+    }
+
+
+class WMT_INSTRUCT(WMT):
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Please", "translate"]
+    COMPLETION_PREFIX = "This is the translation:"
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        super().__init__(num_fewshot)
+        self.stop_sequences: list[str] = ["Please translate"]
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        src_lang, tar_lang = map(self._code_to_language, item["subject"].split("-"))
+        return f"Please translate from {src_lang} to {tar_lang}: {item['source']}"
+
+    def _get_cue(self, item: dict[str, Any]) -> str:
+        return self.COMPLETION_PREFIX
+
+    def _get_fewshot_target_text(self, item: dict[str, Any]) -> str:
+        target = self._get_ground_truth(item)
+        assert target is not None
+        return f" {target}"
+
+    def post_process_generated_completion(self, completion_text: str, sample: Sample | None = None) -> str:
+        completion_text = completion_text.removeprefix(self.COMPLETION_PREFIX)
+        completion_text = completion_text.strip()
+        for stop_sequence in self.stop_sequences:
+            if stop_sequence in completion_text:
+                completion_text = completion_text.split(stop_sequence)[0]
+        return completion_text
+
+
+class WMT14_INSTRUCT(WMT_INSTRUCT):
+    NAME = "WMT14 Instruct"
+    DATASET_PATH = "wmt14"
+    SUBJECTS = ["en-fr", "fr-en"]
+    LANGUAGE = {
+        "en-fr": (Language["ENG"], Language["FRA"]),
+        "fr-en": (Language["FRA"], Language["ENG"]),
+    }
+
+
+class WMT16_INSTRUCT(WMT_INSTRUCT):
+    NAME = "WMT16 Instruct"
+    DATASET_PATH = "wmt16"
+    SUBJECTS = ["de-en", "en-de"]
+    LANGUAGE = {
+        "de-en": (Language["DEU"], Language["ENG"]),
+        "en-de": (Language["ENG"], Language["DEU"]),
+    }
+
+
+class WMT20_INSTRUCT(WMT_INSTRUCT):
+    NAME = "WMT20 Instruct"
+    DATASET_PATH = "wmt20"
+    SUBJECTS = ["de-en", "de-fr", "en-de", "fr-de"]
+    LANGUAGE = {
+        "de-en": (Language["DEU"], Language["ENG"]),
+        "de-fr": (Language["DEU"], Language["FRA"]),
+        "en-de": (Language["ENG"], Language["DEU"]),
+        "fr-de": (Language["FRA"], Language["DEU"]),
+    }

eval_framework/tasks/benchmarks/zero_scrolls.py

@@ -0,0 +1,197 @@
+import re
+from typing import Any
+
+from eval_framework.metrics.completion.exponential_similarity import ExponentialSimilarity
+from eval_framework.metrics.completion.f1 import F1
+from eval_framework.metrics.completion.rouge_geometric_mean import ROUGE_GEOMETRIC_MEAN
+from eval_framework.metrics.loglikelihood.accuracy_loglikelihood import (
+    AccuracyLoglikelihood,
+)
+from eval_framework.tasks.base import BaseTask, Language, ResponseType, Sample
+from eval_framework.tasks.utils import get_n_letters
+
+
+class ZERO_SCROLLS_QUALITY(BaseTask[str]):
+    """ZeroSCROLLS dataset: https://huggingface.co/datasets/tau/zero_scrolls"""
+
+    NAME = "ZeroSCROLLS QuALITY"
+    DATASET_PATH = "tau/zero_scrolls"
+    SAMPLE_SPLIT = "validation"
+    FEWSHOT_SPLIT = "validation"
+    RESPONSE_TYPE = ResponseType.LOGLIKELIHOODS
+    METRICS = [AccuracyLoglikelihood]
+    SUBJECTS = ["quality"]
+
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+    LANGUAGE = Language.ENG
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS QuALITY only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+        self.keys = get_n_letters(4)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\n"
+
+    def _get_cue_text(self, item: dict[str, Any]) -> str:
+        return "Answer:"
+
+    def _get_ground_truth(self, item: dict[str, Any]) -> str | None:
+        return f" {item['output']}"
+
+    def _get_possible_completions(self, item: dict[str, Any]) -> list[str] | None:
+        return [f" {key}" for key in self.keys]
+
+
+class ZERO_SCROLLS_COMPLETION(BaseTask[str]):
+    """ZeroSCROLLS dataset: https://huggingface.co/datasets/tau/zero_scrolls"""
+
+    DATASET_PATH = "tau/zero_scrolls"
+    SAMPLE_SPLIT = "validation"
+    FEWSHOT_SPLIT = "validation"
+    RESPONSE_TYPE = ResponseType.COMPLETION
+
+    def _get_ground_truth(self, item: dict[str, Any]) -> str | None:
+        return item["output"]
+
+
+class ZERO_SCROLLS_GOV_REPORT(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS GovReport"
+    METRICS = [ROUGE_GEOMETRIC_MEAN]
+    SUBJECTS = ["gov_report"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Summary"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS GovReport only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}Summary:"
+
+
+class ZERO_SCROLLS_QMSUM(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS QMSum"
+    METRICS = [ROUGE_GEOMETRIC_MEAN]
+    SUBJECTS = ["qmsum"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS QMSum only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\nAnswer:"
+
+
+class ZERO_SCROLLS_SQUALITY(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS SQuALITY"
+    METRICS = [ROUGE_GEOMETRIC_MEAN]
+    SUBJECTS = ["squality"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS SQuALITY only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\nAnswer:"
+
+
+class ZERO_SCROLLS_QASPER(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS Qasper"
+    METRICS = [F1]
+    SUBJECTS = ["qasper"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS Qasper only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\nAnswer:"
+
+
+class ZERO_SCROLLS_NARRATIVEQA(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS NarrativeQA"
+    METRICS = [F1]
+    SUBJECTS = ["narrative_qa"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS NarrativeQA only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\nAnswer:"
+
+
+class ZERO_SCROLLS_MUSIQUE(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS MuSiQue"
+    METRICS = [F1]
+    SUBJECTS = ["musique"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS MuSiQue only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}\n\nAnswer:"
+
+
+class ZERO_SCROLLS_SPACE_DIGEST(ZERO_SCROLLS_COMPLETION):
+    NAME = "ZeroSCROLLS SpaceDigest"
+    METRICS = [ExponentialSimilarity]
+    SUBJECTS = ["space_digest"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Answer"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        assert num_fewshot == 0, "ZeroSCROLLS SpaceDigest only supports zero fewshot examples"
+        super().__init__(num_fewshot)
+
+    def post_process_generated_completion(self, completion_text: str, sample: Sample | None = None) -> str:
+        # First, try to find patterns like "X%" or "X percent" or "X percentage"
+        percentage_patterns = [
+            r"(\d+(?:\.\d+)?)%",  # Matches: 30%, 30.5%
+            r"(\d+(?:\.\d+)?)\s*percent",  # Matches: 30 percent, 30.5 percent
+            r"(\d+(?:\.\d+)?)\s*percentage",  # Matches: 30 percentage, 30.5 percentage
+            r"percentage\s*(?:is|of|:)?\s*(\d+(?:\.\d+)?)",  # Matches: percentage is 30, percentage: 30.5
+            r"(?:is|equals|equal to|about|approximately|around|roughly)\s*(\d+(?:\.\d+)?)\s*%",
+            # Matches: is 30%, equals 30.5%
+            r"(?:is|equals|equal to|about|approximately|around|roughly)\s*(\d+(?:\.\d+)?)\s*percent",
+            # Matches: is 30 percent
+            r"it'?s\s*(\d+(?:\.\d+)?)",  # Matches: it's 60, its 60
+            r"that'?s\s*(\d+(?:\.\d+)?)",  # Matches: that's 60, thats 60
+        ]
+
+        for pattern in percentage_patterns:
+            match = re.search(pattern, completion_text, re.IGNORECASE)
+            if match:
+                return match.group(1).strip()
+
+        # If no percentage pattern is found, check if the entire text is just a number
+        if re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*", completion_text):
+            return completion_text.strip()
+
+        # If not a standalone number, look for any number in the text
+        # This is a fallback and might be less accurate
+        number_match = re.search(r"(\d+(?:\.\d+)?)", completion_text)
+        if number_match:
+            return number_match.group(1).strip()
+
+        # If no number is found, return the original text stripped
+        return completion_text.strip()
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        query_end_index = item["query_end_index"]
+        return f"{item['input'][:query_end_index]}Answer:"
+
+    def _get_ground_truth(self, item: dict[str, Any]) -> str | None:
+        return self.post_process_generated_completion(item["output"])

eval_framework/tasks/eval_config.py

@@ -0,0 +1,112 @@
+import ast
+import json
+from pathlib import Path
+from typing import Annotated, Any
+
+from pydantic import AfterValidator, Field, field_serializer, field_validator, model_validator
+
+from eval_framework.base_config import BaseConfig
+from eval_framework.llm.base import BaseLLM
+from eval_framework.metrics.llm.base import BaseLLMJudgeMetric
+from eval_framework.tasks.base import BaseTask
+from eval_framework.tasks.perturbation import PerturbationConfig
+from eval_framework.tasks.registry import get_task, validate_task_name
+from eval_framework.utils.constants import ROOT_DIR
+
+
+class EvalConfig(BaseConfig):
+    output_dir: Path = ROOT_DIR
+    wandb_project: str | None = None
+    wandb_entity: str | None = None
+    wandb_run_id: str | None = None
+    hf_upload_dir: str | None = None
+    hf_upload_repo: str | None = None
+    num_fewshot: Annotated[int, Field(ge=0)] = 0
+    num_samples: Annotated[int | None, Field(ge=1)] = 10  # Allows None or int
+    max_tokens: int | None = None
+    perturbation_config: PerturbationConfig | None = None
+    task_name: Annotated[str, AfterValidator(validate_task_name)]
+    task_subjects: list[str] | None = None
+    hf_revision: str | None = None
+    llm_class: type[BaseLLM]
+    llm_args: dict[str, Any] = Field(default_factory=dict)
+    llm_judge_class: type[BaseLLM] | None = None
+    judge_model_args: dict[str, Any] = Field(default_factory=dict)
+    batch_size: Annotated[int, Field(ge=1)] = 1
+    description: str | None = None
+    save_intermediate_results: bool = True
+    save_logs: bool = True
+
+    @property
+    def task_class(self) -> type[BaseTask]:
+        return get_task(self.task_name)
+
+    @field_serializer("output_dir")
+    def serialize_output_dir(self, value: Path) -> str:
+        return str(value)
+
+    @field_validator("output_dir", mode="before")
+    @classmethod
+    def validate_output_dir(cls, value: str | Path) -> Path:
+        if isinstance(value, str):
+            return Path(value)
+        return value
+
+    @field_validator("llm_args", mode="before")
+    @classmethod
+    def validate_llm_args(cls, value: dict[str, Any]) -> dict[str, Any]:
+        def convert_value(v: Any) -> Any:
+            if isinstance(v, dict):
+                # Recursively process nested dictionaries (like sampling_params)
+                return {k: convert_value(nested_v) for k, nested_v in v.items()}
+            elif isinstance(v, str):
+                try:
+                    # Try to evaluate as a Python literal (int, float, bool, None, list, dict, etc.)
+                    return ast.literal_eval(v)
+                except (ValueError, SyntaxError):
+                    return v  # keep as string if not a valid literal
+            else:
+                return v  # already proper type
+
+        return convert_value(value)
+
+    @field_validator("judge_model_args", mode="before")
+    @classmethod
+    def validate_judge_model_args(cls, value: dict[str, Any]) -> dict[str, Any]:
+        typed_value = {}
+        for k, v in value.items():
+            try:  # maybe this llm argument is actually a number?
+                if "." in str(v):
+                    v = float(v)
+                else:
+                    v = int(v)
+            except ValueError:
+                pass
+            typed_value[k] = v
+        return typed_value
+
+    @model_validator(mode="after")
+    def validate_llm_judge_defined(self) -> "EvalConfig":
+        task = get_task(self.task_name)
+        for metric_class in task.METRICS:
+            if issubclass(metric_class, BaseLLMJudgeMetric):
+                assert self.llm_judge_class is not None, "The LLM Judge must be defined for this evaluation task."
+        return self
+
+    @field_serializer("llm_class")
+    def serialize_llm_class(self, value: type[BaseLLM] | None) -> str | None:
+        """Serialize the class into its fully qualified name."""
+        if value:
+            return value.__name__
+        return None
+
+    @field_serializer("llm_judge_class")
+    def serialize_llm_judge_class(self, value: type[BaseLLM] | None) -> str | None:
+        """Serialize the class into its fully qualified name."""
+        if value:
+            return value.__name__
+        return None
+
+    def model_json_dump(self) -> str:
+        model_dump = self.model_dump()
+        return json.dumps(model_dump, sort_keys=True)

eval_framework/tasks/perturbation.py

@@ -0,0 +1,83 @@
+import threading
+from enum import Enum
+from typing import Annotated, Any, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from eval_framework.logger import logger
+from eval_framework.tasks.base import RANDOM_SEED, BaseTask, Sample
+from eval_framework.tasks.utils import Editor, HatPaperEditor
+
+
+class PerturbationType(str, Enum):
+    # Editor methods
+    EDITOR = "editor"
+    # Hat paper methods
+    PERMUTE = "permute"
+    REPLACE = "replace"
+    DELETE = "delete"
+    UPPERCASE = "uppercase"
+
+
+class PerturbationConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    type: PerturbationType = PerturbationType.EDITOR
+    probability: Annotated[float, Field(ge=0.0, le=1.0)] = 0.1
+    seed: int = RANDOM_SEED
+    verbose: bool = False
+
+
+_DOCKER_LAUNCH_LOCK = threading.Lock()
+_AUGMENTER_PORT = 0
+
+SomeBaseTask = TypeVar("SomeBaseTask", bound=BaseTask[Any])
+
+
+def create_perturbation_class[T: BaseTask](base_class: type[T], perturbation_config: PerturbationConfig) -> type[T]:
+    # mypy seems to have trouble inferring the type
+    class EditorPerturbation(base_class):  # type: ignore
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.perturbation_config = perturbation_config
+            self.editor = Editor(
+                language="de" if base_class.LANGUAGE == "German" else "en", seed=perturbation_config.seed
+            )
+
+        def _get_instruction_text(self, sample: Sample) -> str:
+            text = super()._get_instruction_text(sample)
+            if self.perturbation_config.verbose:
+                logger.info(f"Perturbating text: {text}")
+            result = self.editor(
+                text, self.perturbation_config.probability, getattr(self, "PERTURBATION_UNMODIFIABLE_WORDS", [])
+            )
+            if self.perturbation_config.verbose:
+                logger.info(f"Perturbed text: {result}")
+            return result
+
+    class HatPaperPerturbation(base_class):  # type: ignore
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.perturbation_config = perturbation_config
+            self.editor = HatPaperEditor(seed=perturbation_config.seed)
+
+        def _get_instruction_text(self, sample: Sample) -> str:
+            text = super()._get_instruction_text(sample)
+            if self.perturbation_config.verbose:
+                logger.info(f"Perturbating text: {text}")
+            words = getattr(self, "PERTURBATION_UNMODIFIABLE_WORDS", [])
+            if self.perturbation_config.type == PerturbationType.PERMUTE:
+                result = self.editor.permute_chars_in_string(text, self.perturbation_config.probability, words)
+            elif self.perturbation_config.type == PerturbationType.REPLACE:
+                result = self.editor.replace_chars_in_string(text, self.perturbation_config.probability, words)
+            elif self.perturbation_config.type == PerturbationType.DELETE:
+                result = self.editor.delete_chars_in_string(text, self.perturbation_config.probability, words)
+            elif self.perturbation_config.type == PerturbationType.UPPERCASE:
+                result = self.editor.upper_case_string(text)
+            if self.perturbation_config.verbose:
+                logger.info(f"Perturbed text: {result}")
+            return result
+
+    if perturbation_config.type == PerturbationType.EDITOR:
+        return EditorPerturbation
+    else:
+        return HatPaperPerturbation