Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Callable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf9 10 11@ModelBase.register("GlmForCausalLM", "ChatGLMModel", "ChatGLMForConditionalGeneration")12@ModelBase.example("THUDM/chatglm3-6b", "zai-org/glm-4-9b-chat-hf")13class ChatGLMModel(TextModel):14    model_arch = gguf.MODEL_ARCH.CHATGLM15 16    def set_vocab_chatglm3(self):17        dir_model = self.dir_model18        hparams = self.hparams19        tokens: list[bytes] = []20        toktypes: list[int] = []21        scores: list[float] = []22 23        from transformers import AutoTokenizer24        tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True)25        vocab_size = hparams.get("padded_vocab_size", len(tokenizer.get_vocab()))  # ty: ignore[unresolved-attribute]26        assert max(tokenizer.get_vocab().values()) < vocab_size  # ty: ignore[unresolved-attribute]27        role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"]28        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens29        for token_id in range(vocab_size):30            piece = tokenizer._convert_id_to_token(token_id)  # ty: ignore[unresolved-attribute]31            if token_id == 0:32                piece = "<unk>"33            elif token_id == 1:34                piece = "<bos>"35            elif token_id == 2:36                piece = "<eos>"37 38            text = piece.encode("utf-8")  # ty: ignore[unresolved-attribute]39            score = 0.040            # Referencing the tokenizer Python implementation(https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py),41            # it is only valid if it is less than tokenizer.tokenizer.sp_model.vocab_size()42            if len(piece) != 0 and token_id < tokenizer.tokenizer.sp_model.vocab_size():  # ty: ignore[unresolved-attribute, invalid-argument-type]43                score = tokenizer.tokenizer.sp_model.get_score(token_id)  # ty: ignore[unresolved-attribute]44 45            if token_id >= tokenizer.tokenizer.sp_model.vocab_size():  # ty: ignore[unresolved-attribute]46                if piece in special_tokens:47                    toktype = SentencePieceTokenTypes.CONTROL48                elif len(piece) == 0:  # ty: ignore[invalid-argument-type]49                    text = f"[PAD{token_id}]".encode("utf-8")50                    toktype = SentencePieceTokenTypes.UNUSED51                else:52                    toktype = SentencePieceTokenTypes.USER_DEFINED53                tokens.append(text)54                scores.append(score)55                toktypes.append(toktype)56                continue57 58            toktype = SentencePieceTokenTypes.NORMAL59            if tokenizer.tokenizer.sp_model.is_unknown(token_id):  # ty: ignore[unresolved-attribute]60                toktype = SentencePieceTokenTypes.UNKNOWN61            elif tokenizer.tokenizer.sp_model.is_control(token_id):  # ty: ignore[unresolved-attribute]62                toktype = SentencePieceTokenTypes.CONTROL63            elif tokenizer.tokenizer.sp_model.is_unused(token_id):  # ty: ignore[unresolved-attribute]64                toktype = SentencePieceTokenTypes.UNUSED65            elif tokenizer.tokenizer.sp_model.is_byte(token_id):  # ty: ignore[unresolved-attribute]66                toktype = SentencePieceTokenTypes.BYTE67 68            tokens.append(text)69            scores.append(score)70            toktypes.append(toktype)71 72        self.gguf_writer.add_tokenizer_model("llama")73        # glm3 needs prefix and suffix formatted as:74        # prompt = "[gMASK]sop<|user|>\n" + prompt + "<|assistant|>"75        self.gguf_writer.add_tokenizer_pre("chatglm-spm")76        self.gguf_writer.add_token_list(tokens)77        self.gguf_writer.add_token_scores(scores)78        self.gguf_writer.add_token_types(toktypes)79 80        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))81        special_vocab.add_to_gguf(self.gguf_writer)82 83    @staticmethod84    def token_bytes_to_string(b):85        from transformers.convert_slow_tokenizer import bytes_to_unicode86        byte_encoder = bytes_to_unicode()87        return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')])88 89    @staticmethod90    def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]:91        parts = [bytes([b]) for b in token]92        while True:93            min_idx = None94            min_rank = None95            for i, pair in enumerate(zip(parts[:-1], parts[1:])):96                rank = mergeable_ranks.get(pair[0] + pair[1])97                if rank is not None and (min_rank is None or rank < min_rank):98                    min_idx = i99                    min_rank = rank100            if min_rank is None or (max_rank is not None and min_rank >= max_rank):101                break102            assert min_idx is not None103            parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:]104        return parts105 106    def set_vocab(self):107        if "THUDM/chatglm3-6b" in self.hparams.get("_name_or_path", ""):108            self.set_vocab_chatglm3()109            return110 111        dir_model = self.dir_model112        hparams = self.hparams113        tokens: list[str] = []114        toktypes: list[int] = []115 116        from transformers import AutoTokenizer117        tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True)118        vocab_size = hparams.get("padded_vocab_size",hparams["vocab_size"])119        assert max(tokenizer.get_vocab().values()) < vocab_size  # ty: ignore[unresolved-attribute]120 121        tokens, toktypes, tokpre = self.get_vocab_base()122        self.gguf_writer.add_tokenizer_model("gpt2")123        self.gguf_writer.add_tokenizer_pre(tokpre)124        self.gguf_writer.add_token_list(tokens)125        self.gguf_writer.add_token_types(toktypes)126        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)127        # only add special tokens when they were not already loaded from config.json128        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]129        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"])  # ty: ignore[unresolved-attribute]130        # this one is usually not in config.json anyway131        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]132        special_vocab.add_to_gguf(self.gguf_writer)133 134    def set_gguf_parameters(self):135        n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed"))136        assert n_embed is not None137        n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads"))138        assert n_head is not None139        n_head_kv = self.hparams.get("multi_query_group_num", self.hparams.get("num_key_value_heads", n_head))140        self.gguf_writer.add_context_length(self.hparams.get("seq_length", n_embed))141        self.gguf_writer.add_embedding_length(n_embed)142        self.gguf_writer.add_feed_forward_length(self.hparams.get("ffn_hidden_size", self.hparams.get("intermediate_size", 4 * n_embed)))143        self.gguf_writer.add_block_count(self.block_count)144        self.gguf_writer.add_head_count(n_head)145        self.gguf_writer.add_head_count_kv(n_head_kv)146        self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon",1e-5))147        self.gguf_writer.add_file_type(self.ftype)148        if "attention_dim" in self.hparams:149            rope_dim = self.hparams["attention_dim"]150        else:151            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]152        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.5)))153        self.gguf_writer.add_add_bos_token(False)154        rope_freq = 10000155        if "rope_ratio" in self.hparams:156            rope_freq = rope_freq * self.hparams["rope_ratio"]157        self.gguf_writer.add_rope_freq_base(rope_freq)158 159    @classmethod160    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:161        name, gen = item162 163        if name.endswith(".rotary_pos_emb.inv_freq"):164            return None165 166        name = name.removeprefix("transformer.")167 168        return super().filter_tensors((name, gen))169