codekingpro/portable-devtools

1from typing import Dict, Iterator, List, Optional, Tuple, Union
2
3from .. import AddedToken, Tokenizer, decoders, pre_tokenizers, trainers
4from ..models import BPE
5from ..normalizers import BertNormalizer, Lowercase, Sequence, unicode_normalizer_from_str
6from .base_tokenizer import BaseTokenizer
7
8
9class CharBPETokenizer(BaseTokenizer):
10    """Original BPE Tokenizer
11
12    Represents the BPE algorithm, as introduced by Rico Sennrich
13    (https://arxiv.org/abs/1508.07909)
14
15    The defaults settings corresponds to OpenAI GPT BPE tokenizers and differs from the original
16    Sennrich subword-nmt implementation by the following options that you can deactivate:
17        - adding a normalizer to clean up the text (deactivate with `bert_normalizer=False`) by:
18            * removing any control characters and replacing all whitespaces by the classic one.
19            * handle chinese chars by putting spaces around them.
20            * strip all accents.
21        - spitting on punctuation in addition to whitespaces (deactivate it with
22          `split_on_whitespace_only=True`)
23    """
24
25    def __init__(
26        self,
27        vocab: Optional[Union[str, Dict[str, int]]] = None,
28        merges: Optional[Union[str, List[Tuple[str, str]]]] = None,
29        unk_token: Union[str, AddedToken] = "<unk>",
30        suffix: str = "</w>",
31        dropout: Optional[float] = None,
32        lowercase: bool = False,
33        unicode_normalizer: Optional[str] = None,
34        bert_normalizer: bool = True,
35        split_on_whitespace_only: bool = False,
36    ):
37        if vocab is not None and merges is not None:
38            tokenizer = Tokenizer(
39                BPE(
40                    vocab,
41                    merges,
42                    dropout=dropout,
43                    unk_token=str(unk_token),
44                    end_of_word_suffix=suffix,
45                )
46            )
47        else:
48            tokenizer = Tokenizer(BPE(unk_token=str(unk_token), dropout=dropout, end_of_word_suffix=suffix))
49
50        if tokenizer.token_to_id(str(unk_token)) is not None:
51            tokenizer.add_special_tokens([str(unk_token)])
52
53        # Check for Unicode normalization first (before everything else)
54        normalizers = []
55
56        if unicode_normalizer:
57            normalizers += [unicode_normalizer_from_str(unicode_normalizer)]
58
59        if bert_normalizer:
60            normalizers += [BertNormalizer(lowercase=False)]
61
62        if lowercase:
63            normalizers += [Lowercase()]
64
65        # Create the normalizer structure
66        if len(normalizers) > 0:
67            if len(normalizers) > 1:
68                tokenizer.normalizer = Sequence(normalizers)
69            else:
70                tokenizer.normalizer = normalizers[0]
71
72        if split_on_whitespace_only:
73            tokenizer.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
74        else:
75            tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
76
77        tokenizer.decoder = decoders.BPEDecoder(suffix=suffix)
78
79        parameters = {
80            "model": "BPE",
81            "unk_token": unk_token,
82            "suffix": suffix,
83            "dropout": dropout,
84            "lowercase": lowercase,
85            "unicode_normalizer": unicode_normalizer,
86            "bert_normalizer": bert_normalizer,
87            "split_on_whitespace_only": split_on_whitespace_only,
88        }
89
90        super().__init__(tokenizer, parameters)
91
92    @staticmethod
93    def from_file(vocab_filename: str, merges_filename: str, **kwargs):
94        vocab, merges = BPE.read_file(vocab_filename, merges_filename)
95        return CharBPETokenizer(vocab, merges, **kwargs)
96
97    def train(
98        self,
99        files: Union[str, List[str]],
100        vocab_size: int = 30000,
101        min_frequency: int = 2,
102        special_tokens: List[Union[str, AddedToken]] = ["<unk>"],
103        limit_alphabet: int = 1000,
104        initial_alphabet: List[str] = [],
105        suffix: Optional[str] = "</w>",
106        show_progress: bool = True,
107    ):
108        """Train the model using the given files"""
109
110        trainer = trainers.BpeTrainer(
111            vocab_size=vocab_size,
112            min_frequency=min_frequency,
113            special_tokens=special_tokens,
114            limit_alphabet=limit_alphabet,
115            initial_alphabet=initial_alphabet,
116            end_of_word_suffix=suffix,
117            show_progress=show_progress,
118        )
119        if isinstance(files, str):
120            files = [files]
121        self._tokenizer.train(files, trainer=trainer)
122
123    def train_from_iterator(
124        self,
125        iterator: Union[Iterator[str], Iterator[Iterator[str]]],
126        vocab_size: int = 30000,
127        min_frequency: int = 2,
128        special_tokens: List[Union[str, AddedToken]] = ["<unk>"],
129        limit_alphabet: int = 1000,
130        initial_alphabet: List[str] = [],
131        suffix: Optional[str] = "</w>",
132        show_progress: bool = True,
133        length: Optional[int] = None,
134    ):
135        """Train the model using the given iterator"""
136
137        trainer = trainers.BpeTrainer(
138            vocab_size=vocab_size,
139            min_frequency=min_frequency,
140            special_tokens=special_tokens,
141            limit_alphabet=limit_alphabet,
142            initial_alphabet=initial_alphabet,
143            end_of_word_suffix=suffix,
144            show_progress=show_progress,
145        )
146        self._tokenizer.train_from_iterator(
147            iterator,
148            trainer=trainer,
149            length=length,
150        )
151