Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Any, Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf9 10 11@ModelBase.register("OpenELMForCausalLM")12@ModelBase.example("apple/OpenELM-270M")13class OpenELMModel(TextModel):14    model_arch = gguf.MODEL_ARCH.OPENELM15 16    @staticmethod17    def _make_divisible(v: float | int, divisor: int) -> int:18        # ref: https://huggingface.co/apple/OpenELM-270M-Instruct/blob/eb111ff2e6724348e5b905984063d4064d4bc579/configuration_openelm.py#L34-L3819        new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)20        # Make sure that round down does not go down by more than 10%.21        if new_v < 0.9 * v:22            new_v += divisor23        return new_v24 25    def __init__(self, *args, **kwargs):26        super().__init__(*args, **kwargs)27 28        ffn_multipliers: list[float] = self.hparams["ffn_multipliers"]29        ffn_dim_divisor: int = self.hparams["ffn_dim_divisor"]30        self._n_embd: int = self.hparams["model_dim"]31        self._num_kv_heads: list[int] = self.hparams["num_kv_heads"]32        self._num_query_heads: list[int] = self.hparams["num_query_heads"]33        self._ffn_dims: list[int] = [34            OpenELMModel._make_divisible(multiplier * self._n_embd, ffn_dim_divisor)35            for multiplier in ffn_multipliers36        ]37        assert isinstance(self._num_kv_heads, list) and isinstance(self._num_kv_heads[0], int)38        assert isinstance(self._num_query_heads, list) and isinstance(self._num_query_heads[0], int)39 40    # Uses the tokenizer from meta-llama/Llama-2-7b-hf41    def set_vocab(self):42        try:43            self._set_vocab_sentencepiece()44        except FileNotFoundError:45            self._set_vocab_builtin("llama-spm", self.hparams["vocab_size"])46 47    def set_gguf_parameters(self):48        n_embd = self._n_embd49        head_dim = self.hparams["head_dim"]50        rot_pct = 1.051        assert self.block_count == len(self._num_kv_heads)52        assert self.block_count == len(self._num_query_heads)53        assert self.block_count == len(self._ffn_dims)54 55        self.gguf_writer.add_block_count(self.block_count)56        self.gguf_writer.add_context_length(self.hparams["max_context_length"])57        self.gguf_writer.add_embedding_length(n_embd)58        self.gguf_writer.add_feed_forward_length(self._ffn_dims)59        self.gguf_writer.add_head_count(self._num_query_heads)60        self.gguf_writer.add_head_count_kv(self._num_kv_heads)61        self.gguf_writer.add_rope_freq_base(self.hparams["rope_freq_constant"])62        # https://huggingface.co/apple/OpenELM-270M-Instruct/blob/c401df2/modeling_openelm.py#L3063        self.gguf_writer.add_layer_norm_rms_eps(1e-6)64        self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim))65        self.gguf_writer.add_key_length(head_dim)66        self.gguf_writer.add_value_length(head_dim)67        self.gguf_writer.add_file_type(self.ftype)68 69    def find_hparam(self, keys: Iterable[str], optional: bool = False) -> Any:70        if "n_layers" in keys:71            return self.hparams["num_transformer_layers"]72 73        return super().find_hparam(keys, optional)74 75    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:76 77        # split ff78        if bid is not None and name == f"transformer.layers.{bid}.ffn.proj_1.weight":79            ff_dim = self._ffn_dims[bid]80            yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), data_torch[:ff_dim])81            yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), data_torch[ff_dim:])82            return83 84        yield (self.map_tensor_name(name), data_torch)85