tokenizer.py

import json

PAD, BOS, EOS, UNK = 0, 1, 2, 3
SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>"]

class CharTokenizer:
    def __init__(self, vocab=None):
        self.vocab = vocab or {}
        self.inv = {i: c for c, i in self.vocab.items()}

    @classmethod
    def build(cls, texts, min_freq=20):
        from collections import Counter
        cnt = Counter()
        for t in texts:
            cnt.update(t)
        vocab = {s: i for i, s in enumerate(SPECIALS)}
        for ch, freq in sorted(cnt.items()):
            if freq >= min_freq:
                vocab[ch] = len(vocab)
        return cls(vocab)

    def encode(self, text, add_special=True):
        ids = [self.vocab.get(c, UNK) for c in text]
        return [BOS] + ids + [EOS] if add_special else ids

    def decode(self, ids):
        return "".join(self.inv.get(i, "") for i in ids
                       if i not in (PAD, BOS, EOS, UNK))

    def save(self, path):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.vocab, f, ensure_ascii=False)

    @classmethod
    def load(cls, path):
        with open(path, encoding="utf-8") as f:
            return cls(json.load(f))

    def __len__(self):
        return len(self.vocab)