#!/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 已修复")