Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, gguf9 10from .qwen import Qwen2MoeModel11 12 13@ModelBase.register("Dots1ForCausalLM")14class Dots1Model(Qwen2MoeModel):15    model_arch = gguf.MODEL_ARCH.DOTS116 17    def __init__(self, *args, **kwargs):18        super().__init__(*args, **kwargs)19        self.hparams["num_experts"] = self.hparams["n_routed_experts"]20 21    def set_gguf_parameters(self):22        super().set_gguf_parameters()23        self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"])24        self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])25        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])26        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])27 28    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):29        if "shared_experts" in name:30            yield from ModelBase.modify_tensors(self, data_torch, name, bid)31        else:32            yield from super().modify_tensors(data_torch, name, bid)33