Make it compatible with multilingual-lighteval.

swiss-ai · Sep 8, 2024 · e9fb021 · e9fb021
1 parent 60f3047
commit e9fb021
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/nanotron/config/lighteval_config.py b/src/nanotron/config/lighteval_config.py
@@ -51,6 +51,7 @@ def __post_init__(self):
 class LightEvalTasksArgs:
     """Arguments related to tasks for LightEval"""
 
+    langs: Optional[str] = None
     tasks: Optional[str] = None
     custom_tasks: Optional[str] = None
     max_samples: Optional[int] = None

diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py
@@ -757,7 +757,7 @@ def forward(
         self,
         input_ids: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
         input_mask: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
-        lang_code: Union[torch.Tensor, TensorPointer],  # [batch_size, 1]
+        lang_code: Union[torch.Tensor, TensorPointer]=None,  # [batch_size, 1]
     ):
         return self.forward_with_hidden_states(input_ids=input_ids, input_mask=input_mask, lang_code=lang_code)[0]