Fix remote tokenizer loading for transformers inference

inference_mask.py

@@ -2,9 +2,6 @@
 import argparse
 import json
 import os
-import shutil
-import tempfile
-from pathlib import Path
 
 os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
 os.environ.setdefault("TRANSFORMERS_NO_FLAX", "1")
@@ -14,52 +11,10 @@ os.environ["USE_FLAX"] = "0"
 os.environ["USE_TORCH"] = "1"
 
 import torch
-from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
+from transformers import AutoModelForTokenClassification, pipeline
 
 from irish_core_decoder import repair_irish_core_spans
-
-
-def load_tokenizer(model_ref: str):
-    tokenizer_ref = model_ref
-    tokenizer_path = Path(model_ref)
-    if tokenizer_path.exists():
-        tokenizer_cfg_path = tokenizer_path / "tokenizer_config.json"
-        if tokenizer_cfg_path.exists():
-            data = json.loads(tokenizer_cfg_path.read_text(encoding="utf-8"))
-            if "fix_mistral_regex" in data:
-                tmpdir = Path(tempfile.mkdtemp(prefix="openmed_tokenizer_"))
-                keep = {
-                    "tokenizer_config.json",
-                    "tokenizer.json",
-                    "special_tokens_map.json",
-                    "vocab.txt",
-                    "vocab.json",
-                    "merges.txt",
-                    "added_tokens.json",
-                    "sentencepiece.bpe.model",
-                    "spiece.model",
-                }
-                for child in tokenizer_path.iterdir():
-                    if child.is_file() and child.name in keep:
-                        shutil.copy2(child, tmpdir / child.name)
-                data.pop("fix_mistral_regex", None)
-                (tmpdir / "tokenizer_config.json").write_text(
-                    json.dumps(data, ensure_ascii=False, indent=2) + "\n",
-                    encoding="utf-8",
-                )
-                tokenizer_ref = str(tmpdir)
-    try:
-        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=True)
-    except Exception:
-        pass
-    try:
-        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=False)
-    except TypeError:
-        pass
-    try:
-        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True)
-    except Exception:
-        return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=False)
+from onnx_token_classifier import safe_auto_tokenizer
 
 
 def mask_text(text: str, spans: list[dict]) -> str:
@@ -79,7 +34,7 @@ def main() -> None:
     parser.add_argument("--json", action="store_true")
     args = parser.parse_args()
 
-    tokenizer = load_tokenizer(args.model)
+    tokenizer = safe_auto_tokenizer(args.model)
     model = AutoModelForTokenClassification.from_pretrained(args.model)
     if args.device == "auto":
         device = "cuda" if torch.cuda.is_available() else "cpu"