File size: 2,506 Bytes
dd55cba 60a3383 dd55cba 60a3383 dd55cba 60a3383 dd55cba 60a3383 dd55cba | 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 | """BERTc 发布包自带的字级 tokenizer。
这份代码会**随模型一起发到 HF**,所以只能依赖 piece_tokenizer 本身,
不能 import 仓库里的任何东西。
词表文件 BERTc-Tokenizer.pt 与 PieceTokenizer 仓库 save/ 下的那份逐字节相同 ——
同名是为了让来源一目了然。
装 tokenizer:
pip install git+https://github.com/Ismantic/PieceTokenizer
"""
from pathlib import Path
import piece_tokenizer as _pt
class PieceCharTokenizer:
"""字级 tokenizer。
必须用 dict="no" 加载(字模式,不挂分词词典)—— 挂了词典编码结果会跟
训练时不一致,而且不会报错。
"""
MODEL_NAME = "BERTc-Tokenizer.pt"
def __init__(self, model_dir="."):
model_dir = Path(model_dir)
self._tok = _pt.Tokenizer()
self._tok.load(str(model_dir / self.MODEL_NAME), dict="no")
self.pad_token_id = self._tok.piece_to_id("<pad>")
self.unk_token_id = 0
# [MASK] 追加在 piece 词表之后,id 就等于词表大小 —— 不需要单独存一个文件
self.mask_token_id = 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
|