| """
|
| 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
|
| self._vocab_size = 0
|
| self._special_token_to_id = {}
|
| self._id_to_special_token = {}
|
|
|
| @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,
|
| normalization_rule_name="nfkc",
|
| pad_id=3,
|
| input_sentence_size=1_000_000,
|
| 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)
|
|
|
|
|
| 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."""
|
|
|
| unk_id = self.sp.unk_id() if self.sp else 0
|
| ids = [i for i in ids if i != unk_id]
|
|
|
| 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
|
|
|
