Nanotron, Multilingual tasks update + misc (#756)

## Nanotron
- Logprobs working
- Generation works with https://github.com/huggingface/nanotron/tree/nouamane/lighteval-fix, however sometimes there is issue with ordering.
- Only relevant part of config are fetched to prevent some dataclass checks which are not needed in nanotron
- Smol changes to configs

## Metrics
- Probability metrics now works with Char normalization
- Token normalization is fixed when used with transformers

## Multlilingual tasks
- New squads + few new mcf benchmarks + smol fixes to existing tasks

## Misc
- Qa template now only uses unique golds, which speed-upds probs calculation.

---------

Co-authored-by: anton <[email protected]>
Co-authored-by: nouamanetazi <[email protected]>
Co-authored-by: Hynek Kydlicek <[email protected]>

Author: 4 people

src/lighteval/config/lighteval_config.py

@@ -23,11 +23,13 @@
 from dataclasses import dataclass
 from typing import Dict, Optional, Union
 
+from pydantic import BaseModel
+
 from lighteval.utils.imports import is_nanotron_available
 
 
 if is_nanotron_available():
-    from nanotron.config import Config
+    from nanotron.config import GeneralArgs, ModelArgs, TokenizerArgs
     from nanotron.config.parallelism_config import ParallelismArgs
     from nanotron.generation.sampler import SamplerType
     from nanotron.logging import get_logger
@@ -37,9 +39,8 @@
 DEFAULT_GENERATION_SEED = 42
 
 
-@dataclass
-class GenerationArgs:
-    sampler: Optional[Union[str, "SamplerType"]] = None
+class GenerationArgs(BaseModel):
+    sampler: Optional["SamplerType"] = None
     temperature: Optional[float] = None
     top_k: Optional[int] = None
     top_p: Optional[float] = None
@@ -49,8 +50,6 @@ class GenerationArgs:
     use_cache: Optional[bool] = False
 
     def __post_init__(self):
-        if isinstance(self.sampler, str):
-            self.sampler = SamplerType[self.sampler.upper()]
         if self.seed is None:
             self.seed = DEFAULT_GENERATION_SEED
 
@@ -101,4 +100,14 @@ class LightEvalConfig:
 @dataclass
 class FullNanotronConfig:
     lighteval_config: LightEvalConfig
-    nanotron_config: "Config"
+    nanotron_model: "ModelArgs"
+    nanotron_tokenizer: "TokenizerArgs"
+    nanotron_general: "GeneralArgs"
+
+    @property
+    def generation_parameters(self):
+        # Return the generation parameters from the lighteval config
+        # or create default generation parameters if none are set
+        if self.lighteval_config.generation:
+            return self.lighteval_config.generation
+        return GenerationArgs()

src/lighteval/main_nanotron.py

@@ -23,12 +23,12 @@
 # flake8: noqa: C901
 import os
 
+import yaml
 from typer import Option
 from typing_extensions import Annotated
+from yaml import SafeLoader
 
 
-CACHE_DIR: str = os.getenv("HF_HOME", "/scratch")
-
 HELP_PANEL_NAME_1 = "Common Parameters"
 HELP_PANEL_NAME_2 = "Logging Parameters"
 HELP_PANEL_NAME_3 = "Debug Parameters"
@@ -43,41 +43,43 @@ def nanotron(
         str, Option(help="Path to the nanotron checkpoint YAML or python config file, potentially on s3.")
     ],
     lighteval_config_path: Annotated[str, Option(help="Path to a YAML config to be used for the evaluation.")],
-    cache_dir: Annotated[str, Option(help="Cache directory for datasets and models.")] = CACHE_DIR,
 ):
     """
     Evaluate models using nanotron as backend.
     """
-    from nanotron.config import Config, get_config_from_file
+    from nanotron.config import GeneralArgs, ModelArgs, TokenizerArgs, get_config_from_dict, get_config_from_file
 
-    from lighteval.config.lighteval_config import FullNanotronConfig, LightEvalConfig
+    from lighteval.config.lighteval_config import (
+        FullNanotronConfig,
+        LightEvalConfig,
+    )
     from lighteval.logging.evaluation_tracker import EvaluationTracker
-    from lighteval.logging.hierarchical_logger import htrack_block
     from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
     from lighteval.utils.imports import NO_NANOTRON_ERROR_MSG, is_nanotron_available
-    from lighteval.utils.utils import EnvConfig
-
-    env_config = EnvConfig(token=os.getenv("HF_TOKEN"), cache_dir=cache_dir)
 
     if not is_nanotron_available():
         raise ImportError(NO_NANOTRON_ERROR_MSG)
 
-    with htrack_block("Load nanotron config"):
-        # Create nanotron config
-        if not checkpoint_config_path.endswith(".yaml"):
-            raise ValueError("The checkpoint path should point to a YAML file")
+    # Create nanotron config
+    if not checkpoint_config_path.endswith(".yaml"):
+        raise ValueError("The checkpoint path should point to a YAML file")
+
+    with open(checkpoint_config_path) as f:
+        nanotron_yaml = yaml.load(f, Loader=SafeLoader)
 
-        model_config = get_config_from_file(
-            checkpoint_config_path,
-            config_class=Config,
-            model_config_class=None,
+    model_config, tokenizer_config, general_config = [
+        get_config_from_dict(
+            nanotron_yaml[key],
+            config_class=config_class,
             skip_unused_config_keys=True,
             skip_null_keys=True,
         )
+        for key, config_class in [("model", ModelArgs), ("tokenizer", TokenizerArgs), ("general", GeneralArgs)]
+    ]
 
-        # We are getting an type error, because the get_config_from_file is not correctly typed,
-        lighteval_config: LightEvalConfig = get_config_from_file(lighteval_config_path, config_class=LightEvalConfig)  # type: ignore
-        nanotron_config = FullNanotronConfig(lighteval_config, model_config)
+    # Load lighteval config
+    lighteval_config: LightEvalConfig = get_config_from_file(lighteval_config_path, config_class=LightEvalConfig)  # type: ignore
+    nanotron_config = FullNanotronConfig(lighteval_config, model_config, tokenizer_config, general_config)
 
     evaluation_tracker = EvaluationTracker(
         output_dir=lighteval_config.logging.output_dir,
@@ -88,17 +90,15 @@ def nanotron(
         push_to_tensorboard=lighteval_config.logging.push_to_tensorboard,
         save_details=lighteval_config.logging.save_details,
         tensorboard_metric_prefix=lighteval_config.logging.tensorboard_metric_prefix,
-        nanotron_run_info=nanotron_config.nanotron_config.general,
+        nanotron_run_info=nanotron_config.nanotron_general,
     )
 
     pipeline_parameters = PipelineParameters(
         launcher_type=ParallelismManager.NANOTRON,
-        env_config=env_config,
         job_id=os.environ.get("SLURM_JOB_ID", 0),
         nanotron_checkpoint_path=checkpoint_config_path,
         dataset_loading_processes=lighteval_config.tasks.dataset_loading_processes,
         custom_tasks_directory=lighteval_config.tasks.custom_tasks,
-        override_batch_size=lighteval_config.batch_size,
         num_fewshot_seeds=1,
         max_samples=lighteval_config.tasks.max_samples,
         use_chat_template=False,

src/lighteval/metrics/metrics_sample.py

@@ -342,7 +342,8 @@ def __init__(
     def compute(
         self,
         logprobs: list[float],
-        target_tokens: list[list[int]],
+        target_tokens: list[list[int]] | None = None,
+        reference_texts: list[str] | None = None,
         **kwargs,
     ) -> float:
         """Computes the log likelihood probability: chance of choosing the best choice.
@@ -352,8 +353,7 @@ def compute(
             choices_logprob (list[float]): Summed log-probabilities of all the possible choices for the model, ordered as the choices.
             unconditioned_logprob (list[float] | None): Unconditioned log-probabilities for PMI normalization, ordered as the choices.
             choices_tokens (list[list[int]] | None): Tokenized choices for token normalization, ordered as the choices.
-            formatted_doc (Doc): Original document for the sample.
-                Used to get the original choices' length for possible normalization
+            reference_texts (list[str] | None): Reference texts for token normalization, ordered as the choices.
 
         Returns:
             float: The probability of the best log-prob choice being a gold choice.
@@ -364,7 +364,7 @@ def compute(
                 normalization=self.log_prob_normalization,
                 choices_tokens=target_tokens,
                 choices_logprob=logprobs,
-                choices_text=None,
+                choices_text=reference_texts,
                 unconditioned_logprob=None,
             )
             if self.log_prob_normalization