Brunobkr/llama.cpp_AlgMor24_github

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")12class ChatGLMModel(TextModel):13    model_arch = gguf.MODEL_ARCH.CHATGLM14 15    def set_vocab_chatglm3(self):16        dir_model = self.dir_model17        hparams = self.hparams18        tokens: list[bytes] = []19        toktypes: list[int] = []20        scores: list[float] = []21 22        from transformers import AutoTokenizer23        tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True)24        vocab_size = hparams.get("padded_vocab_size", len(tokenizer.get_vocab()))  # ty: ignore[unresolved-attribute]25        assert max(tokenizer.get_vocab().values()) < vocab_size  # ty: ignore[unresolved-attribute]26        role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"]27        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens28        for token_id in range(vocab_size):29            piece = tokenizer._convert_id_to_token(token_id)  # ty: ignore[unresolved-attribute]30            if token_id == 0:31                piece = "<unk>"32            elif token_id == 1:33                piece = "<bos>"34            elif token_id == 2:35                piece = "<eos>"36 37            text = piece.encode("utf-8")  # ty: ignore[unresolved-attribute]38            score = 0.039            # Referencing the tokenizer Python implementation(https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py),40            # it is only valid if it is less than tokenizer.tokenizer.sp_model.vocab_size()41            if len(piece) != 0 and token_id < tokenizer.tokenizer.sp_model.vocab_size():  # ty: ignore[unresolved-attribute, invalid-argument-type]42                score = tokenizer.tokenizer.sp_model.get_score(token_id)  # ty: ignore[unresolved-attribute]43 44            if token_id >= tokenizer.tokenizer.sp_model.vocab_size():  # ty: ignore[unresolved-attribute]45                if piece in special_tokens:46                    toktype = SentencePieceTokenTypes.CONTROL47                elif len(piece) == 0:  # ty: ignore[invalid-argument-type]48                    text = f"[PAD{token_id}]".encode("utf-8")49                    toktype = SentencePieceTokenTypes.UNUSED50                else:51                    toktype = SentencePieceTokenTypes.USER_DEFINED52                tokens.append(text)53                scores.append(score)54                toktypes.append(toktype)55                continue56 57            toktype = SentencePieceTokenTypes.NORMAL58            if tokenizer.tokenizer.sp_model.is_unknown(token_id):  # ty: ignore[unresolved-attribute]59                toktype = SentencePieceTokenTypes.UNKNOWN60            elif tokenizer.tokenizer.sp_model.is_control(token_id):  # ty: ignore[unresolved-attribute]61                toktype = SentencePieceTokenTypes.CONTROL62            elif tokenizer.tokenizer.sp_model.is_unused(token_id):  # ty: ignore[unresolved-attribute]63                toktype = SentencePieceTokenTypes.UNUSED64            elif tokenizer.tokenizer.sp_model.is_byte(token_id):  # ty: ignore[unresolved-attribute]65                toktype = SentencePieceTokenTypes.BYTE66 67            tokens.append(text)68            scores.append(score)69            toktypes.append(toktype)70 71        self.gguf_writer.add_tokenizer_model("llama")72        # glm3 needs prefix and suffix formatted as:73        # prompt = "[gMASK]sop<|user|>\n" + prompt + "<|assistant|>"74        self.gguf_writer.add_tokenizer_pre("chatglm-spm")75        self.gguf_writer.add_token_list(tokens)76        self.gguf_writer.add_token_scores(scores)77        self.gguf_writer.add_token_types(toktypes)78 79        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))80        special_vocab.add_to_gguf(self.gguf_writer)81 82    @staticmethod83    def token_bytes_to_string(b):84        from transformers.convert_slow_tokenizer import bytes_to_unicode85        byte_encoder = bytes_to_unicode()86        return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')])87 88    @staticmethod89    def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]:90        parts = [bytes([b]) for b in token]91        while True:92            min_idx = None93            min_rank = None94            for i, pair in enumerate(zip(parts[:-1], parts[1:])):95                rank = mergeable_ranks.get(pair[0] + pair[1])96                if rank is not None and (min_rank is None or rank < min_rank):97                    min_idx = i98                    min_rank = rank99            if min_rank is None or (max_rank is not None and min_rank >= max_rank):100                break101            assert min_idx is not None102            parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:]103        return parts104 105    def set_vocab(self):106        if "THUDM/chatglm3-6b" in self.hparams.get("_name_or_path", ""):107            self.set_vocab_chatglm3()108            return109 110        dir_model = self.dir_model111        hparams = self.hparams112        tokens: list[str] = []113        toktypes: list[int] = []114 115        from transformers import AutoTokenizer116        tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True)117        vocab_size = hparams.get("padded_vocab_size",hparams["vocab_size"])118        assert max(tokenizer.get_vocab().values()) < vocab_size  # ty: ignore[unresolved-attribute]119 120        tokens, toktypes, tokpre = self.get_vocab_base()121        self.gguf_writer.add_tokenizer_model("gpt2")122        self.gguf_writer.add_tokenizer_pre(tokpre)123        self.gguf_writer.add_token_list(tokens)124        self.gguf_writer.add_token_types(toktypes)125        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)126        # only add special tokens when they were not already loaded from config.json127        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]128        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"])  # ty: ignore[unresolved-attribute]129        # this one is usually not in config.json anyway130        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]131        special_vocab.add_to_gguf(self.gguf_writer)132 133    def set_gguf_parameters(self):134        n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed"))135        assert n_embed is not None136        n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads"))137        assert n_head is not None138        n_head_kv = self.hparams.get("multi_query_group_num", self.hparams.get("num_key_value_heads", n_head))139        self.gguf_writer.add_context_length(self.hparams.get("seq_length", n_embed))140        self.gguf_writer.add_embedding_length(n_embed)141        self.gguf_writer.add_feed_forward_length(self.hparams.get("ffn_hidden_size", self.hparams.get("intermediate_size", 4 * n_embed)))142        self.gguf_writer.add_block_count(self.block_count)143        self.gguf_writer.add_head_count(n_head)144        self.gguf_writer.add_head_count_kv(n_head_kv)145        self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon",1e-5))146        self.gguf_writer.add_file_type(self.ftype)147        if "attention_dim" in self.hparams:148            rope_dim = self.hparams["attention_dim"]149        else:150            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]151        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.5)))152        self.gguf_writer.add_add_bos_token(False)153        rope_freq = 10000154        if "rope_ratio" in self.hparams:155            rope_freq = rope_freq * self.hparams["rope_ratio"]156        self.gguf_writer.add_rope_freq_base(rope_freq)157 158    @classmethod159    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:160        name, gen = item161 162        if name.endswith(".rotary_pos_emb.inv_freq"):163            return None164 165        name = name.removeprefix("transformer.")166 167        return super().filter_tensors((name, gen))168