| from typing import Dict, Iterator, List, Optional, Union
|
|
|
| from tokenizers import AddedToken, Tokenizer, decoders, trainers
|
| from tokenizers.models import WordPiece
|
| from tokenizers.normalizers import BertNormalizer
|
| from tokenizers.pre_tokenizers import BertPreTokenizer
|
| from tokenizers.processors import BertProcessing
|
|
|
| from .base_tokenizer import BaseTokenizer
|
|
|
|
|
| class BertWordPieceTokenizer(BaseTokenizer):
|
| """Bert WordPiece Tokenizer"""
|
|
|
| def __init__(
|
| self,
|
| vocab: Optional[Union[str, Dict[str, int]]] = None,
|
| unk_token: Union[str, AddedToken] = "[UNK]",
|
| sep_token: Union[str, AddedToken] = "[SEP]",
|
| cls_token: Union[str, AddedToken] = "[CLS]",
|
| pad_token: Union[str, AddedToken] = "[PAD]",
|
| mask_token: Union[str, AddedToken] = "[MASK]",
|
| clean_text: bool = True,
|
| handle_chinese_chars: bool = True,
|
| strip_accents: Optional[bool] = None,
|
| lowercase: bool = True,
|
| wordpieces_prefix: str = "##",
|
| ):
|
| if vocab is not None:
|
| tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(unk_token)))
|
| else:
|
| tokenizer = Tokenizer(WordPiece(unk_token=str(unk_token)))
|
|
|
|
|
| if tokenizer.token_to_id(str(unk_token)) is not None:
|
| tokenizer.add_special_tokens([str(unk_token)])
|
| if tokenizer.token_to_id(str(sep_token)) is not None:
|
| tokenizer.add_special_tokens([str(sep_token)])
|
| if tokenizer.token_to_id(str(cls_token)) is not None:
|
| tokenizer.add_special_tokens([str(cls_token)])
|
| if tokenizer.token_to_id(str(pad_token)) is not None:
|
| tokenizer.add_special_tokens([str(pad_token)])
|
| if tokenizer.token_to_id(str(mask_token)) is not None:
|
| tokenizer.add_special_tokens([str(mask_token)])
|
|
|
| tokenizer.normalizer = BertNormalizer(
|
| clean_text=clean_text,
|
| handle_chinese_chars=handle_chinese_chars,
|
| strip_accents=strip_accents,
|
| lowercase=lowercase,
|
| )
|
| tokenizer.pre_tokenizer = BertPreTokenizer()
|
|
|
| if vocab is not None:
|
| sep_token_id = tokenizer.token_to_id(str(sep_token))
|
| if sep_token_id is None:
|
| raise TypeError("sep_token not found in the vocabulary")
|
| cls_token_id = tokenizer.token_to_id(str(cls_token))
|
| if cls_token_id is None:
|
| raise TypeError("cls_token not found in the vocabulary")
|
|
|
| tokenizer.post_processor = BertProcessing((str(sep_token), sep_token_id), (str(cls_token), cls_token_id))
|
| tokenizer.decoder = decoders.WordPiece(prefix=wordpieces_prefix)
|
|
|
| parameters = {
|
| "model": "BertWordPiece",
|
| "unk_token": unk_token,
|
| "sep_token": sep_token,
|
| "cls_token": cls_token,
|
| "pad_token": pad_token,
|
| "mask_token": mask_token,
|
| "clean_text": clean_text,
|
| "handle_chinese_chars": handle_chinese_chars,
|
| "strip_accents": strip_accents,
|
| "lowercase": lowercase,
|
| "wordpieces_prefix": wordpieces_prefix,
|
| }
|
|
|
| super().__init__(tokenizer, parameters)
|
|
|
| @staticmethod
|
| def from_file(vocab: str, **kwargs):
|
| vocab = WordPiece.read_file(vocab)
|
| return BertWordPieceTokenizer(vocab, **kwargs)
|
|
|
| def train(
|
| self,
|
| files: Union[str, List[str]],
|
| vocab_size: int = 30000,
|
| min_frequency: int = 2,
|
| limit_alphabet: int = 1000,
|
| initial_alphabet: List[str] = [],
|
| special_tokens: List[Union[str, AddedToken]] = [
|
| "[PAD]",
|
| "[UNK]",
|
| "[CLS]",
|
| "[SEP]",
|
| "[MASK]",
|
| ],
|
| show_progress: bool = True,
|
| wordpieces_prefix: str = "##",
|
| ):
|
| """Train the model using the given files"""
|
|
|
| trainer = trainers.WordPieceTrainer(
|
| vocab_size=vocab_size,
|
| min_frequency=min_frequency,
|
| limit_alphabet=limit_alphabet,
|
| initial_alphabet=initial_alphabet,
|
| special_tokens=special_tokens,
|
| show_progress=show_progress,
|
| continuing_subword_prefix=wordpieces_prefix,
|
| )
|
| if isinstance(files, str):
|
| files = [files]
|
| self._tokenizer.train(files, trainer=trainer)
|
|
|
| def train_from_iterator(
|
| self,
|
| iterator: Union[Iterator[str], Iterator[Iterator[str]]],
|
| vocab_size: int = 30000,
|
| min_frequency: int = 2,
|
| limit_alphabet: int = 1000,
|
| initial_alphabet: List[str] = [],
|
| special_tokens: List[Union[str, AddedToken]] = [
|
| "[PAD]",
|
| "[UNK]",
|
| "[CLS]",
|
| "[SEP]",
|
| "[MASK]",
|
| ],
|
| show_progress: bool = True,
|
| wordpieces_prefix: str = "##",
|
| length: Optional[int] = None,
|
| ):
|
| """Train the model using the given iterator"""
|
|
|
| trainer = trainers.WordPieceTrainer(
|
| vocab_size=vocab_size,
|
| min_frequency=min_frequency,
|
| limit_alphabet=limit_alphabet,
|
| initial_alphabet=initial_alphabet,
|
| special_tokens=special_tokens,
|
| show_progress=show_progress,
|
| continuing_subword_prefix=wordpieces_prefix,
|
| )
|
| self._tokenizer.train_from_iterator(
|
| iterator,
|
| trainer=trainer,
|
| length=length,
|
| )
|
|
|
