Fix for nllb (#64)

xhluca · Sep 2, 2024 · d9f1fcd · d9f1fcd
1 parent 27c5e9d
commit d9f1fcd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dl_translate/_translation_model.py b/dl_translate/_translation_model.py
@@ -170,7 +170,7 @@ def translate(
             batch_size = len(text)
 
         generation_options.setdefault(
-            "forced_bos_token_id", self._tokenizer.lang_code_to_id[target]
+            "forced_bos_token_id", self._tokenizer.convert_tokens_to_ids(target)
         )
         generation_options.setdefault("max_new_tokens", 512)