Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("FalconForCausalLM", "RWForCausalLM")14@ModelBase.example("tiiuae/falcon-7b")15class FalconModel(TextModel):16    model_arch = gguf.MODEL_ARCH.FALCON17 18    def set_gguf_parameters(self):19        n_head = self.hparams.get("num_attention_heads")20        if n_head is None:21            n_head = self.hparams["n_head"]  # old name22 23        n_head_kv = self.hparams.get("num_kv_heads")24        if n_head_kv is None:25            n_head_kv = self.hparams.get("n_head_kv", 1)  # old name26 27        self.gguf_writer.add_context_length(2048)  # not in config.json28        self.gguf_writer.add_tensor_data_layout("jploski")  # qkv tensor transform29        self.gguf_writer.add_embedding_length(self.hparams["hidden_size"])30        self.gguf_writer.add_feed_forward_length(4 * self.hparams["hidden_size"])31        self.gguf_writer.add_block_count(self.block_count)32        self.gguf_writer.add_head_count(n_head)33        self.gguf_writer.add_head_count_kv(n_head_kv)34        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])35        self.gguf_writer.add_file_type(self.ftype)36 37    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:38        # QKV tensor transform39        # The original query_key_value tensor contains n_head_kv "kv groups",40        # each consisting of n_head/n_head_kv query weights followed by one key41        # and one value weight (shared by all query heads in the kv group).42        # This layout makes it a big pain to work with in GGML.43        # So we rearrange them here,, so that we have n_head query weights44        # followed by n_head_kv key weights followed by n_head_kv value weights,45        # in contiguous fashion.46        # ref: https://github.com/jploski/ggml/blob/falcon40b/examples/falcon/convert-hf-to-ggml.py47 48        if "query_key_value" in name:49            n_head = self.find_hparam(["num_attention_heads", "n_head"])50            n_head_kv = self.find_hparam(["num_kv_heads", "n_head_kv"], optional=True) or 151            head_dim = self.hparams["hidden_size"] // n_head52 53            qkv = data_torch.view(n_head_kv, n_head // n_head_kv + 2, head_dim, head_dim * n_head)54            q = qkv[:, :-2].reshape(n_head * head_dim, head_dim * n_head)55            k = qkv[:, [-2]].reshape(n_head_kv * head_dim, head_dim * n_head)56            v = qkv[:, [-1]].reshape(n_head_kv * head_dim, head_dim * n_head)57            data_torch = torch.cat((q, k, v)).reshape_as(data_torch)58 59        yield from super().modify_tensors(data_torch, name, bid)60