Support MetaX and Vastai tokenizer compatibility images

engine_compat.py

@@ -43,10 +43,11 @@ def _key(value: Any) -> str:
 
 
 def repair_tokenizer_config(
-    model_dir: Path, tokenizer: dict[str, Any]
+    model_dir: Path, tokenizer: dict[str, Any], *, tokenizer_dir: Path | None = None
 ) -> tuple[dict[str, Any], list[str]]:
     patched = dict(tokenizer)
     repairs: list[str] = []
+    effective_tokenizer_dir = tokenizer_dir or model_dir
 
     extra = patched.get("extra_special_tokens")
     if isinstance(extra, list):
@@ -76,7 +77,7 @@ def repair_tokenizer_config(
     # Only replace the known class mismatch from existing ModelHub failures.
     # Arbitrary custom tokenizer classes may require their own Python code.
     if patched.get("tokenizer_class") == "TokenizersBackend":
-        if not (model_dir / "tokenizer.json").is_file():
+        if not (effective_tokenizer_dir / "tokenizer.json").is_file():
             raise CompatibilityError("TokenizersBackend requires tokenizer.json")
         patched["tokenizer_class"] = "PreTrainedTokenizerFast"
         repairs.append("TokenizersBackend -> PreTrainedTokenizerFast")
@@ -85,7 +86,8 @@ def repair_tokenizer_config(
     # tokenizer_config.json. Do not synthesize one when the repository already
     # supplies a template in that format.
     has_standalone_template = any(
-        (model_dir / name).is_file()
+        (directory / name).is_file()
+        for directory in (effective_tokenizer_dir, model_dir)
         for name in ("chat_template.jinja", "chat_template.json")
     )
     if not patched.get("chat_template") and not has_standalone_template:
@@ -102,15 +104,37 @@ def repair_tokenizer_config(
     return patched, repairs
 
 
-def prepare_model(model_dir: str | Path, *, work_root: str | Path | None = None) -> str | None:
-    source = Path(model_dir).resolve()
+def prepare_model(
+    model_dir: str | Path,
+    *,
+    tokenizer_dir: str | Path | None = None,
+    work_root: str | Path | None = None,
+) -> str | None:
+    model_source = Path(model_dir).resolve()
+    if not model_source.is_dir():
+        raise CompatibilityError(f"model directory is missing: {model_source}")
+    source = Path(tokenizer_dir).resolve() if tokenizer_dir is not None else model_source
     if not source.is_dir():
-        raise CompatibilityError(f"model directory is missing: {source}")
+        raise CompatibilityError(f"tokenizer directory is missing: {source}")
     config_path = source / "tokenizer_config.json"
     if not config_path.is_file():
         return None
 
-    patched, repairs = repair_tokenizer_config(source, _json_object(config_path))
+    patched, repairs = repair_tokenizer_config(
+        model_source, _json_object(config_path), tokenizer_dir=source
+    )
+    # Some vendor entrypoints copy tokenizer files to /tmp and pass that path
+    # explicitly with --tokenizer, omitting a standalone template in the copy.
+    # Keep the original template visible to the effective tokenizer directory.
+    extra_templates = [
+        name
+        for name in ("chat_template.jinja", "chat_template.json")
+        if source != model_source
+        and not (source / name).is_file()
+        and (model_source / name).is_file()
+    ]
+    if extra_templates:
+        repairs.append("preserved standalone chat template")
     if not repairs:
         return None
 
@@ -137,8 +161,17 @@ def prepare_model(model_dir: str | Path, *, work_root: str | Path | None = None)
                         raise CompatibilityError(
                             f"cannot link model item: {item.name}"
                         ) from link_exc
+        for name in extra_templates:
+            (view / name).symlink_to(model_source / name)
         (view / ".modelhub-compat-report.json").write_text(
-            json.dumps({"source": str(source), "repairs": repairs}, ensure_ascii=False)
+            json.dumps(
+                {
+                    "model_source": str(model_source),
+                    "tokenizer_source": str(source),
+                    "repairs": repairs,
+                },
+                ensure_ascii=False,
+            )
             + "\n",
             encoding="utf-8",
         )
@@ -153,11 +186,13 @@ def prepare_model(model_dir: str | Path, *, work_root: str | Path | None = None)
 
 
 def main() -> int:
-    if len(sys.argv) != 2:
-        print("usage: engine_compat.py MODEL_DIR", file=sys.stderr)
+    if len(sys.argv) not in (2, 3):
+        print("usage: engine_compat.py MODEL_DIR [TOKENIZER_DIR]", file=sys.stderr)
         return 2
     try:
-        view = prepare_model(sys.argv[1])
+        view = prepare_model(
+            sys.argv[1], tokenizer_dir=sys.argv[2] if len(sys.argv) == 3 else None
+        )
     except CompatibilityError as exc:
         print(f"model compatibility error: {exc}", file=sys.stderr)
         return 2