Update app.py

app.py

@@ -1,13 +1,17 @@
 import gradio as gr
 import re
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 # -------------------------
-# Load Model
+# Load model
 # -------------------------
-translator = pipeline(
-    "text2text-generation",
-    model="facebook/nllb-200-distilled-600M"
+MODEL_NAME = "facebook/nllb-200-distilled-600M"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    MODEL_NAME
 )
 
 # -------------------------
@@ -56,19 +60,25 @@ def shorten_text(text, max_len):
 # -------------------------
 def translate_line(text, max_len):
 
-    prompt = f"""
-Translate this English text to spoken Hindi.
+    tokenizer.src_lang = "eng_Latn"
 
-English:
-{text}
-"""
+    encoded = tokenizer(
+        text,
+        return_tensors="pt"
+    )
 
-    result = translator(
-        prompt,
-        max_new_tokens=128
+    generated_tokens = model.generate(
+        **encoded,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids("hin_Deva"),
+        max_length=128
     )
 
-    hindi = result[0]["generated_text"].strip()
+    hindi = tokenizer.batch_decode(
+        generated_tokens,
+        skip_special_tokens=True
+    )[0]
+
+    hindi = hindi.strip()
 
     # Length control
     if len(hindi) > max_len:
@@ -78,7 +88,7 @@ English:
 
 
 # -------------------------
-# Main Translation
+# Main translation
 # -------------------------
 def translate_srt(srt_text):
 
@@ -92,7 +102,7 @@ def translate_srt(srt_text):
 
         english_len = len(english)
 
-        # Hindi can be max 130%
+        # Hindi max 130%
         max_hindi_len = int(english_len * 1.3)
 
         try:
@@ -115,7 +125,7 @@ def translate_srt(srt_text):
 
 
 # -------------------------
-# UI
+# Gradio UI
 # -------------------------
 demo = gr.Interface(
     fn=translate_srt,
@@ -128,7 +138,7 @@ demo = gr.Interface(
         label="Hindi SRT"
     ),
     title="English → Hindi SRT Translator",
-    description="Translate English subtitles to Hindi with timestamp preservation and length control."
+    description="Translate English SRT subtitles to Hindi with timestamp preservation and subtitle length control."
 )
 
 demo.launch()