forked from facebook/MobileLLM-R1-360M
33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# 自动修复 tokenizer 配置,供 entrypoint.sh 调用
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
|
||
|
|
MODEL_DIR = os.environ.get("MODEL_DIR", "/model")
|
||
|
|
FIX_DIR = "/tmp/fixed_tokenizer"
|
||
|
|
|
||
|
|
os.makedirs(FIX_DIR, exist_ok=True)
|
||
|
|
for name in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt", "tokenizer.model"]:
|
||
|
|
src = os.path.join(MODEL_DIR, name)
|
||
|
|
if os.path.exists(src):
|
||
|
|
shutil.copy(src, os.path.join(FIX_DIR, name))
|
||
|
|
|
||
|
|
cfg_path = os.path.join(FIX_DIR, "tokenizer_config.json")
|
||
|
|
if os.path.exists(cfg_path):
|
||
|
|
with open(cfg_path, "r", encoding="utf-8") as f:
|
||
|
|
cfg = json.load(f)
|
||
|
|
changed = False
|
||
|
|
bad = {"TokenizersBackend", "TiktokenTokenizer"}
|
||
|
|
if cfg.get("tokenizer_class") in bad:
|
||
|
|
cfg["tokenizer_class"] = "PreTrainedTokenizerFast"
|
||
|
|
changed = True
|
||
|
|
extra = cfg.get("extra_special_tokens")
|
||
|
|
if isinstance(extra, list):
|
||
|
|
cfg["extra_special_tokens"] = {t: t for t in extra}
|
||
|
|
changed = True
|
||
|
|
if changed:
|
||
|
|
with open(cfg_path, "w", encoding="utf-8") as f:
|
||
|
|
json.dump(cfg, f, ensure_ascii=False, indent=2)
|
||
|
|
print("[fix] tokenizer_config.json 已修复")
|