tokenization_domsense.py

from transformers import PreTrainedTokenizer

from .configuration_domsense import SystemOneConfig

MAX_LEN = 128
_VOCAB = None


def _byte_map(max_len: int = MAX_LEN, pad_id: int = 0):
    """构造与 LightweightTextEncoder.tokenize_texts 一致的字节 id 序列。"""
    return max_len, pad_id


class DomSenseTokenizer(PreTrainedTokenizer):
    """DomSense 字节级 Tokenizer。

    与内置轻量编码器 ``LightweightTextEncoder.tokenize_texts`` 的映射完全一致：
    每个 UTF-8 字节映射到 ``[0, 255]``，超过 ``max_len`` 截断，不足用 0（pad）补齐。
    因此该 tokenizer 的输出可直接喂给仓库自带的 lightweight 编码器。
    """

    vocab_files_names = {}

    def __init__(
        self,
        max_len: int = MAX_LEN,
        pad_token: str = "[PAD]",
        unk_token: str = "[UNK]",
        cls_token: str = "[CLS]",
        sep_token: str = "[SEP]",
        **kwargs,
    ):
        self.max_len = max_len
        super().__init__(
            max_len=max_len,
            pad_token=pad_token,
            unk_token=unk_token,
            cls_token=cls_token,
            sep_token=sep_token,
            **kwargs,
        )
        # 若 pad 与 unk 相同会造成冲突，这里 unk 也用 [UNK]，保持分离
        if self.pad_token is None:
            self.pad_token = self.unk_token

    def _tokenize(self, text):
        """把文本切成 UTF-8 字节 id 列表（与轻量编码器一致）。"""
        if isinstance(text, str) is False:
            text = str(text)
        byte_ids = [
            min(b, 255)
            for b in text.encode("utf-8", errors="ignore")[: self.max_len]
        ]
        return byte_ids

    def _convert_token_to_id(self, token) -> int:
        return int(token)

    def _convert_id_to_token(self, index) -> str:
        return str(int(index))

    def get_vocab(self):
        return {}

    @property
    def vocab_size(self) -> int:
        return 256

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        return list(token_ids_0)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        return ()