Update translator.py

translator.py

@@ -5,16 +5,20 @@ Fallback chain:
   1. Helsinki-NLP  — dedicated per-language model (best for te/hi/ta/kn)
   2. NLLB-1.3B     — covers all other languages
   3. Google Translate — last resort fallback
+
 LANGUAGE ACCURACY (after upgrade):
   Telugu  (en→te): 85% (was 82% with NLLB)
   Hindi   (en→hi): 87% (was 84% with NLLB)
   Tamil   (en→ta): 84% (was 81% with NLLB)
   Kannada (en→kn): 83% (was 80% with NLLB)
   Others         : NLLB handles (unchanged)
+
 FIXES IN THIS VERSION:
+  - FIX: Removed "translation" / "text2text-generation" pipeline guessing.
+    Newer transformers dropped both for Helsinki models.
+    Now uses AutoTokenizer + MarianMTModel directly (always works).
+  - FIX: Removed duplicate INDIC_TO_EN block in translate()
   - Pre-loads Telugu + Hindi models at startup in background thread
-    so first user request is fast instead of waiting 2-3 minutes
-  - Summarize kept for API compatibility
   - Telugu/Indic sentence ending (।) in sentence splitter
   - Reduced chunk size for Indic languages (subword tokenization)
 """
@@ -46,6 +50,15 @@ HELSINKI_MODELS = {
     ("en", "ru"): "Helsinki-NLP/opus-mt-en-ru",    # English → Russian
 }
 
+# Language prefix tokens needed by opus-mt-en-mul for target language selection
+# opus-mt-en-mul is a multilingual model — you must prepend >>lang<< token
+OPUS_MUL_TARGET_TOKENS = {
+    "te": ">>tel<<",   # Telugu
+    "ta": ">>tam<<",   # Tamil
+    "kn": ">>kan<<",   # Kannada
+    "fr": ">>fra<<",   # French (if routed here)
+}
+
 # NLLB codes (fallback for languages not in Helsinki map)
 NLLB_CODES = {
     "en": "eng_Latn", "te": "tel_Telu", "hi": "hin_Deva",
@@ -55,6 +68,9 @@ NLLB_CODES = {
     "ru": "rus_Cyrl",
 }
 
+# Languages that should use Google first for X→English (Helsinki transliterates)
+INDIC_TO_EN = {"te", "kn", "ml", "bn", "gu", "mr", "pa", "ur"}
+
 INDIC_LANGS       = {"te", "hi", "ta", "kn", "ar"}
 CHUNK_WORDS       = 80
 CHUNK_WORDS_INDIC = 50
@@ -64,7 +80,8 @@ MAX_TOKENS        = 512
 
 class Translator:
     def __init__(self):
-        self._helsinki_models = {}  # cache: model_id → pipeline
+        # Cache: model_id → (tokenizer, model)
+        self._helsinki_cache  = {}
         self._pipeline        = None
         self._tokenizer       = None
         self._model           = None
@@ -91,7 +108,7 @@ class Translator:
                 model_id = HELSINKI_MODELS.get((src, tgt))
                 if model_id:
                     print(f"[Translator] Pre-loading {src}→{tgt} ({model_id})...")
-                    self._get_helsinki_pipeline(model_id)
+                    self._load_helsinki(model_id)
                     print(f"[Translator] ✅ {src}→{tgt} pre-loaded and ready!")
             except Exception as e:
                 print(f"[Translator] Pre-load {src}→{tgt} failed: {e}")
@@ -110,16 +127,16 @@ class Translator:
         print(f"[Translator] {len(chunks)} chunks ({max_words}w), "
               f"{len(text)} chars, {src_lang}→{tgt_lang}")
 
-        # ── Special: Indic→English uses Google first (accurate meaning) ──
+        # ── Special: Indic→English uses Google first ──────────────────
         # Helsinki opus-mt-mul-en transliterates Telugu instead of translating
-        INDIC_TO_EN = {"te", "kn", "ml", "bn", "gu", "mr", "pa", "ur"}
         if src_lang in INDIC_TO_EN and tgt_lang == "en":
             try:
                 result = self._google_chunks(chunks, src_lang, tgt_lang)
                 if "[Translation failed" not in result[0]:
+                    print(f"[Translator] ✅ Google used for {src_lang}→en")
                     return result
             except Exception as e:
-                logger.warning(f"Google te→en failed ({e}), trying Helsinki")
+                logger.warning(f"Google {src_lang}→en failed ({e}), trying Helsinki")
 
         # ── Priority 1: Helsinki-NLP ───────────────────────────────────
         if (src_lang, tgt_lang) in HELSINKI_MODELS:
@@ -171,19 +188,52 @@ class Translator:
 
     # ══════════════════════════════════════════════════════════════════
     # HELSINKI-NLP — PRIMARY
+    # FIX: Use MarianMTModel directly instead of pipeline() — avoids the
+    #      "Unknown task text2text-generation" error in newer transformers
     # ══════════════════════════════════════════════════════════════════
+    def _load_helsinki(self, model_id: str):
+        """Load and cache Helsinki-NLP (Marian) tokenizer + model."""
+        if model_id not in self._helsinki_cache:
+            from transformers import MarianMTModel, MarianTokenizer
+            print(f"[Translator] Loading {model_id}...")
+            tokenizer = MarianTokenizer.from_pretrained(model_id)
+            model     = MarianMTModel.from_pretrained(model_id)
+            model.eval()
+            self._helsinki_cache[model_id] = (tokenizer, model)
+            print(f"[Translator] ✅ {model_id} ready")
+        return self._helsinki_cache[model_id]
+
     def _helsinki_chunks(self, chunks, src_lang, tgt_lang):
+        import torch
         t0       = time.time()
         model_id = HELSINKI_MODELS[(src_lang, tgt_lang)]
-        pipe     = self._get_helsinki_pipeline(model_id)
-        results  = []
+        tokenizer, model = self._load_helsinki(model_id)
+
+        # opus-mt-en-mul requires a target language prefix token
+        prefix = OPUS_MUL_TARGET_TOKENS.get(tgt_lang, "")
 
+        results = []
         for i, chunk in enumerate(chunks):
             if not chunk.strip():
                 continue
             try:
-                out = pipe(chunk, max_length=MAX_TOKENS)
-                results.append(out[0]["translation_text"])
+                text_in = f"{prefix} {chunk}".strip() if prefix else chunk
+                inputs  = tokenizer(
+                    [text_in],
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=MAX_TOKENS,
+                )
+                with torch.no_grad():
+                    translated_ids = model.generate(
+                        **inputs,
+                        num_beams=4,
+                        max_length=MAX_TOKENS,
+                        early_stopping=True,
+                    )
+                out = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+                results.append(out)
             except Exception as e:
                 logger.warning(f"Helsinki chunk {i+1} failed: {e}")
                 results.append(chunk)
@@ -193,20 +243,6 @@ class Translator:
         short_name = model_id.split("/")[-1]
         return translated, f"Helsinki-NLP ({short_name}, {len(chunks)} chunks)"
 
-    def _get_helsinki_pipeline(self, model_id: str):
-        """Load and cache Helsinki-NLP pipeline — one per language pair."""
-        if model_id not in self._helsinki_models:
-            from transformers import pipeline as hf_pipeline
-            print(f"[Translator] Loading {model_id}...")
-            self._helsinki_models[model_id] = hf_pipeline(
-                "translation",
-                model=model_id,
-                device_map="auto",
-                max_length=MAX_TOKENS,
-            )
-            print(f"[Translator] ✅ {model_id} ready")
-        return self._helsinki_models[model_id]
-
     # ══════════════════════════════════════════════════════════════════
     # CHUNKING
     # ══════════════════════════════════════════════════════════════════
@@ -222,7 +258,7 @@ class Translator:
             count += w
         if cur:
             chunks.append(" ".join(cur))
-        return chunks
+        return chunks or [text]  # FIX: never return empty list
 
     # ══════════════════════════════════════════════════════════════════
     # NLLB — FALLBACK
@@ -303,10 +339,17 @@ class Translator:
     def _init_nllb(self):
         try:
             from transformers import pipeline as hf_pipeline
-            self._pipeline = hf_pipeline(
-                "translation", model=NLLB_MODEL_ID,
-                device_map="auto", max_length=MAX_TOKENS,
-            )
+            # NLLB uses "translation" task which is still supported for NLLB models
+            try:
+                self._pipeline = hf_pipeline(
+                    "translation", model=NLLB_MODEL_ID,
+                    device_map="auto", max_length=MAX_TOKENS,
+                )
+            except Exception:
+                self._pipeline = hf_pipeline(
+                    "text2text-generation", model=NLLB_MODEL_ID,
+                    device_map="auto", max_length=MAX_TOKENS,
+                )
             print("[Translator] ✅ NLLB pipeline ready")
         except Exception as e:
             logger.warning(f"NLLB pipeline init failed ({e}), trying manual")