Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf, logger11 12 13@ModelBase.register("MellumForCausalLM")14@ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base")15class MellumModel(TextModel):16    model_arch = gguf.MODEL_ARCH.MELLUM17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:21            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)22            logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")23 24        use_sliding_window = self.hparams.get("use_sliding_window")25        sliding_window = self.hparams.get("sliding_window")26        if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None:27            self.gguf_writer.add_sliding_window(sliding_window)28            logger.info(f"gguf: sliding window = {sliding_window}")29            self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]])30            logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}")31 32    _experts: list[dict[str, Tensor]] | None = None33 34    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:35        if name.find("experts") != -1:36            n_experts = self.find_hparam(["num_local_experts", "num_experts"])37            assert bid is not None38 39            if self._experts is None:40                self._experts = [{} for _ in range(self.block_count)]41 42            self._experts[bid][name] = data_torch43 44            if len(self._experts[bid]) >= n_experts * 3:45                for w_name in ["down_proj", "gate_proj", "up_proj"]:46                    datas: list[Tensor] = []47 48                    for xid in range(n_experts):49                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"50                        datas.append(self._experts[bid][ename])51                        del self._experts[bid][ename]52 53                    data_torch = torch.stack(datas, dim=0)54 55                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"56 57                    yield from super().modify_tensors(data_torch, merged_name, bid)58                return59            else:60                return61 62        yield from super().modify_tensors(data_torch, name, bid)63