Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import re4 5from typing import Iterable, TYPE_CHECKING6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("HrmTextForCausalLM")14@ModelBase.example("danish-foundation-models/DFM-Mimir")15class HrmTextModel(TextModel):16    model_arch = gguf.MODEL_ARCH.HRM_TEXT17 18    def __init__(self, *args, **kwargs):19        super().__init__(*args, **kwargs)20 21        # training-style configs store the per-stack count in num_hidden_layers,22        # transformers-style configs keep it in num_layers_per_stack23        self.layers_per_stack = self.hparams.get("num_layers_per_stack") or self.hparams["num_hidden_layers"]24        self.h_cycles = self.hparams["H_cycles"]25        self.l_cycles = self.hparams["L_cycles"]26 27        # block_count is the expanded cache-slot count; the file only holds28        # 2 * layers_per_stack physical blocks29        self.block_count = self.layers_per_stack * self.h_cycles * (self.l_cycles + 1)30        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, 2 * self.layers_per_stack)31 32    def set_vocab(self):33        self._set_vocab_gpt2()34 35    def set_gguf_parameters(self):36        super().set_gguf_parameters()37 38        head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]39        self.gguf_writer.add_rope_dimension_count(head_dim)40        self.gguf_writer.add_embedding_scale(self.hparams["embedding_scale"])41        self.gguf_writer.add_hrm_layers_per_stack(self.layers_per_stack)42        self.gguf_writer.add_hrm_h_cycles(self.h_cycles)43        self.gguf_writer.add_hrm_l_cycles(self.l_cycles)44        self.gguf_writer.add_hrm_prefix_lm(bool(self.hparams.get("prefix_lm", False)))45 46    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:47        if name == "model.embed_tokens.weight":48            yield self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), data_torch49            return50        if name == "lm_head.weight":51            yield self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT), data_torch52            return53        if name == "model.z_L_init":54            yield self.format_tensor_name(gguf.MODEL_TENSOR.HRM_Z_L_INIT, suffix=""), data_torch55            return56 57        match = re.fullmatch(r"model\.([LH])_module\.layers\.(\d+)\.(.+)", name)58        if match is None:59            raise ValueError(f"can not map tensor: {name}")60 61        stack, layer_s, tensor_name = match.groups()62        # the L stack occupies blocks [0, layers_per_stack), the H stack follows it63        layer_idx = int(layer_s) + (self.layers_per_stack if stack == "H" else 0)64 65        if tensor_name == "attn.gqkv_proj.weight":66            gate, q, k, v = data_torch.chunk(4, dim=0)67            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, layer_idx), gate.contiguous()68            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, layer_idx), q.contiguous()69            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, layer_idx), k.contiguous()70            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, layer_idx), v.contiguous()71        elif tensor_name == "mlp.gate_up_proj.weight":72            gate, up = data_torch.chunk(2, dim=0)73            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, layer_idx), gate.contiguous()74            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, layer_idx), up.contiguous()75        else:76            if tensor_name.startswith("attn."):77                tensor_name = "self_attn." + tensor_name[len("attn."):]78            tensor_name = "model.layers.{bid}." + tensor_name79            yield from super().modify_tensors(data_torch, tensor_name.format(bid=layer_idx), layer_idx)80