fixed bug releated to NLLB language token being duplicated

CENIA-DEV · Dec 17, 2024 · 2d205bd · 2d205bd
1 parent 4ab60f5
commit 2d205bd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/tokenizer_wrapper.py b/src/tokenizer_wrapper.py
@@ -21,7 +21,7 @@ def __call__(
     ) -> dict:
         if isinstance(self.tokenizer, NllbTokenizerFast):
             input_ids = self.tokenizer.encode(text, **self.kwargs)
-            input_ids[0] = self.tokenizer.convert_tokens_to_ids(src_lang)
+            input_ids[:, 0] = self.tokenizer.convert_tokens_to_ids(src_lang)
         elif isinstance(self.tokenizer, T5TokenizerFast):
             input_ids = self.tokenizer.encode(src_lang + text, **self.kwargs)