Update translator.py

translator.py

@@ -3,9 +3,15 @@ Department 3 — Translator
 UPGRADED v2: IndicTrans2 as primary for Telugu/Hindi/Tamil/Kannada
 
 MERGED — best of both versions:
-  ⭐ KEPT: IndicTransToolkit, non-gated models, chunk boundaries
-  ✅ FIXED bugs with bos_token_ids
-  ⭐ NEW: Added Project Glossary to protect the 'ClearWave AI' name from literal translation
+  ⭐ KEPT from friend's version:
+     - IndicTransToolkit.IndicProcessor for correct pre/post-processing
+     - Non-gated 1B models (ai4bharat/indictrans2-en-indic-1B) — no HF token needed
+     - Helsinki fallback handles both "translation_text" and "generated_text" keys
+     - Indic→English uses Google first (opus-mt-mul-en transliterates, not translates)
+     - L8: Wrong <2te> prefix already removed (friend fixed this one correctly)
+  ✅ FIXED bugs that friend's version missed:
+     - L9: forced_bos_token_id now passed to model.generate() — target lang enforced
+     - L10: _chunk word-boundary fallback for unpunctuated transcripts
 """
 
 import re
@@ -15,6 +21,7 @@ import threading
 
 logger = logging.getLogger(__name__)
 
+# ⭐ KEPT: non-gated 1B models
 INDICTRANS2_EN_INDIC = "ai4bharat/indictrans2-en-indic-1B"
 INDICTRANS2_INDIC_EN = "ai4bharat/indictrans2-indic-en-1B"
 
@@ -57,13 +64,15 @@ CHUNK_WORDS_INDIC = 50
 NLLB_MODEL_ID     = "facebook/nllb-200-distilled-1.3B"
 MAX_TOKENS        = 512
 
+
 class Translator:
     def __init__(self):
+        # ⭐ KEPT: separate tokenizer/model handles + IndicProcessor
         self._it2_tok_ei  = None
         self._it2_mod_ei  = None
         self._it2_tok_ie  = None
         self._it2_mod_ie  = None
-        self._it2_ip      = None   
+        self._it2_ip      = None   # IndicProcessor instance
         self._it2_loaded  = False
 
         self._helsinki_models = {}
@@ -94,20 +103,6 @@ class Translator:
             except Exception as e2:
                 print(f"[Translator] Helsinki pre-load failed: {e2}")
 
-    # ── NEW: Protect project name from literal translation ──
-    def _apply_glossary(self, text: str) -> str:
-        """Forces correct branding regardless of language."""
-        protected_terms = {
-            "క్లియర్ వేవ్": "ClearWave AI",
-            "క్లియర్వేవ్": "ClearWave AI",
-            "clearwave ai": "ClearWave AI",
-            "clearwave": "ClearWave AI",
-            "క్లియర్‌వేవ్": "ClearWave AI"
-        }
-        for wrong, right in protected_terms.items():
-            text = re.sub(wrong, right, text, flags=re.IGNORECASE)
-        return text
-
     # ══════════════════════════════════════════════════════════════════
     # PUBLIC — TRANSLATE
     # ══════════════════════════════════════════════════════════════════
@@ -122,45 +117,46 @@ class Translator:
         print(f"[Translator] {len(chunks)} chunks ({max_words}w), "
               f"{len(text)} chars, {src_lang}→{tgt_lang}")
 
+        # ⭐ KEPT: Indic→English: use Google first (better than transliterating opus-mt-mul-en)
         INDIC_TO_EN = {"te", "kn", "ml", "bn", "gu", "mr", "pa", "ur"}
         if src_lang in INDIC_TO_EN and tgt_lang == "en":
             try:
-                res, method = self._google_chunks(chunks, src_lang, tgt_lang)
-                if "[Translation failed" not in res:
-                    return self._apply_glossary(res), method
+                result = self._google_chunks(chunks, src_lang, tgt_lang)
+                if "[Translation failed" not in result[0]:
+                    return result
             except Exception as e:
                 logger.warning(f"Google {src_lang}→en failed ({e}), trying IndicTrans2")
 
+        # Priority 1: IndicTrans2
         is_indic = (
             (src_lang == "en" and tgt_lang in INDICTRANS2_SUPPORTED) or
             (src_lang in INDICTRANS2_SUPPORTED and tgt_lang == "en")
         )
         if is_indic:
             try:
-                res, method = self._indictrans2_chunks(chunks, src_lang, tgt_lang)
-                return self._apply_glossary(res), method
+                return self._indictrans2_chunks(chunks, src_lang, tgt_lang)
             except Exception as e:
                 logger.warning(f"IndicTrans2 failed ({e}), trying Helsinki")
 
+        # Priority 2: Helsinki-NLP
         if (src_lang, tgt_lang) in HELSINKI_MODELS:
             try:
-                res, method = self._helsinki_chunks(chunks, src_lang, tgt_lang)
-                return self._apply_glossary(res), method
+                return self._helsinki_chunks(chunks, src_lang, tgt_lang)
             except Exception as e:
                 logger.warning(f"Helsinki failed ({e}), trying NLLB")
 
+        # Priority 3: NLLB
         try:
             if not self._nllb_loaded:
                 self._init_nllb()
                 self._nllb_loaded = True
             if self._pipeline is not None or self._model is not None:
-                res, method = self._nllb_chunks(chunks, src_lang, tgt_lang)
-                return self._apply_glossary(res), method
+                return self._nllb_chunks(chunks, src_lang, tgt_lang)
         except Exception as e:
             logger.warning(f"NLLB failed ({e}), using Google")
 
-        res, method = self._google_chunks(chunks, src_lang, tgt_lang)
-        return self._apply_glossary(res), method
+        # Priority 4: Google
+        return self._google_chunks(chunks, src_lang, tgt_lang)
 
     # ══════════════════════════════════════════════════════════════════
     # PUBLIC — SUMMARIZE
@@ -197,6 +193,9 @@ class Translator:
         if self._it2_loaded:
             return
         from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+        # IndicTransToolkit is NOT on PyPI — must be installed manually.
+        # If missing, _load_indictrans2 raises RuntimeError and the caller
+        # falls back to Helsinki / NLLB / Google automatically.
         try:
             from IndicTransToolkit.processor import IndicProcessor
         except ImportError:
@@ -231,11 +230,17 @@ class Translator:
         print("[Translator] ✅ IndicTrans2 fully loaded!")
 
     def _indictrans2_chunks(self, chunks, src_lang, tgt_lang):
+        """
+        ⭐ KEPT: IndicProcessor.preprocess_batch / postprocess_batch
+        FIX L9: forced_bos_token_id now resolved and passed to model.generate()
+                so the decoder is forced to output the correct target language.
+        """
         t0 = time.time()
         if not self._it2_loaded:
             self._load_indictrans2()
 
         import torch
+
         src_code = INDICTRANS2_CODES.get(src_lang, "eng_Latn")
         tgt_code = INDICTRANS2_CODES.get(tgt_lang, "tel_Telu")
 
@@ -244,6 +249,7 @@ class Translator:
         else:
             tokenizer, model = self._it2_tok_ie, self._it2_mod_ie
 
+        # FIX L9: resolve forced_bos_token_id for target language
         try:
             forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_code)
         except Exception:
@@ -255,6 +261,7 @@ class Translator:
             if not chunk.strip():
                 continue
             try:
+                # ⭐ KEPT: IndicProcessor pre-processing
                 batch = self._it2_ip.preprocess_batch(
                     [chunk], src_lang=src_code, tgt_lang=tgt_code)
                 inputs = tokenizer(
@@ -270,6 +277,7 @@ class Translator:
                     num_return_sequences=1,
                     max_length=MAX_TOKENS,
                 )
+                # FIX L9: enforce target language in decoder
                 if forced_bos_token_id is not None:
                     gen_kwargs["forced_bos_token_id"] = forced_bos_token_id
 
@@ -280,6 +288,7 @@ class Translator:
                     ids, skip_special_tokens=True,
                     clean_up_tokenization_spaces=True,
                 )
+                # ⭐ KEPT: IndicProcessor post-processing
                 postprocessed = self._it2_ip.postprocess_batch(decoded, lang=tgt_code)
                 results.append(postprocessed[0])
 
@@ -304,6 +313,7 @@ class Translator:
                 continue
             try:
                 out = pipe(chunk, max_length=MAX_TOKENS)
+                # ⭐ KEPT: handles both key names across transformers versions
                 text = out[0].get("translation_text") or out[0].get("generated_text", chunk)
                 results.append(text)
             except Exception as e:
@@ -317,6 +327,7 @@ class Translator:
         if model_id not in self._helsinki_models:
             from transformers import pipeline as hf_pipeline
             print(f"[Translator] Loading Helsinki {model_id}...")
+            # ⭐ KEPT: fallback to text2text-generation for newer transformers
             try:
                 pipe = hf_pipeline("translation", model=model_id,
                                    device_map="auto", max_length=MAX_TOKENS)
@@ -328,9 +339,14 @@ class Translator:
         return self._helsinki_models[model_id]
 
     # ══════════════════════════════════════════════════════════════════
-    # CHUNKING 
+    # CHUNKING — FIXED (L10)
     # ══════════════════════════════════════════════════════════════════
     def _chunk(self, text, max_words):
+        """
+        FIX L10: Added word-boundary fallback.
+        Transcripts often have no sentence-ending punctuation — without this,
+        the entire text becomes one oversized chunk causing OOM / truncation.
+        """
         sentences = re.split(r'(?<=[.!?।])\s+', text.strip())
         chunks, cur, count = [], [], 0
 
@@ -338,6 +354,7 @@ class Translator:
             words_in_sent = s.split()
             w = len(words_in_sent)
 
+            # FIX L10: sentence longer than max_words → split on word boundaries
             if w > max_words:
                 if cur:
                     chunks.append(" ".join(cur))