tokenizer.json compatibility fixup for TrOCR on transformers 5.x

trocr_compat.py

@@ -0,0 +1,134 @@
+"""Compatibility fixup: make TrOCRProcessor.from_pretrained work on transformers
+5.x for microsoft/trocr-base-printed with a cold/legacy HF cache.
+
+transformers 5.x dropped support for slow RoBERTa tokenizers, so TrOCR
+checkpoints (vocab.json + merges.txt, no tokenizer.json) fail to load. Fix:
+locate the cached snapshot dir; if tokenizer.json is missing, fetch
+vocab.json + merges.txt via hf_hub_download, build an equivalent fast-tokenizer
+tokenizer.json (GPT2-style BPE with the RoBERTa special tokens and a
+TemplateProcessing <s> $A </s> post-processor), write it into the snapshot dir,
+and clear any .no_exist marker for tokenizer.json. Idempotent: if
+tokenizer.json already exists, the function returns immediately (no-op).
+"""
+import glob
+import json
+import os
+
+
+def _hub_cache_dir():
+    for var in ("HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"):
+        if os.environ.get(var):
+            return os.path.expanduser(os.environ[var])
+    hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")
+    return os.path.join(hf_home, "hub")
+
+
+def _repo_dirname(model_id):
+    return "models--" + model_id.replace("/", "--")
+
+
+def _find_snapshot_dir(repo_dir, revision="main"):
+    """Pick the snapshot matching refs/<revision>, else the newest one."""
+    ref_file = os.path.join(repo_dir, "refs", revision)
+    if os.path.isfile(ref_file):
+        with open(ref_file) as f:
+            commit = f.read().strip()
+        cand = os.path.join(repo_dir, "snapshots", commit)
+        if os.path.isdir(cand):
+            return cand
+    snaps = sorted(
+        glob.glob(os.path.join(repo_dir, "snapshots", "*")), key=os.path.getmtime
+    )
+    return snaps[-1] if snaps else None
+
+
+def ensure_trocr_tokenizer(model_id="microsoft/trocr-base-printed"):
+    """Patch the HF cache so TrOCRProcessor.from_pretrained(model_id) works on
+    transformers 5.x. Idempotent; raises RuntimeError if it cannot fix the
+    cache. Returns the path to the (existing or newly built) tokenizer.json."""
+    try:
+        from huggingface_hub import hf_hub_download, snapshot_download
+        from tokenizers import Tokenizer, decoders, pre_tokenizers, processors
+        from tokenizers.models import BPE
+    except ImportError as e:
+        raise RuntimeError(
+            "ensure_trocr_tokenizer needs the 'huggingface_hub' and "
+            f"'tokenizers' packages installed: {e}"
+        )
+
+    repo_dir = os.path.join(_hub_cache_dir(), _repo_dirname(model_id))
+    snapshot_dir = None
+    if os.path.isdir(repo_dir):
+        snapshot_dir = _find_snapshot_dir(repo_dir)
+
+    if snapshot_dir is None:
+        # Cold cache: fetch only the small config/tokenizer files (not the
+        # weights); this also creates the standard snapshot layout.
+        snapshot_dir = snapshot_download(model_id, allow_patterns=["*.json", "*.txt"])
+
+    tok_path = os.path.join(snapshot_dir, "tokenizer.json")
+    if os.path.exists(tok_path):
+        return tok_path  # already fixed -> idempotent no-op
+
+    vocab_file = hf_hub_download(model_id, "vocab.json")
+    merges_file = hf_hub_download(model_id, "merges.txt")
+    if not (os.path.isfile(vocab_file) and os.path.isfile(merges_file)):
+        raise RuntimeError(
+            "could not fetch vocab.json/merges.txt for %s; cannot build " % model_id
+            + "tokenizer.json"
+        )
+
+    with open(vocab_file, encoding="utf-8") as f:
+        vocab = json.load(f)
+    merges = []
+    with open(merges_file, encoding="utf-8") as f:
+        for line in f:
+            line = line.rstrip("\n")
+            if not line or line.startswith("#version"):
+                continue
+            parts = line.split(" ")
+            if len(parts) == 2:
+                merges.append(tuple(parts))
+
+    specials = ["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
+    missing = [t for t in specials if t not in vocab]
+    if missing:
+        raise RuntimeError(
+            "vocab.json for %s lacks special tokens %s; cannot build a "
+            "compatible tokenizer.json" % (model_id, missing)
+        )
+
+    bpe = BPE(vocab=vocab, merges=merges, unk_token="<unk>")
+    tok = Tokenizer(bpe)
+    # <s> <pad> </s> <unk> <mask> are already plain vocab entries in TrOCR's
+    # vocab.json, so no add_special_tokens()/added_tokens are needed
+    # (byte-identical behavior to the manual build).
+    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
+    tok.decoder = decoders.ByteLevel()
+    tok.post_processor = processors.TemplateProcessing(
+        single="<s> $A </s>",
+        pair="<s> $A </s> </s> $B",
+        special_tokens=[("<s>", vocab["<s>"]), ("</s>", vocab["</s>"])],
+    )
+
+    # Write atomically so a crash never leaves a partial tokenizer.json.
+    tmp_path = tok_path + ".tmp"
+    with open(tmp_path, "w", encoding="utf-8") as f:
+        f.write(tok.to_str())
+    os.replace(tmp_path, tok_path)
+
+    # Clear cached .no_exist markers so transformers does not conclude
+    # tokenizer.json is absent in the repo.
+    no_exist = os.path.join(repo_dir, ".no_exist")
+    if os.path.isdir(no_exist):
+        for marker in glob.glob(os.path.join(no_exist, "*", "tokenizer.json")):
+            try:
+                os.remove(marker)
+            except OSError:
+                pass
+
+    return tok_path
+
+
+if __name__ == "__main__":
+    print(ensure_trocr_tokenizer())