Upload indic_text.py with huggingface_hub

indic_text.py

@@ -0,0 +1,80 @@
+"""English -> Kannada translation via AI4Bharat IndicTrans2 — local GPU for HF Spaces."""
+
+from __future__ import annotations
+
+import os
+import re
+
+import torch
+
+from config import TRANSLATION_MODEL
+_TRANS_HUB_ID = TRANSLATION_MODEL.hub_id
+
+_tok = None
+_model = None
+
+
+def _get_model():
+    global _tok, _model
+    if _tok is None or _model is None:
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+        token = os.environ.get("HF_TOKEN") or None
+        _tok = AutoTokenizer.from_pretrained(
+            _TRANS_HUB_ID, trust_remote_code=True, token=token)
+        _model = AutoModelForSeq2SeqLM.from_pretrained(
+            _TRANS_HUB_ID, trust_remote_code=True, token=token)
+        _model = _model.to("cuda")
+        _model.eval()
+    return _tok, _model
+
+
+# Load at module level for ZeroGPU (CUDA emulation outside @spaces.GPU)
+try:
+    _get_model()
+except Exception:
+    pass
+
+
+def _split_sentences(text: str, max_chars: int = 180):
+    parts = re.split(r"(?<=[.!?।])\s+|\n+", text.strip())
+    out = []
+    for p in parts:
+        p = p.strip()
+        if not p:
+            continue
+        while len(p) > max_chars:
+            cut = p.rfind(" ", 0, max_chars)
+            cut = cut if cut > 0 else max_chars
+            out.append(p[:cut].strip())
+            p = p[cut:].strip()
+        out.append(p)
+    return out or [text.strip()]
+
+
+def translate_to_kannada(en_text: str) -> str:
+    """Translate English story text to Kannada (Kannada script)."""
+    text = (en_text or "").strip()
+    if not text:
+        raise ValueError("Nothing to translate.")
+
+    from IndicTransToolkit.processor import IndicProcessor
+
+    tok, model = _get_model()
+    ip = IndicProcessor(inference=True)
+
+    sents = _split_sentences(text, max_chars=180)
+    batch = ip.preprocess_batch(sents, src_lang="eng_Latn", tgt_lang="kan_Knda")
+    inputs = tok(batch, truncation=True, padding="longest", return_tensors="pt").to(model.device)
+
+    with torch.inference_mode():
+        generated = model.generate(
+            **inputs, max_length=512, num_beams=5,
+            num_return_sequences=1, length_penalty=1.0,
+        )
+    decoded = tok.batch_decode(generated, skip_special_tokens=True)
+    translations = ip.postprocess_batch(decoded, lang="kan_Knda")
+
+    kn = " ".join(t.strip() for t in translations if t.strip())
+    if not kn:
+        raise RuntimeError("Translation returned empty Kannada text.")
+    return kn