| """加载骨干权重。只依赖 torch。 |
| |
| 支持两种目录: |
| |
| 预训练产出 model.pt {"model": state_dict, "config": {...}, ...} |
| HF 发布包 model.safetensors + config.json |
| |
| 后者不走 safetensors 库 —— 那个格式简单到不值得为它加一个依赖: |
| |
| 8 字节小端 uint64:JSON 头的长度 |
| JSON 头:{张量名: {dtype, shape, data_offsets:[起, 止]}, "__metadata__": ...} |
| 其余:裸张量数据,offsets 相对于头之后的位置 |
| |
| 这样 --ckpt_dir 可以直接指向从 HF 下下来的目录,不用先转格式。 |
| """ |
| import json |
| import struct |
| from pathlib import Path |
|
|
| import torch |
|
|
| |
| _ST_DTYPE = { |
| "BOOL": torch.bool, "U8": torch.uint8, "I8": torch.int8, |
| "I16": torch.int16, "U16": torch.uint16, |
| "I32": torch.int32, "U32": torch.uint32, "I64": torch.int64, |
| "F16": torch.float16, "BF16": torch.bfloat16, |
| "F32": torch.float32, "F64": torch.float64, |
| } |
|
|
|
|
| def load_safetensors(path, device=None) -> dict: |
| """读 .safetensors → state_dict。纯 torch 实现。 |
| |
| device 给了就把张量搬过去,省得调用方再遍历一遍(发布包里的推理代码 |
| 直接这么用)。 |
| """ |
| path = Path(path) |
| with open(path, "rb") as f: |
| header_len = struct.unpack("<Q", f.read(8))[0] |
| header = json.loads(f.read(header_len)) |
| blob = f.read() |
|
|
| out = {} |
| for name, spec in header.items(): |
| if name == "__metadata__": |
| continue |
| dtype = _ST_DTYPE.get(spec["dtype"]) |
| if dtype is None: |
| raise ValueError(f"{path} 里未知的 dtype {spec['dtype']}(张量 {name})") |
| lo, hi = spec["data_offsets"] |
| |
| |
| |
| buf = bytearray(blob[lo:hi]) |
| t = torch.frombuffer(buf, dtype=dtype, count=(hi - lo) // dtype.itemsize) |
| t = t.view(*spec["shape"]) if spec["shape"] else t |
| out[name] = t.to(device) if device is not None else t |
| return out |
|
|
|
|
| def load_backbone(ckpt_dir) -> tuple[dict, dict]: |
| """返回 (骨干 state_dict, config dict)。 |
| |
| 骨干权重在两种格式里都以 bert.* 为前缀(HF 发布包是从 ModernBertForMLM |
| 存的,预训练 ckpt 同理),这里统一剥掉前缀返回,调用方直接喂给 |
| ModernBertModel.load_state_dict。 |
| """ |
| ckpt_dir = Path(ckpt_dir) |
| cfg_path = ckpt_dir / "config.json" |
| if not cfg_path.exists(): |
| raise SystemExit(f"{ckpt_dir} 下没有 config.json") |
| config = json.loads(cfg_path.read_text()) |
|
|
| pt, st = ckpt_dir / "model.pt", ckpt_dir / "model.safetensors" |
| if pt.exists(): |
| ckpt = torch.load(pt, map_location="cpu", weights_only=False) |
| |
| state = ckpt.get("ema") or ckpt["model"] |
| elif st.exists(): |
| state = load_safetensors(st) |
| else: |
| raise SystemExit(f"{ckpt_dir} 下既没有 model.pt 也没有 model.safetensors") |
|
|
| bert = {k[len("bert."):]: v for k, v in state.items() if k.startswith("bert.")} |
| if not bert: |
| raise SystemExit( |
| f"{ckpt_dir} 里没有 bert.* 前缀的张量 —— 这是骨干吗?" |
| f"(微调要从骨干开始,不能从另一个微调结果开始)") |
| return bert, config |
|
|
