Update app.py

app.py

@@ -1,24 +1,27 @@
-import os, torch
+import os, traceback, torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from IndicTransToolkit import IndicProcessor  # https://github.com/VarunGumma/IndicTransToolkit
 
-# --------- Config (override via Space Variables if you like) ----------
+# --- Robust import for IndicProcessor (fallback path per toolkit README) ---
+try:
+    from IndicTransToolkit import IndicProcessor  # preferred
+except Exception:
+    from IndicTransToolkit.IndicTransToolkit import IndicProcessor  # fallback
+
+# ------------------- Config -------------------
 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"
 TE_CODE  = "tel_Telu"
 
-# -------------------- Load model & tokenizer --------------------------
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 dtype  = torch.float16 if torch.cuda.is_available() else torch.float32
 
+# ------------------- Load model/tokenizer -------------------
 tok_kwargs = dict(trust_remote_code=True, use_fast=True)
 if TOKENIZER_REV: tok_kwargs["revision"] = TOKENIZER_REV
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID, **tok_kwargs)
@@ -27,20 +30,29 @@ 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
+
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, **mdl_kwargs).to(device)
 model.eval()
 
 ip = IndicProcessor(inference=True)
 
-# -------------------- Inference helpers -------------------------------
+# Ensure pad/eos ids are set to avoid edge-case crashes
+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
+
+# ------------------- Inference -------------------
 @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
     batch = ip.preprocess_batch([text], src_lang=SRC_CODE, tgt_lang=tgt_code)
 
+    # Tokenize
     enc = tokenizer(
         batch,
         max_length=256,
@@ -50,8 +62,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 +76,13 @@ 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 (no deprecated as_target_tokenizer)
+    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
+    # Postprocess
     final = ip.postprocess_batch(decoded, lang=tgt_code)
     return final[0].strip()
 
@@ -79,15 +90,25 @@ 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:
+        # Surface a friendly error instead of Gradio's generic "Error"
+        msg = f"⚠️ Hindi translation failed: {type(e).__name__}: {str(e).splitlines()[-1]}"
+        print("HI ERROR:\n", traceback.format_exc())
+        hi = msg
+
+    try:
+        te = _translate_to_lang(text, TE_CODE, num_beams, max_new_tokens, temperature, top_p, top_k)
+    except Exception as e:
+        msg = f"⚠️ Telugu translation failed: {type(e).__name__}: {str(e).splitlines()[-1]}"
+        print("TE ERROR:\n", traceback.format_exc())
+        te = msg
+
     return hi, te
 
-# -------------------- UI (professional, clean) ------------------------
-THEME = gr.themes.Soft(
-    primary_hue="blue",
-    neutral_hue="slate",
-).set(
+# ------------------- UI -------------------
+THEME = gr.themes.Soft(primary_hue="blue", neutral_hue="slate").set(
     body_background_fill="#0b1220",
     body_text_color_subdued="#cbd5e1",
     block_background_fill="#0f172a",
@@ -114,12 +135,7 @@ with gr.Blocks(theme=THEME, css=CUSTOM_CSS, title="EN→HI / EN→TE Translator"
 
     with gr.Row():
         with gr.Column(scale=1):
-            src = gr.Textbox(
-                label="English Text",
-                placeholder="Type English here…",
-                lines=8,
-                autofocus=True,
-            )
+            src = gr.Textbox(label="English Text", placeholder="Type English here…", lines=8, autofocus=True)
             with gr.Accordion("Advanced settings", open=False):
                 num_beams = gr.Slider(1, 8, value=4, step=1, label="Beam search: num_beams")
                 max_new   = gr.Slider(16, 512, value=128, step=8, label="Max new tokens")
@@ -135,11 +151,9 @@ with gr.Blocks(theme=THEME, css=CUSTOM_CSS, title="EN→HI / EN→TE Translator"
             te_out = gr.Textbox(label="Telugu (tel_Telu)", lines=8, show_copy_button=True)
 
     examples = gr.Examples(
-        examples=[
-            ["The Constitution guarantees fundamental rights to every citizen of India."],
-            ["Maintenance proceedings shall commence within thirty days from the date of application."],
-            ["The agreement shall remain in force unless terminated by mutual consent in writing."],
-        ],
+        examples=[["The Constitution guarantees fundamental rights to every citizen of India."],
+                  ["Maintenance proceedings shall commence within thirty days from the date of application."],
+                  ["The agreement shall remain in force unless terminated by mutual consent in writing."]],
         inputs=[src],
         label="Quick examples",
     )
@@ -154,5 +168,4 @@ 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
 demo.queue(max_size=48).launch()