Update inference code and model card

README.md

@@ -3,62 +3,64 @@ license: apache-2.0
 language:
 - zh
 tags:
+- bert
 - chinese
+- text2text-generation
 - spelling-correction
-- csc
-- bert
-- text-correction
+pipeline_tag: fill-mask
 library_name: pytorch
 ---
 
 # BERTc-165M-CSC
 
-BERTc-165M-CSC is a Chinese spelling correction model fine-tuned from
-`Ismantic/BERTc-165M`. It uses a Modern BERTc encoder with two heads:
+中文拼写纠错。基于 [BERTc-165M](https://huggingface.co/Ismantic/BERTc-165M) 微调。
 
-- correction head: tied to the input embedding matrix
-- detection head: binary error detection
+双头:cor 逐位置预测正确的字(权重与词嵌入绑定),det 判断该位置有没有错(focal loss)。**只做等长替换**,不处理多字少字。
 
-## Metrics
+## 指标
 
-SIGHAN-15 sentence-level evaluation using the pycorrector-style 707-sample protocol:
+| 指标 | 值 |
+|---|---|
+| 句级 F1 | 0.8308 |
 
-- Sentence F1: **0.8308**
-- Accuracy: **0.8416**
-- Precision: **0.9516**
-- Recall: **0.7373**
-- TP/FP/FN/TN: 275 / 14 / 98 / 320
+## 训练
 
-Training recipe:
+- 配方:5 epoch,batch 64,lr 5e-5,其余同 315M-CSC
+- 数据:同 315M-CSC
 
-- backbone: `Ismantic/BERTc-165M`
-- epochs: 5
-- batch size: 64
-- learning rate: 5e-5
-- warmup ratio: 0.1
-- detection loss weight: 0.3
-- inference threshold: 0.7
-- max length: 128
+## 用法
 
-## Files
+```python
+from csc_model import BERTcForCSC
 
-- `model.safetensors`: CSC state dict. `cor_head.weight` is intentionally omitted and tied to `bert.embed.weight` by `csc_model.py`.
-- `config.json`: BERTc backbone architecture.
-- `csc_config.json`: task and metric metadata.
-- `model.py`: Modern BERTc implementation.
-- `csc_model.py`: CSC wrapper and batch correction helper.
-- `piece.model`: tokenizer model; load with `piece_tokenizer` using `cn_dict="no"`.
+model = BERTcForCSC.from_pretrained(".")
+print(model.correct("我今天很稿兴"))     # 我今天很高兴
+```
 
-## Usage
+## 阈值
 
-```python
-from csc_model import BERTcForCSC, PieceCharTokenizer
+`correct(..., threshold=0.7)`:纠错置信度低于阈值就保留原字。调低提召回、
+调高提精确率。0.7 是与 MacBERT4CSC 对齐的默认值,报告的指标都基于它。
 
-tok = PieceCharTokenizer(".")
-model = BERTcForCSC.from_pretrained(".")
-texts = ["少先队员因该为老人让坐。"]
-print(model.correct(texts, tok, threshold=0.7))
+## Tokenizer
+
+字级 SentencePiece,词表 12536(pad=12531,mask=12535)。**必须用 `dict="no"`
+加载**(字模式,不挂分词词典)——挂了词典编码结果会跟训练时不一致,而且不报错。
+
+```bash
+pip install git+https://github.com/Ismantic/PieceTokenizer
 ```
 
-This is not a generative model. It performs same-length character replacement for
-Chinese spelling correction.
+## 文件
+
+| 文件 | 说明 |
+|---|---|
+| `model.safetensors` | 骨干 + 双头 |
+| `csc_model.py` | 推理入口 `BERTcForCSC` |
+| `model.py` | 骨干定义 |
+| `tokenizer.py` | 字级 tokenizer |
+| `example_correct.py` | 示例 |
+
+## 许可
+
+Apache-2.0。训练语料各自的许可见对应数据集卡。

csc_config.json

@@ -1,34 +1,11 @@
 {
+  "task": "csc",
   "base_model": "Ismantic/BERTc-165M",
-  "task": "Chinese Spelling Correction",
-  "threshold": 0.7,
-  "max_len": 128,
-  "epoch": 5,
   "metrics": {
-    "acc": 0.8415841584158416,
-    "precision": 0.9515570934256056,
-    "recall": 0.7372654155495979,
-    "f1": 0.8308157099697886,
-    "TP": 275,
-    "FP": 14,
-    "FN": 98,
-    "TN": 320,
-    "n": 707
+    "句级 F1": "0.8308"
   },
-  "args": {
-    "backbone_path": "/home/tfbao/Shiyu/BERTc/pretrain/modern_bertc/output_v4_mid/checkpoint-8500",
-    "tokenizer_dir": "/home/tfbao/Shiyu/BERTc/pretrain/modern_bertc/tokenizer",
-    "train_pkl": "/home/tfbao/Shiyu/BERTc/csc/data/sighan_wang271k_pairs.pkl",
-    "test_tsv": "/home/tfbao/Shiyu/BERTc/csc/data/test/sighan2015_test_official.tsv",
-    "output_dir": "/home/tfbao/Shiyu/BERTc/csc/output_v4_mid_csc_v3_tied",
-    "epochs": 5,
-    "batch_size": 64,
-    "lr": 5e-05,
-    "warmup_ratio": 0.1,
-    "max_len": 128,
-    "det_weight": 0.3,
-    "threshold": 0.7,
-    "log_every": 200
-  },
-  "source_checkpoint": "finetune/sota/sota_csc_v4mid_5ep_best.pt"
+  "recipe": "5 epoch,batch 64,lr 5e-5,其余同 315M-CSC",
+  "data": "同 315M-CSC",
+  "source_checkpoint": "finetune/sota/sota_csc_v4mid_5ep_best.pt",
+  "epoch": 5
 }

csc_model.py

@@ -1,39 +1,30 @@
+"""BERTc-CSC 推理:中文拼写纠错。
+
+这份代码随模型一起发到 HF,只依赖同目录的 model.py / tokenizer.py
+和 torch + safetensors。
+
+    from csc_model import BERTcForCSC
+    model = BERTcForCSC.from_pretrained(".")
+    print(model.correct("我今天很稿兴"))   # 我今天很高兴
+"""
 import json
 from pathlib import Path
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from safetensors.torch import load_file
 
 from model import ModernBertConfig, ModernBertModel
+from tokenizer import PieceCharTokenizer
 
 
-class PieceCharTokenizer:
-    def __init__(self, model_dir):
-        import piece_tokenizer as pt
-        model_dir = Path(model_dir)
-        self._tok = pt.Tokenizer()
-        self._tok.load(str(model_dir / "piece.model"), cn_dict="no")
-        mask_path = model_dir / "mask_token_id.txt"
-        self.mask_token_id = int(mask_path.read_text().strip()) if mask_path.exists() else self._tok.vocab_size()
-        self.vocab_size = self._tok.vocab_size() + 1
-        self.pad_token_id = self._tok.piece_to_id("<pad>")
-        self.unk_token_id = 0
-        self.cache = {}
-        self.inv_cache = {}
-
-    def char_to_id(self, char):
-        if char in self.cache:
-            return self.cache[char]
-        ids = self._tok.encode_as_ids(char)
-        tid = ids[0] if ids else self.unk_token_id
-        self.cache[char] = tid
-        self.inv_cache.setdefault(tid, char)
-        return tid
+class BERTcForCSC(nn.Module):
+    """骨干 + 双头:cor 逐位置预测正确的字,det 判断该位置有没有错。
 
+    cor_head 的权重与词嵌入绑定 —— 预训练的 MLM 头就是 logits = h @ embed.weightᵀ,
+    换成独立的 Linear 会废掉这层对齐,纠错效果掉一大截。
+    """
 
-class BERTcForCSC(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
@@ -43,55 +34,54 @@ class BERTcForCSC(nn.Module):
         self.det_head = nn.Linear(config.hidden_size, 1)
 
     @classmethod
-    def from_pretrained(cls, model_dir, map_location="cpu"):
+    def from_pretrained(cls, model_dir=".", map_location="cpu"):
         model_dir = Path(model_dir)
         cfg = ModernBertConfig(**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_file(str(model_dir / "model.safetensors"),
+                          device=str(map_location))
         missing, unexpected = model.load_state_dict(state, strict=False)
-        allowed_missing = {"cor_head.weight"}
-        if set(missing) != allowed_missing or unexpected:
-            raise RuntimeError(f"Bad state dict: missing={missing}, unexpected={unexpected}")
+        # cor_head.weight 跟 embed 共享,safetensors 不重复存,所以它一定"缺"
+        if set(missing) != {"cor_head.weight"} or unexpected:
+            raise RuntimeError(f"权重不匹配:缺 {missing},多 {unexpected}")
         model.cor_head.weight = model.bert.embed.weight
         model.eval()
         return model
 
     def forward(self, input_ids, attention_mask=None):
         h = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cor_logits = self.cor_head(h)
-        det_logits = self.det_head(h).squeeze(-1)
-        return cor_logits, det_logits
+        return self.cor_head(h), self.det_head(h).squeeze(-1)
 
     @torch.no_grad()
-    def correct(self, texts, tokenizer, threshold=0.7, max_len=128, device=None):
-        if isinstance(texts, str):
-            single = True
-            texts = [texts]
-        else:
-            single = False
+    def correct(self, texts, tokenizer=None, threshold=0.7, max_len=128,
+                device=None):
+        """纠错。返回修正后的句子(输入是单句就返回单句)。
+
+        threshold 是纠错置信度下限:softmax 概率低于它就保留原字。
+        低置信度的"纠正"绝大多数是误伤,这个阈值是精确率的主要来源。
+        """
+        if tokenizer is None:
+            tokenizer = PieceCharTokenizer(Path(__file__).resolve().parent)
+        single = isinstance(texts, str)
+        texts = [texts] if single else list(texts)
+
         device = device or next(self.parameters()).device
         self.eval()
-        lengths = [min(len(t), max_len) for t in texts]
-        max_l = max(lengths) if lengths else 0
-        input_ids = torch.full((len(texts), max_l), tokenizer.pad_token_id, dtype=torch.long, device=device)
-        attn = torch.zeros((len(texts), max_l), dtype=torch.long, device=device)
-        for i, text in enumerate(texts):
-            ids = [tokenizer.char_to_id(c) for c in text[:lengths[i]]]
-            if ids:
-                input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device)
-                attn[i, :len(ids)] = 1
+        input_ids, attn, lengths = tokenizer.batch(texts, max_len, device)
+
         cor_logits, _ = self(input_ids, attn)
-        probs = F.softmax(cor_logits, dim=-1)
-        top_probs, top_ids = probs.max(dim=-1)
+        probs = torch.softmax(cor_logits.float(), dim=-1)
+        top_prob, top_id = probs.max(dim=-1)
+
         out = []
         for i, text in enumerate(texts):
-            chars = list(text[:lengths[i]])
-            pred = []
-            for j, orig in enumerate(chars):
-                tid = int(top_ids[i, j].item())
-                prob = float(top_probs[i, j].item())
-                pred.append(tokenizer.inv_cache.get(tid, orig) if prob >= threshold else orig)
-            if len(text) > lengths[i]:
-                pred.extend(list(text[lengths[i]:]))
-            out.append("".join(pred))
+            n = lengths[i]
+            chars = list(text[:n])
+            for j in range(n):
+                if float(top_prob[i, j]) < threshold:
+                    continue
+                c = tokenizer.id_to_char(int(top_id[i, j]))
+                if len(c) == 1:            # 只接受单字替换,保证不改变长度
+                    chars[j] = c
+            out.append("".join(chars) + text[n:])   # 超长部分原样接回
         return out[0] if single else out

example_correct.py

@@ -1,5 +1,6 @@
-from csc_model import BERTcForCSC, PieceCharTokenizer
+"""BERTc-CSC:中文拼写纠错。"""
+from csc_model import BERTcForCSC
 
-tok = PieceCharTokenizer(".")
 model = BERTcForCSC.from_pretrained(".")
-print(model.correct(["少先队员因该为老人让坐。"], tok, threshold=0.7))
+for s in ["我今天很稿兴", "他的身体健康状况很不错,平时喜欢锻练"]:
+    print(f"{s}  →  {model.correct(s)}")

model.py

@@ -1,7 +1,15 @@
-"""Modern BERTc v3 — ModernBERT release-aligned + Cramming-style ScaledSinusoidal PE.
+"""Modern BERTc 骨干网络。ModernBERT release 对齐 + Cramming 式 ScaledSinusoidal PE。
+
+只依赖 torch。**state_dict 的 key 不能动** —— 改任何模块名或嵌套层级都会让
+HF 上已发布的六个模型权重全部失配,而模型照样能随机初始化跑起来、不报错。
+改动后跑 test/test_reproduce_sota.py 验证。
+
+两个已发布规格(都用同一份代码,只是 config 不同):
+  BERTc-165M (v4-Mid)   12L / 1024H / 2752I / 16 heads
+  BERTc-315M (v4-Large) 24L / 1024H / 2752I / 16 heads
 
 主要按 release `modernbert-base-pretrain.yaml` 对齐(除 Alt Attn 和 PE):
-  - 架构: 22L / 768H / 1152I (GLU) / 12 heads,head_dim=64
+  - 默认 config: 22L / 768H / 1152I (GLU) / 12 heads,head_dim=64
   - **ScaledSinusoidal 位置编码**(Hua et al. 2022 FLASH;Cramming 实测短 seq 比
     RoPE 更值:计算几乎免费,RoPE 收益被 5-10% 速度损失抵消)
   - GeGLU FFN(glu + gelu)
@@ -33,7 +41,7 @@
   head_bias (V,)          : 12.5K
   total ≈ 130M
 """
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Optional
 import math
 
@@ -334,114 +342,3 @@ class ModernBertForMLM(nn.Module):
         return sum(p.numel() for p in self.parameters() if p.requires_grad)
 
 
-# ============ Quick sanity test ============
-
-if __name__ == "__main__":
-    print("=== Test 1: default config (22L/768H release-aligned) ===")
-    cfg = ModernBertConfig()
-    print(f"Config: V={cfg.vocab_size} H={cfg.hidden_size} L={cfg.num_hidden_layers} "
-          f"head={cfg.num_attention_heads} d={cfg.head_dim} I={cfg.intermediate_size} "
-          f"pe_theta={cfg.pe_theta} ln_eps={cfg.layer_norm_eps}")
-    print(f"embed_norm={cfg.embed_norm} skip_first_prenorm={cfg.skip_first_prenorm} "
-          f"final_norm={cfg.final_norm} init={cfg.init_method}")
-    print(f"dropout: embed={cfg.embed_dropout} mlp={cfg.mlp_dropout} "
-          f"attn_out={cfg.attn_out_dropout} attn_probs={cfg.attn_probs_dropout}")
-    model = ModernBertForMLM(cfg)
-    n = model.num_parameters()
-    print(f"params: {n:,} = {n / 1e6:.1f}M")
-
-    # forward smoke
-    B, L = 2, 128
-    ids = torch.randint(0, cfg.vocab_size, (B, L))
-    mask = torch.ones(B, L, dtype=torch.long)
-    labels = torch.randint(0, cfg.vocab_size, (B, L))
-    out = model(ids, attention_mask=mask, labels=labels)
-    print(f"logits: {out['logits'].shape}  loss: {out['loss'].item():.4f}")
-
-    # backward smoke
-    out["loss"].backward()
-    # 检查所有 trainable param 都有 grad
-    no_grad = [n for n, p in model.named_parameters() if p.requires_grad and p.grad is None]
-    if no_grad:
-        print(f"ERROR: 缺 grad 的参数: {no_grad}")
-    else:
-        print("backward OK (所有参数都有 grad)")
-
-    # 验证 skip_first_prenorm 生效
-    assert model.bert.layers[0].skip_norm1, "layer[0].skip_norm1 应为 True"
-    assert not model.bert.layers[1].skip_norm1, "layer[1].skip_norm1 应为 False"
-    assert isinstance(model.bert.layers[0].norm1, nn.Identity), "layer[0].norm1 应为 Identity"
-    assert isinstance(model.bert.layers[1].norm1, LayerNormNoBias), "layer[1].norm1 应为 LayerNormNoBias"
-    print("skip_first_prenorm 配置正确")
-
-    # 验证 embed_norm + final_norm
-    assert isinstance(model.bert.embed_norm, LayerNormNoBias), "embed_norm 应为 LayerNormNoBias"
-    assert isinstance(model.bert.final_norm, LayerNormNoBias), "final_norm 应为 LayerNormNoBias"
-    print("embed_norm + final_norm 配置正确")
-
-    # 验证 Megatron init
-    L_layers = cfg.num_hidden_layers
-    expected_scale = (2.0 * L_layers) ** -0.5
-    # 一个未缩的初始值 std ≈ 0.02,缩后 std ≈ 0.02 * expected_scale
-    o_std = model.bert.layers[5].attn.o.weight.std().item()
-    expected_std = cfg.initializer_range * expected_scale
-    # 允许 ±50% 容差(单 tensor std 估计有 noise)
-    assert 0.5 * expected_std < o_std < 1.5 * expected_std, \
-        f"Megatron init: attn.o.weight.std={o_std:.6f}, expected≈{expected_std:.6f}"
-    print(f"Megatron init OK: attn.o.weight.std={o_std:.6f} ≈ {expected_std:.6f}")
-
-    # 验证 no bias
-    for name, p in model.named_parameters():
-        if "bias" in name and name != "head_bias":
-            print(f"ERROR: 不该有的 bias: {name}")
-    print("Linear/Norm no-bias 配置正确")
-
-    print("\n=== Test 1b: seg_ids → flex_attention 路径(cross-doc 隔离)===")
-    if torch.cuda.is_available():
-        # flex_attention 需要 CUDA
-        model_cuda = ModernBertForMLM(cfg).cuda().to(torch.bfloat16)
-        ids_c = ids.cuda()
-        labels_c = labels.cuda()
-        # 构造 seg_ids:chunk 内 0 0 0 ... 0 | 1 1 ... 1 (前半 doc0, 后半 doc1)
-        seg_ids = torch.zeros(B, L, dtype=torch.int32, device="cuda")
-        seg_ids[:, L // 2:] = 1
-        out_flex = model_cuda(ids_c, seg_ids=seg_ids, labels=labels_c)
-        print(f"  flex_attention forward OK  loss={out_flex['loss'].item():.4f}  "
-              f"logits={out_flex['logits'].shape} dtype={out_flex['logits'].dtype}")
-        out_flex["loss"].backward()
-        no_grad = [n for n, p in model_cuda.named_parameters() if p.requires_grad and p.grad is None]
-        if no_grad:
-            print(f"  ERROR: 缺 grad: {no_grad[:5]}")
-        else:
-            print(f"  flex_attention backward OK")
-
-        # 验证 attention 真正隔离了:用相同 ids 但 seg_ids 不同,前半 token 应该输出不同
-        # (因为 doc0 token 在 seg=0 时只 attend 同 seg=0;改成 seg=2 后 attend 不同的 kv 集)
-        with torch.no_grad():
-            seg_v1 = torch.zeros(B, L, dtype=torch.int32, device="cuda")
-            seg_v2 = torch.zeros(B, L, dtype=torch.int32, device="cuda")
-            seg_v2[:, :L // 2] = 0
-            seg_v2[:, L // 2:] = 1  # 前半 vs 后半 隔离
-            h1 = model_cuda.bert(ids_c, seg_ids=seg_v1)  # 全 attend(单 doc)
-            h2 = model_cuda.bert(ids_c, seg_ids=seg_v2)  # 前后半隔离
-            # 前半 token 在 v1 attend 全 chunk,在 v2 只 attend 前半 → 输出应不同
-            front_diff = (h1[:, :L // 2] - h2[:, :L // 2]).abs().mean().item()
-            print(f"  前半 token 输出差异(应 > 0,验证隔离生效):{front_diff:.6f}")
-            assert front_diff > 1e-3, f"隔离未生效:front_diff={front_diff}"
-            print(f"  跨 doc 隔离工作正常 ✓")
-    else:
-        print("  跳过(无 CUDA)")
-
-    print("\n=== Test 2: legacy v1 config (12L/1024H) — 验证向后兼容 ===")
-    cfg_v1 = ModernBertConfig(
-        hidden_size=1024, num_hidden_layers=12, num_attention_heads=16,
-        intermediate_size=2752,
-        embed_norm=False, skip_first_prenorm=False, init_method="normal",
-    )
-    model_v1 = ModernBertForMLM(cfg_v1)
-    n_v1 = model_v1.num_parameters()
-    print(f"v1-style params: {n_v1 / 1e6:.1f}M")
-    out_v1 = model_v1(ids, attention_mask=mask, labels=labels)
-    print(f"v1 forward OK  loss={out_v1['loss'].item():.4f}")
-
-    print("\n=== All smoke tests passed ===")

tokenizer.py

@@ -0,0 +1,64 @@
+"""BERTc 发布包自带的字级 tokenizer。
+
+这份代码会**随模型一起发到 HF**,所以只能依赖 piece_tokenizer 本身,
+不能 import 仓库里的任何东西。
+
+装 tokenizer:
+    pip install git+https://github.com/Ismantic/PieceTokenizer
+"""
+from pathlib import Path
+
+import piece_tokenizer as _pt
+
+
+class PieceCharTokenizer:
+    """字级 tokenizer。
+
+    必须用 dict="no" 加载(字模式,不挂分词词典)—— 挂了词典编码结果会跟
+    训练时不一致,而且不会报错。
+    """
+
+    def __init__(self, model_dir="."):
+        model_dir = Path(model_dir)
+        self._tok = _pt.Tokenizer()
+        self._tok.load(str(model_dir / "piece.model"), dict="no")
+
+        self.pad_token_id = self._tok.piece_to_id("<pad>")
+        self.unk_token_id = 0
+        mask_path = model_dir / "mask_token_id.txt"
+        self.mask_token_id = (int(mask_path.read_text().strip())
+                              if mask_path.exists() else self._tok.vocab_size())
+        self.vocab_size = self._tok.vocab_size() + 1
+        self._cache = {}
+
+    def char_to_id(self, char: str) -> int:
+        tid = self._cache.get(char)
+        if tid is None:
+            ids = self._tok.encode_as_ids(char)
+            tid = ids[0] if ids else self.unk_token_id
+            self._cache[char] = tid
+        return tid
+
+    def id_to_char(self, tid: int) -> str:
+        piece = self._tok.id_to_piece(int(tid))
+        return piece.replace("▁", "")
+
+    def encode(self, text: str) -> list:
+        return [self.char_to_id(c) for c in text]
+
+    def batch(self, texts, max_len, device=None):
+        """一批文本 → (input_ids, attention_mask, 每条的有效长度)。"""
+        import torch
+
+        lengths = [min(len(t), max_len) for t in texts]
+        width = max(lengths) if lengths else 0
+        input_ids = torch.full((len(texts), width), self.pad_token_id,
+                               dtype=torch.long, device=device)
+        attn = torch.zeros((len(texts), width), dtype=torch.long, device=device)
+        for i, t in enumerate(texts):
+            ids = self.encode(t[:lengths[i]])
+            if ids:
+                input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long,
+                                                       device=device)
+                attn[i, :len(ids)] = 1
+        return input_ids, attn, lengths