Tiny improvements to endpoint_model.py, base_model.py,...

src/lighteval/models/endpoints/endpoint_model.py

@@ -395,7 +395,7 @@ def _async_process_request(
             grammar=grammar,
             max_new_tokens=max_tokens,
             stop_sequences=stop_tokens,
-            # truncate=,
+            truncate=self.max_length,
         )
 
         return generated_text
@@ -416,7 +416,7 @@ def _process_request(
             grammar=grammar,
             max_new_tokens=max_tokens,
             stop_sequences=stop_tokens,
-            # truncate=,
+            truncate=self.max_length,
         )
 
         return generated_text
@@ -492,7 +492,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,
@@ -514,12 +514,15 @@ 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(
                         GenerativeResponse(
                             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,
                         )
                     )
@@ -538,7 +541,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,
@@ -589,7 +592,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,7 +116,9 @@ def load_openai_model(config: OpenAIModelConfig, env_config: EnvConfig):
     return model
 
 
-def load_model_with_inference_endpoints(config: InferenceEndpointModelConfig, env_config: EnvConfig):
+def load_model_with_inference_endpoints(
+    config: Union[InferenceEndpointModelConfig, ServerlessEndpointModelConfig], env_config: EnvConfig
+):
     logger.info("Spin up model using inference endpoint.")
     model = InferenceEndpointModel(config=config, env_config=env_config)
     return model

src/lighteval/models/transformers/base_model.py

@@ -826,10 +826,7 @@ def greedy_until(
                     input_ids=tokenized["input_ids"],
                     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
-                        for c in context
-                    ],
+                    truncated=[max(len(c) - tokenized["input_ids"].shape[1], 0) for c in context],
                     padded=[sum(mask == 0) for mask in tokenized["attention_mask"]],
                 )