Update src/translate.py

src/translate.py

@@ -2,110 +2,157 @@
 from __future__ import annotations
 import os
 import re
-from typing import List, Optional
+from typing import Optional, List
 
-import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from IndicTransToolkit.processor import IndicProcessor
+# Public constants
+ENGLISH = "en"
+HINDI = "hi"
 
-# -------- Model choices (CPU-friendly distilled) ----------
-# You can switch to 1B by replacing the *_CKPT names below with:
-#   ai4bharat/indictrans2-indic-en-1B
-#   ai4bharat/indictrans2-en-indic-1B
-INDIC_EN_CKPT = os.getenv("INDIC_EN_MODEL", "ai4bharat/indictrans2-indic-en-dist-200M")
-EN_INDIC_CKPT = os.getenv("EN_INDIC_MODEL", "ai4bharat/indictrans2-en-indic-dist-200M")
+# ENV knobs
+ENABLE_TRANSLATION = os.getenv("ENABLE_TRANSLATION", "1") == "1"
+MODEL_ID_EN2INDIC = os.getenv(
+    "INDICTRANS2_EN2INDIC_MODEL",
+    "ai4bharat/indictrans2-en-indic-distilled"
+)
+# If later you add Indic→English, you can add the reverse distilled model:
+MODEL_ID_INDIC2EN = os.getenv(
+    "INDICTRANS2_INDIC2EN_MODEL",
+    "ai4bharat/indictrans2-indic-en-distilled"
+)
 
-_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Globals (loaded once)
+_MODEL_EN2INDIC = None
+_TOKENIZER_EN2INDIC = None
+_MODEL_INDIC2EN = None
+_TOKENIZER_INDIC2EN = None
+_IPROCESSOR = None  # Indic pre/post processor
 
-# We keep singletons so the large models load once per process
-_ip: Optional[IndicProcessor] = None
-_tok_indic_en: Optional[AutoTokenizer] = None
-_mod_indic_en: Optional[AutoModelForSeq2SeqLM] = None
-_tok_en_indic: Optional[AutoTokenizer] = None
-_mod_en_indic: Optional[AutoModelForSeq2SeqLM] = None
+# Light Hindi detection (Devanagari range)
+_RE_DEVANAGARI = re.compile(r"[\u0900-\u097F]")
 
-# Hindi (Devanagari) target code for IndicTrans2
-HINDI = "hin_Deva"
-ENGLISH = "eng_Latn"
+def _likely_hindi(text: str) -> bool:
+    return bool(_RE_DEVANAGARI.search(text or ""))
 
-def _iproc() -> IndicProcessor:
-    global _ip
-    if _ip is None:
-        _ip = IndicProcessor(inference=True)
-    return _ip
+def _try_imports():
+    """Import heavy libs lazily."""
+    global transformers, torch, IndicProcessor
+    import transformers  # type: ignore
+    import torch  # type: ignore
+    from IndicTransToolkit.processor import IndicProcessor  # type: ignore
+    return transformers, torch, IndicProcessor
 
-def _load_indic_en():
-    global _tok_indic_en, _mod_indic_en
-    if _tok_indic_en is None or _mod_indic_en is None:
-        _tok_indic_en = AutoTokenizer.from_pretrained(INDIC_EN_CKPT, trust_remote_code=True)
-        _mod_indic_en = AutoModelForSeq2SeqLM.from_pretrained(
-            INDIC_EN_CKPT, trust_remote_code=True
-        ).to(_DEVICE)
-    return _tok_indic_en, _mod_indic_en
+def _device():
+    # Force CPU on Spaces (safe default)
+    return "cpu"
 
-def _load_en_indic():
-    global _tok_en_indic, _mod_en_indic
-    if _tok_en_indic is None or _mod_en_indic is None:
-        _tok_en_indic = AutoTokenizer.from_pretrained(EN_INDIC_CKPT, trust_remote_code=True)
-        _mod_en_indic = AutoModelForSeq2SeqLM.from_pretrained(
-            EN_INDIC_CKPT, trust_remote_code=True
-        ).to(_DEVICE)
-    return _tok_en_indic, _mod_en_indic
+def _load_iprocessor():
+    global _IPROCESSOR
+    if _IPROCESSOR is not None:
+        return _IPROCESSOR
+    try:
+        _, _, IndicProcessor = _try_imports()
+        _IPROCESSOR = IndicProcessor(inference=True)
+    except Exception:
+        _IPROCESSOR = None
+    return _IPROCESSOR
 
-_DEVANAGARI_RE = re.compile(r"[\u0900-\u097F]")
+def _load_en2indic():
+    """Load the distilled en→indic model once."""
+    global _MODEL_EN2INDIC, _TOKENIZER_EN2INDIC
+    if _MODEL_EN2INDIC is not None:
+        return _MODEL_EN2INDIC, _TOKENIZER_EN2INDIC
+    try:
+        transformers, torch, _ = _try_imports()
+        tok = transformers.AutoTokenizer.from_pretrained(MODEL_ID_EN2INDIC, trust_remote_code=True)
+        model = transformers.AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID_EN2INDIC, trust_remote_code=True)
+        model.to(_device())
+        model.eval()
+        _MODEL_EN2INDIC, _TOKENIZER_EN2INDIC = model, tok
+    except Exception:
+        _MODEL_EN2INDIC, _TOKENIZER_EN2INDIC = None, None
+    return _MODEL_EN2INDIC, _TOKENIZER_EN2INDIC
 
-def looks_devanagari(text: str) -> bool:
-    """Heuristic: any Devanagari char ⇒ treat as Hindi."""
-    return bool(_DEVANAGARI_RE.search(text or ""))
+def _load_indic2en():
+    """Load the distilled indic→en model once (only if needed)."""
+    global _MODEL_INDIC2EN, _TOKENIZER_INDIC2EN
+    if _MODEL_INDIC2EN is not None:
+        return _MODEL_INDIC2EN, _TOKENIZER_INDIC2EN
+    try:
+        transformers, torch, _ = _try_imports()
+        tok = transformers.AutoTokenizer.from_pretrained(MODEL_ID_INDIC2EN, trust_remote_code=True)
+        model = transformers.AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID_INDIC2EN, trust_remote_code=True)
+        model.to(_device())
+        model.eval()
+        _MODEL_INDIC2EN, _TOKENIZER_INDIC2EN = model, tok
+    except Exception:
+        _MODEL_INDIC2EN, _TOKENIZER_INDIC2EN = None, None
+    return _MODEL_INDIC2EN, _TOKENIZER_INDIC2EN
 
-def _batch_decode(
-    model, tok, inputs: List[str], src_lang: str, tgt_lang: str, max_new_tokens=256
-) -> List[str]:
-    ip = _iproc()
-    # Pre-process
-    sents, srcl = ip.preprocess_batch(inputs, src_lang=src_lang, tgt_lang=tgt_lang)
-    enc = tok(sents, return_tensors="pt", padding=True, truncation=True).to(_DEVICE)
-    with torch.no_grad():
-        gen = model.generate(
-            **enc,
-            max_new_tokens=max_new_tokens,
-            num_beams=4,
-            length_penalty=1.0,
+def _generate(model, tokenizer, inputs: List[str], max_new_tokens=256) -> List[str]:
+    """Run generation on a small batch of strings."""
+    if model is None or tokenizer is None:
+        return inputs  # graceful fallback
+    try:
+        import torch  # local import
+        enc = tokenizer(
+            inputs,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512,
         )
-    out = tok.batch_decode(gen, skip_special_tokens=True)
-    # Post-process
-    return ip.postprocess_batch(out, lang=tgt_lang)
+        enc = {k: v.to(_device()) for k, v in enc.items()}
+        with torch.no_grad():
+            outs = model.generate(
+                **enc,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+            )
+        return tokenizer.batch_decode(outs, skip_special_tokens=True)
+    except Exception:
+        return inputs
 
-def indic_to_en(text: str, src_lang: str = HINDI) -> str:
-    """Translate Indic→English; default assumes Hindi (hin_Deva)."""
-    if not text:
+def ensure_english(text: str) -> str:
+    """
+    If input text looks Hindi, translate to English. Otherwise return as is.
+    We keep this very light: only detect Devanagari → hi→en.
+    """
+    if not ENABLE_TRANSLATION:
         return text
-    tok, mod = _load_indic_en()
-    return _batch_decode(mod, tok, [text], src_lang=src_lang, tgt_lang=ENGLISH)[0]
-
-def en_to_lang(text: str, tgt_lang: str = HINDI) -> str:
-    """Translate English→Indic (Hindi by default)."""
-    if not text:
+    try:
+        if _likely_hindi(text):
+            model, tok = _load_indic2en()
+            ip = _load_iprocessor()
+            src = text
+            if ip:
+                # Normalize/romanize as the toolkit suggests (safe to skip if None)
+                src = ip.preprocess_batch([src], src_lang=HINDI, tgt_lang=ENGLISH)[0]
+            out = _generate(model, tok, [src])[0]
+            if ip:
+                out = ip.postprocess_batch([out], lang=ENGLISH)[0]
+            return out
+        return text
+    except Exception:
         return text
-    tok, mod = _load_en_indic()
-    return _batch_decode(mod, tok, [text], src_lang=ENGLISH, tgt_lang=tgt_lang)[0]
 
-def ensure_english(text: str, src_hint: Optional[str] = None) -> tuple[str, Optional[str]]:
+def en_to_lang(text: str, tgt_lang: str = HINDI) -> str:
     """
-    If text seems Hindi (or src_hint given), translate to English.
-    Returns (english_text, original_lang_code_or_None).
+    Translate English → target Indic language (default: Hindi).
+    If translation stack is unavailable, returns original text.
     """
-    orig_lang = None
-    # If the caller knows it's Hindi, pass src_hint="hin_Deva"
-    if src_hint:
-        if src_hint != ENGLISH:
-            return indic_to_en(text, src_lang=src_hint), src_hint
-        return text, None
-
-    if looks_devanagari(text):
-        try:
-            return indic_to_en(text, src_lang=HINDI), HINDI
-        except Exception:
-            # Fallback: return original text if translation fails
-            return text, None
-    return text, None
+    if not ENABLE_TRANSLATION:
+        return text
+    if not text:
+        return text
+    try:
+        model, tok = _load_en2indic()
+        ip = _load_iprocessor()
+        src = text
+        if ip:
+            src = ip.preprocess_batch([src], src_lang=ENGLISH, tgt_lang=tgt_lang)[0]
+        out = _generate(model, tok, [src])[0]
+        if ip:
+            out = ip.postprocess_batch([out], lang=tgt_lang)[0]
+        return out
+    except Exception:
+        return text