ack5

app.py

@@ -6,7 +6,10 @@ from tokenizers import Tokenizer
 
 ROOT = Path(__file__).parent
 
+# ✅ match your actual repo folder name (case-sensitive on Spaces)
 MODEL_ROOT = ROOT / "vMAX_512_8L_32k_OpenSubs"
+
+# ✅ adjust this to what your model was trained on
 PREFIX = "translate English to Indonesian: "
 
 
@@ -19,7 +22,6 @@ def latest_checkpoint(root: Path) -> Path | None:
         ckpts.sort(key=lambda p: int(p.name.split("-")[-1]))
         return ckpts[-1]
 
-    # fallback: root itself contains model files
     if (root / "config.json").exists() or (root / "model.safetensors").exists():
         return root
 
@@ -27,6 +29,9 @@ def latest_checkpoint(root: Path) -> Path | None:
 
 
 CKPT = latest_checkpoint(MODEL_ROOT)
+print("[BOOT] MODEL_ROOT =", MODEL_ROOT)
+print("[BOOT] CKPT =", CKPT)
+
 _pipe = {"tok": None, "model": None, "device": None, "ckpt": None}
 
 
@@ -36,15 +41,11 @@ def model_ready():
 
 class SimpleTokenizerWrapper:
     """
-    Minimal tokenizer wrapper around `tokenizers.Tokenizer` that matches what this app needs:
-    - callable returning PyTorch tensors (input_ids + attention_mask)
-    - decode(ids)
-    - eos_token_id attribute (optional)
+    Wrapper over tokenizers.Tokenizer so we avoid the TokenizersBackend crash on Spaces.
+    Needs: __call__ -> tensors, decode(), eos_token_id (optional)
     """
     def __init__(self, base_tok: Tokenizer):
         self.base_tok = base_tok
-
-        # Try to discover EOS/PAD ids if they exist in the vocab; safe if missing
         self.eos_token_id = self._find_id(["</s>", "<eos>", "[EOS]"])
         self.pad_token_id = self._find_id(["<pad>", "[PAD]"])
 
@@ -60,14 +61,9 @@ class SimpleTokenizerWrapper:
 
     def __call__(self, text, return_tensors="pt", truncation=True, max_length=256):
         enc = self.base_tok.encode(text)
-        ids = enc.ids
-
-        if truncation and max_length is not None:
-            ids = ids[:max_length]
-
+        ids = enc.ids[:max_length] if (truncation and max_length is not None) else enc.ids
         input_ids = torch.tensor([ids], dtype=torch.long)
         attention_mask = torch.ones_like(input_ids)
-
         return {"input_ids": input_ids, "attention_mask": attention_mask}
 
     def decode(self, ids, skip_special_tokens=True):
@@ -84,98 +80,104 @@ def load_tokenizer(ckpt: Path) -> SimpleTokenizerWrapper:
     return SimpleTokenizerWrapper(base)
 
 
+def load_eos_meme_b64() -> str:
+    b64_path = ROOT / "eos_ack.b64"
+    if not b64_path.exists():
+        return ""
+    return b64_path.read_text(encoding="utf-8").strip()
+
+
+EOS_ACK_B64 = load_eos_meme_b64()
+
+
 def get_model():
-    if _pipe["model"] is None:
-        if CKPT is None:
-            raise RuntimeError(f"No checkpoint found under: {MODEL_ROOT}")
+    if _pipe["model"] is not None:
+        return _pipe["tok"], _pipe["model"], _pipe["device"]
 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if CKPT is None:
+        raise RuntimeError(f"No checkpoint found under: {MODEL_ROOT}")
 
-        # ✅ bypass broken tokenizer_config.json (TokenizersBackend error on Spaces)
-        tok = load_tokenizer(CKPT)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("[LOAD] device =", device)
+    print("[LOAD] loading from =", CKPT)
 
-        model = AutoModelForSeq2SeqLM.from_pretrained(CKPT).to(device)
-        model.eval()
+    tok = load_tokenizer(CKPT)
 
-        _pipe.update({"tok": tok, "model": model, "device": device, "ckpt": str(CKPT)})
+    # low-memory load
+    kwargs = dict(low_cpu_mem_usage=True)
+    if device == "cuda":
+        kwargs["torch_dtype"] = torch.float16  # comment out if you don’t want fp16
 
-        # quick startup self-test (will run on first translate click)
-        try:
-            test = translate_text("I am happy")
-            print(f"[LOAD OK] ckpt={CKPT} device={device} test_out={repr(test)}")
-        except Exception as e:
-            print(f"[LOAD FAIL] ckpt={CKPT} device={device} err={e}")
-            raise
+    model = AutoModelForSeq2SeqLM.from_pretrained(CKPT, **kwargs).to(device)
+    model.eval()
+
+    _pipe.update({"tok": tok, "model": model, "device": device, "ckpt": str(CKPT)})
+    print("[LOAD] OK")
 
-    return _pipe["tok"], _pipe["model"], _pipe["device"]
+    return tok, model, device
 
 
 def translate_text(text: str):
     text = (text or "").strip()
     if not text:
         return ""
+
     if not model_ready():
         return "[Model not ready — checkpoint folder not found.]"
 
-    tok, model, device = get_model()
-
-    prompt = PREFIX + text
-    inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=256)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    try:
+        tok, model, device = get_model()
 
-    input_len = inputs["input_ids"].shape[1]
-    max_new = int(input_len * 1.3)
-    max_new = min(max_new, 40)
+        prompt = PREFIX + text
+        inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=256)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
 
-    gen_kwargs = dict(
-        num_beams=2,
-        length_penalty=1.15,
-        no_repeat_ngram_size=4,
-        repetition_penalty=1.30,
-        early_stopping=True,
-        min_new_tokens=4,
-        max_new_tokens=max_new,
-    )
+        input_len = inputs["input_ids"].shape[1]
+        max_new = min(int(input_len * 1.3), 40)
 
-    # only pass eos_token_id if we found one
-    if getattr(tok, "eos_token_id", None) is not None:
-        gen_kwargs["eos_token_id"] = tok.eos_token_id
+        gen_kwargs = dict(
+            num_beams=2,
+            length_penalty=1.15,
+            no_repeat_ngram_size=4,
+            repetition_penalty=1.30,
+            early_stopping=True,
+            min_new_tokens=4,
+            max_new_tokens=max_new,
+        )
 
-    with torch.no_grad():
-        out_ids = model.generate(**inputs, **gen_kwargs)
+        if getattr(tok, "eos_token_id", None) is not None:
+            gen_kwargs["eos_token_id"] = tok.eos_token_id
 
-    result = tok.decode(out_ids[0], skip_special_tokens=True).strip()
+        with torch.no_grad():
+            out_ids = model.generate(**inputs, **gen_kwargs)
 
-    if not result:
-        raw = tok.decode(out_ids[0])
-        return f"[Blank output] raw={raw!r} ckpt={_pipe['ckpt']} device={device}"
+        return tok.decode(out_ids[0], skip_special_tokens=True).strip()
 
-    return result
+    except Exception as e:
+        return f"[ERROR] {type(e).__name__}: {e}"
 
 
-title = "Translator Demo"
-status = "✅ Model found" if model_ready() else "⏳ Model not found"
-
 with gr.Blocks() as demo:
-    gr.Markdown(f"# {title}")
-    gr.Markdown(f"**Status:** {status}")
+    gr.Markdown("# Translator Demo")
+    gr.Markdown(f"**Status:** {'✅ Model found' if model_ready() else '⏳ Model not found'}")
     gr.Markdown(f"**Model root:** `{MODEL_ROOT}`")
-    gr.Markdown(f"**Loaded checkpoint:** `{str(CKPT) if CKPT else 'None'}`")
+    gr.Markdown(f"**Checkpoint:** `{str(CKPT) if CKPT else 'None'}`")
     gr.Markdown(f"**Prefix:** `{PREFIX}`")
 
-    # ✅ EOS note + image
+    # ✅ EOS note + embedded meme (no png file, so HF won’t block the push)
     gr.Markdown("### ⚠️ Note: EOS wasn’t properly learnt by the model")
-    gr.Image(
-        value=str(ROOT / "eos_ack.png"),
-        show_label=False,
-        interactive=False,
-        height=260,
-    )
-
-    inp = gr.Textbox(label="Input", lines=4, placeholder="Type here...")
+    if EOS_ACK_B64:
+        gr.HTML(f"""
+        <div style="max-width:520px;">
+          <img src="data:image/png;base64,{EOS_ACK_B64}"
+               style="width:100%; border-radius:12px;" />
+        </div>
+        """)
+    else:
+        gr.Markdown("_Missing `eos_ack.b64` — meme not loaded._")
+
+    inp = gr.Textbox(label="Input", lines=4)
     out = gr.Textbox(label="Output", lines=4)
-
-    btn = gr.Button("Translate")
-    btn.click(translate_text, inp, out)
+    gr.Button("Translate").click(translate_text, inp, out)
 
 demo.launch()