Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import json4import sys5 6from typing import Callable, Iterable, TYPE_CHECKING7 8if TYPE_CHECKING:9    from torch import Tensor10 11from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger12 13from .llama import LlamaModel14 15 16@ModelBase.register("InternLM2ForCausalLM")17@ModelBase.example("internlm/internlm2-chat-7b")18class InternLM2Model(TextModel):19    model_arch = gguf.MODEL_ARCH.INTERNLM220 21    def set_vocab(self):22        # (TODO): Is there a better way?23        # Copy from _set_vocab_sentencepiece, The only difference is that we will treat the character24        # \x00 specially and convert it into an emoji character to prevent it from being mistakenly25        # recognized as an empty string in C++.26        from sentencepiece import SentencePieceProcessor27        from sentencepiece import sentencepiece_model_pb2 as model28 29        tokenizer_path = self.dir_model / 'tokenizer.model'30 31        tokens: list[bytes] = []32        scores: list[float] = []33        toktypes: list[int] = []34 35        if not tokenizer_path.is_file():36            logger.error(f'Error: Missing {tokenizer_path}')37            sys.exit(1)38 39        sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]40        sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())41        add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix42 43        tokenizer = SentencePieceProcessor()44        tokenizer.LoadFromFile(str(tokenizer_path))45 46        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())47 48        for token_id in range(vocab_size):49            piece = tokenizer.IdToPiece(token_id)50            text = piece.encode("utf-8")51            score = tokenizer.GetScore(token_id)52            if text == b"\x00":53                # (TODO): fixme54                # Hack here and replace the \x00 characters.55                logger.warning(f"InternLM2 convert token '{text}' to '🐉'!")56                text = "🐉".encode("utf-8")57 58            toktype = SentencePieceTokenTypes.NORMAL59            if tokenizer.IsUnknown(token_id):60                toktype = SentencePieceTokenTypes.UNKNOWN61            elif tokenizer.IsControl(token_id):62                toktype = SentencePieceTokenTypes.CONTROL63            elif tokenizer.IsUnused(token_id):64                toktype = SentencePieceTokenTypes.UNUSED65            elif tokenizer.IsByte(token_id):66                toktype = SentencePieceTokenTypes.BYTE67            # take care of ununsed raw token68            if piece.startswith('[UNUSED'):69                toktype = SentencePieceTokenTypes.UNUSED70 71            tokens.append(text)72            scores.append(score)73            toktypes.append(toktype)74 75        added_tokens_file = self.dir_model / 'added_tokens.json'76        if added_tokens_file.is_file():77            with open(added_tokens_file, "r", encoding="utf-8") as f:78                added_tokens_json = json.load(f)79 80                for key in added_tokens_json:81                    tokens.append(key.encode("utf-8"))82                    scores.append(-1000.0)83                    toktypes.append(SentencePieceTokenTypes.USER_DEFINED)84 85        chat_eos_token = '<|im_end|>'86        chat_eos_token_id = None87 88        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'89        if tokenizer_config_file.is_file():90            with open(tokenizer_config_file, "r", encoding="utf-8") as f:91                tokenizer_config_json = json.load(f)92                added_tokens_decoder = tokenizer_config_json.get("added_tokens_decoder", {})93                for token_id, foken_data in added_tokens_decoder.items():94                    token_id = int(token_id)95                    token = foken_data["content"]96                    if token == chat_eos_token:97                        chat_eos_token_id = token_id98                    token = token.encode("utf-8")99                    if toktypes[token_id] != SentencePieceTokenTypes.UNUSED:100                        if tokens[token_id] != token:101                            logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}')102                    tokens[token_id] = token103                    scores[token_id] = -1000.0104                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED105                    if foken_data.get("special"):106                        toktypes[token_id] = SentencePieceTokenTypes.CONTROL107 108        tokenizer_file = self.dir_model / 'tokenizer.json'109        if tokenizer_file.is_file():110            with open(tokenizer_file, "r", encoding="utf-8") as f:111                tokenizer_json = json.load(f)112                added_tokens = tokenizer_json.get("added_tokens", [])113                for foken_data in added_tokens:114                    token_id = int(foken_data["id"])115                    token = foken_data["content"]116                    if token == chat_eos_token:117                        chat_eos_token_id = token_id118                    token = token.encode("utf-8")119                    if toktypes[token_id] != SentencePieceTokenTypes.UNUSED:120                        if tokens[token_id] != token:121                            logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}')122                    tokens[token_id] = token123                    scores[token_id] = -1000.0124                    toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED125                    if foken_data.get("special"):126                        toktypes[token_id] = SentencePieceTokenTypes.CONTROL127 128        self.gguf_writer.add_tokenizer_model("llama")129        self.gguf_writer.add_tokenizer_pre("default")130        self.gguf_writer.add_token_list(tokens)131        self.gguf_writer.add_token_scores(scores)132        self.gguf_writer.add_token_types(toktypes)133        self.gguf_writer.add_add_space_prefix(add_prefix)134 135        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))136        old_eos = special_vocab.special_token_ids["eos"]137        if chat_eos_token_id is not None:138            # For the chat model, we replace the eos with '<|im_end|>'.139            # TODO: this is a hack, should be fixed140            #       https://github.com/ggml-org/llama.cpp/pull/6745#issuecomment-2067687048141            special_vocab.special_token_ids["eos"] = chat_eos_token_id142            logger.warning(f"Replace eos:{old_eos} with a special token:{chat_eos_token_id}"143                           " in chat mode so that the conversation can end normally.")144 145        special_vocab.add_to_gguf(self.gguf_writer)146 147    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:148        num_heads = self.hparams["num_attention_heads"]149        num_kv_heads = self.hparams["num_key_value_heads"]150        n_embd = self.hparams["hidden_size"]151        q_per_kv = num_heads // num_kv_heads152        head_dim = n_embd // num_heads153        num_groups = num_heads // q_per_kv154 155        if bid is not None and f"model.layers.{bid}.attention.wqkv" in name:156            qkv = data_torch157 158            qkv = qkv.reshape((num_groups, q_per_kv + 2, head_dim, n_embd))159            q, k, v = qkv[:, : q_per_kv], qkv[:, -2], qkv[:, -1]160 161            # The model weights of q and k equire additional reshape.162            q = LlamaModel.permute(q.reshape((-1, q.shape[-1])), num_heads, num_heads)163            k = LlamaModel.permute(k.reshape((-1, k.shape[-1])), num_heads, num_kv_heads)164            v = v.reshape((-1, v.shape[-1]))165 166            yield from super().modify_tensors(q, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), bid)167            yield from super().modify_tensors(k, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), bid)168            yield from super().modify_tensors(v, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), bid)169        else:170            yield from super().modify_tensors(data_torch, name, bid)171 172 173@ModelBase.register("InternLM3ForCausalLM")174@ModelBase.example("internlm/internlm3-8b-instruct")175class InternLM3Model(TextModel):176    model_arch = gguf.MODEL_ARCH.LLAMA177 178    def set_vocab(self):179        tokens, scores, toktypes = self._create_vocab_sentencepiece()180 181        self.gguf_writer.add_tokenizer_model("llama")182        self.gguf_writer.add_tokenizer_pre("default")183        self.gguf_writer.add_token_list(tokens)184        self.gguf_writer.add_token_scores(scores)185        self.gguf_writer.add_token_types(toktypes)186 187        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))188 189        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'190        if tokenizer_config_file.is_file():191            with open(tokenizer_config_file, "r", encoding="utf-8") as f:192                tokenizer_config_json = json.load(f)193                if "add_prefix_space" in tokenizer_config_json:194                    self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"])195 196                if "added_tokens_decoder" in tokenizer_config_json:197                    for token_id, token_data in tokenizer_config_json["added_tokens_decoder"].items():198                        if token_data.get("special"):199                            token_id = int(token_id)200                            token = token_data["content"]201                            special_vocab._set_special_token(token, token_id)202                            # update eos token203                            if token == '<|im_end|>' and "eos" in special_vocab.special_token_ids:204                                special_vocab.special_token_ids["eos"] = token_id205 206        special_vocab.add_to_gguf(self.gguf_writer)207 208    def set_gguf_parameters(self):209        super().set_gguf_parameters()210        hparams = self.hparams211        self.gguf_writer.add_vocab_size(hparams["vocab_size"])212 213        if (rope_dim := hparams.get("head_dim")) is None:214            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]215        self.gguf_writer.add_rope_dimension_count(rope_dim)216 217    @classmethod218    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:219        name, gen = item220 221        if name.startswith(("mlp", "vision_model")):222            # skip visual tensors223            return None224 225        return super().filter_tensors(item)226 227    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:228        n_head = self.hparams["num_attention_heads"]229        n_kv_head = self.hparams.get("num_key_value_heads")230        if name.endswith(("q_proj.weight", "q_proj.bias")):231            data_torch = LlamaModel.permute(data_torch, n_head, n_head)232        if name.endswith(("k_proj.weight", "k_proj.bias")):233            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)234        yield from super().modify_tensors(data_torch, name, bid)235