Update app.py

app.py

@@ -84,13 +84,18 @@ def load_model(model_name: str):
     if model_name in _model_cache:
         return _model_cache[model_name]
 
-    tok = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True, use_fast=True)
+    token = os.getenv("hf_token")
+
+    tok = AutoTokenizer.from_pretrained(
+        "ai4bharat/indictrans2-en-indic-1B",
+        trust_remote_code=True, use_fast=True
+    )
     mdl = AutoModelForSeq2SeqLM.from_pretrained(
         model_name, trust_remote_code=True,
-        low_cpu_mem_usage=True, dtype=dtype
+        low_cpu_mem_usage=True, dtype=dtype, token = token
     ).to(device).eval()
 
-    # Fix vocab
+    # Fix vocab (some HF models have mismatched config.vocab_size)
     try:
         mdl.config.vocab_size = mdl.get_output_embeddings().weight.shape[0]
     except Exception:
@@ -109,56 +114,88 @@ def build_bad_words_ids_from_vocab(tok):
     ] + [f"<ID{i}>" for i in range(10)]
     out = []
     for c in candidates:
-        if c in vocab: out.append([vocab[c]]); continue
+        if c in vocab:
+            out.append([vocab[c]])
+            continue
         sp_c = "▁" + c
-        if sp_c in vocab: out.append([vocab[sp_c]])
+        if sp_c in vocab:
+            out.append([vocab[sp_c]])
     return out
 
-# --------------------- Translation ---------------------
+
+# --------------------- Streaming Translation ---------------------
+BATCH_SIZE = 6
+
 @torch.inference_mode()
-def _translate(text: str, tgt_lang: str, model_choice: str,
-               num_beams=4, max_new=128, batch_size=3) -> str:
-    
+def translate_dual_stream(text, model_choice, num_beams, max_new):
+    """
+    Generator that yields (hindi_accumulated_text, telugu_accumulated_text)
+    after each processed batch so the UI updates progressively.
+    """
+    if not text or not text.strip():
+        yield "", ""
+        return
+
+    # Prepare once
     tok, mdl = load_model(MODELS[model_choice])
     BAD_WORDS_IDS = build_bad_words_ids_from_vocab(tok)
-
     sentences = split_into_sentences(text)
-    full_trans = []
-
-    for i in range(0, len(sentences), batch_size):
-        batch = sentences[i:i+batch_size]
-        proc = ip.preprocess_batch(batch, src_lang=SRC_CODE, tgt_lang=tgt_lang)
-        enc = tok(proc, padding=True, truncation=True, max_length=256, return_tensors="pt").to(device)
-
-        out = mdl.generate(
-            **enc, max_length=max_new, num_beams=num_beams,
-            early_stopping=True, no_repeat_ngram_size=3, use_cache=False,
-            bad_words_ids=BAD_WORDS_IDS if BAD_WORDS_IDS else None
-        )
 
-        decoded = tok.batch_decode(out, skip_special_tokens=True)
-        decoded = [strip_lang_tags(t) for t in decoded]
-        post = ip.postprocess_batch(decoded, lang=tgt_lang)
-
-        if tgt_lang == HI_CODE:
-            post = [ensure_hindi_danda(x) for x in post]
-
-        full_trans.extend(p.strip() for p in post)
-
-    return " ".join(full_trans)
-
-
-def translate_dual(text, model_choice, num_beams, max_new):
-    if not text.strip(): return "", ""
-    try:
-        hi = _translate(text, HI_CODE, model_choice, num_beams=num_beams, max_new=max_new)
-    except Exception as e:
-        hi = f"⚠️ Hindi failed: {e}"
-    try:
-        te = _translate(text, TE_CODE, model_choice, num_beams=num_beams, max_new=max_new)
-    except Exception as e:
-        te = f"⚠️ Telugu failed: {e}"
-    return hi, te
+    hi_acc, te_acc = [], []
+
+    # Clear outputs immediately for a snappy feel
+    yield "", ""
+
+    for i in range(0, len(sentences), BATCH_SIZE):
+        batch = sentences[i:i + BATCH_SIZE]
+
+        # --- Hindi batch ---
+        try:
+            proc_hi = ip.preprocess_batch(batch, src_lang=SRC_CODE, tgt_lang=HI_CODE)
+            enc_hi  = tok(
+                proc_hi, padding=True, truncation=True, max_length=256, return_tensors="pt"
+            ).to(device)
+            out_hi  = mdl.generate(
+                **enc_hi,
+                max_length=max_new,              # keep semantics same as your original
+                num_beams=int(num_beams),
+                early_stopping=True,
+                no_repeat_ngram_size=3,
+                use_cache=False,
+                bad_words_ids=BAD_WORDS_IDS if BAD_WORDS_IDS else None
+            )
+            dec_hi  = tok.batch_decode(out_hi, skip_special_tokens=True)
+            dec_hi  = [strip_lang_tags(t) for t in dec_hi]
+            post_hi = ip.postprocess_batch(dec_hi, lang=HI_CODE)
+            post_hi = [ensure_hindi_danda(x) for x in post_hi]
+            hi_acc.extend(p.strip() for p in post_hi)
+        except Exception as e:
+            hi_acc.append(f"⚠️ Hindi failed (batch {i//BATCH_SIZE+1}): {e}")
+
+        # --- Telugu batch ---
+        try:
+            proc_te = ip.preprocess_batch(batch, src_lang=SRC_CODE, tgt_lang=TE_CODE)
+            enc_te  = tok(
+                proc_te, padding=True, truncation=True, max_length=256, return_tensors="pt"
+            ).to(device)
+            out_te  = mdl.generate(
+                **enc_te,
+                max_length=max_new,
+                num_beams=int(num_beams),
+                early_stopping=True,
+                no_repeat_ngram_size=3,
+                use_cache=False,
+                bad_words_ids=BAD_WORDS_IDS if BAD_WORDS_IDS else None
+            )
+            dec_te  = tok.batch_decode(out_te, skip_special_tokens=True)
+            dec_te  = [strip_lang_tags(t) for t in dec_te]
+            post_te = ip.postprocess_batch(dec_te, lang=TE_CODE)
+            te_acc.extend(p.strip() for p in post_te)
+        except Exception as e:
+            te_acc.append(f"⚠️ Telugu failed (batch {i//BATCH_SIZE+1}): {e}")
+
+        # Stream the accumulators so far
+        yield (" ".join(hi_acc), " ".join(te_acc))
 
 
 # --------------------- Dark Theme ---------------------
@@ -169,7 +206,7 @@ THEME = gr.themes.Soft(
     body_text_color="#f3f4f6",
     block_background_fill="#111827",
     block_border_color="#1f2937",
-    block_title_text_color="#e5e7eb",
+    block_title_text_color="#123456",
     button_primary_background_fill="#2563eb",
     button_primary_text_color="#ffffff",
 )
@@ -262,11 +299,13 @@ with gr.Blocks(theme=THEME, css=CUSTOM_CSS, title="EN → HI/TE Translator") as
                 num_beams   = gr.Slider(1, 8, value=4, step=1, label="Beam Search", elem_id="model_dd")
                 max_new     = gr.Slider(32, 512, value=128, step=16, label="Max New Tokens", elem_id="model_dd")
 
+    # Use streaming generator
     translate_btn.click(
-        translate_dual,
+        translate_dual_stream,
         inputs=[src, model_choice, num_beams, max_new],
         outputs=[hi_out, te_out]
     )
     clear_btn.click(lambda: ("", "", ""), outputs=[src, hi_out, te_out])
 
+# Enable queue for streaming
 demo.queue(max_size=48).launch()