File size: 1,154 Bytes
caf0502 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
"""
This module supplies `transformers`-compatible wrappers for
`GPTXTokenizer`s.
The tokenizers in this do not conform to the `PreTrainedTokenizer` API,
but allow for better practical usage.
"""
from typing import List
from gptx_tokenizer.hf_wrappers import (
HFTokenizer as _HFTokenizer,
SPTokenizer as _SPTokenizer,
)
class HFTokenizer(_HFTokenizer):
# The tokenizer is ridiculously slow without this; however, this
# doesn't implement all APIs of `PreTrainedTokenizer`.
def encode(self, text: str, **kwargs) -> List[int]:
return_tokens = kwargs.pop('return_tokens', False)
return self._tok.encode(text, return_tokens=return_tokens)
class SPTokenizer(_SPTokenizer):
# `is_continuation` does not work without this, but it doesn't
# implement all APIs of `PreTrainedTokenizer`.
def encode(self, text: str, **kwargs) -> List[int]:
return_tokens = kwargs.pop('return_tokens', False)
is_continuation = kwargs.pop('is_continuation', False)
return self._tok.encode(
text,
return_tokens=return_tokens,
is_continuation=is_continuation,
)
|