codekingpro/portable-devtools

1from typing import Dict, Iterator, List, Optional, Tuple, Union
2
3from tokenizers import AddedToken, Tokenizer, decoders, pre_tokenizers, trainers
4from tokenizers.models import BPE
5from tokenizers.normalizers import NFKC
6
7from .base_tokenizer import BaseTokenizer
8
9
10class SentencePieceBPETokenizer(BaseTokenizer):
11    """SentencePiece BPE Tokenizer
12
13    Represents the BPE algorithm, with the pretokenization used by SentencePiece
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        unk_token: Union[str, AddedToken] = "<unk>",
21        replacement: str = "▁",
22        add_prefix_space: bool = True,
23        dropout: Optional[float] = None,
24        fuse_unk: Optional[bool] = False,
25    ):
26        if vocab is not None and merges is not None:
27            tokenizer = Tokenizer(BPE(vocab, merges, dropout=dropout, unk_token=unk_token, fuse_unk=fuse_unk))
28        else:
29            tokenizer = Tokenizer(BPE(dropout=dropout, unk_token=unk_token, fuse_unk=fuse_unk))
30
31        if tokenizer.token_to_id(str(unk_token)) is not None:
32            tokenizer.add_special_tokens([str(unk_token)])
33
34        tokenizer.normalizer = NFKC()
35        prepend_scheme = "always" if add_prefix_space else "never"
36        tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)
37        tokenizer.decoder = decoders.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)
38
39        parameters = {
40            "model": "SentencePieceBPE",
41            "unk_token": unk_token,
42            "replacement": replacement,
43            "add_prefix_space": add_prefix_space,
44            "dropout": dropout,
45        }
46
47        super().__init__(tokenizer, parameters)
48
49    @staticmethod
50    def from_file(vocab_filename: str, merges_filename: str, **kwargs):
51        vocab, merges = BPE.read_file(vocab_filename, merges_filename)
52        return SentencePieceBPETokenizer(vocab, merges, **kwargs)
53
54    def train(
55        self,
56        files: Union[str, List[str]],
57        vocab_size: int = 30000,
58        min_frequency: int = 2,
59        special_tokens: List[Union[str, AddedToken]] = ["<unk>"],
60        limit_alphabet: int = 1000,
61        initial_alphabet: List[str] = [],
62        show_progress: bool = True,
63    ):
64        """Train the model using the given files"""
65
66        trainer = trainers.BpeTrainer(
67            vocab_size=vocab_size,
68            min_frequency=min_frequency,
69            special_tokens=special_tokens,
70            limit_alphabet=limit_alphabet,
71            initial_alphabet=initial_alphabet,
72            show_progress=show_progress,
73        )
74        if isinstance(files, str):
75            files = [files]
76        self._tokenizer.train(files, trainer=trainer)
77
78    def train_from_iterator(
79        self,
80        iterator: Union[Iterator[str], Iterator[Iterator[str]]],
81        vocab_size: int = 30000,
82        min_frequency: int = 2,
83        special_tokens: List[Union[str, AddedToken]] = ["<unk>"],
84        limit_alphabet: int = 1000,
85        initial_alphabet: List[str] = [],
86        show_progress: bool = True,
87        length: Optional[int] = None,
88    ):
89        """Train the model using the given iterator"""
90
91        trainer = trainers.BpeTrainer(
92            vocab_size=vocab_size,
93            min_frequency=min_frequency,
94            special_tokens=special_tokens,
95            limit_alphabet=limit_alphabet,
96            initial_alphabet=initial_alphabet,
97            show_progress=show_progress,
98        )
99        self._tokenizer.train_from_iterator(
100            iterator,
101            trainer=trainer,
102            length=length,
103        )
104