Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import math4 5from typing import Callable, Iterable, TYPE_CHECKING6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("Jais2ForCausalLM")14class Jais2Model(TextModel):15    model_arch = gguf.MODEL_ARCH.JAIS216 17    def set_gguf_parameters(self):18        super().set_gguf_parameters()19        hparams = self.hparams20        head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])21        self.gguf_writer.add_rope_dimension_count(head_dim)22 23 24@ModelBase.register("JAISLMHeadModel")25class JaisModel(TextModel):26    model_arch = gguf.MODEL_ARCH.JAIS27 28    def __init__(self, *args, **kwargs):29        super().__init__(*args, **kwargs)30 31        # SwigLU activation32        assert self.hparams["activation_function"] == "swiglu"33        # ALiBi position embedding34        assert self.hparams["position_embedding_type"] == "alibi"35 36        # Embeddings scale37        self.embeddings_scale = 1.038        if 'mup_embeddings_scale' in self.hparams:39            self.embeddings_scale = self.hparams['mup_embeddings_scale']40        elif 'embeddings_scale' in self.hparams:41            self.embeddings_scale = self.hparams['embeddings_scale']42        else:43            assert False44 45        self.width_scale = 1.046        if 'mup_output_alpha' in self.hparams:47            assert 'mup_width_scale' in self.hparams48            self.width_scale = self.hparams['mup_output_alpha'] * self.hparams['mup_width_scale']49        elif 'width_scale' in self.hparams:50            self.width_scale = self.hparams['width_scale']51        else:52            assert False53 54        self.max_alibi_bias = 8.055 56    def set_vocab(self):57        self._set_vocab_gpt2()58 59    def set_gguf_parameters(self):60        self.gguf_writer.add_block_count(self.block_count)61        self.gguf_writer.add_context_length(self.hparams["n_positions"])62        self.gguf_writer.add_embedding_length(self.hparams["n_embd"])63        self.gguf_writer.add_feed_forward_length(self.hparams["n_inner"])64        self.gguf_writer.add_head_count(self.hparams["n_head"])65        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])66        self.gguf_writer.add_file_type(self.ftype)67 68    @classmethod69    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:70        name, gen = item71 72        # we don't need these73        if name.endswith((".attn.bias")):74            return None75 76        return super().filter_tensors(item)77 78    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:79        if name.endswith(("relative_pe.slopes")):80            # Calculate max ALiBi bias (this is the inverse of the ALiBi calculation)81            # Some other models has max_alibi_bias spelled out explicitly in the hyperparams,82            # but Jais's PyTorch model simply precalculates the slope values and places them83            # in relative_pes.slopes84            n_head_closest_log2 = 2 ** math.floor(math.log2(self.hparams["n_head"]))85            first_val = float(data_torch[0].item())86            self.max_alibi_bias = -round(math.log2(first_val) * n_head_closest_log2)87 88            return89 90        if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight", ".c_fc2.weight")):91            data_torch = data_torch.transpose(1, 0)92 93        new_name = self.map_tensor_name(name)94 95        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD):96            yield from super().modify_tensors(data_torch * self.embeddings_scale, new_name, bid)97        elif new_name == self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT):98            yield from super().modify_tensors(data_torch * self.width_scale, new_name, bid)99        else:100            yield from super().modify_tensors(data_torch, new_name, bid)101 102    def prepare_tensors(self):103        super().prepare_tensors()104        self.gguf_writer.add_max_alibi_bias(self.max_alibi_bias)105