Do the improvements

src/lighteval/data.py

@@ -22,6 +22,7 @@
 
 from typing import Iterator
 
+import math
 import torch
 from torch.utils.data import Dataset
 from torch.utils.data.distributed import DistributedSampler, T_co
@@ -80,7 +81,7 @@ def init_split_limits(self, num_dataset_splits):
             )
             num_dataset_splits = 1
 
-        split_size = self.total_size // num_dataset_splits + 1
+        split_size = math.ceil(self.total_size / num_dataset_splits)
         splits_indices = [
             (ix * split_size, min((ix + 1) * split_size, self.total_size)) for ix in range(num_dataset_splits)
         ]

src/lighteval/models/base_model.py

@@ -565,7 +565,7 @@ def greedy_until(
                     input_lengths=[len(item == 1) for item in tokenized["attention_mask"]],
                     input_mask=tokenized["attention_mask"],
                     truncated=[
-                        len(c) - tokenized["input_ids"].shape[1] if len(c) > tokenized["input_ids"].shape[1] else 0
+                        max(len(c) - tokenized["input_ids"].shape[1], 0)
                         for c in context
                     ],
                     padded=[sum(mask == 0) for mask in tokenized["attention_mask"]],

src/lighteval/models/endpoint_model.py

@@ -86,7 +86,7 @@ def __init__(
                     custom_image={
                         "health_route": "/health",
                         "env": {
-                            # Documentaiton: https://huggingface.co/docs/text-generation-inference/en/basic_tutorials/launcher
+                            # Documentation: https://huggingface.co/docs/text-generation-inference/en/basic_tutorials/launcher
                             "MAX_BATCH_PREFILL_TOKENS": "2048",
                             "MAX_INPUT_LENGTH": "2047",
                             "MAX_TOTAL_TOKENS": "2048",
@@ -162,7 +162,7 @@ def _async_process_request(
             decoder_input_details=True,
             max_new_tokens=max_tokens,
             stop_sequences=stop_tokens,
-            # truncate=,
+            truncate=self.max_length,
         )
 
         return generated_text
@@ -176,7 +176,7 @@ def _process_request(self, context: str, stop_tokens: list[str], max_tokens: int
             decoder_input_details=True,
             max_new_tokens=max_tokens,
             stop_sequences=stop_tokens,
-            # truncate=,
+            truncate=self.max_length,
         )
 
         return generated_text
@@ -250,7 +250,7 @@ def greedy_until(
 
         for _, _ in tqdm(
             dataset.splits_start_end_iterator(),
-            total=self.DATASET_SPLITS,
+            total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
@@ -272,12 +272,13 @@ def greedy_until(
                     responses = asyncio.run(self._async_process_batch_generate(batch))
                 else:
                     responses = self._process_batch_generate(batch)
-                for response in responses:
+                for i, response in enumerate(responses):
                     results.append(
                         GenerateReturn(
                             result=response.generated_text,
                             logits=[item.logprob for item in response.details.prefill] if returns_logits else None,
-                            truncated_tokens_count=-1,
+                            generated_tokens=[token.id for token in response.details.tokens],
+                            truncated_tokens_count=max(len(self.tokenizer.encode(batch[i].context))-self.max_length, 0),
                             padded_tokens_count=-1,
                         )
                     )
@@ -296,7 +297,7 @@ def loglikelihood(
 
         for _, _ in tqdm(
             dataset.splits_start_end_iterator(),
-            total=self.DATASET_SPLITS,
+            total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
@@ -342,7 +343,7 @@ def loglikelihood_rolling(
 
         for _, _ in tqdm(
             dataset.splits_start_end_iterator(),
-            total=self.DATASET_SPLITS,
+            total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,

src/lighteval/models/model_loader.py

@@ -116,13 +116,14 @@ def load_model_with_tgi(config: TGIModelConfig):
     return model, model_info
 
 
-def load_model_with_inference_endpoints(config: InferenceEndpointModelConfig, env_config: EnvConfig):
+def load_model_with_inference_endpoints(config: Union[InferenceEndpointModelConfig, InferenceModelConfig],
+                                        env_config: EnvConfig):
     hlog("Spin up model using inference endpoint.")
     model = InferenceEndpointModel(config=config, env_config=env_config)
     model_info = ModelInfo(
         model_name=model.name,
         model_sha=model.revision,
-        model_dtype=config.model_dtype or "default",
+        model_dtype=getattr(config, "model_dtype", "default"),
         model_size=-1,
     )
     return model, model_info

src/lighteval/tasks/lighteval_task.py

@@ -468,7 +468,7 @@ def construct_requests(
                     example_index=document_id_seed,
                     request_index=0,
                     context=context,
-                    stop_sequence=self.stop_sequence,
+                    stop_sequence=self.stop_sequence or [],
                     generation_size=self.generation_size,
                     num_samples=max(self.num_samples),  # If we have several samplings to apply, we use the max
                     use_logits=use_logits,
@@ -502,7 +502,7 @@ def construct_requests(
                     example_index=document_id_seed,
                     request_index=0,
                     context=context,
-                    stop_sequence=self.stop_sequence,
+                    stop_sequence=self.stop_sequence or [],
                     generation_size=self.generation_size,
                 )
             ]