Update app.py

app.py

@@ -1,15 +1,18 @@
-import os, torch
+import os, traceback, types, torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from IndicTransToolkit import IndicProcessor  # https://github.com/VarunGumma/IndicTransToolkit
+
+# Robust import for IndicProcessor (fallback path)
+try:
+    from IndicTransToolkit import IndicProcessor  # preferred
+except Exception:
+    from IndicTransToolkit.IndicTransToolkit import IndicProcessor  # fallback
 
 # --------- Config (override via Space Variables if you like) ----------
 TOKENIZER_ID = os.getenv("TOKENIZER_ID", "ai4bharat/indictrans2-en-indic-1B")
 MODEL_ID     = os.getenv("MODEL_ID",     "law-ai/InLegalTrans-En2Indic-1B")
-
-# (Optional) pin revisions to avoid surprise upstream changes
-TOKENIZER_REV = os.getenv("TOKENIZER_REV", None)  # e.g., "b1a2c3d"
-MODEL_REV     = os.getenv("MODEL_REV",     None)  # e.g., "e4f5a6b"
+TOKENIZER_REV = os.getenv("TOKENIZER_REV", None)  # optional pin
+MODEL_REV     = os.getenv("MODEL_REV",     None)  # optional pin
 
 SRC_CODE = "eng_Latn"
 HI_CODE  = "hin_Deva"
@@ -20,27 +23,60 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 dtype  = torch.float16 if torch.cuda.is_available() else torch.float32
 
 tok_kwargs = dict(trust_remote_code=True, use_fast=True)
-if TOKENIZER_REV: tok_kwargs["revision"] = TOKENIZER_REV
+if TOKENIZER_REV:
+    tok_kwargs["revision"] = TOKENIZER_REV
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID, **tok_kwargs)
 
 mdl_kwargs = dict(
     trust_remote_code=True,
     attn_implementation="eager",
     low_cpu_mem_usage=True,
-    dtype=dtype,        # <- fixes the torch_dtype deprecation warning
+    dtype=dtype,  # modern kw (no deprecation warning)
 )
-if MODEL_REV: mdl_kwargs["revision"] = MODEL_REV
+if MODEL_REV:
+    mdl_kwargs["revision"] = MODEL_REV
+
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, **mdl_kwargs).to(device)
 model.eval()
 
+# Patch generation config safety
+if getattr(model.generation_config, "pad_token_id", None) is None:
+    model.generation_config.pad_token_id = (
+        getattr(tokenizer, "pad_token_id", None) or getattr(tokenizer, "eos_token_id", 0)
+    )
+if getattr(model.generation_config, "eos_token_id", None) is None and getattr(tokenizer, "eos_token_id", None) is not None:
+    model.generation_config.eos_token_id = tokenizer.eos_token_id
+
+# ---- Runtime compatibility patch for newer Transformers beam search ----
+# Newer versions call self.config.get_text_config().vocab_size.
+# Some custom configs (IndicTransConfig) don't define these.
+if not hasattr(model.config, "vocab_size") or model.config.vocab_size is None:
+    try:
+        model.config.vocab_size = getattr(tokenizer, "vocab_size", None) or len(tokenizer)
+    except Exception:
+        # Fallback to a safe default if tokenizer doesn't expose size
+        model.config.vocab_size = 64000
+if not hasattr(model.config, "get_text_config") or not callable(getattr(model.config, "get_text_config", None)):
+    def _get_text_config(self):
+        return self
+    model.config.get_text_config = types.MethodType(_get_text_config, model.config)
+
+# Mirror into generation_config as well (some codepaths read there)
+try:
+    model.generation_config.vocab_size = model.config.vocab_size
+except Exception:
+    pass
+
 ip = IndicProcessor(inference=True)
 
 # -------------------- Inference helpers -------------------------------
 @torch.inference_mode()
 def _translate_to_lang(text: str, tgt_code: str, num_beams: int, max_new_tokens: int,
                        temperature: float, top_p: float, top_k: int):
+    # Preprocess via IndicTransToolkit
     batch = ip.preprocess_batch([text], src_lang=SRC_CODE, tgt_lang=tgt_code)
 
+    # Tokenize
     enc = tokenizer(
         batch,
         max_length=256,
@@ -50,8 +86,10 @@ def _translate_to_lang(text: str, tgt_code: str, num_beams: int, max_new_tokens:
         return_attention_mask=True,
     ).to(device)
 
+    # Sampling toggles
     do_sample = (temperature is not None) and (float(temperature) > 0)
 
+    # Generate
     outputs = model.generate(
         **enc,
         max_new_tokens=int(max_new_tokens),
@@ -62,16 +100,15 @@ def _translate_to_lang(text: str, tgt_code: str, num_beams: int, max_new_tokens:
         top_k=int(top_k) if do_sample else None,
         use_cache=True,
         early_stopping=False,
-        pad_token_id=tokenizer.pad_token_id or 0,
+        pad_token_id=model.generation_config.pad_token_id,
     )
 
-    with tokenizer.as_target_tokenizer():
-        decoded = tokenizer.batch_decode(
-            outputs.detach().cpu().tolist(),
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
+    # Decode
+    decoded = tokenizer.batch_decode(
+        outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )
 
+    # Postprocess back to target script
     final = ip.postprocess_batch(decoded, lang=tgt_code)
     return final[0].strip()
 
@@ -79,15 +116,23 @@ def translate_dual(text, num_beams, max_new_tokens, temperature, top_p, top_k):
     text = (text or "").strip()
     if not text:
         return "", ""
-    hi = _translate_to_lang(text, HI_CODE, num_beams, max_new_tokens, temperature, top_p, top_k)
-    te = _translate_to_lang(text, TE_CODE, num_beams, max_new_tokens, temperature, top_p, top_k)
+
+    try:
+        hi = _translate_to_lang(text, HI_CODE, num_beams, max_new_tokens, temperature, top_p, top_k)
+    except Exception as e:
+        print("HI ERROR:\n", traceback.format_exc())
+        hi = f"⚠️ Hindi translation failed: {type(e).__name__}: {str(e).splitlines()[-1]}"
+
+    try:
+        te = _translate_to_lang(text, TE_CODE, num_beams, max_new_tokens, temperature, top_p, top_k)
+    except Exception as e:
+        print("TE ERROR:\n", traceback.format_exc())
+        te = f"⚠️ Telugu translation failed: {type(e).__name__}: {str(e).splitlines()[-1]}"
+
     return hi, te
 
 # -------------------- UI (professional, clean) ------------------------
-THEME = gr.themes.Soft(
-    primary_hue="blue",
-    neutral_hue="slate",
-).set(
+THEME = gr.themes.Soft(primary_hue="blue", neutral_hue="slate").set(
     body_background_fill="#0b1220",
     body_text_color_subdued="#cbd5e1",
     block_background_fill="#0f172a",
@@ -154,5 +199,5 @@ with gr.Blocks(theme=THEME, css=CUSTOM_CSS, title="EN→HI / EN→TE Translator"
 
     gr.Markdown('<div class="footer">Model: law-ai/InLegalTrans-En2Indic-1B · Tokenizer: ai4bharat/indictrans2-en-indic-1B</div>')
 
-# IMPORTANT: remove unsupported arg; keep queue to enable request buffering
+# Keep queue to enable buffering; omit unsupported args on older Gradio
 demo.queue(max_size=48).launch()