Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf9 10 11@ModelBase.register("MPTForCausalLM")12@ModelBase.example("anas-awadalla/mpt-7b")13class MPTModel(TextModel):14    model_arch = gguf.MODEL_ARCH.MPT15 16    def set_vocab(self):17        try:18            self._set_vocab_gpt2()19        except Exception:20            # Fallback for SEA-LION model21            self._set_vocab_sentencepiece()22            self.gguf_writer.add_add_bos_token(False)23            self.gguf_writer.add_pad_token_id(3)24            self.gguf_writer.add_eos_token_id(1)25            self.gguf_writer.add_unk_token_id(0)26 27    def set_gguf_parameters(self):28        self.gguf_writer.add_context_length(self.hparams["max_seq_len"])29        self.gguf_writer.add_embedding_length(self.hparams["d_model"])30        self.gguf_writer.add_block_count(self.block_count)31        self.gguf_writer.add_feed_forward_length(4 * self.hparams["d_model"])32        self.gguf_writer.add_head_count(self.hparams["n_heads"])33        if kv_n_heads := self.hparams["attn_config"].get("kv_n_heads"):34            self.gguf_writer.add_head_count_kv(kv_n_heads)35        self.gguf_writer.add_layer_norm_eps(1e-5)36        if self.hparams["attn_config"]["clip_qkv"] is not None:37            self.gguf_writer.add_clamp_kqv(self.hparams["attn_config"]["clip_qkv"])38        if self.hparams["attn_config"]["alibi"]:39            self.gguf_writer.add_max_alibi_bias(self.hparams["attn_config"]["alibi_bias_max"])40        else:41            self.gguf_writer.add_max_alibi_bias(0.0)42 43    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:44        if "scales" in name:45            new_name = self.map_tensor_name(name, try_suffixes=(".weight", ".bias", ".scales"))46            new_name = new_name.replace("scales", "act.scales")47        else:48            new_name = self.map_tensor_name(name, try_suffixes=(".weight", ".bias"))49 50        yield from super().modify_tensors(data_torch, new_name, bid)51