""" Byte-Pair Encoding (BPE) Tokenizer for English-Malay Translation ================================================================= We support two modes: 1. **Shared tokenizer** (preferred for 10+2 Tied Transformer): A single BPE tokenizer trained on the concatenated en+ms corpus. Both encoder and decoder share the same vocabulary. 2. **Separate tokenizers** (legacy): Two independent BPE tokenizers, one per language. Why BPE? • Handles subword units, so rare / unseen words are decomposed into known subword pieces instead of mapping to [UNK]. • Malay is morphologically rich (prefixes: me-, ber-, di-; suffixes: -kan, -an, -i). BPE naturally learns these affixes as subword units, giving much better coverage than a word-level tokenizer. • Keeps vocabulary compact while still reaching high coverage on both English and Malay. Why shared vocabulary for en-ms? • Both languages use the Latin script with significant lexical overlap (loanwords: "teknologi", "matematik", "universiti"; numbers; proper nouns). • A joint BPE captures cross-lingual subword patterns and enables tied embeddings in the model (Press & Wolf, 2017), saving ~26M params. Design choices: • NFKC normalisation + lowercase – ensures consistent encoding of Unicode characters and removes casing noise. • Whitespace pre-tokeniser – splits on spaces before BPE merges; simple and effective for Latin-script languages. • Special tokens: [PAD] – padding for uniform sequence lengths in batches [UNK] – fallback for unknown characters [CLS] – beginning-of-sequence / classification token [SEP] – separator (unused in basic seq2seq but reserved) [MASK] – reserved for masked-LM pretraining objectives [BOS] – beginning of sentence (fed to decoder at step 0) [EOS] – end of sentence (signals the decoder to stop) """ from __future__ import annotations import os import tempfile from pathlib import Path from typing import Iterator, List, Optional, Union from tokenizers import Tokenizer from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer from tokenizers.pre_tokenizers import Whitespace from tokenizers.normalizers import Sequence, NFKC, Lowercase from tokenizers.processors import TemplateProcessing # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- SPECIAL_TOKENS: List[str] = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[BOS]", "[EOS]"] PAD_TOKEN = "[PAD]" UNK_TOKEN = "[UNK]" CLS_TOKEN = "[CLS]" SEP_TOKEN = "[SEP]" MASK_TOKEN = "[MASK]" BOS_TOKEN = "[BOS]" EOS_TOKEN = "[EOS]" DEFAULT_VOCAB_SIZE = 50_000 DEFAULT_MIN_FREQUENCY = 2 # --------------------------------------------------------------------------- # Helper: write an iterator of strings to a temporary file (needed by the # HuggingFace `tokenizers` training API which expects file paths). # --------------------------------------------------------------------------- def _write_texts_to_tmpfile(texts: Iterator[str]) -> str: """Write an iterable of strings to a temp file, one per line. Returns path.""" tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8") for line in texts: line = line.strip() if line: tmp.write(line + "\n") tmp.close() return tmp.name # --------------------------------------------------------------------------- # Core: build & train a BPE tokenizer # --------------------------------------------------------------------------- def build_tokenizer( vocab_size: int = DEFAULT_VOCAB_SIZE, min_frequency: int = DEFAULT_MIN_FREQUENCY, ) -> tuple[Tokenizer, BpeTrainer]: """ Create an *untrained* BPE tokenizer and its trainer. Returns ------- tokenizer : Tokenizer Ready to call ``tokenizer.train(files, trainer)``. trainer : BpeTrainer Configured trainer instance. """ tokenizer = Tokenizer(BPE(unk_token=UNK_TOKEN)) # --- Normalisation: NFKC (canonical Unicode) + lowercase ------------- tokenizer.normalizer = Sequence([NFKC(), Lowercase()]) # --- Pre-tokenisation: split on whitespace --------------------------- tokenizer.pre_tokenizer = Whitespace() # --- Trainer --------------------------------------------------------- trainer = BpeTrainer( vocab_size=vocab_size, min_frequency=min_frequency, special_tokens=SPECIAL_TOKENS, show_progress=True, ) return tokenizer, trainer def train_tokenizer( texts: Union[List[str], Iterator[str]], vocab_size: int = DEFAULT_VOCAB_SIZE, min_frequency: int = DEFAULT_MIN_FREQUENCY, files: Optional[List[str]] = None, ) -> Tokenizer: """ Train a BPE tokenizer on the given texts **or** files. Parameters ---------- texts : list[str] or iterator of str, optional Raw sentences. Ignored when *files* is provided. vocab_size : int Target vocabulary size (default 30 000). min_frequency : int Minimum frequency for a pair to be merged. files : list[str], optional Paths to plain-text files (one sentence per line). Returns ------- Tokenizer Trained tokenizer ready for encoding / decoding. """ tokenizer, trainer = build_tokenizer(vocab_size, min_frequency) if files is not None: tokenizer.train(files, trainer) else: # Write texts to a temporary file so we can use the fast Rust trainer tmp_path = _write_texts_to_tmpfile(iter(texts)) try: tokenizer.train([tmp_path], trainer) finally: os.remove(tmp_path) # --- Post-processing: wrap every encoded sequence with [BOS] … [EOS] - bos_id = tokenizer.token_to_id(BOS_TOKEN) eos_id = tokenizer.token_to_id(EOS_TOKEN) tokenizer.post_processor = TemplateProcessing( single=f"[BOS]:0 $A:0 [EOS]:0", pair=f"[BOS]:0 $A:0 [EOS]:0 [BOS]:1 $B:1 [EOS]:1", special_tokens=[ ("[BOS]", bos_id), ("[EOS]", eos_id), ], ) return tokenizer # --------------------------------------------------------------------------- # Convenience wrappers for saving / loading # --------------------------------------------------------------------------- def save_tokenizer(tokenizer: Tokenizer, path: Union[str, Path]) -> None: """Save a trained tokenizer to a JSON file.""" path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) tokenizer.save(str(path)) print(f"[✓] Tokenizer saved → {path}") def load_tokenizer(path: Union[str, Path]) -> Tokenizer: """Load a previously saved tokenizer from a JSON file.""" tokenizer = Tokenizer.from_file(str(path)) print(f"[✓] Tokenizer loaded ← {path}") return tokenizer # --------------------------------------------------------------------------- # Encoding / decoding helpers # --------------------------------------------------------------------------- def encode(tokenizer: Tokenizer, text: str) -> List[int]: """Encode a single string and return token IDs (includes [BOS]/[EOS]).""" return tokenizer.encode(text).ids def decode(tokenizer: Tokenizer, ids: List[int]) -> str: """Decode token IDs back to a string, skipping special tokens.""" return tokenizer.decode(ids, skip_special_tokens=True) def get_vocab_size(tokenizer: Tokenizer) -> int: """Return the size of the tokenizer's vocabulary.""" return tokenizer.get_vocab_size() def token_to_id(tokenizer: Tokenizer, token: str) -> Optional[int]: """Look up the integer ID for a single token string.""" return tokenizer.token_to_id(token) def id_to_token(tokenizer: Tokenizer, id: int) -> Optional[str]: """Look up the token string for a single integer ID.""" return tokenizer.id_to_token(id) # --------------------------------------------------------------------------- # High-level: train a SHARED tokenizer on both languages (for tied embeddings) # --------------------------------------------------------------------------- def train_shared_tokenizer_from_dataset( dataset, src_lang: str = "en", tgt_lang: str = "ms", vocab_size: int = DEFAULT_VOCAB_SIZE, save_dir: Union[str, Path] = "tokenizer", ) -> Tokenizer: """ Train a single shared BPE tokenizer on the concatenated en+ms corpus. This is used with the 10+2 Tied Transformer architecture, where both encoder and decoder share the same vocabulary and embedding matrix. Parameters ---------- dataset : datasets.Dataset A HuggingFace dataset split where each example has a ``'translation'`` dict with keys for each language code. src_lang : str Source language code (default ``'en'``). tgt_lang : str Target language code (default ``'ms'``). vocab_size : int Vocabulary size for the shared tokenizer. save_dir : str or Path Directory to save the trained tokenizer JSON file. Returns ------- Tokenizer A single shared tokenizer for both languages. """ save_dir = Path(save_dir) # Concatenate all source and target sentences into one corpus src_texts = [example["translation"][src_lang] for example in dataset] tgt_texts = [example["translation"][tgt_lang] for example in dataset] all_texts = src_texts + tgt_texts print(f"Training shared BPE tokenizer on {len(all_texts):,} sentences " f"({len(src_texts):,} {src_lang} + {len(tgt_texts):,} {tgt_lang}) …") shared_tokenizer = train_tokenizer(all_texts, vocab_size=vocab_size) save_tokenizer(shared_tokenizer, save_dir / "tokenizer_shared.json") # Sanity check for name, sample in [(src_lang, src_texts[0]), (tgt_lang, tgt_texts[0])]: enc = shared_tokenizer.encode(sample) print(f"\n[{name}] Sample: {sample[:80]}…") print(f" Tokens : {enc.tokens[:15]}…") print(f" IDs : {enc.ids[:15]}…") print(f" Decoded: {shared_tokenizer.decode(enc.ids, skip_special_tokens=True)[:80]}…") print(f"\n[✓] Shared tokenizer trained and saved to {save_dir}/tokenizer_shared.json") return shared_tokenizer # --------------------------------------------------------------------------- # High-level: train source (English) & target (Malay) tokenizers from a # HuggingFace dataset split. # --------------------------------------------------------------------------- def train_tokenizers_from_dataset( dataset, src_lang: str = "en", tgt_lang: str = "ms", vocab_size: int = DEFAULT_VOCAB_SIZE, save_dir: Union[str, Path] = "tokenizer", ) -> tuple[Tokenizer, Tokenizer]: """ Train separate BPE tokenizers for source and target languages. Parameters ---------- dataset : datasets.Dataset A HuggingFace dataset split (e.g. ``dataset['train']``) where each example has a ``'translation'`` dict with keys for each language code. src_lang : str Source language code (default ``'en'``). tgt_lang : str Target language code (default ``'ms'``). vocab_size : int Vocabulary size for each tokenizer. save_dir : str or Path Directory to save the trained tokenizer JSON files. Returns ------- (src_tokenizer, tgt_tokenizer) """ save_dir = Path(save_dir) # Extract raw sentences from the dataset src_texts = [example["translation"][src_lang] for example in dataset] tgt_texts = [example["translation"][tgt_lang] for example in dataset] print(f"Training source tokenizer ({src_lang}) on {len(src_texts):,} sentences …") src_tokenizer = train_tokenizer(src_texts, vocab_size=vocab_size) save_tokenizer(src_tokenizer, save_dir / f"tokenizer_{src_lang}.json") print(f"Training target tokenizer ({tgt_lang}) on {len(tgt_texts):,} sentences …") tgt_tokenizer = train_tokenizer(tgt_texts, vocab_size=vocab_size) save_tokenizer(tgt_tokenizer, save_dir / f"tokenizer_{tgt_lang}.json") # Quick sanity check for name, tok, sample in [ (src_lang, src_tokenizer, src_texts[0]), (tgt_lang, tgt_tokenizer, tgt_texts[0]), ]: enc = tok.encode(sample) print(f"\n[{name}] Sample: {sample[:80]}…") print(f" Tokens : {enc.tokens[:15]}…") print(f" IDs : {enc.ids[:15]}…") print(f" Decoded: {tok.decode(enc.ids, skip_special_tokens=True)[:80]}…") print(f"\n[✓] Both tokenizers trained and saved to {save_dir}/") return src_tokenizer, tgt_tokenizer # --------------------------------------------------------------------------- # Standalone usage # --------------------------------------------------------------------------- if __name__ == "__main__": from datasets import load_from_disk print("Loading TED Talks IWSLT dataset (en ↔ ms, 2016) …") ds = load_from_disk("dataset/en_ms_2016") src_tok, tgt_tok = train_tokenizers_from_dataset( ds, src_lang="en", tgt_lang="ms", vocab_size=DEFAULT_VOCAB_SIZE, save_dir="tokenizer", ) print(f"\nEnglish vocab size : {get_vocab_size(src_tok):,}") print(f"Malay vocab size : {get_vocab_size(tgt_tok):,}") print(f"[PAD] id (en) : {token_to_id(src_tok, PAD_TOKEN)}") print(f"[EOS] id (ms) : {token_to_id(tgt_tok, EOS_TOKEN)}")