| from os import PathLike |
| from typing import Dict, List, Union |
| from wenet.text.base_tokenizer import BaseTokenizer, T as Type |
|
|
|
|
| class HuggingFaceTokenizer(BaseTokenizer): |
|
|
| def __init__(self, model: Union[str, PathLike], *args, **kwargs) -> None: |
| |
| self.model = model |
| self.tokenizer = None |
|
|
| self.args = args |
| self.kwargs = kwargs |
|
|
| def __getstate__(self): |
| state = self.__dict__.copy() |
| del state['tokenizer'] |
| return state |
|
|
| def __setstate__(self, state): |
| self.__dict__.update(state) |
| recovery = {'tokenizer': None} |
| self.__dict__.update(recovery) |
|
|
| def _build_hugging_face(self): |
| from transformers import AutoTokenizer |
| if self.tokenizer is None: |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| self.model, **self.kwargs) |
| self.t2i = self.tokenizer.get_vocab() |
|
|
| def text2tokens(self, line: str) -> List[Type]: |
| self._build_hugging_face() |
| return self.tokenizer.tokenize(line) |
|
|
| def tokens2text(self, tokens: List[Type]) -> str: |
| self._build_hugging_face() |
| ids = self.tokens2ids(tokens) |
| return self.tokenizer.decode(ids) |
|
|
| def tokens2ids(self, tokens: List[Type]) -> List[int]: |
| self._build_hugging_face() |
| return self.tokenizer.convert_tokens_to_ids(tokens) |
|
|
| def ids2tokens(self, ids: List[int]) -> List[Type]: |
| self._build_hugging_face() |
| return self.tokenizer.convert_ids_to_tokens(ids) |
|
|
| def vocab_size(self) -> int: |
| self._build_hugging_face() |
| |
| return len(self.tokenizer) |
|
|
| @property |
| def symbol_table(self) -> Dict[Type, int]: |
| self._build_hugging_face() |
| return self.t2i |
|
|
