Update app.py

app.py

@@ -1,18 +1,15 @@
 import gradio as gr
 import re
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+from transformers import MarianMTModel, MarianTokenizer
 
 # -------------------------
-# Load model
+# Fast Model
 # -------------------------
-MODEL_NAME = "facebook/nllb-200-distilled-600M"
+MODEL_NAME = "Helsinki-NLP/opus-mt-en-hi"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
 
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    MODEL_NAME
-)
+model = MarianMTModel.from_pretrained(MODEL_NAME)
 
 # -------------------------
 # Parse SRT
@@ -42,76 +39,78 @@ def parse_srt(srt_text):
 
 
 # -------------------------
-# Shorten text
+# Shorten Hindi
 # -------------------------
 def shorten_text(text, max_len):
 
+    if len(text) <= max_len:
+        return text
+
     words = text.split()
 
     while len(text) > max_len and len(words) > 1:
-        words.pop(-1)
+        words.pop()
         text = " ".join(words)
 
     return text
 
 
 # -------------------------
-# Translate line
+# Batch Translate
 # -------------------------
-def translate_line(text, max_len):
-
-    tokenizer.src_lang = "eng_Latn"
+def batch_translate(texts):
 
-    encoded = tokenizer(
-        text,
-        return_tensors="pt"
+    inputs = tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True
     )
 
-    generated_tokens = model.generate(
-        **encoded,
-        forced_bos_token_id=tokenizer.convert_tokens_to_ids("hin_Deva"),
-        max_length=128
+    translated = model.generate(
+        **inputs,
+        max_new_tokens=64
     )
 
-    hindi = tokenizer.batch_decode(
-        generated_tokens,
+    outputs = tokenizer.batch_decode(
+        translated,
         skip_special_tokens=True
-    )[0]
-
-    hindi = hindi.strip()
-
-    # Length control
-    if len(hindi) > max_len:
-        hindi = shorten_text(hindi, max_len)
+    )
 
-    return hindi
+    return outputs
 
 
 # -------------------------
-# Main translation
+# Main Function
 # -------------------------
 def translate_srt(srt_text):
 
     subtitles = parse_srt(srt_text)
 
-    output = []
+    english_texts = [
+        sub["text"] for sub in subtitles
+    ]
+
+    # FAST batch translation
+    hindi_texts = batch_translate(
+        english_texts
+    )
 
-    for sub in subtitles:
+    output = []
 
-        english = sub["text"]
+    for sub, hindi in zip(subtitles, hindi_texts):
 
-        english_len = len(english)
+        english_len = len(sub["text"])
 
-        # Hindi max 130%
-        max_hindi_len = int(english_len * 1.3)
+        # 130% rule
+        max_hindi_len = int(
+            english_len * 1.3
+        )
 
-        try:
-            hindi = translate_line(
-                english,
-                max_hindi_len
-            )
-        except:
-            hindi = english
+        hindi = shorten_text(
+            hindi.strip(),
+            max_hindi_len
+        )
 
         block = (
             f'{sub["index"]}\n'
@@ -125,7 +124,7 @@ def translate_srt(srt_text):
 
 
 # -------------------------
-# Gradio UI
+# UI
 # -------------------------
 demo = gr.Interface(
     fn=translate_srt,
@@ -137,8 +136,8 @@ demo = gr.Interface(
         lines=20,
         label="Hindi SRT"
     ),
-    title="English → Hindi SRT Translator",
-    description="Translate English SRT subtitles to Hindi with timestamp preservation and subtitle length control."
+    title="Fast English → Hindi SRT Translator",
+    description="Ultra-fast subtitle translation with timestamp preservation and subtitle length control."
 )
 
 demo.launch()