codekingpro/portable-devtools

1from typing import Dict, Iterator, List, Optional, Tuple, Union
2
3from tokenizers import AddedToken, Tokenizer, decoders, pre_tokenizers, processors, trainers
4from tokenizers.models import BPE
5from tokenizers.normalizers import Lowercase, Sequence, unicode_normalizer_from_str
6
7from .base_tokenizer import BaseTokenizer
8
9
10class ByteLevelBPETokenizer(BaseTokenizer):
11    """ByteLevelBPETokenizer
12
13    Represents a Byte-level BPE as introduced by OpenAI with their GPT-2 model
14    """
15
16    def __init__(
17        self,
18        vocab: Optional[Union[str, Dict[str, int]]] = None,
19        merges: Optional[Union[str, List[Tuple[str, str]]]] = None,
20        add_prefix_space: bool = False,
21        lowercase: bool = False,
22        dropout: Optional[float] = None,
23        unicode_normalizer: Optional[str] = None,
24        continuing_subword_prefix: Optional[str] = None,
25        end_of_word_suffix: Optional[str] = None,
26        trim_offsets: bool = False,
27    ):
28        if vocab is not None and merges is not None:
29            tokenizer = Tokenizer(
30                BPE(
31                    vocab,
32                    merges,
33                    dropout=dropout,
34                    continuing_subword_prefix=continuing_subword_prefix or "",
35                    end_of_word_suffix=end_of_word_suffix or "",
36                )
37            )
38        else:
39            tokenizer = Tokenizer(BPE())
40
41        # Check for Unicode normalization first (before everything else)
42        normalizers = []
43
44        if unicode_normalizer:
45            normalizers += [unicode_normalizer_from_str(unicode_normalizer)]
46
47        if lowercase:
48            normalizers += [Lowercase()]
49
50        # Create the normalizer structure
51        if len(normalizers) > 0:
52            if len(normalizers) > 1:
53                tokenizer.normalizer = Sequence(normalizers)
54            else:
55                tokenizer.normalizer = normalizers[0]
56
57        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space)
58        tokenizer.decoder = decoders.ByteLevel()
59        tokenizer.post_processor = processors.ByteLevel(trim_offsets=trim_offsets)
60
61        parameters = {
62            "model": "ByteLevelBPE",
63            "add_prefix_space": add_prefix_space,
64            "lowercase": lowercase,
65            "dropout": dropout,
66            "unicode_normalizer": unicode_normalizer,
67            "continuing_subword_prefix": continuing_subword_prefix,
68            "end_of_word_suffix": end_of_word_suffix,
69            "trim_offsets": trim_offsets,
70        }
71
72        super().__init__(tokenizer, parameters)
73
74    @staticmethod
75    def from_file(vocab_filename: str, merges_filename: str, **kwargs):
76        vocab, merges = BPE.read_file(vocab_filename, merges_filename)
77        return ByteLevelBPETokenizer(vocab, merges, **kwargs)
78
79    def train(
80        self,
81        files: Union[str, List[str]],
82        vocab_size: int = 30000,
83        min_frequency: int = 2,
84        show_progress: bool = True,
85        special_tokens: List[Union[str, AddedToken]] = [],
86    ):
87        """Train the model using the given files"""
88
89        trainer = trainers.BpeTrainer(
90            vocab_size=vocab_size,
91            min_frequency=min_frequency,
92            show_progress=show_progress,
93            special_tokens=special_tokens,
94            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
95        )
96        if isinstance(files, str):
97            files = [files]
98        self._tokenizer.train(files, trainer=trainer)
99
100    def train_from_iterator(
101        self,
102        iterator: Union[Iterator[str], Iterator[Iterator[str]]],
103        vocab_size: int = 30000,
104        min_frequency: int = 2,
105        show_progress: bool = True,
106        special_tokens: List[Union[str, AddedToken]] = [],
107        length: Optional[int] = None,
108    ):
109        """Train the model using the given iterator"""
110
111        trainer = trainers.BpeTrainer(
112            vocab_size=vocab_size,
113            min_frequency=min_frequency,
114            show_progress=show_progress,
115            special_tokens=special_tokens,
116            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
117        )
118        self._tokenizer.train_from_iterator(
119            iterator,
120            trainer=trainer,
121            length=length,
122        )
123