Update app.py

app.py

@@ -3,10 +3,14 @@ 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) ----------
+# --------- 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"
+
 SRC_CODE = "eng_Latn"
 HI_CODE  = "hin_Deva"
 TE_CODE  = "tel_Telu"
@@ -15,18 +19,18 @@ TE_CODE  = "tel_Telu"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 dtype  = torch.float16 if torch.cuda.is_available() else torch.float32
 
-tokenizer = AutoTokenizer.from_pretrained(
-    TOKENIZER_ID, trust_remote_code=True, use_fast=True
-)
+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)
 
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    MODEL_ID,
+mdl_kwargs = dict(
     trust_remote_code=True,
     attn_implementation="eager",
     low_cpu_mem_usage=True,
-    torch_dtype=dtype,
+    dtype=dtype,        # <- fixes the torch_dtype deprecation warning
 )
-model = model.to(device)
+if MODEL_REV: mdl_kwargs["revision"] = MODEL_REV
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, **mdl_kwargs).to(device)
 model.eval()
 
 ip = IndicProcessor(inference=True)
@@ -35,7 +39,6 @@ ip = IndicProcessor(inference=True)
 @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):
-    """Runs IndicTrans2-style preprocess -> generate -> postprocess for a single target language."""
     batch = ip.preprocess_batch([text], src_lang=SRC_CODE, tgt_lang=tgt_code)
 
     enc = tokenizer(
@@ -62,7 +65,6 @@ def _translate_to_lang(text: str, tgt_code: str, num_beams: int, max_new_tokens:
         pad_token_id=tokenizer.pad_token_id or 0,
     )
 
-    # decode
     with tokenizer.as_target_tokenizer():
         decoded = tokenizer.batch_decode(
             outputs.detach().cpu().tolist(),
@@ -70,7 +72,6 @@ def _translate_to_lang(text: str, tgt_code: str, num_beams: int, max_new_tokens:
             clean_up_tokenization_spaces=True,
         )
 
-    # postprocess
     final = ip.postprocess_batch(decoded, lang=tgt_code)
     return final[0].strip()
 
@@ -78,7 +79,6 @@ 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)
     return hi, te
@@ -154,4 +154,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>')
 
-demo.queue(concurrency_count=4, max_size=48).launch()
+# IMPORTANT: remove unsupported arg; keep queue to enable request buffering
+demo.queue(max_size=48).launch()