File size: 3,289 Bytes
5e6d9f5 | 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 | """分词器：优先用 HuggingFace tokenizers 训 BPE，装不上就退化成字节级分词。
字节级方案零依赖、永远不会 OOV，缺点是序列变长；小语料上其实够用。
"""
import json
import os
from typing import List, Optional
try:
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
_HAS_TOKENIZERS = True
except Exception: # pragma: no cover
_HAS_TOKENIZERS = False
PAD, BOS, EOS, UNK = "<pad>", "<bos>", "<eos>", "<unk>"
SPECIALS = [PAD, BOS, EOS, UNK]
class ByteTokenizer:
"""UTF-8 字节级分词器：vocab = 4 个特殊 token + 256 个字节。"""
kind = "byte"
def __init__(self):
self.vocab_size = 256 + len(SPECIALS)
self.pad_id, self.bos_id, self.eos_id, self.unk_id = 0, 1, 2, 3
self.offset = len(SPECIALS)
def encode(self, text: str, bos: bool = False, eos: bool = False) -> List[int]:
ids = [b + self.offset for b in text.encode("utf-8")]
if bos:
ids = [self.bos_id] + ids
if eos:
ids = ids + [self.eos_id]
return ids
def decode(self, ids: List[int]) -> str:
buf = bytes(i - self.offset for i in ids if i >= self.offset)
return buf.decode("utf-8", errors="replace")
def save(self, path: str):
with open(path, "w", encoding="utf-8") as f:
json.dump({"kind": "byte"}, f)
class BPETokenizer:
kind = "bpe"
def __init__(self, tok: "Tokenizer"):
self.tok = tok
self.vocab_size = tok.get_vocab_size()
self.pad_id = tok.token_to_id(PAD)
self.bos_id = tok.token_to_id(BOS)
self.eos_id = tok.token_to_id(EOS)
self.unk_id = tok.token_to_id(UNK)
def encode(self, text: str, bos: bool = False, eos: bool = False) -> List[int]:
ids = self.tok.encode(text).ids
if bos:
ids = [self.bos_id] + ids
if eos:
ids = ids + [self.eos_id]
return ids
def decode(self, ids: List[int]) -> str:
return self.tok.decode([i for i in ids if i not in (self.pad_id, self.bos_id)])
def save(self, path: str):
self.tok.save(path)
def train_bpe(texts: List[str], vocab_size: int, out_path: str) -> "BPETokenizer":
if not _HAS_TOKENIZERS:
raise RuntimeError("未安装 tokenizers 库，无法训练 BPE")
tok = Tokenizer(models.BPE(unk_token=UNK))
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=SPECIALS,
show_progress=False,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(texts, trainer=trainer)
os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
tok.save(out_path)
return BPETokenizer(tok)
def load_tokenizer(path: str):
with open(path, "r", encoding="utf-8") as f:
head = f.read(200)
if '"kind": "byte"' in head or '"kind":"byte"' in head:
return ByteTokenizer()
if not _HAS_TOKENIZERS:
raise RuntimeError("该分词器需要 tokenizers 库")
return BPETokenizer(Tokenizer.from_file(path))
|