189 lines
6.8 KiB
Python
Executable File
189 lines
6.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""把 R2 chat_template 与 R4 architecture 白名单应用到模型目录。
|
||
|
||
不访问平台，不启动 vLLM。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
FORBIDDEN_FRAMEWORK = "vllm_fix_tokenizer"
|
||
DEFAULT_FRAMEWORK = "vllm_bi150_r2r4"
|
||
TOKENIZER_PATCH_KEY = "modelhub_tokenizer_patch"
|
||
|
||
|
||
class PatchError(Exception):
|
||
"""模型目录不满足本引擎的修补前提。"""
|
||
|
||
|
||
def engine_root() -> Path:
|
||
return Path(__file__).resolve().parents[1]
|
||
|
||
|
||
def assert_framework(framework: str) -> None:
|
||
if not framework or framework == FORBIDDEN_FRAMEWORK:
|
||
raise PatchError(
|
||
f"FRAMEWORK 必须是 {DEFAULT_FRAMEWORK}，"
|
||
"不能使用 vllm_fix_tokenizer"
|
||
)
|
||
|
||
|
||
def load_whitelist(path: Path) -> set[str]:
|
||
names: set[str] = set()
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
item = line.split("#", 1)[0].strip()
|
||
if item:
|
||
names.add(item)
|
||
if not names:
|
||
raise PatchError(f"架构白名单为空: {path}")
|
||
return names
|
||
|
||
|
||
def load_aliases(path: Path) -> dict[str, str]:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
if not isinstance(data, dict):
|
||
raise PatchError(f"别名文件必须是对象: {path}")
|
||
return {str(key): str(value) for key, value in data.items()}
|
||
|
||
|
||
def preview_default_template(messages: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
|
||
"""与 default_chat_template.jinja 同一套角色标记，供无 Jinja 依赖的本地预览。"""
|
||
parts: list[str] = []
|
||
for message in messages:
|
||
role = message["role"]
|
||
content = message["content"]
|
||
parts.append(f"<|{role}|>\n{content}<|end|>\n")
|
||
if add_generation_prompt:
|
||
parts.append("<|assistant|>\n")
|
||
return "".join(parts)
|
||
|
||
|
||
def apply_chat_template(model_dir: Path, template_text: str, force: bool = False) -> dict:
|
||
path = model_dir / "tokenizer_config.json"
|
||
if not path.is_file():
|
||
raise PatchError(f"缺少 tokenizer_config.json: {path}")
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
if not isinstance(data, dict):
|
||
raise PatchError("tokenizer_config.json 必须是对象")
|
||
current = data.get("chat_template")
|
||
if isinstance(current, str) and current.strip() and not force:
|
||
return {"injected": False, "reason": "已有 chat_template"}
|
||
data["chat_template"] = template_text
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
return {"injected": True, "reason": "已写入 chat_template"}
|
||
|
||
|
||
def apply_architectures(
|
||
model_dir: Path,
|
||
whitelist: set[str],
|
||
aliases: dict[str, str],
|
||
unknown_policy: str = "fail",
|
||
) -> dict:
|
||
if unknown_policy not in {"fail", "skip"}:
|
||
raise PatchError("ARCH_UNKNOWN_POLICY 只能是 fail 或 skip")
|
||
path = model_dir / "config.json"
|
||
if not path.is_file():
|
||
raise PatchError(f"缺少 config.json: {path}")
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
if not isinstance(data, dict):
|
||
raise PatchError("config.json 必须是对象")
|
||
architectures = data.get("architectures")
|
||
if not isinstance(architectures, list) or not architectures:
|
||
raise PatchError("config.json 缺少 architectures")
|
||
|
||
rewritten: list[str] = []
|
||
skipped: list[str] = []
|
||
updated: list[str] = []
|
||
for name in architectures:
|
||
if not isinstance(name, str) or not name:
|
||
raise PatchError("architectures 含有空名称")
|
||
if name in whitelist:
|
||
updated.append(name)
|
||
continue
|
||
target = aliases.get(name)
|
||
if target in whitelist:
|
||
updated.append(target)
|
||
rewritten.append(f"{name}->{target}")
|
||
continue
|
||
if unknown_policy == "skip":
|
||
skipped.append(name)
|
||
updated.append(name)
|
||
continue
|
||
raise PatchError(f"未知架构: {name}")
|
||
|
||
changed = updated != architectures
|
||
if changed:
|
||
data["architectures"] = updated
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
return {
|
||
"architectures": updated,
|
||
"rewritten": rewritten,
|
||
"skipped": skipped,
|
||
"changed": changed,
|
||
}
|
||
|
||
|
||
def apply_tokenizer_patch(model_dir: Path, enabled: bool) -> dict:
|
||
if not enabled:
|
||
return {"applied": False}
|
||
path = model_dir / "tokenizer_config.json"
|
||
if not path.is_file():
|
||
raise PatchError(f"缺少 tokenizer_config.json: {path}")
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
data[TOKENIZER_PATCH_KEY] = "optional-v1"
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
return {"applied": True}
|
||
|
||
|
||
def apply_all(
|
||
model_dir: Path,
|
||
template_text: str,
|
||
whitelist: set[str],
|
||
aliases: dict[str, str],
|
||
unknown_policy: str = "fail",
|
||
force_chat_template: bool = False,
|
||
tokenizer_patch: bool = False,
|
||
) -> dict:
|
||
chat = apply_chat_template(model_dir, template_text, force=force_chat_template)
|
||
architectures = apply_architectures(model_dir, whitelist, aliases, unknown_policy)
|
||
tokenizer = apply_tokenizer_patch(model_dir, tokenizer_patch)
|
||
return {"chat_template": chat, "architectures": architectures, "tokenizer": tokenizer}
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
root = engine_root()
|
||
parser = argparse.ArgumentParser(description="应用 R2/R4 模型目录修补")
|
||
parser.add_argument("--model-dir", required=True, type=Path)
|
||
parser.add_argument("--framework", default=DEFAULT_FRAMEWORK)
|
||
parser.add_argument("--unknown-policy", default="fail", choices=("fail", "skip"))
|
||
parser.add_argument("--force-chat-template", action="store_true")
|
||
parser.add_argument("--tokenizer-patch", action="store_true")
|
||
parser.add_argument("--template", type=Path, default=root / "patches/chat_template/default_chat_template.jinja")
|
||
parser.add_argument("--whitelist", type=Path, default=root / "patches/architectures/whitelist.txt")
|
||
parser.add_argument("--aliases", type=Path, default=root / "patches/architectures/aliases.json")
|
||
args = parser.parse_args(argv)
|
||
try:
|
||
assert_framework(args.framework)
|
||
result = apply_all(
|
||
args.model_dir,
|
||
args.template.read_text(encoding="utf-8"),
|
||
load_whitelist(args.whitelist),
|
||
load_aliases(args.aliases),
|
||
unknown_policy=args.unknown_policy,
|
||
force_chat_template=args.force_chat_template,
|
||
tokenizer_patch=args.tokenizer_patch,
|
||
)
|
||
except (PatchError, OSError, json.JSONDecodeError) as exc:
|
||
print(f"修补失败: {exc}", file=sys.stderr)
|
||
return 1
|
||
print(json.dumps(result, ensure_ascii=False))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|