Update app.py

app.py

@@ -13,17 +13,19 @@ import subprocess
 import torch 
 import bitsandbytes
 import scipy
+from googletrans import Translator
 
 ZipFile("ffmpeg.zip").extractall()
 st = os.stat('ffmpeg')
 os.chmod('ffmpeg', st.st_mode | stat.S_IEXEC)
 
-with open('language_codes.json', 'r') as f:
-    lang_codes = json.load(f)
+with open('google_lang_codes.json', 'r') as f:
+    google_lang_codes = json.load(f)
 
+translator = Translator()
 
-tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-3.3B")
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-3.3B")
+#tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-3.3B")
+#model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-3.3B")
 whisper_model = WhisperModel("large-v2", device="cuda", compute_type="float16")
 
 print("cwd", os.getcwd())
@@ -75,21 +77,16 @@ def process_video(Video, target_language):
         f.seek(0)
     
         # Translating the SRT from Whisper with NLLB.
-        flores_code = lang_codes.get(target_language, "eng_Latn")
+        target_language_code = google_lang_codes.get(target_language, "en")
         paragraph = ""
         for line in f:
             if line.strip().isnumeric() or "-->" in line:
-                if paragraph:
-                    inputs = tokenizer(paragraph, return_tensors="pt")
-                    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[flores_code], max_length=100)
-                    translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-                    translated_lines.append(translated_text + "\n")
-                    paragraph = ""
                 translated_lines.append(line)
             elif line.strip() != "":
-                paragraph += " " + line.strip()
+                translated_text = translator.translate(line.strip(), dest=target_language_code).text
+                translated_lines.append(translated_text + "\n")
             else:
-                translated_lines.append("\n")
+                translated_lines.append("\n"))
     
         # Move the file pointer to the beginning of the file and truncate it.
         f.seek(0)