File size: 2,686 Bytes
71dabc7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | #!/usr/bin/env python3
"""从 Laya 完整 checkpoint 中抽取 mmBERT(ModernBERT) encoder 底座。
Laya 的 safetensors 里 encoder 权重带 `encoder.` 前缀（其余是决策头）。
本脚本剥掉前缀后存成独立目录，使其可直接用 AutoModelForSequenceClassification
做标准分类微调；分类头（score.*）是新增的，训练时随机初始化属正常现象。
用法:
python scripts/extract_encoder.py # 默认 models/laya-multilingual -> models/mmbert-ad
python scripts/extract_encoder.py --inspect # 只打印键前缀分布，不写文件
"""
import argparse
import shutil
import sys
from pathlib import Path
from safetensors.torch import load_file, save_file
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--laya_dir", default="models/laya-multilingual")
ap.add_argument("--out_dir", default="models/mmbert-ad")
ap.add_argument("--inspect", action="store_true")
args = ap.parse_args()
ckpt = Path(args.laya_dir) / "model.safetensors"
if not ckpt.exists():
sys.exit(f"[ERR] 找不到 {ckpt}，请先下载模型（见 README 快速开始）")
state = load_file(str(ckpt))
prefixes = {}
for k in state:
top = ".".join(k.split(".")[:2])
prefixes[top] = prefixes.get(top, 0) + 1
print("键前缀分布（top2 级）:")
for p, n in sorted(prefixes.items()):
print(f" {p:<40} {n} tensors")
if args.inspect:
return
enc_keys = [k for k in state if k.startswith("encoder.")]
if not enc_keys:
sys.exit(f"[ERR] checkpoint 中没有 encoder.* 前缀的键，请根据上面的分布手动调整")
stripped = {k[len("encoder."):]: state[k] for k in enc_keys}
# ModernBertForSequenceClassification 的主体键以 model. 开头；
# 若 encoder 是裸 ModernBertModel（键为 embeddings./layers. 开头），补上 model. 前缀
if not any(k.startswith("model.") for k in stripped):
stripped = {f"model.{k}": v for k, v in stripped.items()}
out = Path(args.out_dir)
out.mkdir(parents=True, exist_ok=True)
save_file(stripped, str(out / "model.safetensors"), metadata={"format": "pt"})
shutil.copy(Path(args.laya_dir) / "encoder" / "config.json", out / "config.json")
tok_dir = Path(args.laya_dir) / "tokenizer"
for f in ("tokenizer.json", "tokenizer_config.json"):
if (tok_dir / f).exists():
shutil.copy(tok_dir / f, out / f)
n_params = sum(v.numel() for v in stripped.values())
print(f"[OK] encoder 已抽取 -> {out} (tensors={len(stripped)}, params={n_params / 1e6:.1f}M)")
if __name__ == "__main__":
main()
|