superlillm/tokenizer.py

import json
import re
from collections import Counter


TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)


class WordTokenizer:
    def __init__(self, token_to_id=None):
        self.specials = ["<pad>", "<unk>", "<bos>", "<eos>"]
        self.token_to_id = token_to_id or {token: i for i, token in enumerate(self.specials)}
        self.id_to_token = {i: token for token, i in self.token_to_id.items()}

    def tokenize(self, text):
        text = text.replace("’", "'").replace("“", '"').replace("”", '"')
        return TOKEN_RE.findall(text.lower())

    def build(self, texts, min_freq=1):
        counts = Counter()
        for text in texts:
            counts.update(self.tokenize(text))
        for token, freq in counts.most_common():
            if freq >= min_freq and token not in self.token_to_id:
                self.token_to_id[token] = len(self.token_to_id)
        self.id_to_token = {i: token for token, i in self.token_to_id.items()}

    def encode(self, text, add_bos=False, add_eos=False):
        ids = []
        if add_bos:
            ids.append(self.token_to_id["<bos>"])
        ids.extend(self.token_to_id.get(token, self.token_to_id["<unk>"]) for token in self.tokenize(text))
        if add_eos:
            ids.append(self.token_to_id["<eos>"])
        return ids

    def decode(self, ids):
        tokens = []
        for idx in ids:
            token = self.id_to_token.get(int(idx), "<unk>")
            if token in self.specials:
                continue
            tokens.append(token)
        text = " ".join(tokens)
        text = re.sub(r"\s+([,.!?;:])", r"\1", text)
        text = re.sub(r"([(\[{])\s+", r"\1", text)
        text = re.sub(r"\s+([)\]}])", r"\1", text)
        return text.strip()

    def save(self, path):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.token_to_id, f, indent=2)

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