Update off_topic.py

off_topic.py

@@ -21,7 +21,15 @@ class Translator:
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_id)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(self.device)
-        self.bos_token_map = self.tokenizer.get_lang_id if hasattr(self.tokenizer, "get_lang_id") else self.tokenizer.lang_code_to_id
+    
+    @property
+    def _bos_token_attr(self):
+        if hasattr(self.tokenizer, "get_lang_id"):
+            return self.tokenizer.get_lang_id
+        elif hasattr(self.tokenizer, "lang_code_to_id"):
+            return self.tokenizer.lang_code_to_id
+        else:
+            return
 
     @property
     def _language_code_mapper(self):
@@ -33,12 +41,20 @@ class Translator:
             return {"en": "en",
                     "es": "es",
                     "pt": "pt"}
+        else:
+            return {"en": "eng",
+                    "es": "spa",
+                    "pt": "por"}
 
     def translate(self, texts: List[str], src_lang: str, dest_lang: str = "en", max_length: int = 100):
         self.tokenizer.src_lang = self._language_code_mapper[src_lang]
         inputs = self.tokenizer(texts, return_tensors="pt").to(self.device)
+        if "opus" in self.model_id.lower():
+            forced_bos_token_id = None
+        else:
+            forced_bos_token_id = self._bos_token_attr[self._language_code_mapper["en"]]
         translated_tokens = self.model.generate(
-            **inputs, forced_bos_token_id=self.bos_token_map["eng_Latn"], max_length=max_length
+            **inputs, forced_bos_token_id=forced_bos_token_id, max_length=max_length
         )
         return self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)