char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
13.9 KiB288 linesraw
1from __future__ import annotations
2
3import json
4import os
5
6from typing import Iterable, TYPE_CHECKING
7
8if TYPE_CHECKING:
9    from torch import Tensor
10
11from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
12
13
14@ModelBase.register("T5WithLMHeadModel")
15@ModelBase.register("T5ForConditionalGeneration")
16@ModelBase.register("MT5ForConditionalGeneration")
17@ModelBase.register("UMT5ForConditionalGeneration")
18@ModelBase.register("UMT5Model")
19@ModelBase.example("google-t5/t5-small", "google/flan-t5-small", "google/umt5-small")
20class T5Model(TextModel):
21    model_arch = gguf.MODEL_ARCH.T5
22
23    def __init__(self, *args, **kwargs):
24        super().__init__(*args, **kwargs)
25        self.shared_token_embeddings_found = False
26
27    def set_vocab(self):
28        # to avoid TypeError: Descriptors cannot be created directly
29        # exception when importing sentencepiece_model_pb2
30        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
31        from sentencepiece import SentencePieceProcessor
32        from sentencepiece import sentencepiece_model_pb2 as model
33
34        tokenizer_path = self.dir_model / 'tokenizer.model'
35
36        # many older models use spiece.model tokenizer model filename
37        if not tokenizer_path.is_file():
38            tokenizer_path = self.dir_model / 'spiece.model'
39
40        if not tokenizer_path.is_file():
41            raise FileNotFoundError(f"File not found: {tokenizer_path}")
42
43        sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
44        sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())
45
46        # some models like Pile-T5 family use BPE tokenizer instead of Unigram
47        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE
48            # assure the tokenizer model file name is correct
49            assert tokenizer_path.name == 'tokenizer.model'
50            return self._set_vocab_sentencepiece()
51        else:
52            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM
53
54        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix
55        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces
56        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap
57
58        tokenizer = SentencePieceProcessor()
59        tokenizer.LoadFromFile(str(tokenizer_path))
60
61        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())
62
63        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]
64        scores: list[float] = [-10000.0] * vocab_size
65        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size
66
67        for token_id in range(tokenizer.vocab_size()):
68            piece = tokenizer.IdToPiece(token_id)
69            text = piece.encode("utf-8")
70            score = tokenizer.GetScore(token_id)
71
72            toktype = SentencePieceTokenTypes.NORMAL
73            if tokenizer.IsUnknown(token_id):
74                toktype = SentencePieceTokenTypes.UNKNOWN
75            elif tokenizer.IsControl(token_id):
76                toktype = SentencePieceTokenTypes.CONTROL
77            elif tokenizer.IsUnused(token_id):
78                toktype = SentencePieceTokenTypes.UNUSED
79            elif tokenizer.IsByte(token_id):
80                toktype = SentencePieceTokenTypes.BYTE
81
82            tokens[token_id] = text
83            scores[token_id] = score
84            toktypes[token_id] = toktype
85
86        added_tokens_file = self.dir_model / 'added_tokens.json'
87        if added_tokens_file.is_file():
88            with open(added_tokens_file, "r", encoding="utf-8") as f:
89                added_tokens_json = json.load(f)
90                for key in added_tokens_json:
91                    token_id = added_tokens_json[key]
92                    if token_id >= vocab_size:
93                        logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}')
94                        continue
95
96                    tokens[token_id] = key.encode("utf-8")
97                    scores[token_id] = -1000.0
98                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED
99
100        if vocab_size > len(tokens):
101            pad_count = vocab_size - len(tokens)
102            logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]")
103            for i in range(1, pad_count + 1):
104                tokens.append(bytes(f"[PAD{i}]", encoding="utf-8"))
105                scores.append(-1000.0)
106                toktypes.append(SentencePieceTokenTypes.UNUSED)
107
108        self.gguf_writer.add_tokenizer_model("t5")
109        self.gguf_writer.add_tokenizer_pre("default")
110        self.gguf_writer.add_token_list(tokens)
111        self.gguf_writer.add_token_scores(scores)
112        self.gguf_writer.add_token_types(toktypes)
113        self.gguf_writer.add_add_space_prefix(add_prefix)
114        self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces)
115        if precompiled_charsmap:
116            self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap)
117
118        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))
119        special_vocab.add_to_gguf(self.gguf_writer)
120
121    def set_gguf_parameters(self):
122        if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None:
123            logger.warning("Couldn't find context length in config.json, assuming default value of 512")
124            n_ctx = 512
125        self.gguf_writer.add_context_length(n_ctx)
126        self.gguf_writer.add_embedding_length(self.hparams["d_model"])
127        self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"])
128        self.gguf_writer.add_block_count(self.block_count)
129        if (dec_n_layer := self.hparams.get("num_decoder_layers")) is not None:
130            self.gguf_writer.add_decoder_block_count(dec_n_layer)
131        self.gguf_writer.add_head_count(self.hparams["num_heads"])
132        self.gguf_writer.add_key_length(self.hparams["d_kv"])
133        self.gguf_writer.add_value_length(self.hparams["d_kv"])
134        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])
135        self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"])
136        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"])
137        self.gguf_writer.add_decoder_start_token_id(self.hparams["decoder_start_token_id"])
138        self.gguf_writer.add_file_type(self.ftype)
139
140    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
141        # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight",
142        # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored
143        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder
144        # and decoder and ignore the remaining ones.
145        if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]:
146            if not self.shared_token_embeddings_found:
147                name = "shared.weight"
148                self.shared_token_embeddings_found = True
149            else:
150                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")
151                return
152
153        yield from super().modify_tensors(data_torch, name, bid)
154
155
156@ModelBase.register("T5EncoderModel")
157@ModelBase.example("sentence-transformers/sentence-t5-base")
158class T5EncoderModel(TextModel):
159    model_arch = gguf.MODEL_ARCH.T5ENCODER
160
161    def __init__(self, *args, **kwargs):
162        super().__init__(*args, **kwargs)
163        self.shared_token_embeddings_found = False
164
165    def set_vocab(self):
166        # to avoid TypeError: Descriptors cannot be created directly
167        # exception when importing sentencepiece_model_pb2
168        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
169        from sentencepiece import SentencePieceProcessor
170        from sentencepiece import sentencepiece_model_pb2 as model
171
172        tokenizer_path = self.dir_model / 'tokenizer.model'
173
174        # many older models use spiece.model tokenizer model filename
175        if not tokenizer_path.is_file():
176            tokenizer_path = self.dir_model / 'spiece.model'
177
178        if not tokenizer_path.is_file():
179            raise FileNotFoundError(f"File not found: {tokenizer_path}")
180
181        sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
182        sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())
183
184        # some models like Pile-T5 family use BPE tokenizer instead of Unigram
185        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE
186            # assure the tokenizer model file name is correct
187            assert tokenizer_path.name == 'tokenizer.model'
188            return self._set_vocab_sentencepiece()
189        else:
190            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM
191
192        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix
193        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces
194        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap
195
196        tokenizer = SentencePieceProcessor()
197        tokenizer.LoadFromFile(str(tokenizer_path))
198
199        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())
200
201        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]
202        scores: list[float] = [-10000.0] * vocab_size
203        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size
204
205        for token_id in range(tokenizer.vocab_size()):
206            piece = tokenizer.IdToPiece(token_id)
207            text = piece.encode("utf-8")
208            score = tokenizer.GetScore(token_id)
209
210            toktype = SentencePieceTokenTypes.NORMAL
211            if tokenizer.IsUnknown(token_id):
212                toktype = SentencePieceTokenTypes.UNKNOWN
213            elif tokenizer.IsControl(token_id):
214                toktype = SentencePieceTokenTypes.CONTROL
215            elif tokenizer.IsUnused(token_id):
216                toktype = SentencePieceTokenTypes.UNUSED
217            elif tokenizer.IsByte(token_id):
218                toktype = SentencePieceTokenTypes.BYTE
219
220            tokens[token_id] = text
221            scores[token_id] = score
222            toktypes[token_id] = toktype
223
224        added_tokens_file = self.dir_model / 'added_tokens.json'
225        if added_tokens_file.is_file():
226            with open(added_tokens_file, "r", encoding="utf-8") as f:
227                added_tokens_json = json.load(f)
228                for key in added_tokens_json:
229                    token_id = added_tokens_json[key]
230                    if token_id >= vocab_size:
231                        logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}')
232                        continue
233
234                    tokens[token_id] = key.encode("utf-8")
235                    scores[token_id] = -1000.0
236                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED
237
238        if vocab_size > len(tokens):
239            pad_count = vocab_size - len(tokens)
240            logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]")
241            for i in range(1, pad_count + 1):
242                tokens.append(bytes(f"[PAD{i}]", encoding="utf-8"))
243                scores.append(-1000.0)
244                toktypes.append(SentencePieceTokenTypes.UNUSED)
245
246        self.gguf_writer.add_tokenizer_model("t5")
247        self.gguf_writer.add_tokenizer_pre("default")
248        self.gguf_writer.add_token_list(tokens)
249        self.gguf_writer.add_token_scores(scores)
250        self.gguf_writer.add_token_types(toktypes)
251        self.gguf_writer.add_add_space_prefix(add_prefix)
252        self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces)
253        if precompiled_charsmap:
254            self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap)
255
256        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))
257        special_vocab.add_to_gguf(self.gguf_writer)
258
259    def set_gguf_parameters(self):
260        if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None:
261            logger.warning("Couldn't find context length in config.json, assuming default value of 512")
262            n_ctx = 512
263        self.gguf_writer.add_context_length(n_ctx)
264        self.gguf_writer.add_embedding_length(self.hparams["d_model"])
265        self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"])
266        self.gguf_writer.add_block_count(self.block_count)
267        self.gguf_writer.add_head_count(self.hparams["num_heads"])
268        self.gguf_writer.add_key_length(self.hparams["d_kv"])
269        self.gguf_writer.add_value_length(self.hparams["d_kv"])
270        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])
271        self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"])
272        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"])
273        self.gguf_writer.add_file_type(self.ftype)
274
275    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
276        # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight",
277        # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored
278        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder
279        # and decoder and ignore the remaining ones.
280        if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]:
281            if not self.shared_token_embeddings_found:
282                name = "shared.weight"
283                self.shared_token_embeddings_found = True
284            else:
285                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")
286                return
287
288        yield from super().modify_tensors(data_torch, name, bid)