Felipe97/llama-cpp-compiled

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")19@ModelBase.example("google-t5/t5-small", "google/flan-t5-small", "google/umt5-small")20class T5Model(TextModel):21    model_arch = gguf.MODEL_ARCH.T522 23    def __init__(self, *args, **kwargs):24        super().__init__(*args, **kwargs)25        self.shared_token_embeddings_found = False26 27    def set_vocab(self):28        # to avoid TypeError: Descriptors cannot be created directly29        # exception when importing sentencepiece_model_pb230        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"31        from sentencepiece import SentencePieceProcessor32        from sentencepiece import sentencepiece_model_pb2 as model33 34        tokenizer_path = self.dir_model / 'tokenizer.model'35 36        # many older models use spiece.model tokenizer model filename37        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 Unigram47        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE48            # assure the tokenizer model file name is correct49            assert tokenizer_path.name == 'tokenizer.model'50            return self._set_vocab_sentencepiece()51        else:52            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM53 54        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix55        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces56        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap57 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_size65        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size66 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.NORMAL73            if tokenizer.IsUnknown(token_id):74                toktype = SentencePieceTokenTypes.UNKNOWN75            elif tokenizer.IsControl(token_id):76                toktype = SentencePieceTokenTypes.CONTROL77            elif tokenizer.IsUnused(token_id):78                toktype = SentencePieceTokenTypes.UNUSED79            elif tokenizer.IsByte(token_id):80                toktype = SentencePieceTokenTypes.BYTE81 82            tokens[token_id] = text83            scores[token_id] = score84            toktypes[token_id] = toktype85 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                        continue95 96                    tokens[token_id] = key.encode("utf-8")97                    scores[token_id] = -1000.098                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED99 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 = 512125        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 stored143        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder144        # 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 = True149            else:150                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")151                return152 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.T5ENCODER160 161    def __init__(self, *args, **kwargs):162        super().__init__(*args, **kwargs)163        self.shared_token_embeddings_found = False164 165    def set_vocab(self):166        # to avoid TypeError: Descriptors cannot be created directly167        # exception when importing sentencepiece_model_pb2168        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"169        from sentencepiece import SentencePieceProcessor170        from sentencepiece import sentencepiece_model_pb2 as model171 172        tokenizer_path = self.dir_model / 'tokenizer.model'173 174        # many older models use spiece.model tokenizer model filename175        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 Unigram185        if sentencepiece_model.trainer_spec.model_type == 2:  # BPE186            # assure the tokenizer model file name is correct187            assert tokenizer_path.name == 'tokenizer.model'188            return self._set_vocab_sentencepiece()189        else:190            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM191 192        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix193        remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces194        precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap195 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_size203        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size204 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.NORMAL211            if tokenizer.IsUnknown(token_id):212                toktype = SentencePieceTokenTypes.UNKNOWN213            elif tokenizer.IsControl(token_id):214                toktype = SentencePieceTokenTypes.CONTROL215            elif tokenizer.IsUnused(token_id):216                toktype = SentencePieceTokenTypes.UNUSED217            elif tokenizer.IsByte(token_id):218                toktype = SentencePieceTokenTypes.BYTE219 220            tokens[token_id] = text221            scores[token_id] = score222            toktypes[token_id] = toktype223 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                        continue233 234                    tokens[token_id] = key.encode("utf-8")235                    scores[token_id] = -1000.0236                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED237 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 = 512263        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 stored278        # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder279        # 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 = True284            else:285                logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.")286                return287 288        yield from super().modify_tensors(data_torch, name, bid)289