Upload code/tokenizer.py with huggingface_hub

code/tokenizer.py

@@ -0,0 +1,63 @@
+import json
+from pathlib import Path
+
+from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
+
+VOCAB_SIZE = 32000
+SPECIAL_TOKENS = ["<pad>", "<unk>", "<bos>", "<eos>"]
+
+
+class BPETokenizer:
+    def __init__(self, tokenizer: Tokenizer):
+        self._tok = tokenizer
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tok.get_vocab_size()
+
+    def encode(self, text: str) -> list[int]:
+        return self._tok.encode(text).ids
+
+    def decode(self, ids) -> str:
+        return self._tok.decode(list(ids))
+
+    def save(self, path: Path):
+        self._tok.save(str(path))
+
+    @classmethod
+    def load(cls, path: Path) -> "BPETokenizer":
+        return cls(Tokenizer.from_file(str(path)))
+
+    @classmethod
+    def build_from_text(cls, text: str, vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
+        tok = Tokenizer(models.BPE(unk_token="<unk>"))
+        tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
+        tok.decoder = decoders.ByteLevel()
+        trainer = trainers.BpeTrainer(
+            vocab_size=vocab_size,
+            special_tokens=SPECIAL_TOKENS,
+            min_frequency=2,
+        )
+        tok.train_from_iterator(_chunk(text), trainer=trainer)
+        return cls(tok)
+
+    @classmethod
+    def build_from_files(cls, paths: list[Path], vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
+        tok = Tokenizer(models.BPE(unk_token="<unk>"))
+        tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
+        tok.decoder = decoders.ByteLevel()
+        trainer = trainers.BpeTrainer(
+            vocab_size=vocab_size,
+            special_tokens=SPECIAL_TOKENS,
+            min_frequency=2,
+        )
+        tok.train([str(p) for p in paths], trainer=trainer)
+        return cls(tok)
+
+
+def _chunk(text: str, size: int = 1_000_000):
+    for i in range(0, len(text), size):
+        yield text[i:i + size]
+
+
+CharTokenizer = BPETokenizer