Refactor translator to use a more accurate language model

Author: andrewwkimm

examples/subtitles/sample_srt_file.srt

@@ -1,15 +1,15 @@
 1
 00:00:00,000 --> 00:00:06,360
-When we got out of the long tunnel of the border, it was an eyebled.
+It was a snowflake when we came out of the long tunnel.
 
 2
 00:00:07,800 --> 00:00:10,740
-The bottom of the night has been changed.
+The bottom of the night has been hallucinating.
 
 3
 00:00:12,320 --> 00:00:15,360
-The train stopped at the signal station.
+The train stopped in front of the signal station.
 
 4
 00:00:16,980 --> 00:00:23,540
-On the other side, a virgin approached and opened a window in front of the box.
+A virgin came from the other side of the seat and opened the window in front of the movie theater.

examples/subtitles/sample_vtt_file.vtt

@@ -1,19 +1,19 @@
 WEBVTT
 
 00:00:00.000 --> 00:00:04.880
-It's okay, son.
+I'm all right, son.
 
 00:00:04.880 --> 00:00:09.240
-There's no fighting.
+I'm not going to fight.
 
 00:00:09.240 --> 00:00:16.440
-No, I don't think you should be a brother.
+No, I should call him brother.
 
 00:00:16.440 --> 00:00:18.060
 What?
 
 00:00:18.060 --> 00:00:21.280
-It's over.
+It's over now.
 
 00:00:21.280 --> 00:00:25.140
-You can throw away your gun and live.
+You can throw your guns away and live.

examples/videos/sample_japanese_video_translated.mp4

@@ -2,7 +2,7 @@
 
 import logging
 
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 logging.getLogger("transformers").setLevel(logging.ERROR)
 log = logging.getLogger(__name__)
@@ -28,13 +28,29 @@ def translate_transcript(transcript: dict, target_language: str) -> list:
 
 def translate_text(text: str, source_language: str, target_language: str) -> str:
     """Translates source language to target language."""
-    model_name = f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
-    model = MarianMTModel.from_pretrained(model_name)
-
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    tokenized_text = tokenizer([text], return_tensors="pt")
-
-    translation = model.generate(**tokenized_text)
+    model_name = "facebook/nllb-200-distilled-600M"
+
+    languages = {
+        "ko": "kor_Hang",
+        "ja": "jpn_Jpan",
+        "en": "eng_Latn",
+        "es": "spa_Latn",
+        "fr": "fra_Latn",
+        "de": "deu_Latn",
+        "zh": "zho_Hans",
+    }
+
+    source_language_code = languages.get(source_language, source_language)
+    target_language_code = languages.get(target_language, target_language)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=source_language_code)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    tokenized_text = tokenizer(text, return_tensors="pt")
+    translation = model.generate(
+        **tokenized_text,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids(target_language_code),
+    )
 
     translated_text = tokenizer.decode(translation[0], skip_special_tokens=True)
     return translated_text

examples/videos/sample_korean_video_translated.mp4

@@ -7,15 +7,15 @@ def test_translate_transcript() -> None:
     """Tests that the transcript was translated properly."""
     sample_transcript = {
         "segments": [
-            {"text": "안녕하세요."},
+            {"text": "시대를 초월한 마음."},
             {"text": "음식."},
         ],
         "language": "ko",
     }
 
     actual = translate_transcript(sample_transcript, target_language="en")
     expected = [
-        {"text": "Hello."},
+        {"text": "A mind beyond its time."},
         {"text": "Food."},
     ]