| """把 DomSense 训练 checkpoint 转换为 Hugging Face 标准格式并导出。 |
| |
| 用法: |
| python scripts/convert_to_hf.py |
| --ckpt ../outputs/rlcd_mdp_run1/best_model.pt |
| --out ./saved_model |
| |
| 产出（HF 标准目录）: |
| saved_model/ |
| config.json # SystemOneConfig 序列化 |
| model.safetensors # 转换后的权重（自动去掉 base_model_prefix "core."） |
| tokenizer.json # DomSenseTokenizer |
| |
| 导出后，可这样加载: |
| import sys; sys.path.insert(0, "repo_root") |
| from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer |
| m = SystemOneModelForDecision.from_pretrained("saved_model_dir") |
| """ |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| REPO_ROOT = Path(__file__).resolve().parent.parent.parent |
| sys.path.insert(0, str(REPO_ROOT)) |
|
|
| from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer |
| from huggingface_repo.configuration_domsense import SystemOneConfig |
|
|
|
|
| def convert(ckpt_path: str, out_dir: str): |
| print(f"[1/4] 读取 checkpoint: {ckpt_path}") |
| ckpt = torch.load(ckpt_path, map_location="cpu") |
| state_dict = ckpt["model_state_dict"] |
|
|
| |
| raw_cfg = ckpt.get("config", {}) |
| if isinstance(raw_cfg, dict) and "model" in raw_cfg: |
| raw_cfg = raw_cfg["model"] |
| cfg_kwargs = { |
| k: v |
| for k, v in raw_cfg.items() |
| if v is not None and not isinstance(v, dict) |
| } |
| |
| if "mdp_head.action_embed_table.weight" in state_dict: |
| cfg_kwargs["max_actions"] = state_dict[ |
| "mdp_head.action_embed_table.weight" |
| ].shape[0] |
| |
| if "domain_moe.gate.2.weight" in state_dict: |
| n_slots = state_dict["domain_moe.gate.2.weight"].shape[0] |
| n_domains = n_slots - 1 |
| cfg_kwargs["use_domain_experts"] = n_domains > 0 |
| if "domains" not in cfg_kwargs or not cfg_kwargs["domains"]: |
| cfg_kwargs["domains"] = [f"domain_{i}" for i in range(n_domains)] |
| |
| if len(cfg_kwargs["domains"]) != n_domains: |
| cfg_kwargs["domains"] = [f"domain_{i}" for i in range(n_domains)] |
| cfg = SystemOneConfig(**cfg_kwargs) |
|
|
| print(f"[2/4] 构建 HF 模型；hidden_dim={cfg.hidden_dim}, " |
| f"num_bellman_steps={cfg.num_bellman_steps}, " |
| f"use_domain_experts={cfg.use_domain_experts}") |
| model = SystemOneModelForDecision(cfg) |
|
|
| print("[3/4] 加载权重到内核 core.*（strict=False）") |
| missing, unexpected = model.core.load_state_dict(state_dict, strict=False) |
| if missing: |
| print(f" - missing {len(missing)} keys: {missing[:10]}") |
| if unexpected: |
| print(f" - unexpected {len(unexpected)} keys: {unexpected[:10]}") |
| total, trainable = model.count_parameters() |
| print(f" 总参数量: {total:,} | 可训练: {trainable:,}") |
|
|
| out = Path(out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| print(f"[4/4] 导出 HF 格式到: {out}") |
| |
| model.save_pretrained(str(out)) |
| DomSenseTokenizer().save_pretrained(str(out)) |
|
|
| print("完成！现在可以从该目录加载:") |
| print(f' from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer') |
| print(f' m = SystemOneModelForDecision.from_pretrained("{out}")') |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--ckpt", default="../outputs/rlcd_mdp_run1/best_model.pt") |
| parser.add_argument("--out", default="./saved_model") |
| args = parser.parse_args() |
| convert(args.ckpt, args.out) |
|
|
