Moved some files (#905)

* fixed nanotron config file placed in a very random place

* also put litellm model config file with other endpoint models

* fix docs

Author: clefourrier

docs/source/package_reference/models.mdx

@@ -37,7 +37,7 @@ set in the `model-args` or in the model yaml file (see example
 [[autodoc]] models.endpoints.tgi_model.TGIModelConfig
 
 ### Litellm Model
-[[autodoc]] models.litellm_model.LiteLLMModelConfig
+[[autodoc]] models.endpoints.litellm_model.LiteLLMModelConfig
 
 ## Custom Model
 [[autodoc]] models.custom.custom_model.CustomModelConfig

src/lighteval/config/lighteval_config.py

@@ -377,7 +377,7 @@ def litellm(
     import yaml
 
     from lighteval.logging.evaluation_tracker import EvaluationTracker
-    from lighteval.models.litellm_model import LiteLLMModelConfig
+    from lighteval.models.endpoints.litellm_model import LiteLLMModelConfig
     from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
 
     evaluation_tracker = EvaluationTracker(

src/lighteval/main_endpoint.py

@@ -57,18 +57,19 @@ def nanotron(
     """
     Evaluate models using nanotron as backend.
     """
+    from lighteval.utils.imports import NO_NANOTRON_ERROR_MSG, is_nanotron_available
+
+    if not is_nanotron_available():
+        raise ImportError(NO_NANOTRON_ERROR_MSG)
+
     from nanotron.config import GeneralArgs, ModelArgs, TokenizerArgs, get_config_from_dict, get_config_from_file
 
-    from lighteval.config.lighteval_config import (
+    from lighteval.logging.evaluation_tracker import EvaluationTracker
+    from lighteval.models.nanotron import (
         FullNanotronConfig,
         LightEvalConfig,
     )
-    from lighteval.logging.evaluation_tracker import EvaluationTracker
     from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
-    from lighteval.utils.imports import NO_NANOTRON_ERROR_MSG, is_nanotron_available
-
-    if not is_nanotron_available():
-        raise ImportError(NO_NANOTRON_ERROR_MSG)
 
     # Create nanotron config
     if not checkpoint_config_path.endswith(".yaml"):

src/lighteval/main_nanotron.py

@@ -35,8 +35,8 @@
     InferenceProvidersClient,
     InferenceProvidersModelConfig,
 )
+from lighteval.models.endpoints.litellm_model import LiteLLMClient, LiteLLMModelConfig
 from lighteval.models.endpoints.tgi_model import ModelClient, TGIModelConfig
-from lighteval.models.litellm_model import LiteLLMClient, LiteLLMModelConfig
 from lighteval.models.sglang.sglang_model import SGLangModel, SGLangModelConfig
 from lighteval.models.transformers.adapter_model import AdapterModel, AdapterModelConfig
 from lighteval.models.transformers.delta_model import DeltaModel, DeltaModelConfig

src/lighteval/models/litellm_model.py renamed to src/lighteval/models/endpoints/litellm_model.py

@@ -19,23 +19,23 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
-
 # ruff: noqa: C901
 import logging
 import os
 import time
-from typing import List, Optional, Tuple, Type, Union
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import torch.nn.functional as F
 import transformers
 from datasets.download.streaming_download_manager import xPath
+from pydantic import BaseModel
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 from tqdm import tqdm
 from transformers import AutoTokenizer, BatchEncoding
 
-from lighteval.config.lighteval_config import FullNanotronConfig
 from lighteval.data import (
     GenDistributedSampler,
     GenerativeTaskDatasetNanotron,
@@ -69,7 +69,10 @@
 if is_nanotron_available():
     from nanotron import distributed as dist
     from nanotron import logging
+    from nanotron.config import GeneralArgs, ModelArgs, TokenizerArgs
+    from nanotron.config.parallelism_config import ParallelismArgs
     from nanotron.generation.decode import decode_tokenized
+    from nanotron.generation.sampler import SamplerType
     from nanotron.logging import human_format, log_rank
     from nanotron.models import build_model
     from nanotron.parallel.context import ParallelContext
@@ -83,6 +86,82 @@
 
     logger = logging.get_logger(__name__)
 
+DEFAULT_GENERATION_SEED = 42
+
+
+class GenerationArgs(BaseModel):
+    sampler: Optional["SamplerType"] = None
+    temperature: Optional[float] = None
+    top_k: Optional[int] = None
+    top_p: Optional[float] = None
+    n_samples: Optional[int] = None
+    eos: Optional[str] = None
+    seed: Optional[int] = None
+    use_cache: Optional[bool] = False
+
+    def __post_init__(self):
+        if self.seed is None:
+            self.seed = DEFAULT_GENERATION_SEED
+
+
+@dataclass
+class LightEvalLoggingArgs:
+    """Arguments related to logging for LightEval"""
+
+    output_dir: str
+    results_path_template: str | None = None
+    save_details: bool = True
+    push_to_hub: bool = False
+    push_to_tensorboard: bool = False
+    public_run: bool = False
+    results_org: str | None = None
+    tensorboard_metric_prefix: str = "eval"
+
+
+@dataclass
+class LightEvalTasksArgs:
+    """Arguments related to tasks for LightEval"""
+
+    tasks: str
+    custom_tasks: Optional[str] = None
+    max_samples: Optional[int] = None
+    num_fewshot_seeds: Optional[int] = None
+
+    dataset_loading_processes: int = 8
+    multichoice_continuations_start_space: Optional[bool] = None
+    pairwise_tokenization: bool = False
+
+
+@dataclass
+class LightEvalConfig:
+    """Arguments related to running LightEval on checkpoints.
+
+    All is optional because you can also use this class to later supply arguments to override
+    the saved config when running LightEval after training.
+    """
+
+    logging: LightEvalLoggingArgs
+    tasks: LightEvalTasksArgs
+    parallelism: "ParallelismArgs"
+    batch_size: int = 0
+    generation: Optional[Union[GenerationArgs, Dict[str, GenerationArgs]]] = None
+
+
+@dataclass
+class FullNanotronConfig:
+    lighteval_config: LightEvalConfig
+    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()
+
 
 class NanotronLightevalModel(LightevalModel):
     # Default max sequence length setting for when no `max_length` is provided