| from pathlib import Path |
| from typing import Any, List, Optional, Union |
|
|
| from transformers import AutoTokenizer, PreTrainedTokenizerBase |
|
|
|
|
| class TokenizerBase(PreTrainedTokenizerBase): |
| ''' This is a protocol for the tokenizer. Users can implement their own tokenizer by inheriting this class. ''' |
|
|
|
|
| class TransformersTokenizer(TokenizerBase): |
| ''' A wrapper for the Transformers' tokenizer. |
| This is the default tokenizer for LLM. ''' |
|
|
| def __init__(self, tokenizer): |
| self.tokenizer = tokenizer |
|
|
| def __call__(self, text: str, *args, **kwargs) -> Any: |
| return self.tokenizer(text, *args, **kwargs) |
|
|
| @property |
| def eos_token_id(self) -> int: |
| return self.tokenizer.eos_token_id |
|
|
| @property |
| def pad_token_id(self) -> int: |
| return self.tokenizer.pad_token_id |
|
|
| def encode(self, text: str, *args, **kwargs) -> List[int]: |
| return self.tokenizer.encode(text, *args, **kwargs) |
|
|
| def decode(self, token_ids: List[int], *args, **kwargs) -> str: |
| return self.tokenizer.decode(token_ids, *args, **kwargs) |
|
|
| def batch_encode_plus(self, texts: List[str], *args, **kwargs) -> dict: |
| return self.tokenizer.batch_encode_plus(texts, *args, **kwargs) |
|
|
| def __repr__(self) -> str: |
| return f"{self.__class__.__name__}({self.tokenizer})" |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_dir: str, **kwargs): |
| tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, |
| **kwargs) |
| return cls(tokenizer) |
|
|
|
|
| def tokenizer_factory(obj: Optional[Union[str, Path, PreTrainedTokenizerBase, |
| TokenizerBase]] = None, |
| **kwargs) -> Optional[TokenizerBase]: |
| if obj is None: |
| return None |
| elif isinstance(obj, (str, Path)): |
| default_kwargs = { |
| 'legacy': False, |
| 'padding_side': 'left', |
| 'truncation_side': 'left', |
| 'trust_remote_code': True, |
| 'use_fast': True, |
| } |
| default_kwargs.update(kwargs) |
| return TransformersTokenizer.from_pretrained(obj, **default_kwargs) |
| elif isinstance(obj, PreTrainedTokenizerBase): |
| return TransformersTokenizer(obj) |
| elif isinstance(obj, TokenizerBase): |
| return obj |
| else: |
| raise TypeError(f"Unrecognized tokenizer {obj}") |
|
|
