char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.7 KiB163 linesraw
1from __future__ import annotations
2
3import json
4import sys
5
6from typing import Iterable, TYPE_CHECKING
7
8import torch
9
10if TYPE_CHECKING:
11    from torch import Tensor
12
13from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
14
15from .llama import LlamaModel
16
17
18@ModelBase.register("ArcticForCausalLM")
19@ModelBase.example("Snowflake/snowflake-arctic-instruct")
20class ArcticModel(TextModel):
21    model_arch = gguf.MODEL_ARCH.ARCTIC
22
23    def set_vocab(self):
24        # The reason for using a custom implementation here is that the
25        # snowflake-arctic-instruct model redefined tokens 31998 and 31999 from
26        # tokenizer.model and used them as BOS and EOS instead of adding new tokens.
27        from sentencepiece import SentencePieceProcessor
28
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 file
36        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_size
43        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size
44
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.NORMAL
52            if tokenizer.IsUnknown(token_id):
53                toktype = SentencePieceTokenTypes.UNKNOWN
54            elif tokenizer.IsControl(token_id):
55                toktype = SentencePieceTokenTypes.CONTROL
56            elif tokenizer.IsUnused(token_id):
57                toktype = SentencePieceTokenTypes.UNUSED
58            elif tokenizer.IsByte(token_id):
59                toktype = SentencePieceTokenTypes.BYTE
60
61            tokens[token_id] = text
62            scores[token_id] = score
63            toktypes[token_id] = toktype
64
65        # Use the added_tokens_decoder field from tokeniser_config.json as the source
66        # 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                            continue
79
80                        token_content = token_json["content"]
81                        token_type = SentencePieceTokenTypes.USER_DEFINED
82                        token_score = -10000.0
83
84                        # Map unk_token to UNKNOWN, other special tokens to CONTROL
85                        # Set the score to 0.0 as in the original tokenizer.model
86                        if ("special" in token_json) and token_json["special"]:
87                            if token_content == tokenizer_config_json["unk_token"]:
88                                token_type = SentencePieceTokenTypes.UNKNOWN
89                            else:
90                                token_type = SentencePieceTokenTypes.CONTROL
91                            token_score = 0.0
92
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_type
96                        scores[token_id] = token_score
97
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.hparams
110        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 = None
114
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 separately
125        if name.find("block_sparse_moe.experts") != -1:
126            n_experts = self.hparams["num_local_experts"]
127
128            assert bid is not None
129
130            if self._experts is None:
131                self._experts = [{} for _ in range(self.block_count)]
132
133            self._experts[bid][name] = data_torch
134
135            if len(self._experts[bid]) >= n_experts * 3:
136                # merge the experts into a single 3d tensor
137                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                return
151            else:
152                return
153
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}")