| import os |
| from tokenizers import Tokenizer |
| from tokenizers.models import BPE |
| from tokenizers.trainers import BpeTrainer |
| from tokenizers.pre_tokenizers import Whitespace |
| from tokenizers.decoders import BPEDecoder |
|
|
| class SocrateXTokenizer: |
| """ |
| Wrapper over HuggingFace Tokenizers to provide a clean interface: |
| .fit(), .encode(), .decode(), .save() |
| """ |
| def __init__(self, tokenizer=None, vocab_size=1000): |
| if tokenizer is None: |
| self._tokenizer = Tokenizer(BPE(unk_token="<unk>")) |
| self._tokenizer.pre_tokenizer = Whitespace() |
| self._tokenizer.decoder = BPEDecoder() |
| else: |
| self._tokenizer = tokenizer |
| if self._tokenizer.decoder is None: |
| self._tokenizer.decoder = BPEDecoder() |
| self.vocab_size = vocab_size |
|
|
| def fit(self, data_source, special_tokens=None): |
| """ |
| Trains the tokenizer. |
| data_source can be a path to a text file (e.g. "data.txt") |
| or a list of strings in memory. |
| """ |
| if special_tokens is None: |
| special_tokens = ["<pad>", "<bos>", "<eos>", "<unk>"] |
| |
| trainer = BpeTrainer(vocab_size=self.vocab_size, special_tokens=special_tokens) |
| |
| if isinstance(data_source, str) and os.path.exists(data_source): |
| self._tokenizer.train(files=[data_source], trainer=trainer) |
| elif isinstance(data_source, list): |
| self._tokenizer.train_from_iterator(data_source, trainer=trainer) |
| else: |
| raise ValueError("data_source must be a list of strings or a path to a text file.") |
| |
| def encode(self, text): |
| return self._tokenizer.encode(text) |
| |
| def decode(self, ids): |
| return self._tokenizer.decode(ids) |
| |
| def get_vocab_size(self): |
| return self._tokenizer.get_vocab_size() |
| |
| def token_to_id(self, token): |
| return self._tokenizer.token_to_id(token) |
| |
| def save(self, path): |
| self._tokenizer.save(path) |
| |
| @classmethod |
| def from_file(cls, path): |
| tk = Tokenizer.from_file(path) |
| return cls(tokenizer=tk, vocab_size=tk.get_vocab_size()) |
|
|
| def init_tokenizer(vocab_size=1000): |
| """ |
| Factory function to quickly instantiate an empty tokenizer, |
| ready to be trained via .fit() |
| """ |
| return SocrateXTokenizer(vocab_size=vocab_size) |
|
|
