Fix for nllb (#64)

Galoist · web-flow · commit d9f1fcd6d7d9 · 2024-09-02T17:08:29.000-04:00
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)
 

Original file line number	Diff line number	Diff line change
`@@ -170,7 +170,7 @@ def translate(`
`170`	`170`	`batch_size = len(text)`
`171`	`171`
`172`	`172`	`generation_options.setdefault(`
`173`		`- "forced_bos_token_id", self._tokenizer.lang_code_to_id[target]`
	`173`	`+ "forced_bos_token_id", self._tokenizer.convert_tokens_to_ids(target)`
`174`	`174`	`)`
`175`	`175`	`generation_options.setdefault("max_new_tokens", 512)`
`176`	`176`