Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf, logger9 10 11@ModelBase.register("BaichuanForCausalLM", "BaiChuanForCausalLM")12class BaichuanModel(TextModel):13    model_arch = gguf.MODEL_ARCH.BAICHUAN14 15    def set_vocab(self):16        self._set_vocab_sentencepiece()17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20 21        self.gguf_writer.add_tensor_data_layout("Meta AI original pth")22        self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"])23 24    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:25        head_count = self.hparams["num_attention_heads"]26        head_count_kv = self.hparams.get("num_key_value_heads", head_count)27 28        if bid is not None and name == f"model.layers.{bid}.self_attn.W_pack.weight":29            logger.info(f"Unpacking and permuting layer {bid}")30            yield from [31                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid),32                    self._reverse_hf_permute_part(data_torch, 0, head_count, head_count)),33                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid),34                    self._reverse_hf_permute_part(data_torch, 1, head_count, head_count_kv)),35                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid),36                    self._reverse_hf_part(data_torch, 2)),37            ]38        else:39            yield from self.modify_tensors(data_torch, self.map_tensor_name(name), bid)40 41    def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor:42        if n_kv_head is not None and n_head != n_kv_head:43            n_head //= n_kv_head44 45        return (46            weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])47            .swapaxes(1, 2)48            .reshape(weights.shape)49        )50 51    def _reverse_hf_permute_part(52        self, weights: Tensor, n_part: int, n_head: int, n_head_kv: int | None = None,53    ) -> Tensor:54        r = weights.shape[0] // 355        return self._reverse_hf_permute(weights[r * n_part:r * n_part + r, ...], n_head, n_head_kv)56 57    def _reverse_hf_part(self, weights: Tensor, n_part: int) -> Tensor:58        r = weights.shape[0] // 359        return weights[r * n_part:r * n_part + r, ...]60