Update inference code and model card

checkpoint.py

@@ -0,0 +1,90 @@
+"""加载骨干权重。只依赖 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
+
+# safetensors 的 dtype 名 → torch dtype
+_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"]
+        # 先 bytearray 拷一份再 frombuffer:直接映射 bytes 会得到只读缓冲区上的
+        # 张量,torch 会警告"可以写但不该写";而且零拷贝视图会把整个 blob
+        # 一直拖在内存里不释放。
+        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)
+        # 有 EMA 就优先用 shadow 权重(更稳),否则用原始权重
+        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

csc_model.py

@@ -1,7 +1,7 @@
 """BERTc-CSC 推理:中文拼写纠错。
 
 这份代码随模型一起发到 HF,只依赖同目录的 model.py / tokenizer.py
-和 torch + safetensors。
+和 torch。
 
     from csc_model import BERTcForCSC
     model = BERTcForCSC.from_pretrained(".")
@@ -12,7 +12,7 @@ from pathlib import Path
 
 import torch
 import torch.nn as nn
-from safetensors.torch import load_file
+from checkpoint import load_safetensors
 
 from model import ModernBertConfig, ModernBertModel
 from tokenizer import PieceCharTokenizer
@@ -36,10 +36,11 @@ class BERTcForCSC(nn.Module):
     @classmethod
     def from_pretrained(cls, model_dir=".", map_location="cpu"):
         model_dir = Path(model_dir)
-        cfg = ModernBertConfig(**json.loads((model_dir / "config.json").read_text()))
+        cfg = ModernBertConfig.from_dict(
+            json.loads((model_dir / "config.json").read_text()))
         model = cls(cfg)
-        state = load_file(str(model_dir / "model.safetensors"),
-                          device=str(map_location))
+        state = load_safetensors(model_dir / "model.safetensors",
+                                 device=str(map_location))
         missing, unexpected = model.load_state_dict(state, strict=False)
         # cor_head.weight 跟 embed 共享,safetensors 不重复存,所以它一定"缺"
         if set(missing) != {"cor_head.weight"} or unexpected:

model.py

@@ -76,11 +76,9 @@ class ModernBertConfig:
     layer_norm_eps: float = 1e-5
     initializer_range: float = 0.02
     tie_word_embeddings: bool = True
-    # Dropout(对齐 ModernBERT release)
-    embed_dropout: float = 0.0
-    mlp_dropout: float = 0.0
-    attn_out_dropout: float = 0.1
-    attn_probs_dropout: float = 0.0
+    # 全程无 dropout:预训练数据量远大于参数量,不存在过拟合,dropout 只拖慢
+    # 收敛。所以没有 dropout 配置项 —— 已发布的 config.json 里那几个 dropout
+    # 字段都是 0,load_backbone 按字段名过滤,多出来的会被忽略。
     # 架构开关(对齐 release)
     embed_norm: bool = True            # embedding 后立刻 LayerNorm
     skip_first_prenorm: bool = True    # 第 1 层不做 pre-norm
@@ -88,6 +86,13 @@ class ModernBertConfig:
     # init
     init_method: str = "megatron"      # "megatron"(残差层 ×1/sqrt(2L))或 "normal"
 
+    @classmethod
+    def from_dict(cls, d: dict) -> "ModernBertConfig":
+        """按字段名过滤后构造。config.json 可能带这个版本不认识的键 ——
+        已发布的六个模型就带着四个 dropout 字段(现在全程无 dropout,字段去掉了),
+        直接 `ModernBertConfig(**d)` 会 TypeError。"""
+        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
+
     @property
     def head_dim(self) -> int:
         assert self.hidden_size % self.num_attention_heads == 0
@@ -143,11 +148,9 @@ class ModernBertAttention(nn.Module):
         self.num_heads = config.num_attention_heads
         self.head_dim = config.head_dim
         self.scale = self.head_dim ** -0.5
-        self.attn_probs_dropout = config.attn_probs_dropout
         # 无 bias
         self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
         self.o = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-        self.out_dropout = nn.Dropout(config.attn_out_dropout)
 
     def forward(self, x: torch.Tensor,
                 block_mask=None,
@@ -175,11 +178,10 @@ class ModernBertAttention(nn.Module):
             out = F.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask=sdpa_mask,
-                dropout_p=self.attn_probs_dropout if self.training else 0.0,
                 is_causal=False,
             )  # [B, h, L, d]
         out = out.transpose(1, 2).reshape(B, L, H)
-        return self.out_dropout(self.o(out))
+        return self.o(out)
 
 
 # ============ GeGLU MLP ============
@@ -193,11 +195,10 @@ class GeGLU(nn.Module):
         I = config.intermediate_size
         self.w_in = nn.Linear(config.hidden_size, 2 * I, bias=False)
         self.w_out = nn.Linear(I, config.hidden_size, bias=False)
-        self.dropout = nn.Dropout(config.mlp_dropout)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         gate, up = self.w_in(x).chunk(2, dim=-1)
-        return self.dropout(self.w_out(F.gelu(gate) * up))
+        return self.w_out(F.gelu(gate) * up)
 
 
 # ============ Layer(pre-norm,支持 skip_first_prenorm)============
@@ -235,7 +236,6 @@ class ModernBertModel(nn.Module):
         )
         self.embed_norm = (LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps)
                            if config.embed_norm else nn.Identity())
-        self.embed_dropout = nn.Dropout(config.embed_dropout)
         self.layers = nn.ModuleList(
             [ModernBertLayer(config, is_first=(i == 0))
              for i in range(config.num_hidden_layers)]
@@ -284,7 +284,6 @@ class ModernBertModel(nn.Module):
         x = self.embed(input_ids)
         x = x + self.pos_emb(L).to(x.dtype)  # 加 scaled sinusoidal PE
         x = self.embed_norm(x)
-        x = self.embed_dropout(x)
 
         block_mask = self._build_block_mask(seg_ids, B, L) if seg_ids is not None else None