Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import re4 5from typing import Iterable, TYPE_CHECKING6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("XverseForCausalLM")14class XverseModel(TextModel):15    model_arch = gguf.MODEL_ARCH.XVERSE16 17    def set_vocab(self):18        assert (self.dir_model / "tokenizer.json").is_file()19        dir_model = self.dir_model20        hparams = self.hparams21 22        tokens: list[bytes] = []23        toktypes: list[int] = []24 25        from transformers import AutoTokenizer26        tokenizer = AutoTokenizer.from_pretrained(dir_model)27        vocab_size = hparams.get("vocab_size", len(tokenizer.vocab))  # ty: ignore[unresolved-attribute]28        # Since we are checking the maximum index, we need to ensure it's strictly less than vocab_size,29        # because vocab_size is the count of items, and indexes start at 0.30        max_vocab_index = max(tokenizer.get_vocab().values())  # ty: ignore[unresolved-attribute]31        if max_vocab_index >= vocab_size:32            raise ValueError("Vocabulary size exceeds expected maximum size.")33 34        reverse_vocab: dict[int, str] = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()}  # ty: ignore[unresolved-attribute]35        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]36 37        for token_id in range(vocab_size):38            token_text = reverse_vocab[token_id].encode('utf-8')39            # replace "\x00" to string with length > 040            if token_text == b"\x00":41                toktype = gguf.TokenType.BYTE  # special42                token_text = f"<{token_text}>".encode('utf-8')43            elif re.fullmatch(br"<0x[0-9A-Fa-f]{2}>", token_text):44                toktype = gguf.TokenType.BYTE  # special45            elif reverse_vocab[token_id] in added_vocab:46                if tokenizer.added_tokens_decoder[token_id].special:  # ty: ignore[unresolved-attribute]47                    toktype = gguf.TokenType.CONTROL48                else:49                    toktype = gguf.TokenType.USER_DEFINED50            else:51                toktype = gguf.TokenType.NORMAL52 53            tokens.append(token_text)54            toktypes.append(toktype)55 56        self.gguf_writer.add_tokenizer_model("llama")57        self.gguf_writer.add_tokenizer_pre("default")58        self.gguf_writer.add_token_list(tokens)59        self.gguf_writer.add_token_types(toktypes)60 61        special_vocab = gguf.SpecialVocab(dir_model, n_vocab=len(tokens))62        special_vocab.add_to_gguf(self.gguf_writer)63 64    def set_gguf_parameters(self):65        super().set_gguf_parameters()66 67        self.gguf_writer.add_tensor_data_layout("Meta AI original pth")68        self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"])69 70    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:71        head_count = self.hparams["num_attention_heads"]72        head_count_kv = self.hparams.get("num_key_value_heads", head_count)73 74        # HF models permute some of the tensors, so we need to undo that75        if name.endswith("q_proj.weight"):76            data_torch = self._reverse_hf_permute(data_torch, head_count, head_count)77        if name.endswith("k_proj.weight"):78            data_torch = self._reverse_hf_permute(data_torch, head_count, head_count_kv)79 80        yield from super().modify_tensors(data_torch, name, bid)81 82    def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor:83        if n_kv_head is not None and n_head != n_kv_head:84            n_head //= n_kv_head85 86        return (87            weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])88            .swapaxes(1, 2)89            .reshape(weights.shape)90        )91