Translate now returns a single string instead of a list

KennethEnevoldsen · Jan 23, 2024 · e1cefd9 · e1cefd9
1 parent c75629a
commit e1cefd9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/seb/registered_models/translate_e5_models.py b/src/seb/registered_models/translate_e5_models.py
@@ -26,7 +26,9 @@ def translate(self, sentence: str, src_lang: str) -> str:
         gen_tokens = self.trans_model.generate(
             **encoded_sent, forced_bos_token_id=self.trans_tokenizer.get_lang_id("en")
         )
-        return self.trans_tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)
+        return self.trans_tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)[
+            0
+        ]
 
     def encode(
         self,