Update translator.py

translator.py

@@ -1,469 +1,36 @@
-"""
-Department 3 — Translator
-UPGRADED v2: IndicTrans2 as primary for Telugu/Hindi/Tamil/Kannada
-
-MERGED — best of both versions:
-  ⭐ 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
-import time
 import logging
-import threading
+from transformers import pipeline
 
 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"
-
-INDICTRANS2_SUPPORTED = {"te", "hi", "ta", "kn", "ml", "bn", "gu", "mr", "pa", "ur"}
-
-INDICTRANS2_CODES = {
-    "en": "eng_Latn", "te": "tel_Telu", "hi": "hin_Deva",
-    "ta": "tam_Taml", "kn": "kan_Knda", "ml": "mal_Mlym",
-    "bn": "ben_Beng", "gu": "guj_Gujr", "mr": "mar_Deva",
-    "pa": "pan_Guru", "ur": "urd_Arab",
-}
-
-HELSINKI_MODELS = {
-    ("en", "te"): "Helsinki-NLP/opus-mt-en-mul",
-    ("en", "hi"): "Helsinki-NLP/opus-mt-en-hi",
-    ("en", "ta"): "Helsinki-NLP/opus-mt-en-mul",
-    ("en", "kn"): "Helsinki-NLP/opus-mt-en-mul",
-    ("hi", "en"): "Helsinki-NLP/opus-mt-hi-en",
-    ("te", "en"): "Helsinki-NLP/opus-mt-mul-en",
-    ("ta", "en"): "Helsinki-NLP/opus-mt-mul-en",
-    ("en", "es"): "Helsinki-NLP/opus-mt-en-es",
-    ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
-    ("en", "de"): "Helsinki-NLP/opus-mt-en-de",
-    ("en", "zh"): "Helsinki-NLP/opus-mt-en-zh",
-    ("en", "ar"): "Helsinki-NLP/opus-mt-en-ar",
-    ("en", "ru"): "Helsinki-NLP/opus-mt-en-ru",
-}
-
-NLLB_CODES = {
-    "en": "eng_Latn", "te": "tel_Telu", "hi": "hin_Deva",
-    "ta": "tam_Taml", "kn": "kan_Knda", "es": "spa_Latn",
-    "fr": "fra_Latn", "de": "deu_Latn", "ja": "jpn_Jpan",
-    "zh": "zho_Hans", "ar": "arb_Arab", "pt": "por_Latn",
-    "ru": "rus_Cyrl",
-}
-
-INDIC_LANGS       = {"te", "hi", "ta", "kn", "ar", "ml", "bn", "gu", "mr", "pa"}
-CHUNK_WORDS       = 80
-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   # IndicProcessor instance
-        self._it2_loaded  = False
-
-        self._helsinki_models = {}
-
-        self._pipeline    = None
-        self._tokenizer   = None
-        self._model       = None
-        self._nllb_loaded = False
-
-        print("[Translator] IndicTrans2 v2 — pre-loading Telugu + Hindi...")
-        threading.Thread(target=self._preload_common_models, daemon=True).start()
-
-    def _preload_common_models(self):
-        time.sleep(5)
-        try:
-            print("[Translator] Pre-loading IndicTrans2...")
-            self._load_indictrans2()
-            print("[Translator] ✅ IndicTrans2 ready! Telugu/Hindi fast from now.")
-        except Exception as e:
-            print(f"[Translator] IndicTrans2 pre-load failed: {e}")
-            print("[Translator] Falling back to Helsinki-NLP pre-load...")
-            try:
-                for src, tgt in [("en", "te"), ("en", "hi")]:
-                    model_id = HELSINKI_MODELS.get((src, tgt))
-                    if model_id:
-                        self._get_helsinki_pipeline(model_id)
-                        print(f"[Translator] ✅ Helsinki {src}→{tgt} ready")
-            except Exception as e2:
-                print(f"[Translator] Helsinki pre-load failed: {e2}")
-
-    # ══════════════════════════════════════════════════════════════════
-    # PUBLIC — TRANSLATE
-    # ══════════════════════════════════════════════════════════════════
-    def translate(self, text: str, src_lang: str, tgt_lang: str):
-        if not text or not text.strip():
-            return "", "skipped (empty)"
-        if src_lang == tgt_lang:
-            return text, "skipped (same language)"
-
-        max_words = CHUNK_WORDS_INDIC if src_lang in INDIC_LANGS else CHUNK_WORDS
-        chunks    = self._chunk(text, max_words)
-        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:
-                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:
-                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:
-                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:
-                return self._nllb_chunks(chunks, src_lang, tgt_lang)
-        except Exception as e:
-            logger.warning(f"NLLB failed ({e}), using Google")
-
-        # Priority 4: Google
-        return self._google_chunks(chunks, src_lang, tgt_lang)
-
-    # ══════════════════════════════════════════════════════════════════
-    # PUBLIC — SUMMARIZE
-    # ══════════════════════════════════════════════════════════════════
-    def summarize(self, text: str, max_sentences: int = 5) -> str:
-        try:
-            sentences = re.split(r'(?<=[.!?।])\s+', text.strip())
-            sentences = [s.strip() for s in sentences if len(s.split()) > 5]
-            if len(sentences) <= max_sentences:
-                return text
-            n = len(sentences)
-
-            def score(idx, sent):
-                if idx == 0:          pos = 1.0
-                elif idx == n - 1:    pos = 0.7
-                elif idx <= n * 0.2:  pos = 0.6
-                else:                 pos = 0.3
-                wc    = len(sent.split())
-                bonus = 0.3 if 10 <= wc <= 30 else (0.0 if wc < 10 else 0.1)
-                return pos + bonus
-
-            scored      = sorted(enumerate(sentences),
-                                 key=lambda x: score(x[0], x[1]), reverse=True)
-            top_indices = sorted([i for i, _ in scored[:max_sentences]])
-            return " ".join(sentences[i] for i in top_indices).strip()
-        except Exception as e:
-            logger.warning(f"Summarize failed: {e}")
-            return text[:800] + "..."
-
-    # ══════════════════════════════════════════════════════════════════
-    # INDICTRANS2
-    # ══════════════════════════════════════════════════════════════════
-    def _load_indictrans2(self):
-        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:
-            raise RuntimeError(
-                "IndicTransToolkit not installed. "
-                "Run: pip install git+https://github.com/AI4Bharat/IndicTransToolkit.git"
-            )
-        import torch
-
-        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-        print(f"[Translator] Loading {INDICTRANS2_EN_INDIC}...")
-        self._it2_tok_ei = AutoTokenizer.from_pretrained(
-            INDICTRANS2_EN_INDIC, trust_remote_code=True)
-        self._it2_mod_ei = AutoModelForSeq2SeqLM.from_pretrained(
-            INDICTRANS2_EN_INDIC, trust_remote_code=True, torch_dtype=dtype)
-        if torch.cuda.is_available():
-            self._it2_mod_ei = self._it2_mod_ei.cuda()
-        self._it2_mod_ei.eval()
-
-        print(f"[Translator] Loading {INDICTRANS2_INDIC_EN}...")
-        self._it2_tok_ie = AutoTokenizer.from_pretrained(
-            INDICTRANS2_INDIC_EN, trust_remote_code=True)
-        self._it2_mod_ie = AutoModelForSeq2SeqLM.from_pretrained(
-            INDICTRANS2_INDIC_EN, trust_remote_code=True, torch_dtype=dtype)
-        if torch.cuda.is_available():
-            self._it2_mod_ie = self._it2_mod_ie.cuda()
-        self._it2_mod_ie.eval()
-
-        self._it2_ip     = IndicProcessor(inference=True)
-        self._it2_loaded = True
-        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")
-
-        if src_lang == "en":
-            tokenizer, model = self._it2_tok_ei, self._it2_mod_ei
-        else:
-            tokenizer, model = self._it2_tok_ie, self._it2_mod_ie
-
-        # FIX L9: resolve forced_bos_token_id for target language
+        self.model_name = "Helsinki-NLP/opus-mt-en-mul"
+        self.translator = None
         try:
-            forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_code)
-        except Exception:
-            forced_bos_token_id = None
-            logger.warning(f"[Translator] Could not resolve token id for {tgt_code}")
-
-        results = []
-        for i, chunk in enumerate(chunks):
-            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(
-                    batch, truncation=True, padding="longest",
-                    return_tensors="pt", return_attention_mask=True,
-                    max_length=MAX_TOKENS,
-                )
-                if torch.cuda.is_available():
-                    inputs = {k: v.cuda() for k, v in inputs.items()}
-
-                gen_kwargs = dict(
-                    num_beams=4,
-                    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
-
-                with torch.no_grad():
-                    ids = model.generate(**inputs, **gen_kwargs)
-
-                decoded = tokenizer.batch_decode(
-                    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])
-
-            except Exception as e:
-                logger.warning(f"IndicTrans2 chunk {i+1} failed: {e}")
-                results.append(chunk)
-
-        translated = " ".join(results)
-        logger.info(f"IndicTrans2 done in {time.time()-t0:.2f}s")
-        return translated, f"IndicTrans2-1B ({src_lang}→{tgt_lang}, {len(chunks)} chunks)"
-
-    # ══════════════════════════════════════════════════════════════════
-    # HELSINKI-NLP
-    # ══════════════════════════════════════════════════════════════════
-    def _helsinki_chunks(self, chunks, src_lang, tgt_lang):
-        t0       = time.time()
-        model_id = HELSINKI_MODELS[(src_lang, tgt_lang)]
-        pipe     = self._get_helsinki_pipeline(model_id)
-        results  = []
-        for i, chunk in enumerate(chunks):
-            if not chunk.strip():
-                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:
-                logger.warning(f"Helsinki chunk {i+1} failed: {e}")
-                results.append(chunk)
-        translated = " ".join(results)
-        logger.info(f"Helsinki done in {time.time()-t0:.2f}s")
-        return translated, f"Helsinki-NLP ({model_id.split('/')[-1]}, {len(chunks)} chunks)"
-
-    def _get_helsinki_pipeline(self, model_id: str):
-        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)
-            except Exception:
-                pipe = hf_pipeline("text2text-generation", model=model_id,
-                                   device_map="auto", max_length=MAX_TOKENS)
-            self._helsinki_models[model_id] = pipe
-            print(f"[Translator] ✅ {model_id} ready")
-        return self._helsinki_models[model_id]
-
-    # ══════════════════════════════════════════════════════════════════
-    # 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
-
-        for s in sentences:
-            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))
-                    cur, count = [], 0
-                for start in range(0, w, max_words):
-                    sub = " ".join(words_in_sent[start:start + max_words])
-                    chunks.append(sub)
-                continue
-
-            if count + w > max_words and cur:
-                chunks.append(" ".join(cur))
-                cur, count = [], 0
-
-            cur.append(s)
-            count += w
-
-        if cur:
-            chunks.append(" ".join(cur))
-        return chunks
-
-    # ══════════════════════════════════════════════════════════════════
-    # NLLB
-    # ══════════════════════════════════════════════════════════════════
-    def _nllb_chunks(self, chunks, src_lang, tgt_lang):
-        t0       = time.time()
-        src_code = NLLB_CODES.get(src_lang, "eng_Latn")
-        tgt_code = NLLB_CODES.get(tgt_lang, "tel_Telu")
-        results  = []
-        for i, chunk in enumerate(chunks):
-            if not chunk.strip():
-                continue
-            try:
-                if self._pipeline is not None:
-                    out = self._pipeline(chunk, src_lang=src_code,
-                                         tgt_lang=tgt_code, max_length=MAX_TOKENS)
-                    results.append(out[0]["translation_text"])
-                else:
-                    import torch
-                    inputs = self._tokenizer(
-                        chunk, return_tensors="pt", padding=True,
-                        truncation=True, max_length=MAX_TOKENS,
-                    )
-                    if torch.cuda.is_available():
-                        inputs = {k: v.cuda() for k, v in inputs.items()}
-                    tid = self._tokenizer.convert_tokens_to_ids(tgt_code)
-                    with torch.no_grad():
-                        ids = self._model.generate(
-                            **inputs, forced_bos_token_id=tid,
-                            max_length=MAX_TOKENS, num_beams=4, early_stopping=True,
-                        )
-                    results.append(
-                        self._tokenizer.batch_decode(ids, skip_special_tokens=True)[0])
-            except Exception as e:
-                logger.warning(f"NLLB chunk {i+1} failed: {e}")
-                results.append(chunk)
-        translated = " ".join(results)
-        logger.info(f"NLLB done in {time.time()-t0:.2f}s")
-        return translated, f"NLLB-200-1.3B ({len(chunks)} chunks)"
-
-    # ══════════════════════════════════════════════════════════════════
-    # GOOGLE
-    # ══════════════════════════════════════════════════════════════════
-    def _google_chunks(self, chunks, src_lang, tgt_lang):
-        t0 = time.time()
-        try:
-            from deep_translator import GoogleTranslator
-            results = []
-            for chunk in chunks:
-                if not chunk.strip():
-                    continue
-                out = GoogleTranslator(
-                    source=src_lang if src_lang != "auto" else "auto",
-                    target=tgt_lang,
-                ).translate(chunk)
-                results.append(out)
-            full = " ".join(results)
-            logger.info(f"Google done in {time.time()-t0:.2f}s")
-            return full, f"Google Translate ({len(chunks)} chunks)"
+            # FIX: Change task to "translation" to match Helsinki requirements
+            self.translator = pipeline("translation", model=self.model_name)
+            logger.info("[Translator] 🌐 Helsinki-NLP loaded successfully")
         except Exception as e:
-            logger.error(f"Google failed: {e}")
-            return f"[Translation failed: {e}]", "error"
-
-    # ══════════════════════════════════════════════════════════════════
-    # NLLB INIT
-    # ══════════════════════════════════════════════════════════════════
-    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,
-            )
-            print("[Translator] ✅ NLLB pipeline ready")
-        except Exception as e:
-            logger.warning(f"NLLB pipeline init failed ({e}), trying manual")
-            self._init_nllb_manual()
-
-    def _init_nllb_manual(self):
-        try:
-            from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-            import torch
-            self._tokenizer = AutoTokenizer.from_pretrained(NLLB_MODEL_ID)
-            self._model     = AutoModelForSeq2SeqLM.from_pretrained(
-                NLLB_MODEL_ID,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            )
-            if torch.cuda.is_available():
-                self._model = self._model.cuda()
-            self._model.eval()
-            print("[Translator] ✅ NLLB manual load ready")
+            logger.warning(f"[Translator] Pre-load waiting for first run: {e}")
+
+    def translate(self, text, target_lang):
+        if not text: return ""
+        
+        lang_map = {
+            "Hindi": ">>hin<<", "Telugu": ">>tel<<", 
+            "Spanish": ">>spa<<", "French": ">>fra<<", "German": ">>deu<<"
+        }
+        
+        prefix = lang_map.get(target_lang, ">>hin<<")
+        input_text = f"{prefix} {text}"
+        
+        try:
+            if self.translator is None:
+                self.translator = pipeline("translation", model=self.model_name)
+            
+            result = self.translator(input_text)
+            return result[0]['translation_text']
         except Exception as e:
-            logger.error(f"NLLB manual load failed: {e}")
+            logger.error(f"Translation failed: {e}")
+            return f"[Translation Error: {e}]"