Felipe97/llama-cpp-compiled

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")14@ModelBase.example("rednote-hilab/dots.llm1.inst")15class Dots1Model(Qwen2MoeModel):16    model_arch = gguf.MODEL_ARCH.DOTS117 18    def __init__(self, *args, **kwargs):19        super().__init__(*args, **kwargs)20        self.hparams["num_experts"] = self.hparams["n_routed_experts"]21 22    def set_gguf_parameters(self):23        super().set_gguf_parameters()24        self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"])25        self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])26        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])27        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])28 29    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):30        if "shared_experts" in name:31            yield from ModelBase.modify_tensors(self, data_torch, name, bid)32        else:33            yield from super().modify_tensors(data_torch, name, bid)34