Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import json4import sys5 6from typing import Iterable, TYPE_CHECKING7 8import torch9 10if TYPE_CHECKING:11    from torch import Tensor12 13from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger14 15from .llama import LlamaModel16 17 18@ModelBase.register("ArcticForCausalLM")19@ModelBase.example("Snowflake/snowflake-arctic-instruct")20class ArcticModel(TextModel):21    model_arch = gguf.MODEL_ARCH.ARCTIC22 23    def set_vocab(self):24        # The reason for using a custom implementation here is that the25        # snowflake-arctic-instruct model redefined tokens 31998 and 31999 from26        # tokenizer.model and used them as BOS and EOS instead of adding new tokens.27        from sentencepiece import SentencePieceProcessor28 29        tokenizer_path = self.dir_model / 'tokenizer.model'30 31        if not tokenizer_path.is_file():32            logger.error(f'Error: Missing {tokenizer_path}')33            sys.exit(1)34 35        # Read the whole vocabulary from the tokenizer.model file36        tokenizer = SentencePieceProcessor()37        tokenizer.LoadFromFile(str(tokenizer_path))38 39        vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size())40 41        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]42        scores: list[float] = [-10000.0] * vocab_size43        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size44 45        for token_id in range(tokenizer.vocab_size()):46 47            piece = tokenizer.IdToPiece(token_id)48            text = piece.encode("utf-8")49            score = tokenizer.GetScore(token_id)50 51            toktype = SentencePieceTokenTypes.NORMAL52            if tokenizer.IsUnknown(token_id):53                toktype = SentencePieceTokenTypes.UNKNOWN54            elif tokenizer.IsControl(token_id):55                toktype = SentencePieceTokenTypes.CONTROL56            elif tokenizer.IsUnused(token_id):57                toktype = SentencePieceTokenTypes.UNUSED58            elif tokenizer.IsByte(token_id):59                toktype = SentencePieceTokenTypes.BYTE60 61            tokens[token_id] = text62            scores[token_id] = score63            toktypes[token_id] = toktype64 65        # Use the added_tokens_decoder field from tokeniser_config.json as the source66        # of information about added/redefined tokens and modify them accordingly.67        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'68        if tokenizer_config_file.is_file():69            with open(tokenizer_config_file, "r", encoding="utf-8") as f:70                tokenizer_config_json = json.load(f)71 72                if "added_tokens_decoder" in tokenizer_config_json:73                    added_tokens_decoder = tokenizer_config_json["added_tokens_decoder"]74                    for token_id, token_json in added_tokens_decoder.items():75                        token_id = int(token_id)76                        if token_id >= vocab_size:77                            logger.debug(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}')78                            continue79 80                        token_content = token_json["content"]81                        token_type = SentencePieceTokenTypes.USER_DEFINED82                        token_score = -10000.083 84                        # Map unk_token to UNKNOWN, other special tokens to CONTROL85                        # Set the score to 0.0 as in the original tokenizer.model86                        if ("special" in token_json) and token_json["special"]:87                            if token_content == tokenizer_config_json["unk_token"]:88                                token_type = SentencePieceTokenTypes.UNKNOWN89                            else:90                                token_type = SentencePieceTokenTypes.CONTROL91                            token_score = 0.092 93                        logger.info(f"Setting added token {token_id} to '{token_content}' (type: {token_type}, score: {token_score:.2f})")94                        tokens[token_id] = token_content.encode("utf-8")95                        toktypes[token_id] = token_type96                        scores[token_id] = token_score97 98        self.gguf_writer.add_tokenizer_model("llama")99        self.gguf_writer.add_tokenizer_pre("default")100        self.gguf_writer.add_token_list(tokens)101        self.gguf_writer.add_token_scores(scores)102        self.gguf_writer.add_token_types(toktypes)103 104        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))105        special_vocab.add_to_gguf(self.gguf_writer)106 107    def set_gguf_parameters(self):108        super().set_gguf_parameters()109        hparams = self.hparams110        self.gguf_writer.add_vocab_size(hparams["vocab_size"])111        self.gguf_writer.add_rope_dimension_count(hparams["hidden_size"] // hparams["num_attention_heads"])112 113    _experts: list[dict[str, Tensor]] | None = None114 115    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:116        n_head = self.hparams["num_attention_heads"]117        n_kv_head = self.hparams.get("num_key_value_heads")118 119        if name.endswith("q_proj.weight"):120            data_torch = LlamaModel.permute(data_torch, n_head, n_head)121        if name.endswith("k_proj.weight"):122            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)123 124        # process the experts separately125        if name.find("block_sparse_moe.experts") != -1:126            n_experts = self.hparams["num_local_experts"]127 128            assert bid is not None129 130            if self._experts is None:131                self._experts = [{} for _ in range(self.block_count)]132 133            self._experts[bid][name] = data_torch134 135            if len(self._experts[bid]) >= n_experts * 3:136                # merge the experts into a single 3d tensor137                for wid in ["w1", "w2", "w3"]:138                    datas: list[Tensor] = []139 140                    for xid in range(n_experts):141                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"142                        datas.append(self._experts[bid][ename])143                        del self._experts[bid][ename]144 145                    data_torch = torch.stack(datas, dim=0)146 147                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"148 149                    yield from super().modify_tensors(data_torch, merged_name, bid)150                return151            else:152                return153 154        yield from super().modify_tensors(data_torch, name, bid)155 156    def prepare_tensors(self):157        super().prepare_tensors()158 159        if self._experts is not None:160            # flatten `list[dict[str, Tensor]]` into `list[str]`161            experts = [k for d in self._experts for k in d.keys()]162            if len(experts) > 0:163                raise ValueError(f"Unprocessed experts: {experts}")164