Merge branch 'add-custom-model' into dev

docs/source/package_reference/models.mdx

@@ -21,7 +21,7 @@
 ## Endpoints-based Models
 ### InferenceEndpointModel
 [[autodoc]] models.endpoints.endpoint_model.InferenceEndpointModelConfig
-[[autodoc]] models.endpoints.endpoint_model.InferenceModelConfig
+[[autodoc]] models.endpoints.endpoint_model.ServerlessEndpointModelConfig
 [[autodoc]] models.endpoints.endpoint_model.InferenceEndpointModel
 
 ### TGI ModelClient

examples/custom_models/local_mt_model.py

@@ -24,6 +24,7 @@
 from typing import Optional
 
 import pycountry
+import torch
 from tqdm import tqdm
 from transformers import (
     AutoModelForSeq2SeqLM,
@@ -86,6 +87,7 @@ def __init__(self, config, env_config) -> None:
         self.model = config.model
         self.model_definition_file_path = config.model_definition_file_path
         self.batch_size = 32
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         self.model_info = ModelInfo(
             model_name=config.model,
@@ -106,6 +108,9 @@ def __init__(self, config, env_config) -> None:
         else:
             raise ValueError(f"Unsupported model: {config.model}")
 
+        self._model.to(self.device)
+        self._model.eval()
+
     def _convert_to_iso3(self, lang_code: str) -> str:
         """Convert 2-letter ISO code to 3-letter ISO code."""
         try:
@@ -166,7 +171,9 @@ def get_langs(task_name: str) -> tuple[str, str]:
             current_requests = dataset.sorted_data[split_start:split_end]
 
             # Process in batches
-            for batch_idx in range(0, len(current_requests), batch_size):
+            for batch_idx in tqdm(
+                range(0, len(current_requests), batch_size), desc="Batches", position=1, disable=False
+            ):
                 batch = current_requests[batch_idx : batch_idx + batch_size]
 
                 # Batch tokenize all inputs together instead of concatenating pre-tokenized inputs
@@ -178,15 +185,19 @@ def get_langs(task_name: str) -> tuple[str, str]:
                 if self.model_type == "seamless-4mt":
                     tokenizer_kwargs["src_lang"] = src_lang
 
-                input_ids, attention_mask = self._tokenizer(**tokenizer_kwargs).values()
+                input_ids, attention_mask = self._tokenizer(**tokenizer_kwargs).to(self.device).values()
 
                 tgt_langs = [get_langs(r.task_name)[1] for r in batch]
                 assert set(tgt_langs) == {tgt_langs[0]}, "All target languages must be the same"
 
+                generation_sizes = [r.generation_size for r in batch]
+                assert set(generation_sizes) == {generation_sizes[0]}, "All generation sizes must be the same"
+
                 # Use unpacked values directly
                 generate_kwargs = {
                     "input_ids": input_ids,
                     "attention_mask": attention_mask,
+                    "max_new_tokens": generation_sizes[0],
                 }
                 if self.model_type == "seamless-4mt":
                     generate_kwargs["tgt_lang"] = tgt_langs[0]
@@ -212,7 +223,7 @@ def tokenizer(self):
         return self._tokenizer
 
     def tok_encode(self, str_to_encode: str | list[str], add_special_tokens: Optional[bool] = None) -> TokenSequence:
-        return self._tokenizer(text=str_to_encode, add_special_tokens=add_special_tokens or False)
+        return self._tokenizer(text=str_to_encode, add_special_tokens=add_special_tokens or False).to(self.device)
 
     @property
     def add_special_tokens(self) -> bool:

src/lighteval/main_endpoint.py

@@ -146,6 +146,13 @@ def inference_endpoint(
         str, Argument(help="Path to model config yaml file. (examples/model_configs/endpoint_model.yaml)")
     ],
     tasks: Annotated[str, Argument(help="Comma-separated list of tasks to evaluate on.")],
+    free_endpoint: Annotated[
+        bool,
+        Option(
+            help="Use serverless free endpoints instead of spinning up your own inference endpoint.",
+            rich_help_panel=HELP_PANEL_NAME_4,
+        ),
+    ] = False,
     # === Common parameters ===
     use_chat_template: Annotated[
         bool, Option(help="Use chat template for evaluation.", rich_help_panel=HELP_PANEL_NAME_4)
@@ -200,9 +207,7 @@ def inference_endpoint(
     """
 
     from lighteval.logging.evaluation_tracker import EvaluationTracker
-    from lighteval.models.endpoints.endpoint_model import (
-        InferenceEndpointModelConfig,
-    )
+    from lighteval.models.endpoints.endpoint_model import InferenceEndpointModelConfig, ServerlessEndpointModelConfig
     from lighteval.pipeline import EnvConfig, ParallelismManager, Pipeline, PipelineParameters
 
     env_config = EnvConfig(token=TOKEN, cache_dir=cache_dir)
@@ -220,10 +225,10 @@ def inference_endpoint(
     parallelism_manager = ParallelismManager.NONE  # since we're using inference endpoints in remote
 
     # Find a way to add this back
-    # if config["base_params"].get("endpoint_name", None):
-    #    return InferenceModelConfig(model=config["base_params"]["endpoint_name"])
-
-    model_config = InferenceEndpointModelConfig.from_path(model_config_path)
+    if free_endpoint:
+        model_config = ServerlessEndpointModelConfig.from_path(model_config_path)
+    else:
+        model_config = InferenceEndpointModelConfig.from_path(model_config_path)
 
     pipeline_params = PipelineParameters(
         launcher_type=parallelism_manager,

src/lighteval/models/endpoints/endpoint_model.py

@@ -75,10 +75,18 @@
 
 
 @dataclass
-class InferenceModelConfig:
-    model: str
+class ServerlessEndpointModelConfig:
+    model_name: str
     add_special_tokens: bool = True
 
+    @classmethod
+    def from_path(cls, path: str) -> "ServerlessEndpointModelConfig":
+        import yaml
+
+        with open(path, "r") as f:
+            config = yaml.safe_load(f)["model"]
+        return cls(**config["base_params"])
+
 
 @dataclass
 class InferenceEndpointModelConfig:
@@ -150,7 +158,7 @@ class InferenceEndpointModel(LightevalModel):
     """
 
     def __init__(  # noqa: C901
-        self, config: Union[InferenceEndpointModelConfig, InferenceModelConfig], env_config: EnvConfig
+        self, config: Union[InferenceEndpointModelConfig, ServerlessEndpointModelConfig], env_config: EnvConfig
     ) -> None:
         self.reuse_existing = getattr(config, "reuse_existing", False)
         self._max_length = None
@@ -280,10 +288,10 @@ def __init__(  # noqa: C901
         else:  # Free inference client
             self.endpoint = None
             self.endpoint_name = None
-            self.name = config.model
+            self.name = config.model_name
             self.revision = "default"
-            self.async_client = AsyncInferenceClient(model=config.model, token=env_config.token)
-            self.client = InferenceClient(model=config.model, token=env_config.token)
+            self.async_client = AsyncInferenceClient(model=config.model_name, token=env_config.token)
+            self.client = InferenceClient(model=config.model_name, token=env_config.token)
 
         self.use_async = True  # set to False for debug - async use is faster
 
@@ -293,7 +301,7 @@ def __init__(  # noqa: C901
         self.model_info = ModelInfo(
             model_name=self.name,
             model_sha=self.revision,
-            model_dtype=config.model_dtype or "default",
+            model_dtype=getattr(config, "model_dtype", "default"),
             model_size=-1,
         )
 
@@ -545,7 +553,12 @@ def loglikelihood(
                     cont_toks = torch.tensor(cur_request.tokenized_continuation)
                     len_choice = len(cont_toks)
 
-                    logits = [t.logprob for t in response.details.prefill[-len_choice:] if t.logprob is not None]
+                    if self.endpoint:  # inference endpoint
+                        logits = [
+                            t.logprob for t in response.details.prefill[-len_choice:] if t.logprob is not None
+                        ]  # to check
+                    else:  # serverless endpoint
+                        logits = [t.logprob for t in response.details.tokens[-len_choice:] if t.logprob is not None]
 
                     greedy_tokens = torch.tensor(logits).argmax(dim=-1)
                     max_equal = (greedy_tokens == cont_toks).all().squeeze(0)

src/lighteval/models/model_loader.py

@@ -29,7 +29,7 @@
 from lighteval.models.endpoints.endpoint_model import (
     InferenceEndpointModel,
     InferenceEndpointModelConfig,
-    InferenceModelConfig,
+    ServerlessEndpointModelConfig,
 )
 from lighteval.models.endpoints.openai_model import OpenAIClient, OpenAIModelConfig
 from lighteval.models.endpoints.tgi_model import ModelClient, TGIModelConfig
@@ -87,7 +87,7 @@ def load_model(  # noqa: C901
     if isinstance(config, TGIModelConfig):
         return load_model_with_tgi(config)
 
-    if isinstance(config, InferenceEndpointModelConfig) or isinstance(config, InferenceModelConfig):
+    if isinstance(config, InferenceEndpointModelConfig) or isinstance(config, ServerlessEndpointModelConfig):
         return load_model_with_inference_endpoints(config, env_config=env_config)
 
     if isinstance(config, BaseModelConfig):

tests/models/endpoints/test_endpoint_model.py

@@ -53,7 +53,7 @@ class TestInferenceEndpointModelConfig:
                 },
             ),
             (
-                "examples/model_configs/endpoint_model_lite.yaml",
+                "examples/model_configs/serverless_model.yaml",
                 {
                     "model_name": "meta-llama/Llama-3.1-8B-Instruct",
                     # Defaults: