| 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 |
|
|
|
|
| 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): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.bert = ModernBertModel(config) |
| self.cor_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.cor_head.weight = self.bert.embed.weight |
| self.det_head = nn.Linear(config.hidden_size, 1) |
|
|
| @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())) |
| model = cls(cfg) |
| 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}") |
| 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 |
|
|
| @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 |
| 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 |
| cor_logits, _ = self(input_ids, attn) |
| probs = F.softmax(cor_logits, dim=-1) |
| top_probs, top_ids = 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)) |
| return out[0] if single else out |
|
|
