add sample for testing

app.py

@@ -1,13 +1,16 @@
 # app.py
-# Minimal Gradio app:
-# - User uploads an image
-# - App loads your private HF repo (best.pt + vocab_char.json) using HF_TOKEN secret
-# - Returns OCR text
+# Minimal Gradio app with Examples:
+# - Loads your private HF model repo using HF_TOKEN (Space secret)
+# - User can upload an image OR click an example image to test quickly
 #
-# Hugging Face Space:
-# 1) Add a Space secret named HF_TOKEN (Settings → Secrets)
-# 2) Make sure your private model repo contains: best.pt, vocab_char.json
-# 3) requirements.txt should include python-multipart>=0.0.9
+# Put sample images in the Space repo root (same folder as app.py):
+#   image.png, image1.png, image2.png, image3.png, image4.png
+#
+# Space Secrets:
+#   HF_TOKEN = <your HF access token with access to the private model repo>
+#
+# Private model repo must contain:
+#   best.pt, vocab_char.json
 
 import os
 import math
@@ -31,6 +34,14 @@ REPO_ID = "SoyVitou/infinity-khmer-ocr-large"
 CKPT_FILENAME = "best.pt"
 VOCAB_FILENAME = "vocab_char.json"
 
+EXAMPLES = [
+    ["./image.png"],
+    ["./image1.png"],
+    ["./image2.png"],
+    ["./image3.png"],
+    ["./image4.png"],
+]
+
 
 @dataclass
 class CFG:
@@ -75,7 +86,6 @@ class CFG:
     UNK_LOGP_PENALTY: float = 1.0
 
     EMA_DECAY: float = 0.999
-    VOCAB_JSON: str = VOCAB_FILENAME
 
     USE_FP16: bool = True
     USE_AUTOCAST: bool = True
@@ -83,13 +93,7 @@ class CFG:
 
 
 class CharTokenizer:
-    def __init__(
-        self,
-        vocab_json: str,
-        unk_token: str = "<unk>",
-        collapse_whitespace: bool = True,
-        unicode_nfc: bool = True
-    ):
+    def __init__(self, vocab_json: str, unk_token: str = "<unk>", collapse_whitespace: bool = True, unicode_nfc: bool = True):
         with open(vocab_json, "r", encoding="utf-8") as f:
             vocab_raw: Dict[str, int] = json.load(f)
 
@@ -105,8 +109,6 @@ class CharTokenizer:
 
         self.unk_token = unk_token
         self.unk_id = self.token_to_id[unk_token]
-        self.collapse_whitespace = collapse_whitespace
-        self.unicode_nfc = unicode_nfc
 
         self.blank_id = 0
         self.pad_id = 1
@@ -218,8 +220,6 @@ class ConvStem(nn.Module):
 class HybridContextOCRV2(nn.Module):
     def __init__(self, cfg: CFG, tok: CharTokenizer):
         super().__init__()
-        self.cfg = cfg
-        self.tok = tok
         d = cfg.DROPOUT
 
         self.stem = ConvStem(cfg.ENC_DIM, d)
@@ -238,35 +238,29 @@ class HybridContextOCRV2(nn.Module):
         self.enc = nn.TransformerEncoder(enc_layer, num_layers=cfg.ENC_LAYERS)
         self.enc_ln = nn.LayerNorm(cfg.ENC_DIM)
 
-        self.use_ctc = cfg.USE_CTC
-        if self.use_ctc:
-            self.ctc_head = nn.Sequential(
-                nn.LayerNorm(cfg.ENC_DIM),
-                nn.Dropout(d),
-                nn.Linear(cfg.ENC_DIM, tok.ctc_classes),
-            )
-
-        self.use_decoder = cfg.USE_DECODER
-        if self.use_decoder:
-            self.mem_proj = nn.Linear(cfg.ENC_DIM, cfg.DEC_DIM, bias=False)
-            self.dec_emb = nn.Embedding(tok.dec_vocab, cfg.DEC_DIM)
-
-            dec_layer = nn.TransformerDecoderLayer(
-                d_model=cfg.DEC_DIM,
-                nhead=cfg.DEC_HEADS,
-                dim_feedforward=cfg.DEC_FF,
-                dropout=d,
-                batch_first=True,
-                activation="gelu",
-                norm_first=True,
-            )
-            self.dec = nn.TransformerDecoder(dec_layer, num_layers=cfg.DEC_LAYERS)
-            self.dec_ln = nn.LayerNorm(cfg.DEC_DIM)
-            self.dec_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)
-
-        self.use_lm = cfg.USE_LM
-        if self.use_lm:
-            self.lm_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)
+        self.ctc_head = nn.Sequential(
+            nn.LayerNorm(cfg.ENC_DIM),
+            nn.Dropout(d),
+            nn.Linear(cfg.ENC_DIM, tok.ctc_classes),
+        )
+
+        self.mem_proj = nn.Linear(cfg.ENC_DIM, cfg.DEC_DIM, bias=False)
+        self.dec_emb = nn.Embedding(tok.dec_vocab, cfg.DEC_DIM)
+
+        dec_layer = nn.TransformerDecoderLayer(
+            d_model=cfg.DEC_DIM,
+            nhead=cfg.DEC_HEADS,
+            dim_feedforward=cfg.DEC_FF,
+            dropout=d,
+            batch_first=True,
+            activation="gelu",
+            norm_first=True,
+        )
+        self.dec = nn.TransformerDecoder(dec_layer, num_layers=cfg.DEC_LAYERS)
+        self.dec_ln = nn.LayerNorm(cfg.DEC_DIM)
+        self.dec_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)
+
+        self.lm_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)
 
     def encode(self, imgs: torch.Tensor) -> torch.Tensor:
         x = self.stem(imgs)
@@ -280,8 +274,7 @@ class HybridContextOCRV2(nn.Module):
 
 
 class EMA:
-    def __init__(self, model: nn.Module, decay: float):
-        self.decay = decay
+    def __init__(self, model: nn.Module):
         self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()}
 
     @torch.no_grad()
@@ -330,9 +323,6 @@ def beam_decode_one_batched(
             max_steps = min(max_steps, max(1, int(target_len * cfg.DEC_MAX_LEN_RATIO) + cfg.DEC_MAX_LEN_PAD))
         else:
             max_steps = min(max_steps, cfg.DEC_MAX_LEN_PAD)
-    else:
-        mem_len = mem_proj_1.size(1)
-        max_steps = min(max_steps, max(1, int(mem_len * cfg.MEM_MAX_LEN_RATIO) + cfg.DEC_MAX_LEN_PAD))
 
     full_causal = torch.triu(
         torch.ones((cfg.MAX_DEC_LEN + 2, cfg.MAX_DEC_LEN + 2), device=device, dtype=torch.bool),
@@ -368,28 +358,15 @@ def beam_decode_one_batched(
             logits = model.dec_head(out)[:, -1, :]
             logp = F.log_softmax(logits, dim=-1)
 
-            if cfg.USE_LM and cfg.USE_LM_FUSION_EVAL:
-                lm_logits = model.lm_head(out)[:, -1, :]
-                logp = logp + cfg.LM_FUSION_ALPHA * F.log_softmax(lm_logits, dim=-1)
+            lm_logits = model.lm_head(out)[:, -1, :]
+            logp = logp + cfg.LM_FUSION_ALPHA * F.log_softmax(lm_logits, dim=-1)
 
         unk_id = tok.unk_id + tok.dec_offset
 
         for i, (_, seq, _) in enumerate(alive):
             cur_len = len(seq) - 1
-
-            if target_len is not None and target_len > 0:
-                min_len = min(cfg.EOS_BIAS_UNTIL_LEN, max(1, int(target_len * 0.7)))
-                if cur_len < min_len:
-                    logp[i, tok.dec_eos] = logp[i, tok.dec_eos] - cfg.EOS_LOGP_BIAS
-                elif cur_len >= target_len:
-                    logp[i, tok.dec_eos] = logp[i, tok.dec_eos] + cfg.EOS_LOGP_BOOST
-            else:
-                if cur_len < cfg.EOS_BIAS_UNTIL_LEN:
-                    logp[i, tok.dec_eos] = logp[i, tok.dec_eos] - cfg.EOS_LOGP_BIAS
-
-            if len(seq) >= 4 and seq[-1] == seq[-2] == seq[-3]:
-                logp[i, seq[-1]] = logp[i, seq[-1]] - cfg.REPEAT_LAST_PENALTY
-
+            if cur_len < cfg.EOS_BIAS_UNTIL_LEN:
+                logp[i, tok.dec_eos] = logp[i, tok.dec_eos] - cfg.EOS_LOGP_BIAS
             logp[i, unk_id] = logp[i, unk_id] - cfg.UNK_LOGP_PENALTY
 
         topv, topi = torch.topk(logp, k=cfg.BEAM, dim=-1)
@@ -410,65 +387,7 @@ def beam_decode_one_batched(
         new_beams.sort(key=lambda x: normed(x[0], x[1]), reverse=True)
         beams = new_beams[:cfg.BEAM]
 
-    def length_norm(score: float, seq: List[int]) -> float:
-        return score / (max(1, len(seq) - 1) ** cfg.BEAM_LENP)
-
-    if ctc_logits_1 is not None and cfg.CTC_FUSION_ALPHA > 0:
-        log_probs = F.log_softmax(ctc_logits_1.squeeze(0), dim=-1)
-
-        def ctc_sequence_log_prob(label_ids: List[int]) -> torch.Tensor:
-            if len(label_ids) == 0:
-                return log_probs[:, tok.blank_id].sum()
-
-            blank = tok.blank_id
-            ext = [blank]
-            for lid in label_ids:
-                ext.append(lid)
-                ext.append(blank)
-
-            s_len = len(ext)
-            alpha = log_probs.new_full((s_len,), float("-inf"))
-            alpha[0] = log_probs[0, blank]
-            alpha[1] = log_probs[0, ext[1]]
-
-            for t in range(1, log_probs.size(0)):
-                next_alpha = log_probs.new_full((s_len,), float("-inf"))
-                for s in range(s_len):
-                    candidates = [alpha[s]]
-                    if s - 1 >= 0:
-                        candidates.append(alpha[s - 1])
-                    if s - 2 >= 0 and ext[s] != blank and ext[s] != ext[s - 2]:
-                        candidates.append(alpha[s - 2])
-                    next_alpha[s] = torch.logsumexp(torch.stack(candidates), dim=0) + log_probs[t, ext[s]]
-                alpha = next_alpha
-
-            if s_len == 1:
-                return alpha[0]
-            return torch.logsumexp(torch.stack([alpha[s_len - 1], alpha[s_len - 2]]), dim=0)
-
-        def seq_to_ctc_labels(seq: List[int]) -> List[int]:
-            labels = []
-            for x in seq[1:]:
-                if x == tok.dec_eos:
-                    break
-                if x in (tok.dec_pad, tok.dec_bos):
-                    continue
-                y = x - tok.dec_offset
-                if 0 <= y < tok.vocab_size:
-                    labels.append(y + tok.ctc_offset)
-                else:
-                    labels.append(tok.unk_id + tok.ctc_offset)
-            return labels
-
-        def combined_score(entry):
-            dec_score = length_norm(entry[0], entry[1])
-            labels = seq_to_ctc_labels(entry[1])
-            ctc_score = ctc_sequence_log_prob(labels) / max(1, len(labels))
-            return dec_score + cfg.CTC_FUSION_ALPHA * float(ctc_score)
-
-        best = max(beams, key=combined_score)[1]
-    else:
-        best = max(beams, key=lambda x: length_norm(x[0], x[1]))[1]
+    best = max(beams, key=lambda x: x[0])[1]
 
     ids = []
     for x in best[1:]:
@@ -515,12 +434,7 @@ def load_model():
             if hasattr(cfg, k):
                 setattr(cfg, k, v)
 
-    tok = CharTokenizer(
-        vocab_path,
-        unk_token=cfg.UNK_TOKEN,
-        collapse_whitespace=cfg.COLLAPSE_WHITESPACE,
-        unicode_nfc=cfg.UNICODE_NFC,
-    )
+    tok = CharTokenizer(vocab_path, unk_token=cfg.UNK_TOKEN)
 
     device = setup_device(cfg)
 
@@ -528,7 +442,7 @@ def load_model():
     model.load_state_dict(ckpt["model"], strict=True)
 
     if isinstance(ckpt, dict) and "ema" in ckpt and isinstance(ckpt["ema"], dict):
-        ema = EMA(model, decay=cfg.EMA_DECAY)
+        ema = EMA(model)
         ema.shadow = {k: v.detach().clone() for k, v in ckpt["ema"].items()}
         ema.copy_to(model)
 
@@ -550,10 +464,9 @@ def predict(img: Image.Image) -> str:
     x = preprocess_pil(CFG_OBJ, img).to(DEVICE)
     if CFG_OBJ.USE_FP16 and DEVICE == "cuda":
         x = x.half()
-
     mem = MODEL.encode(x)
     mem_proj = MODEL.mem_proj(mem)
-    ctc_logits = MODEL.ctc_head(mem) if CFG_OBJ.USE_CTC else None
+    ctc_logits = MODEL.ctc_head(mem)
     return beam_decode_one_batched(MODEL, mem_proj, TOK, CFG_OBJ, ctc_logits_1=ctc_logits)
 
 
@@ -562,6 +475,8 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil", label="Upload image"),
     outputs=gr.Textbox(label="OCR result", lines=6),
     title="Infinity Khmer OCR",
+    examples=EXAMPLES,
+    cache_examples=False,
 )
 
 if __name__ == "__main__":