Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import srt
+import torch
 import os
 
 # --- Configuration ---
@@ -8,30 +9,45 @@ MODEL_CHECKPOINT = "facebook/nllb-200-distilled-600M"
 SRC_LANG = "eng_Latn"
 TGT_LANG = "arb_Arab"
 
-# --- Load Model ---
+# --- Load Model Directly (No Pipeline) ---
 print("Loading model...")
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
-
-translator = pipeline(
-    "translation", 
-    model=model, 
-    tokenizer=tokenizer, 
-    src_lang=SRC_LANG, 
-    tgt_lang=TGT_LANG, 
-    device=-1
-)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT)
+print("Model loaded!")
 
 def batch_translate(texts, batch_size=8):
+    """
+    Directly uses the model to translate without the pipeline abstraction.
+    """
     results = []
+    
+    # 1. Set the source language for the tokenizer
+    tokenizer.src_lang = SRC_LANG
+
     for i in range(0, len(texts), batch_size):
         batch = texts[i : i + batch_size]
-        outputs = translator(batch, max_length=400)
-        results.extend([out['translation_text'] for out in outputs])
+        
+        # 2. Tokenize the batch
+        inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        
+        # 3. Generate translation (Force the target language ID)
+        # NLLB requires forcing the 'bos_token_id' to the target language
+        forced_bos_token_id = tokenizer.lang_code_to_id[TGT_LANG]
+        
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **inputs,
+                forced_bos_token_id=forced_bos_token_id,
+                max_length=512
+            )
+            
+        # 4. Decode the results
+        batch_results = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+        results.extend(batch_results)
+        
     return results
 
 def process_srt(filepath):
-    # Gradio 4 passes the file path as a string
     if filepath is None:
         return None
     
@@ -43,12 +59,15 @@ def process_srt(filepath):
     except Exception as e:
         return f"Error parsing SRT: {str(e)}"
 
+    # Translate content
     texts_to_translate = [sub.content for sub in subtitles]
     translated_texts = batch_translate(texts_to_translate)
     
+    # Update subtitles
     for sub, trans_text in zip(subtitles, translated_texts):
         sub.content = trans_text
 
+    # Save output
     output_path = "translated_subtitles.srt"
     with open(output_path, 'w', encoding='utf-8') as f:
         f.write(srt.compose(subtitles))
@@ -56,7 +75,7 @@ def process_srt(filepath):
     return output_path
 
 # --- Gradio Interface ---
-with gr.Blocks(title="SRT Translator") as demo:
+with gr.Blocks(title="NLLB SRT Translator") as demo:
     gr.Markdown("# 🇬🇧 English to 🇸🇦 Arabic SRT Translator")
     
     with gr.Row():