Class HgfTokenizerWrapper

codegeex/megatron/tokenizer/tokenizer.py:217–258

Wrapper for Hugging Face tokenizer.

class HgfTokenizerWrapper(AbstractTokenizer):
    """Wrapper for Hugging Face tokenizer."""

    def __init__(
            self,
            tokenizer,
            ws_start: int = None,
            ws_len: int = None,
    ):
        super(HgfTokenizerWrapper, self).__init__(tokenizer.__class__.__name__)
        self.tokenizer = tokenizer
        self.ws_start = ws_start
        self.ws_len = ws_len

    def tokenize(self, text):
        if self.ws_start:
            text = encode_whitespaces(text, self.ws_start, self.ws_len)
        input_ids = self.tokenizer(text, is_split_into_words=False).input_ids

        return input_ids

    def detokenize(self, token_ids):
        text = self.tokenizer.decode(token_ids, skip_special_tokens=False)
        if self.ws_start:
            text = decode_whitespaces(text, self.ws_start, self.ws_len)
        return text

    @property
    def eod(self):
        return self.tokenizer.eos_token_id

    @property
    def inv_vocab(self):
        return len(self.tokenizer.decoder)

    @property
    def vocab(self):
        return self.tokenizer.vocab

    @property
    def vocab_size(self):
        return len(self.vocab)