EuroEval 16.3.0__py3-none-any.whl → 16.5.0__py3-none-any.whl

euroeval/finetuning.py

@@ -1,12 +1,12 @@
 """Functions related to the finetuning of models."""
 
+import collections.abc as c
 import logging
 import sys
 import typing as t
 from functools import partial
 
 import torch
-from tqdm.auto import tqdm
 from transformers.trainer_callback import (
     EarlyStoppingCallback,
     PrinterCallback,
@@ -18,13 +18,9 @@ from transformers.training_args import OptimizerNames, TrainingArguments
 from .callbacks import NeverLeaveProgressCallback
 from .enums import DataType
 from .exceptions import InvalidBenchmark, NaNValueInModelOutput
+from .logging_utils import block_terminal_output, get_pbar, log, log_once
 from .model_loading import load_model
-from .utils import (
-    block_terminal_output,
-    clear_memory,
-    enforce_reproducibility,
-    log_once,
-)
+from .utils import clear_memory, enforce_reproducibility
 
 if t.TYPE_CHECKING:
     from datasets import DatasetDict
@@ -32,16 +28,14 @@ if t.TYPE_CHECKING:
     from .benchmark_modules import BenchmarkModule
     from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
-logger = logging.getLogger("euroeval")
-
 
 def finetune(
     model: "BenchmarkModule",
-    datasets: list["DatasetDict"],
+    datasets: c.Sequence["DatasetDict"],
     model_config: "ModelConfig",
     dataset_config: "DatasetConfig",
     benchmark_config: "BenchmarkConfig",
-) -> list[dict[str, float]]:
+) -> c.Sequence[dict[str, float]]:
     """Evaluate a model on a dataset through finetuning.
 
     Args:
@@ -58,6 +52,10 @@ def finetune(
 
     Returns:
         A list of dicts containing the scores for each metric for each iteration.
+
+    Raises:
+        InvalidBenchmark:
+            If the benchmark could not be completed.
     """
     # Set the data type to use for the model weights
     using_cuda = benchmark_config.device == torch.device("cuda")
@@ -70,7 +68,7 @@ def finetune(
 
     bs: int = benchmark_config.batch_size
     scores: list[dict[str, float]] = list()
-    for idx in tqdm(
+    for idx in get_pbar(
         iterable=range(benchmark_config.num_iterations),
         desc="Benchmarking",
         disable=not benchmark_config.progress_bar,
@@ -80,7 +78,7 @@ def finetune(
         model_already_initialized = idx == 0
 
         # Run a loop here to deal with automatic reduction of batch size
-        while True:
+        for _ in range(num_attempts := 10):
             # Clear GPU memory
             if not model_already_initialized:
                 try:
@@ -112,7 +110,10 @@ def finetune(
                 )
 
                 scores.append(itr_scores)
-                logger.debug(f"Test scores for iteration {idx}: {itr_scores}")
+                log(
+                    f"Test scores for iteration {idx}: {itr_scores}",
+                    level=logging.DEBUG,
+                )
 
                 break
 
@@ -123,9 +124,10 @@ def finetune(
                 if dtype != DataType.FP32:
                     dtype = DataType.FP32
                     model_already_initialized = False
-                    logger.debug(
+                    log(
                         "NaN value detected in model outputs while using mixed "
-                        "precision. Retrying with full fp32 precision."
+                        "precision. Retrying with full fp32 precision.",
+                        level=logging.DEBUG,
                     )
                 else:
                     raise InvalidBenchmark(
@@ -151,7 +153,12 @@ def finetune(
                 model_already_initialized = False
 
                 bs //= 2
-                logger.debug(f"Reduced batch size to {bs}")
+                log(f"Reduced batch size to {bs}", level=logging.DEBUG)
+
+        else:
+            raise InvalidBenchmark(
+                f"Could not benchmark the model after {num_attempts} attempts!"
+            )
 
     return scores
 

euroeval/generation.py

@@ -1,5 +1,6 @@
 """Functions related to text generation of models."""
 
+import collections.abc as c
 import logging
 import sys
 import typing as t
@@ -11,12 +12,13 @@ from tqdm.auto import tqdm
 
 from .enums import BatchingPreference, TaskGroup
 from .exceptions import InvalidBenchmark
+from .logging_utils import get_pbar, log, log_once
 from .model_cache import (
     ModelCache,
     load_cached_model_outputs,
     split_dataset_into_cached_and_non_cached,
 )
-from .utils import clear_memory, log_once
+from .utils import clear_memory
 
 if t.TYPE_CHECKING:
     from datasets import DatasetDict
@@ -29,16 +31,14 @@ if t.TYPE_CHECKING:
         ModelConfig,
     )
 
-logger = logging.getLogger("euroeval")
-
 
 def generate(
     model: "BenchmarkModule",
-    datasets: list["DatasetDict"],
+    datasets: c.Sequence["DatasetDict"],
     model_config: "ModelConfig",
     dataset_config: "DatasetConfig",
     benchmark_config: "BenchmarkConfig",
-) -> list[dict[str, float]]:
+) -> c.Sequence[dict[str, float]]:
     """Evaluate a model on a dataset through generation.
 
     Args:
@@ -78,7 +78,7 @@ def generate(
     )
 
     scores: list[dict[str, float]] = list()
-    for idx in tqdm(
+    for idx in get_pbar(
         iterable=range(len(datasets)),
         desc="Benchmarking",
         disable=not benchmark_config.progress_bar,
@@ -90,7 +90,7 @@ def generate(
             dataset_config=dataset_config,
             benchmark_config=benchmark_config,
         )
-        logger.debug(f"Test scores for iteration {idx}: {test_scores}")
+        log(f"Test scores for iteration {idx}: {test_scores}", level=logging.DEBUG)
         scores.append(test_scores)
         clear_memory()
 
@@ -142,14 +142,14 @@ def generate_single_iteration(
         itr: t.Iterable
         match model.batching_preference:
             case BatchingPreference.SINGLE_SAMPLE:
-                itr = tqdm(iterable=non_cached_dataset, leave=False)
+                itr = get_pbar(iterable=non_cached_dataset)
             case BatchingPreference.ALL_AT_ONCE:
                 itr = [non_cached_dataset[:]]
             case _:
                 num_batches = len(non_cached_dataset) // benchmark_config.batch_size
                 if len(non_cached_dataset) % benchmark_config.batch_size != 0:
                     num_batches += 1
-                itr = tqdm(
+                itr = get_pbar(
                     iterable=mit.batched(
                         iterable=non_cached_dataset, n=benchmark_config.batch_size
                     ),
@@ -254,7 +254,7 @@ def generate_single_iteration(
 def debug_log(
     batch: dict[str, t.Any],
     model_output: "GenerativeModelOutput",
-    extracted_labels: list[dict | str | list[str]],
+    extracted_labels: c.Sequence[dict | str | c.Sequence[str]],
     dataset_config: "DatasetConfig",
 ) -> None:
     """Log inputs and outputs for debugging purposes.
@@ -297,7 +297,7 @@ def debug_log(
                         + "\n"
                         + "\t".join(labels)
                     )
-            logger.info("\n\n".join(log_msgs))
+            log("\n\n".join(log_msgs), level=logging.DEBUG)
             return
 
         case (
@@ -332,7 +332,7 @@ def debug_log(
     else:
         input_texts = batch["text"]
 
-    metadata_keys: list[str] = [
+    metadata_keys: c.Sequence[str] = [
         key
         for key in batch.keys()
         if key not in ["text", "messages", "label", "labels", "target_text"]
@@ -347,6 +347,7 @@ def debug_log(
         if labels[idx]:
             data_to_log["Label"] = labels[idx]
         data_to_log |= {key.capitalize(): batch[key][idx] for key in metadata_keys}
-        logger.info(
-            "\n".join(f"{key}: {value!r}" for key, value in data_to_log.items())
+        log(
+            "\n".join(f"{key}: {value!r}" for key, value in data_to_log.items()),
+            level=logging.DEBUG,
         )

euroeval/generation_utils.py

@@ -1,5 +1,6 @@
 """Utility functions related to generative models."""
 
+import collections.abc as c
 import itertools as it
 import json
 import logging
@@ -9,8 +10,9 @@ import typing as t
 
 from .enums import GenerativeType, TaskGroup
 from .exceptions import InvalidBenchmark, InvalidModel
+from .logging_utils import log_once
 from .tokenisation_utils import apply_chat_template
-from .utils import extract_multiple_choice_labels, log_once
+from .utils import extract_multiple_choice_labels
 
 if t.TYPE_CHECKING:
     from datasets import DatasetDict
@@ -18,15 +20,13 @@ if t.TYPE_CHECKING:
 
     from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
-logger = logging.getLogger("euroeval")
-
 
 def extract_few_shot_examples(
     dataset: "DatasetDict",
     dataset_config: "DatasetConfig",
     benchmark_config: "BenchmarkConfig",
     itr_idx: int,
-) -> list[dict[str, t.Any]]:
+) -> c.Sequence[dict[str, t.Any]]:
     """Extract few-shot examples from a dataset.
 
     This will always extract the examples from the training split.
@@ -79,7 +79,7 @@ def extract_few_shot_examples(
                     lambda example: len(example["text"]) < max_num_tokens
                 )
                 num_short_examples = len(train_with_short_examples)
-                if num_short_examples >= dataset_config.num_few_shot_examples:
+                if num_short_examples >= num_few_shots:
                     break
             else:
                 raise InvalidBenchmark(
@@ -144,7 +144,7 @@ def extract_few_shot_examples(
                     lambda example: len(example["context"]) < max_num_tokens
                 )
                 num_short_examples = len(train_with_short_examples)
-                if num_short_examples >= dataset_config.num_few_shot_examples:
+                if num_short_examples >= num_few_shots:
                     break
             else:
                 raise InvalidBenchmark(
@@ -171,7 +171,7 @@ def extract_few_shot_examples(
 
 def apply_prompt(
     examples: dict[str, t.Any],
-    few_shot_examples: list[dict[str, t.Any]],
+    few_shot_examples: c.Sequence[dict[str, t.Any]],
     model_config: "ModelConfig",
     dataset_config: "DatasetConfig",
     generative_type: GenerativeType | None,
@@ -432,7 +432,7 @@ def apply_prompt(
 
 
 def raise_if_wrong_params(
-    model_config: "ModelConfig", allowed_params: dict[re.Pattern, list[str]]
+    model_config: "ModelConfig", allowed_params: dict[re.Pattern, c.Sequence[str]]
 ) -> None:
     """Raise an error if the model configuration has invalid parameters.