File size: 3,933 Bytes
73d3d73 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | 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
|