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 LazyTorchTensor, ModelBase, TextModel, gguf11 12 13@ModelBase.register("TalkieForCausalLM")14@ModelBase.example("lewtun/talkie-1930-13b-it-hf")15class TalkieModel(TextModel):16    model_arch = gguf.MODEL_ARCH.TALKIE17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20        # Talkie used F.rms_norm without an explicit eps21        self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps)22 23    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:24        prefix = f"model.blocks.{bid}." if bid is not None else ""25        suffix = name.removeprefix(prefix)26 27        if suffix == "attn_gain.a_g":28            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid, ".scale"), data_torch29            return30        elif suffix == "mlp_gain.a_g":31            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN, bid, ".scale"), data_torch32            return33        elif suffix == "lm_head_gain.w_g":34            self.gguf_writer.add_logit_scale(LazyTorchTensor.to_eager(data_torch).item())35            return36        elif suffix in ("attn.attn_query.weight", "attn.attn_key.weight"):37            # absorb inverse rope38            head_dim = self.hparams["head_dim"]39            shape = data_torch.shape40            data_torch = torch.reshape(data_torch, (-1, head_dim, shape[-1]))41            signs = torch.ones((1, head_dim, 1), dtype=data_torch.dtype)42            signs[:, head_dim // 2 :, :] = -143            if self.lazy:44                signs = LazyTorchTensor.from_eager(signs)45            # (n_head, head_dim, n_in) -> (n_out, n_in)46            data_torch = torch.reshape(data_torch * signs, shape)47        elif suffix == "attn.head_gain.head_g":48            # allow head gain to broadcast49            data_torch = data_torch.unsqueeze(-1)50 51        if not name.endswith(".weight"):52            name += ".weight"53 54        yield from super().modify_tensors(data_torch, name, bid)55