Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from pathlib import Path4from typing import Callable, TYPE_CHECKING5 6if TYPE_CHECKING:7    from torch import Tensor8 9from .base import MistralTokenizerType, MistralVocab, _mistral_common_installed, _mistral_import_error_msg, gguf, logger10 11from .deepseek import DeepseekV2Model12from .llama import LlamaModel13 14if _mistral_common_installed:15    from mistral_common.tokens.tokenizers.base import TokenizerVersion  # type: ignore[import-not-found, ty:unresolved-import]16    from mistral_common.tokens.tokenizers.tekken import Tekkenizer  # type: ignore[import-not-found, ty:unresolved-import]17    from mistral_common.tokens.tokenizers.sentencepiece import SentencePieceTokenizer  # type: ignore[import-not-found, ty:unresolved-import]18else:19    TokenizerVersion = None  # type: ignore[assignment]20    Tekkenizer = None  # type: ignore[assignment]21    SentencePieceTokenizer = None  # type: ignore[assignment]22 23 24class MistralModel(LlamaModel):25    model_arch = gguf.MODEL_ARCH.MISTRAL326    model_name = "Mistral"27    hf_arch = ""28    is_mistral_format = True29    undo_permute = False30 31    def __init__(self, *args, **kwargs):32        super().__init__(*args, **kwargs)33        # for compatibility, we use LLAMA arch for older models34        # TODO: remove this once everyone migrates to newer version of llama.cpp35        if "llama_4_scaling" not in self.hparams:36            self.model_arch = gguf.MODEL_ARCH.LLAMA37            self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[self.model_arch]38            self.gguf_writer.add_architecture()39            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)40 41    def dequant_model(self):42        # transform quantization config into HF format43        quant_config = self.hparams.get("quantization")44        if quant_config is not None:45            assert quant_config["qformat_weight"] == "fp8_e4m3"46            self.hparams["quantization_config"] = {47                "activation_scheme": "static",48                "quant_method": "fp8",49                "weight_block_size": None,50            }51        return super().dequant_model()52 53    @staticmethod54    def get_community_chat_template(vocab: MistralVocab, templates_dir: Path, is_mistral_format: bool):55        assert TokenizerVersion is not None and Tekkenizer is not None and SentencePieceTokenizer is not None, _mistral_import_error_msg56        assert isinstance(vocab.tokenizer, (Tekkenizer, SentencePieceTokenizer)), (57            f"Expected Tekkenizer or SentencePieceTokenizer, got {type(vocab.tokenizer)}"58        )59 60        if vocab.tokenizer.version == TokenizerVersion.v1:61            return "mistral-v1"62        elif vocab.tokenizer.version == TokenizerVersion.v3 and vocab.tokenizer_type == MistralTokenizerType.spm:63            return "mistral-v3"64        elif vocab.tokenizer.version == TokenizerVersion.v3 and vocab.tokenizer_type == MistralTokenizerType.tekken:65            return "mistral-v3-tekken"66        elif vocab.tokenizer.version == TokenizerVersion.v7 and vocab.tokenizer_type == MistralTokenizerType.spm:67            return "mistral-v7"68        elif vocab.tokenizer.version == TokenizerVersion.v7 and vocab.tokenizer_type == MistralTokenizerType.tekken:69            return "mistral-v7-tekken"70        elif vocab.tokenizer.version == TokenizerVersion.v11:71            template_file = "Mistral-Small-3.2-24B-Instruct-2506.jinja"72        elif vocab.tokenizer.version == TokenizerVersion.v13:73            template_file = "unsloth-mistral-Devstral-Small-2507.jinja"74        else:75            err_message = f"Unknown tokenizer type: {vocab.tokenizer_type} and version {vocab.tokenizer.version}"76            if is_mistral_format:77                err_message += (78                    " . Please pass --disable-mistral-community-chat-template argument to the CLI "79                    "if you want to skip this error and use the Mistral official `mistral-common` pre-processing library."80                )81            raise ValueError(err_message)82 83        template_path = templates_dir / template_file84        if not template_path.exists():85            raise FileNotFoundError(f"Template file not found: {template_path}")86 87        with open(template_path, "r", encoding="utf-8") as f:88            template = f.read()89 90        return template91 92    def set_gguf_parameters(self):93        super().set_gguf_parameters()94        MistralModel.set_mistral_config(self.gguf_writer, self.hparams)95 96    @staticmethod97    def set_mistral_config(gguf_writer: gguf.GGUFWriter, hparams: dict):98        if "yarn" in hparams:99            yarn_params = hparams["yarn"]100            mscale_all_dim = 1.0 if not yarn_params["apply_scale"] else 0.0101            gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN)102            gguf_writer.add_rope_scaling_factor(yarn_params["factor"])103            gguf_writer.add_rope_scaling_yarn_beta_fast(yarn_params["beta"])104            gguf_writer.add_rope_scaling_yarn_beta_slow(yarn_params["alpha"])105            gguf_writer.add_rope_scaling_yarn_log_mul(mscale_all_dim)106            gguf_writer.add_rope_scaling_orig_ctx_len(yarn_params["original_max_position_embeddings"])107 108        llama_4_scaling = hparams.get("llama_4_scaling")109        if llama_4_scaling is not None:110            gguf_writer.add_attn_temperature_scale(llama_4_scaling["beta"])111 112 113class MistralMoeModel(DeepseekV2Model):114    model_arch = gguf.MODEL_ARCH.DEEPSEEK2115    model_name = "Mistral"116    hf_arch = ""117    is_mistral_format = True118 119    def __init__(self, *args, **kwargs):120        super().__init__(*args, **kwargs)121        logger.info("Using MistralMoeModel")122        # remap hparams from Mistral MoE format to DeepseekV2 format123        # we do this way to be able to reuse DeepseekV2Model set_gguf_parameters logic124        # ref: https://github.com/vllm-project/vllm/blob/b294e28db2c5dee61bc25157664edcada8b90b31/vllm/transformers_utils/configs/mistral.py125        config = self.hparams126        # Mistral key -> HF key127        config_mapping = {128            "dim": "hidden_size",129            "norm_eps": "rms_norm_eps",130            "n_kv_heads": "num_key_value_heads",131            "n_layers": "num_hidden_layers",132            "n_heads": "num_attention_heads",133            "hidden_dim": "intermediate_size",134        }135        # HF key -> (Mistral key, default value)136        top_level_mapping_with_default = {137            "model_type": ("model_type", "transformer"),138            "hidden_act": ("activation", "silu"),139            "tie_word_embeddings": ("tied_embeddings", False),140            "max_seq_len": ("max_seq_len", config.get("max_position_embeddings", 128_000)),141            "max_position_embeddings": ("max_position_embeddings", 128_000),142        }143        # mapping top-level keys144        for key, new_key in config_mapping.items():145            if key in config:146                config[new_key] = config[key]147        for new_key, (key, default_value) in top_level_mapping_with_default.items():148            config[new_key] = config.get(key, default_value)149        # mapping MoE-specific keys150        moe_config_map = {151            "route_every_n": "moe_layer_freq",152            "first_k_dense_replace": "first_k_dense_replace",153            "num_experts_per_tok": "num_experts_per_tok",154            "num_experts": "n_routed_experts",155            "expert_hidden_dim": "moe_intermediate_size",156            "routed_scale": "routed_scaling_factor",157            "num_shared_experts": "n_shared_experts",158            "num_expert_groups": "n_group",159            "num_expert_groups_per_tok": "topk_group",160        }161        moe = config["moe"]162        for key, new_key in moe_config_map.items():163            if key in moe:164                config[new_key] = moe[key]165        # provide missing values166        config["topk_method"] = None167        config["norm_topk_prob"] = True168        config["scoring_func"] = "softmax"169 170    def set_vocab(self):171        self._set_vocab_mistral()172 173    def set_gguf_parameters(self):174        super().set_gguf_parameters()175        MistralModel.set_mistral_config(self.gguf_writer, self.hparams)176        yarn_params = self.hparams["yarn"]177        self.gguf_writer.add_attn_temperature_length(yarn_params["original_max_position_embeddings"])178 179        # [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]180        # note: for legacy reasons, this is not consistent with the other usages of self.gguf_writer.add_rope_scaling_yarn_log_mul181        # ref https://github.com/ggml-org/llama.cpp/pull/17945182        self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1) # mscale_all_dim * 0.1183 184    @classmethod185    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:186        name, gen = item187 188        # rename certain tensors so that we can reuse DeepseekV2Model modify_tensors logic189        if name.endswith(".qscale_act"):190            name = name.replace(".qscale_act", ".input_scale")191        if name.endswith(".qscale_weight"):192            name = name.replace(".qscale_weight", ".weight_scale")193        if ".wkv_b." in name:194            name = name.replace(".wkv_b.", ".kv_b_proj.")195        if ".experts." in name:196            name = name.replace(".experts.", ".mlp.experts.")197            name = name.replace(".w1.", ".gate_proj.")198            name = name.replace(".w2.", ".down_proj.")199            name = name.replace(".w3.", ".up_proj.")200            name = "model." + name201 202        return super().filter_tensors((name, gen))203