enginex-qwen25-compat/engine_compat.py

#!/usr/bin/env python3
"""Prepare an isolated model view for known Qwen2.5 tokenizer defects.

The source model is never edited. Large model files are linked into /tmp.
Only tokenizer_config.json is rewritten when a supported repair is needed.
"""

from __future__ import annotations

import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any


QWEN_CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "<|im_start|>{{ message['role'] }}\n"
    "{{ message['content'] }}<|im_end|>\n"
    "{% endfor %}"
    "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)


class CompatibilityError(RuntimeError):
    pass


def _json_object(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        raise CompatibilityError(f"invalid JSON object: {path.name}") from exc
    if not isinstance(value, dict):
        raise CompatibilityError(f"invalid JSON object: {path.name}")
    return value


def _key(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True)


def repair_tokenizer_config(
    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):
        if any(
            not (
                isinstance(item, str)
                or isinstance(item, dict) and isinstance(item.get("content"), str)
            )
            for item in extra
        ):
            raise CompatibilityError("unsupported extra_special_tokens list")
        additional = patched.get("additional_special_tokens", [])
        if not isinstance(additional, list):
            raise CompatibilityError("additional_special_tokens must be a list")
        merged = list(additional)
        seen = {_key(item) for item in merged}
        for item in extra:
            if _key(item) not in seen:
                merged.append(item)
                seen.add(_key(item))
        patched["additional_special_tokens"] = merged
        del patched["extra_special_tokens"]
        repairs.append("extra_special_tokens list -> additional_special_tokens")
    elif extra is not None and not isinstance(extra, dict):
        raise CompatibilityError("extra_special_tokens must be a mapping or list")

    # 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 (effective_tokenizer_dir / "tokenizer.json").is_file():
            raise CompatibilityError("TokenizersBackend requires tokenizer.json")
        patched["tokenizer_class"] = "PreTrainedTokenizerFast"
        repairs.append("TokenizersBackend -> PreTrainedTokenizerFast")

    # Transformers can load a standalone template with higher precedence than
    # tokenizer_config.json. Do not synthesize one when the repository already
    # supplies a template in that format.
    has_standalone_template = any(
        (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:
        config_path = model_dir / "config.json"
        if config_path.is_file():
            model_config = _json_object(config_path)
            if (
                model_config.get("model_type") == "qwen2"
                and patched.get("eos_token") == "<|im_end|>"
            ):
                patched["chat_template"] = QWEN_CHAT_TEMPLATE
                repairs.append("added Qwen2 ChatML chat_template")

    return patched, repairs


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"tokenizer directory is missing: {source}")
    config_path = source / "tokenizer_config.json"
    if not config_path.is_file():
        return None

    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

    view = Path(tempfile.mkdtemp(prefix="modelhub-compat-", dir=work_root))
    try:
        for item in source.iterdir():
            target = view / item.name
            if item.name == "tokenizer_config.json":
                target.write_text(
                    json.dumps(patched, ensure_ascii=False, indent=2) + "\n",
                    encoding="utf-8",
                )
            else:
                # The evaluation container is Linux; symlinks also work when
                # /model is a read-only mount on a different filesystem.
                try:
                    target.symlink_to(item.resolve(), target_is_directory=item.is_dir())
                except OSError as exc:
                    if not item.is_file():
                        raise CompatibilityError(f"cannot link model item: {item.name}") from exc
                    try:
                        os.link(item, target)
                    except OSError as link_exc:
                        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(
                {
                    "model_source": str(model_source),
                    "tokenizer_source": str(source),
                    "repairs": repairs,
                },
                ensure_ascii=False,
            )
            + "\n",
            encoding="utf-8",
        )
    except Exception:
        for item in view.iterdir():
            item.unlink()
        view.rmdir()
        raise

    print(json.dumps({"model_view": str(view), "repairs": repairs}, ensure_ascii=False), file=sys.stderr)
    return str(view)


def main() -> int:
    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], 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
    if view:
        print(view)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())