Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4import os5 6from typing import Iterable, TYPE_CHECKING7 8if TYPE_CHECKING:9    from torch import Tensor10 11from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger12 13 14@ModelBase.register("T5WithLMHeadModel")15@ModelBase.register("T5ForConditionalGeneration")16@ModelBase.register("MT5ForConditionalGeneration")17@ModelBase.register("UMT5ForConditionalGeneration")18@ModelBase.register("UMT5Model")19class T5Model(TextModel):20    model_arch = gguf.MODEL_ARCH.T521 22    def __init__(self, *args, **kwargs):23        super().__init__(*args, **kwargs)24        self.shared_token_embeddings_found = False25 26    def set_vocab(self):27        # to avoid TypeError: Descriptors cannot be created directly28        # exception when importing sentencepiece_model_pb229        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"30        from sentencepiece import SentencePieceProcessor31        from sentencepiece import sentencepiece_model_pb2 as model32 33        tokenizer_path = self.dir_model / 'tokenizer.model'34 35        # many older models use spiece.model tokenizer model filename36        if not tokenizer_path.is_file():37            tokenizer_path = self.dir_model / 'spiece.model'38 39        if not tokenizer_path.is_file():40            raise FileNotFoundError(f"File not found: {tokenizer_path}")41 42        sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]43        sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())44 45        # some models like Pile-T5 family use BPE tokenizer instead of Unigram46        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE47            # assure the tokenizer model file name is correct48            assert tokenizer_path.name == 'tokenizer.model'49            return self._set_vocab_sentencepiece()50        else:51            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM52 53        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix54        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces55        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap56 57        tokenizer = SentencePieceProcessor()58        tokenizer.LoadFromFile(str(tokenizer_path))59 60        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())61 62        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]63        scores: list[float] = [-10000.0] * vocab_size64        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size65 66        for token_id in range(tokenizer.vocab_size()):67            piece = tokenizer.IdToPiece(token_id)68            text = piece.encode("utf-8")69            score = tokenizer.GetScore(token_id)70 71            toktype = SentencePieceTokenTypes.NORMAL72            if tokenizer.IsUnknown(token_id):73                toktype = SentencePieceTokenTypes.UNKNOWN74            elif tokenizer.IsControl(token_id):75                toktype = SentencePieceTokenTypes.CONTROL76            elif tokenizer.IsUnused(token_id):77                toktype = SentencePieceTokenTypes.UNUSED78            elif tokenizer.IsByte(token_id):79                toktype = SentencePieceTokenTypes.BYTE80 81            tokens[token_id] = text82            scores[token_id] = score83            toktypes[token_id] = toktype84 85        added_tokens_file = self.dir_model / 'added_tokens.json'86        if added_tokens_file.is_file():87            with open(added_tokens_file, "r", encoding="utf-8") as f:88                added_tokens_json = json.load(f)89                for key in added_tokens_json:90                    token_id = added_tokens_json[key]91                    if token_id >= vocab_size:92                        logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}')93                        continue94 95                    tokens[token_id] = key.encode("utf-8")96                    scores[token_id] = -1000.097                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED98 99        if vocab_size > len(tokens):100            pad_count = vocab_size - len(tokens)101            logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]")102            for i in range(1, pad_count + 1):103                tokens.append(bytes(f"[PAD{i}]", encoding="utf-8"))104                scores.append(-1000.0)105                toktypes.append(SentencePieceTokenTypes.UNUSED)106 107        self.gguf_writer.add_tokenizer_model("t5")108        self.gguf_writer.add_tokenizer_pre("default")109        self.gguf_writer.add_token_list(tokens)110        self.gguf_writer.add_token_scores(scores)111        self.gguf_writer.add_token_types(toktypes)112        self.gguf_writer.add_add_space_prefix(add_prefix)113        self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces)114        if precompiled_charsmap:115            self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap)116 117        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))118        special_vocab.add_to_gguf(self.gguf_writer)119 120    def set_gguf_parameters(self):121        if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None:122            logger.warning("Couldn't find context length in config.json, assuming default value of 512")123            n_ctx = 512124        self.gguf_writer.add_context_length(n_ctx)125        self.gguf_writer.add_embedding_length(self.hparams["d_model"])126        self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"])127        self.gguf_writer.add_block_count(self.block_count)128        if (dec_n_layer := self.hparams.get("num_decoder_layers")) is not None:129            self.gguf_writer.add_decoder_block_count(dec_n_layer)130        self.gguf_writer.add_head_count(self.hparams["num_heads"])131        self.gguf_writer.add_key_length(self.hparams["d_kv"])132        self.gguf_writer.add_value_length(self.hparams["d_kv"])133        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])134        self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"])135        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"])136        self.gguf_writer.add_decoder_start_token_id(self.hparams["decoder_start_token_id"])137        self.gguf_writer.add_file_type(self.ftype)138 139    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:140        # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight",141        # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored142        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder143        # and decoder and ignore the remaining ones.144        if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]:145            if not self.shared_token_embeddings_found:146                name = "shared.weight"147                self.shared_token_embeddings_found = True148            else:149                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")150                return151 152        yield from super().modify_tensors(data_torch, name, bid)153 154 155@ModelBase.register("T5EncoderModel")156class T5EncoderModel(TextModel):157    model_arch = gguf.MODEL_ARCH.T5ENCODER158 159    def __init__(self, *args, **kwargs):160        super().__init__(*args, **kwargs)161        self.shared_token_embeddings_found = False162 163    def set_vocab(self):164        # to avoid TypeError: Descriptors cannot be created directly165        # exception when importing sentencepiece_model_pb2166        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"167        from sentencepiece import SentencePieceProcessor168        from sentencepiece import sentencepiece_model_pb2 as model169 170        tokenizer_path = self.dir_model / 'tokenizer.model'171 172        # many older models use spiece.model tokenizer model filename173        if not tokenizer_path.is_file():174            tokenizer_path = self.dir_model / 'spiece.model'175 176        if not tokenizer_path.is_file():177            raise FileNotFoundError(f"File not found: {tokenizer_path}")178 179        sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]180        sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())181 182        # some models like Pile-T5 family use BPE tokenizer instead of Unigram183        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE184            # assure the tokenizer model file name is correct185            assert tokenizer_path.name == 'tokenizer.model'186            return self._set_vocab_sentencepiece()187        else:188            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM189 190        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix191        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces192        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap193 194        tokenizer = SentencePieceProcessor()195        tokenizer.LoadFromFile(str(tokenizer_path))196 197        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())198 199        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]200        scores: list[float] = [-10000.0] * vocab_size201        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size202 203        for token_id in range(tokenizer.vocab_size()):204            piece = tokenizer.IdToPiece(token_id)205            text = piece.encode("utf-8")206            score = tokenizer.GetScore(token_id)207 208            toktype = SentencePieceTokenTypes.NORMAL209            if tokenizer.IsUnknown(token_id):210                toktype = SentencePieceTokenTypes.UNKNOWN211            elif tokenizer.IsControl(token_id):212                toktype = SentencePieceTokenTypes.CONTROL213            elif tokenizer.IsUnused(token_id):214                toktype = SentencePieceTokenTypes.UNUSED215            elif tokenizer.IsByte(token_id):216                toktype = SentencePieceTokenTypes.BYTE217 218            tokens[token_id] = text219            scores[token_id] = score220            toktypes[token_id] = toktype221 222        added_tokens_file = self.dir_model / 'added_tokens.json'223        if added_tokens_file.is_file():224            with open(added_tokens_file, "r", encoding="utf-8") as f:225                added_tokens_json = json.load(f)226                for key in added_tokens_json:227                    token_id = added_tokens_json[key]228                    if token_id >= vocab_size:229                        logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}')230                        continue231 232                    tokens[token_id] = key.encode("utf-8")233                    scores[token_id] = -1000.0234                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED235 236        if vocab_size > len(tokens):237            pad_count = vocab_size - len(tokens)238            logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]")239            for i in range(1, pad_count + 1):240                tokens.append(bytes(f"[PAD{i}]", encoding="utf-8"))241                scores.append(-1000.0)242                toktypes.append(SentencePieceTokenTypes.UNUSED)243 244        self.gguf_writer.add_tokenizer_model("t5")245        self.gguf_writer.add_tokenizer_pre("default")246        self.gguf_writer.add_token_list(tokens)247        self.gguf_writer.add_token_scores(scores)248        self.gguf_writer.add_token_types(toktypes)249        self.gguf_writer.add_add_space_prefix(add_prefix)250        self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces)251        if precompiled_charsmap:252            self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap)253 254        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))255        special_vocab.add_to_gguf(self.gguf_writer)256 257    def set_gguf_parameters(self):258        if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None:259            logger.warning("Couldn't find context length in config.json, assuming default value of 512")260            n_ctx = 512261        self.gguf_writer.add_context_length(n_ctx)262        self.gguf_writer.add_embedding_length(self.hparams["d_model"])263        self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"])264        self.gguf_writer.add_block_count(self.block_count)265        self.gguf_writer.add_head_count(self.hparams["num_heads"])266        self.gguf_writer.add_key_length(self.hparams["d_kv"])267        self.gguf_writer.add_value_length(self.hparams["d_kv"])268        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])269        self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"])270        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"])271        self.gguf_writer.add_file_type(self.ftype)272 273    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:274        # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight",275        # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored276        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder277        # and decoder and ignore the remaining ones.278        if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]:279            if not self.shared_token_embeddings_found:280                name = "shared.weight"281                self.shared_token_embeddings_found = True282            else:283                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")284                return285 286        yield from super().modify_tensors(data_torch, name, bid)287