File size: 5,767 Bytes
d993048 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
Level 2 (Advanced): BPE Tokenizer using SentencePiece
Instead of one token per character, BPE groups common character sequences
into "subwords". For example, the common Armenian word "Հայաստան" might
become just 1-2 tokens instead of 8 characters.
This gives better results but requires an extra library:
pip install sentencepiece
How it works:
1. Train: learn common character groups from Armenian text
2. Encode: split text into subword tokens
3. Decode: join subword tokens back into text
"""
import os
import json
class BPETokenizer:
"""Subword tokenizer using SentencePiece BPE."""
def __init__(self):
self.sp = None # SentencePiece processor
self._vocab_size = 0
self._special_token_to_id = {} # e.g. {"<|user|>": 8000}
self._id_to_special_token = {} # e.g. {8000: "<|user|>"}
@property
def vocab_size(self):
base = self.sp.get_piece_size() if self.sp is not None else self._vocab_size
return base + len(self._special_token_to_id)
def train(self, text_file, model_prefix="data/bpe_model", vocab_size=16000):
"""
Train a BPE model on Armenian text.
Args:
text_file: path to a .txt file with training text
model_prefix: where to save the model (creates .model and .vocab files)
vocab_size: number of subword tokens to learn (8000 is good for Armenian)
"""
try:
import sentencepiece as spm
except ImportError:
print("Error: sentencepiece not installed!")
print("Install it with: pip install sentencepiece")
raise
print(f"Training BPE tokenizer (vocab_size={vocab_size})...")
spm.SentencePieceTrainer.train(
input=text_file,
model_prefix=model_prefix,
vocab_size=vocab_size,
model_type="bpe",
character_coverage=0.9999, # cover almost all Armenian characters
normalization_rule_name="nfkc",
pad_id=3,
input_sentence_size=1_000_000, # sample 1M sentences for large files
shuffle_input_sentence=True,
num_threads=16,
)
self.sp = spm.SentencePieceProcessor(model_file=f"{model_prefix}.model")
print(f"BPE tokenizer trained! Vocab size: {self.vocab_size}")
def add_special_tokens(self, tokens):
"""
Register multi-character special tokens.
SentencePiece doesn't natively add tokens after training, so we
map them to IDs beyond the existing vocab.
"""
for token in tokens:
if token not in self._special_token_to_id:
idx = self.vocab_size
self._special_token_to_id[token] = idx
self._id_to_special_token[idx] = token
return self
def encode(self, text):
"""Convert text to a list of integer token IDs."""
if not self._special_token_to_id:
return self.sp.encode(text)
# Split text around special tokens, encode each segment, insert special IDs
import re
pattern = re.compile("(" + "|".join(re.escape(t) for t in self._special_token_to_id) + ")")
parts = pattern.split(text)
ids = []
for part in parts:
if part in self._special_token_to_id:
ids.append(self._special_token_to_id[part])
elif part:
ids.extend(self.sp.encode(part))
return ids
def decode(self, ids):
"""Convert a list of integer token IDs back to text."""
# Filter out unk tokens (id 0) to avoid ⁇ in output
unk_id = self.sp.unk_id() if self.sp else 0
ids = [i for i in ids if i != unk_id]
# Decode in segments, replacing special token IDs with their strings
result = []
sp_ids = []
for i in ids:
if i in self._id_to_special_token:
if sp_ids:
result.append(self.sp.decode(sp_ids))
sp_ids = []
result.append(self._id_to_special_token[i])
else:
sp_ids.append(i)
if sp_ids:
result.append(self.sp.decode(sp_ids))
return "".join(result)
def save(self, path):
"""Save tokenizer metadata (the .model file is saved during training)."""
data = {
"type": "bpe",
"vocab_size": self.vocab_size,
"model_file": self.sp.serialized_model_proto().hex()
if self.sp else None,
"special_tokens": self._special_token_to_id,
}
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f)
@classmethod
def load(cls, path):
"""Load BPE tokenizer from saved metadata."""
try:
import sentencepiece as spm
except ImportError:
print("Error: sentencepiece not installed!")
print("Install it with: pip install sentencepiece")
raise
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
tok = cls()
if data.get("model_file"):
tok.sp = spm.SentencePieceProcessor()
tok.sp.load_from_serialized_proto(bytes.fromhex(data["model_file"]))
else:
tok._vocab_size = data["vocab_size"]
if data.get("special_tokens"):
tok._special_token_to_id = {k: int(v) for k, v in data["special_tokens"].items()}
tok._id_to_special_token = {v: k for k, v in tok._special_token_to_id.items()}
return tok
|