Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import re4from typing import Iterable, TYPE_CHECKING5 6import torch7 8if TYPE_CHECKING:9    from torch import Tensor10 11from .base import ModelBase, TextModel, gguf, logger12 13 14@ModelBase.register("CohereForCausalLM")15# [TAG_HF_EXAMPLE_GATED] CohereLabs/c4ai-command-r-v01 is gated16# [TAG_HF_EXAMPLE_MISSING]17class CommandR2Model(TextModel):18    model_arch = gguf.MODEL_ARCH.COMMAND_R19 20    def __init__(self, *args, **kwargs):21        super().__init__(*args, **kwargs)22 23        # max_position_embeddings = 8192 in config.json but model was actually24        # trained on 128k context length25        # aya-23 models don't have model_max_length specified26        self.hparams["max_position_embeddings"] = self.find_hparam(["model_max_length", "max_position_embeddings"])27 28    def set_gguf_parameters(self):29        super().set_gguf_parameters()30        self.gguf_writer.add_logit_scale(self.hparams["logit_scale"])31        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)32 33 34@ModelBase.register("Cohere2ForCausalLM")35# [TAG_HF_EXAMPLE_GATED] CohereLabs/c4ai-command-r7b-12-2024 is gated36@ModelBase.example("hf-tiny-v2/tiny-random-Cohere2ForCausalLM")37class Cohere2Model(TextModel):38    model_arch = gguf.MODEL_ARCH.COHERE239 40    def set_gguf_parameters(self):41        super().set_gguf_parameters()42 43        self.gguf_writer.add_logit_scale(self.hparams["logit_scale"])44        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])45        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])46 47        rotary_pct = self.hparams["rotary_pct"]48        hidden_size = self.hparams["hidden_size"]49        num_attention_heads = self.hparams["num_attention_heads"]50        self.gguf_writer.add_rope_dimension_count(int(rotary_pct * (hidden_size // num_attention_heads)))51        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)52 53    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:54        # Cohere2 runtime in llama.cpp expects no bias tensors;55        # the actual weight only contains 0-value tensors as bias, we can skip them56        if name.endswith(".bias"):57            if torch.any(data_torch != 0):58                raise ValueError(f"Bias tensor {name!r} is not zero.")59            logger.debug(f"Skipping bias tensor {name!r} for Cohere2 conversion.")60            return61 62        yield from super().modify_tensors(data_torch, name, bid)63 64 65@ModelBase.register("Cohere2MoeForCausalLM")66@ModelBase.example("CohereLabs/North-Mini-Code-1.0")67class Cohere2MoeModel(TextModel):68    model_arch = gguf.MODEL_ARCH.COHERE2MOE69    _n_main_layers: int | None = None70    _expert_tensor_re = re.compile(71        r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(down_proj|gate_proj|up_proj)\.weight"72    )73 74    def __init__(self, *args, **kwargs):75        super().__init__(*args, **kwargs)76        if (n_nextn := int(self.hparams.get("num_nextn_predict_layers", 0) or 0)) > 0 and not self.no_mtp:77            self.block_count += n_nextn78            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)79        self._experts: list[dict[str, Tensor]] = [{} for _ in range(self.block_count)]80 81    def _set_vocab_gpt2(self) -> None:82        tokens, toktypes, tokpre = self.get_vocab_base()83        self.gguf_writer.add_tokenizer_model("gpt2")84        self.gguf_writer.add_tokenizer_pre(tokpre)85        self.gguf_writer.add_token_list(tokens)86        self.gguf_writer.add_token_types(toktypes)87 88        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)89        special_vocab.add_to_gguf(self.gguf_writer)90 91    def set_gguf_parameters(self):92        hparams = self.hparams93        expert_intermediate_size = hparams["intermediate_size"]94        mlp_layer_types = hparams.get("mlp_layer_types")95        n_dense_lead = hparams.get("first_k_dense_replace", 0)96        if mlp_layer_types is not None:97            n_dense_lead = next((i for i, t in enumerate(mlp_layer_types) if t != "dense"), len(mlp_layer_types))98 99        super().set_gguf_parameters()100 101        self.gguf_writer.add_logit_scale(hparams["logit_scale"])102        self.gguf_writer.add_sliding_window(hparams["sliding_window"])103        self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]])104        self.gguf_writer.add_vocab_size(hparams["vocab_size"])105        self.gguf_writer.add_expert_feed_forward_length(expert_intermediate_size)106        self.gguf_writer.add_leading_dense_block_count(n_dense_lead)107        self.gguf_writer.add_expert_weights_norm(hparams.get("norm_topk_prob", False))108        if (num_shared_experts := hparams.get("num_shared_experts", 0)) > 0:109            if hparams.get("shared_expert_combination_strategy", "average") != "average":110                raise ValueError("Cohere2 MoE only supports average shared expert combination")111            self.gguf_writer.add_expert_shared_count(num_shared_experts)112            self.gguf_writer.add_expert_shared_feed_forward_length(expert_intermediate_size * num_shared_experts)113        if (n_nextn := hparams.get("num_nextn_predict_layers", 0)) > 0 and not self.no_mtp:114            self.gguf_writer.add_nextn_predict_layers(n_nextn)115        self.gguf_writer.add_rope_dimension_count(hparams["head_dim"])116        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)117 118    def index_tensors(self, remote_hf_model_id: str | None = None):119        hparams = {**self.hparams, **self.hparams.get("text_config", {})}120        self._n_main_layers = hparams.get("num_hidden_layers")121        type(self)._n_main_layers = self._n_main_layers122        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)123 124    @classmethod125    def filter_tensors(cls, item):126        if (titem := super().filter_tensors(item)) is None:127            return None128        name, gen = titem129 130        if cls._n_main_layers is not None:131            is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers132            if is_mtp and cls.no_mtp:133                return None134            if cls.mtp_only and not is_mtp and name not in (135                "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",136            ):137                return None138 139        return name, gen140 141    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:142        if name.endswith(".bias"):143            if torch.any(data_torch != 0):144                raise ValueError(f"Bias tensor {name!r} is not zero.")145            logger.debug(f"Skipping bias tensor {name!r}.")146            return147 148        if (m := self._expert_tensor_re.fullmatch(name)) is not None:149            n_experts = self.hparams["num_experts"]150            layer_idx = int(m.group(1))151            assert bid is None or bid == layer_idx152 153            self._experts[layer_idx][name] = data_torch154 155            expected = {156                f"model.layers.{layer_idx}.mlp.experts.{xid}.{w_name}.weight"157                for xid in range(n_experts)158                for w_name in ("down_proj", "gate_proj", "up_proj")159            }160            if expected.issubset(self._experts[layer_idx]):161                for w_name in ["down_proj", "gate_proj", "up_proj"]:162                    datas: list[Tensor] = []163 164                    for xid in range(n_experts):165                        ename = f"model.layers.{layer_idx}.mlp.experts.{xid}.{w_name}.weight"166                        datas.append(self._experts[layer_idx][ename])167                        del self._experts[layer_idx][ename]168 169                    data_torch = torch.stack(datas, dim=0)170                    merged_name = f"model.layers.{layer_idx}.mlp.experts.{w_name}.weight"171 172                    yield from super().modify_tensors(data_torch, merged_name, layer_idx)173            return174 175        yield from super().modify_tensors(data_torch, name, bid)176 177    def prepare_tensors(self):178        super().prepare_tensors()179 180        experts = [k for d in self._experts for k in d.keys()]181        if len(experts) > 0:182            raise ValueError(f"Unprocessed experts: {experts}")183