From 488cdee96f911420fe9343714f6316c7ab705ddc Mon Sep 17 00:00:00 2001 From: suhyun-hwang Date: Fri, 9 Jan 2026 22:12:16 +0900 Subject: [PATCH] model : add VAETKI architecture support --- convert_hf_to_gguf.py | 233 +++++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 33 +++++ gguf-py/gguf/tensor_mapping.py | 5 + src/CMakeLists.txt | 1 + src/llama-arch.cpp | 31 +++++ src/llama-arch.h | 2 + src/llama-model.cpp | 116 ++++++++++++++++ src/llama-model.h | 1 + src/llama-vocab.cpp | 72 +++++++++- src/llama-vocab.h | 1 + src/models/models.h | 4 + src/models/vaetki.cpp | 194 +++++++++++++++++++++++++++ src/unicode.cpp | 5 +- src/unicode.h | 2 +- 14 files changed, 693 insertions(+), 7 deletions(-) create mode 100644 src/models/vaetki.cpp diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 464ecbaab9..e662f39c29 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -1120,6 +1120,9 @@ class TextModel(ModelBase): if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea": # ref: https://huggingface.co/Qwen/Qwen1.5-7B res = "qwen2" + if chkhsh == "f5f8b79793693cfcca1c36aac854ab481ae887cf7dde234b889f8f4bf009891a": + # ref: https://huggingface.co/nc-ai-consortium/VAETKI-VL-7B-A1B + res = "vaetki" if chkhsh == "b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166": # ref: https://huggingface.co/allenai/OLMo-1.7-7B-hf res = "olmo" @@ -7664,6 +7667,236 @@ class DeepseekV2Model(TextModel): raise ValueError(f"Unprocessed experts: {experts}") +@ModelBase.register("VaetkiForCausalLM") +@ModelBase.register("VaetkiVLForCausalLM") +class VaetkiModel(TextModel): + """VAETKI MoE model with MLA attention and 4-norm layer structure""" + model_arch = gguf.MODEL_ARCH.VAETKI + + _experts: list[dict[str, Tensor]] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Flatten text_config parameters to top level + if "text_config" in self.hparams: + text_config = self.hparams["text_config"] + for key, value in text_config.items(): + if key not in self.hparams: + self.hparams[key] = value + + def set_vocab(self): + # VAETKI uses Metaspace-based BPE tokenizer, load vocab from tokenizer.json + import json + import re + from transformers import AutoTokenizer + + dir_model = self.dir_model + hparams = self.hparams + + tokenizer_json_path = dir_model / "tokenizer.json" + if not tokenizer_json_path.is_file(): + raise FileNotFoundError(f"VAETKI tokenizer.json not found: {tokenizer_json_path}") + + with open(tokenizer_json_path, "r", encoding="utf-8") as f: + tokenizer_json = json.load(f) + + # Get vocab from tokenizer.json + vocab = tokenizer_json["model"]["vocab"] + merges = tokenizer_json["model"].get("merges", []) + + vocab_size = hparams.get("vocab_size", len(vocab)) + + # Build reverse vocab + reverse_vocab = {v: k for k, v in vocab.items()} + + # Get added tokens from tokenizer.json + added_tokens = {} + for token_info in tokenizer_json.get("added_tokens", []): + added_tokens[token_info["id"]] = { + "content": token_info["content"], + "special": token_info.get("special", False) + } + + tokens: list[str] = [] + toktypes: list[int] = [] + + for i in range(vocab_size): + if i in added_tokens: + token = added_tokens[i]["content"] + if added_tokens[i]["special"]: + toktypes.append(gguf.TokenType.CONTROL) + else: + # pre-normalize user-defined spaces (Metaspace → space) + token = token.replace("\xe2\x96\x81", " ") + toktypes.append(gguf.TokenType.USER_DEFINED) + tokens.append(token) + elif i in reverse_vocab: + token = reverse_vocab[i] + # Check for byte tokens (format: <0xXX>) + if re.fullmatch(r"<0x[0-9A-Fa-f]{2}>", token): + toktypes.append(gguf.TokenType.BYTE) + else: + toktypes.append(gguf.TokenType.NORMAL) + tokens.append(token) + else: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + + # Get pre-tokenizer type + tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True) + tokpre = self.get_vocab_base_pre(tokenizer) + + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + # Add merges (convert from [['a', 'b'], ...] to ['a b', ...] format) + if merges: + # tokenizer.json stores merges as list of pairs, GGUF expects space-separated strings + if isinstance(merges[0], list): + merges = [' '.join(pair) for pair in merges] + self.gguf_writer.add_token_merges(merges) + + # Add special tokens + special_vocab = gguf.SpecialVocab(dir_model, load_merges=False) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + hparams = self.hparams + self.gguf_writer.add_block_count(hparams["num_hidden_layers"]) + self.gguf_writer.add_context_length(hparams.get("max_position_embeddings", 32768)) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + # For MLA without absorption, n_head_kv = n_head (full MHA after decompression) + self.gguf_writer.add_head_count_kv(hparams["num_attention_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(hparams.get("rms_norm_eps", 1e-5)) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + # MLA parameters (like DeepSeek2) + self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) + + # For MLA without absorption, key_length/value_length are the full (MHA) dimensions + # key = qk_nope + qk_rope, value = v_head_dim + self.gguf_writer.add_key_length(hparams["qk_head_dim"]) + self.gguf_writer.add_value_length(hparams["v_head_dim"]) + + # key_length_mla/value_length_mla are the MLA head dimensions (same as key/value for non-absorption) + self.gguf_writer.add_key_length_mla(hparams["qk_head_dim"]) + self.gguf_writer.add_value_length_mla(hparams["v_head_dim"]) + self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) + + # VAETKI uses hybrid attention with different rope_theta per layer type: + # - sliding_attention layers use rope_theta (local, default 10000.0) + # - full_attention layers use rope_theta_global (global, default 1000000.0) + # In llama.cpp: rope_freq_base is for non-SWA (full), rope_freq_base_swa is for SWA (sliding) + rope_theta_local = hparams.get("rope_theta", 10000.0) + rope_theta_global = hparams.get("rope_theta_global", 1000000.0) + self.gguf_writer.add_rope_freq_base(rope_theta_global) # for full_attention layers + self.gguf_writer.add_rope_freq_base_swa(rope_theta_local) # for sliding_attention layers + + # MoE parameters + self.gguf_writer.add_leading_dense_block_count(hparams.get("first_k_dense_replace", 1)) + self.gguf_writer.add_expert_count(hparams["n_routed_experts"]) + self.gguf_writer.add_expert_used_count(hparams["num_experts_per_tok"]) + self.gguf_writer.add_expert_shared_count(hparams.get("n_shared_experts", 1)) + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_weights_scale(hparams.get("routed_scaling_factor", 1.0)) + # VAETKI uses sigmoid gating function (WBLTopkRouter uses router_logits.sigmoid()) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + # Normalize top-k probabilities (norm_topk_prob=true in config) + if hparams.get("norm_topk_prob", False): + self.gguf_writer.add_expert_weights_norm(True) + + # Sliding window and hybrid attention pattern + if "sliding_window" in hparams: + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + + # Add sliding window pattern from layer_types + if "layer_types" in hparams: + # Convert layer_types to sliding_window_pattern (1 = sliding, 0 = full) + # Store as uint32 array to match llama.cpp hparams.swa_layers type + sliding_window_pattern = [1 if t == "sliding_attention" else 0 for t in hparams["layer_types"]] + self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # Skip vision encoder tensors + if "vision_tower" in name or "vision_model" in name or "visual" in name: + return [] + if name.startswith("model.vision_model.") or name.startswith("vision_model."): + return [] + + # Handle lm_head.weight (VAETKI does not use tied embeddings) + if name == "lm_head.weight": + return [(self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT), data_torch)] + + # Remove language_model prefix + if name.startswith("model.language_model."): + name = name.replace("model.language_model.", "model.") + elif name.startswith("language_model."): + name = name.replace("language_model.", "model.") + + # VAETKI WBLRMSNorm: add 1 to weights for standard RMSNorm compatibility + norm_weight_patterns = [ + "input_layernorm.weight", + "post_attention_layernorm.weight", + "pre_mlp_layernorm.weight", + "post_mlp_layernorm.weight", + "q_a_layernorm.weight", + "kv_a_layernorm.weight", + "model.norm.weight", + ] + if any(pattern in name for pattern in norm_weight_patterns): + data_torch = data_torch + 1.0 + + # Handle MoE expert tensors + if ".mlp.experts." in name and ".shared_experts." not in name: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + # Check if all experts for this layer are collected (n_experts * 3 tensors: down/gate/up) + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # Merge experts into 3D tensors + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + new_name = self.map_tensor_name(merged_name) + tensors.append((new_name, data_torch)) + + return tensors + else: + return [] + + return super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # Check for unprocessed experts + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + @ModelBase.register("MiniMaxM2ForCausalLM") class MiniMaxM2Model(TextModel): model_arch = gguf.MODEL_ARCH.MINIMAXM2 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 31273b2b5a..3f35e6c09a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -459,6 +459,7 @@ class MODEL_ARCH(IntEnum): MIMO2 = auto() LLAMA_EMBED = auto() MAINCODER = auto() + VAETKI = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -655,6 +656,7 @@ class MODEL_TENSOR(IntEnum): V_MMPROJ_MLP = auto() V_MMPROJ_PEG = auto() V_ENC_EMBD_CLS = auto() + V_ENC_EMBD_CLS_POS = auto() V_ENC_EMBD_PATCH = auto() V_ENC_EMBD_NORM = auto() V_ENC_EMBD_POS = auto() @@ -880,6 +882,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.MIMO2: "mimo2", MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", + MODEL_ARCH.VAETKI: "vaetki", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1073,6 +1076,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.V_MMPROJ_MLP: "mm.model.mlp.{bid}", MODEL_TENSOR.V_MMPROJ_PEG: "mm.model.peg.{bid}", MODEL_TENSOR.V_ENC_EMBD_CLS: "v.class_embd", + MODEL_TENSOR.V_ENC_EMBD_CLS_POS: "v.class_pos_embd", MODEL_TENSOR.V_ENC_EMBD_PATCH: "v.patch_embd", MODEL_TENSOR.V_ENC_EMBD_NORM: "v.norm_embd", MODEL_TENSOR.V_ENC_EMBD_POS: "v.position_embd", @@ -1191,6 +1195,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.V_MMPROJ_MLP, MODEL_TENSOR.V_MMPROJ_PEG, MODEL_TENSOR.V_ENC_EMBD_CLS, + MODEL_TENSOR.V_ENC_EMBD_CLS_POS, MODEL_TENSOR.V_ENC_EMBD_PATCH, MODEL_TENSOR.V_ENC_EMBD_NORM, MODEL_TENSOR.V_ENC_EMBD_POS, @@ -3377,6 +3382,34 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.VAETKI: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + ], # TODO } diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 84aa868809..0b182ef995 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1281,6 +1281,11 @@ class TensorNameMap: "model.vision_tower.embeddings.cls_token", # Intern-S1 "vision_model.class_embedding", # llama 4 "model.vision.patch_embedding.cls_embedding", # cogvlm + "model.visual.class_embedding", # vaetki + ), + + MODEL_TENSOR.V_ENC_EMBD_CLS_POS: ( + "model.visual.class_pos_emb", # vaetki ), MODEL_TENSOR.V_ENC_EMBD_PATCH: ( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c15c281a5e..7f39b4f342 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -136,6 +136,7 @@ add_library(llama models/t5-dec.cpp models/t5-enc.cpp models/wavtokenizer-dec.cpp + models/vaetki.cpp models/xverse.cpp models/mistral3.cpp models/graph-context-mamba.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index a54bc1956a..040dd8f529 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -120,6 +120,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MIMO2, "mimo2" }, { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, + { LLM_ARCH_VAETKI, "vaetki" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -339,6 +340,7 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + { LLM_TENSOR_FFN_PRE_NORM, "blk.%d.ffn_norm" }, { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, @@ -2289,6 +2291,35 @@ static std::set llm_get_tensor_names(llm_arch arch) { LLM_TENSOR_FFN_DOWN, LLM_TENSOR_FFN_UP, }; + case LLM_ARCH_VAETKI: + return { + LLM_TENSOR_TOKEN_EMBD, + LLM_TENSOR_OUTPUT_NORM, + LLM_TENSOR_OUTPUT, + LLM_TENSOR_ATTN_NORM, + LLM_TENSOR_ATTN_Q_A_NORM, + LLM_TENSOR_ATTN_KV_A_NORM, + LLM_TENSOR_ATTN_Q_A, + LLM_TENSOR_ATTN_Q_B, + LLM_TENSOR_ATTN_KV_A_MQA, + LLM_TENSOR_ATTN_KV_B, + LLM_TENSOR_ATTN_K_B, + LLM_TENSOR_ATTN_V_B, + LLM_TENSOR_ATTN_OUT, + LLM_TENSOR_ATTN_POST_NORM, + LLM_TENSOR_FFN_NORM, + LLM_TENSOR_FFN_POST_NORM, + LLM_TENSOR_FFN_GATE, + LLM_TENSOR_FFN_UP, + LLM_TENSOR_FFN_DOWN, + LLM_TENSOR_FFN_GATE_INP, + LLM_TENSOR_FFN_GATE_EXPS, + LLM_TENSOR_FFN_DOWN_EXPS, + LLM_TENSOR_FFN_UP_EXPS, + LLM_TENSOR_FFN_GATE_SHEXP, + LLM_TENSOR_FFN_DOWN_SHEXP, + LLM_TENSOR_FFN_UP_SHEXP, + }; default: GGML_ABORT("unknown architecture for tensor mapping"); } diff --git a/src/llama-arch.h b/src/llama-arch.h index 270d28b16a..f784558dfc 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -124,6 +124,7 @@ enum llm_arch { LLM_ARCH_MIMO2, LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, + LLM_ARCH_VAETKI, LLM_ARCH_UNKNOWN, }; @@ -345,6 +346,7 @@ enum llm_tensor { LLM_TENSOR_FFN_GATE_INP, LLM_TENSOR_FFN_GATE_INP_SHEXP, LLM_TENSOR_FFN_NORM, + LLM_TENSOR_FFN_PRE_NORM, LLM_TENSOR_FFN_POST_NORM, LLM_TENSOR_FFN_GATE, LLM_TENSOR_FFN_DOWN, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 94c47dc248..fe4c4e0979 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1128,6 +1128,47 @@ void llama_model::load_hparams(llama_model_loader & ml) { default: type = LLM_TYPE_UNKNOWN; } } break; + case LLM_ARCH_VAETKI: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla, false); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla, false); + + if (hparams.n_embd_head_k_mla != 0 && hparams.n_embd_head_v_mla != 0) { + hparams.n_embd_head_k = hparams.n_embd_head_k_mla; + hparams.n_embd_head_v = hparams.n_embd_head_v_mla; + hparams.n_head_kv_arr = hparams.n_head_arr; + } + + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID; + } + + { + uint32_t n_swa_temp = 0; + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, n_swa_temp, false); + if (n_swa_temp > 0) { + hparams.n_swa = n_swa_temp; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + } + } + + switch (hparams.n_layer) { + case 24: type = LLM_TYPE_7B; break; + case 48: type = LLM_TYPE_109B; break; + default: type = LLM_TYPE_UNKNOWN; + } + } break; case LLM_ARCH_QWEN3VL: { ml.get_key(LLM_KV_NUM_DEEPSTACK_LAYERS, hparams.n_deepstack_layers, false); @@ -6966,6 +7007,64 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); } } break; + case LLM_ARCH_VAETKI: + { + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla; + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla; + const int64_t n_embd_head_qk_rope = hparams.n_rot; + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_layer_dense = hparams.n_layer_dense_lead; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (!output) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_qk_rope)}, 0); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, 0); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v_mla)}, 0); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0); + + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + if (i < n_layer_dense) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + } + } + } break; default: throw std::runtime_error("unknown architecture"); } @@ -7319,6 +7418,18 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: expert_gating_func = %s\n", __func__, llama_expert_gating_func_name((llama_expert_gating_func_type) hparams.expert_gating_func)); } + if (arch == LLM_ARCH_VAETKI) { + LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); + LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv); + LLAMA_LOG_INFO("%s: n_embd_head_k_mla = %d\n", __func__, hparams.n_embd_head_k_mla); + LLAMA_LOG_INFO("%s: n_embd_head_v_mla = %d\n", __func__, hparams.n_embd_head_v_mla); + LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp); + LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp); + LLAMA_LOG_INFO("%s: n_expert_shared = %d\n", __func__, hparams.n_expert_shared); + LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale); + LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm); + } + if (arch == LLM_ARCH_QWEN2MOE) { LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp); LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp); @@ -7619,6 +7730,10 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { { llm = std::make_unique(*this, params); } break; + case LLM_ARCH_VAETKI: + { + llm = std::make_unique(*this, params); + } break; case LLM_ARCH_DECI: { llm = std::make_unique(*this, params); @@ -8242,6 +8357,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MISTRAL3: case LLM_ARCH_LLAMA_EMBED: case LLM_ARCH_MAINCODER: + case LLM_ARCH_VAETKI: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 diff --git a/src/llama-model.h b/src/llama-model.h index d1de16e3f2..8edf932217 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -261,6 +261,7 @@ struct llama_layer { struct ggml_tensor * ffn_norm = nullptr; struct ggml_tensor * ffn_norm_b = nullptr; struct ggml_tensor * ffn_post_norm = nullptr; + struct ggml_tensor * ffn_pre_norm = nullptr; struct ggml_tensor * layer_out_norm = nullptr; struct ggml_tensor * layer_out_norm_b = nullptr; struct ggml_tensor * ffn_norm_exps = nullptr; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index a23950d007..5a7d17ece6 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -468,6 +468,12 @@ struct llm_tokenizer_bpe : llm_tokenizer { "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}\\p{M}*(?: \\p{L}\\p{M}*)*)+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]?|\\s*[\\r\\n]|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_VAETKI: + regex_exprs = { + "[^\r\n]+", + "[\r\n]+", + }; + break; default: // default regex for BPE tokenization pre-processing regex_exprs = { @@ -525,7 +531,23 @@ struct llm_tokenizer_bpe_session { void tokenize(const std::string & text, std::vector & output) { int final_prev_index = -1; - const auto word_collection = unicode_regex_split(text, tokenizer.regex_exprs); + const bool skip_byte_encoding = (vocab.get_pre_type() == LLAMA_VOCAB_PRE_TYPE_VAETKI); + + std::string normalized; + const std::string * input = &text; + if (skip_byte_encoding) { + normalized.reserve(text.size() * 3); + for (char c : text) { + if (c == ' ') { + normalized += "\xe2\x96\x81"; + } else { + normalized += c; + } + } + input = &normalized; + } + + const auto word_collection = unicode_regex_split(*input, tokenizer.regex_exprs, skip_byte_encoding); symbols_final.clear(); @@ -615,10 +637,15 @@ struct llm_tokenizer_bpe_session { if (token == LLAMA_TOKEN_NULL) { for (auto j = str.begin(); j != str.end(); ++j) { - std::string byte_str(1, *j); - auto token_multibyte = vocab.text_to_token(byte_str); - if (token_multibyte != LLAMA_TOKEN_NULL) { - output.push_back(token_multibyte); + llama_token token_byte; + if (skip_byte_encoding) { + token_byte = vocab.byte_to_token(static_cast(*j)); + } else { + std::string byte_str(1, *j); + token_byte = vocab.text_to_token(byte_str); + } + if (token_byte != LLAMA_TOKEN_NULL) { + output.push_back(token_byte); } } } else { @@ -2041,6 +2068,11 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "solar-open") { pre_type = LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN; clean_spaces = false; + } else if ( + tokenizer_pre == "vaetki") { + pre_type = LLAMA_VOCAB_PRE_TYPE_VAETKI; + clean_spaces = false; + add_space_prefix = false; } else { throw std::runtime_error(format("unknown pre-tokenizer type: '%s'", tokenizer_pre.c_str())); } @@ -2675,6 +2707,11 @@ uint8_t llama_vocab::impl::token_to_byte(llama_token id) const { return strtol(buf.c_str(), NULL, 16); } case LLAMA_VOCAB_TYPE_BPE: { + // VAETKI uses <0xXX> format for byte tokens + if (pre_type == LLAMA_VOCAB_PRE_TYPE_VAETKI) { + auto buf = token_data.text.substr(3, 2); + return strtol(buf.c_str(), NULL, 16); + } GGML_ABORT("fatal error"); } case LLAMA_VOCAB_TYPE_WPM: { @@ -3143,9 +3180,21 @@ int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t return _try_copy(token_text.data(), token_text.size()); } if (attr & LLAMA_TOKEN_ATTR_NORMAL) { + if (pre_type == LLAMA_VOCAB_PRE_TYPE_VAETKI) { + std::string result = token_text; + llama_unescape_whitespace(result); + return _try_copy(result.data(), result.size()); + } std::string result = llama_decode_text(token_text); return _try_copy(result.data(), result.size()); } + if (attr & LLAMA_TOKEN_ATTR_BYTE) { + // VAETKI uses <0xXX> format for byte tokens + if (pre_type == LLAMA_VOCAB_PRE_TYPE_VAETKI) { + char byte = (char) token_to_byte(token); + return _try_copy(&byte, 1); + } + } break; } case LLAMA_VOCAB_TYPE_RWKV: { @@ -3418,6 +3467,19 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const { } case LLAMA_VOCAB_TYPE_WPM: case LLAMA_VOCAB_TYPE_BPE: { + if (pimpl->pre_type == LLAMA_VOCAB_PRE_TYPE_VAETKI) { + const char buf[7] = { '<', '0', 'x', hex[ch >> 4], hex[ch & 15], '>', 0 }; + auto token = pimpl->token_to_id.find(buf); + if (token != pimpl->token_to_id.end()) { + return (*token).second; + } + const char buf2[2] = { (char)ch, 0 }; + auto token2 = pimpl->token_to_id.find(buf2); + if (token2 != pimpl->token_to_id.end()) { + return (*token2).second; + } + return LLAMA_TOKEN_NULL; + } return pimpl->token_to_id.at(unicode_byte_to_utf8(ch)); } case LLAMA_VOCAB_TYPE_PLAMO2: { diff --git a/src/llama-vocab.h b/src/llama-vocab.h index 28c3a82b91..8ac8a6036e 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -54,6 +54,7 @@ enum llama_vocab_pre_type { LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN = 43, LLAMA_VOCAB_PRE_TYPE_YOUTU = 44, LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE = 45, + LLAMA_VOCAB_PRE_TYPE_VAETKI = 46, }; struct LLM_KV; diff --git a/src/models/models.h b/src/models/models.h index 3a44f7f140..e5ce8141e4 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -568,6 +568,10 @@ struct llm_build_wavtokenizer_dec : public llm_graph_context { llm_build_wavtokenizer_dec(const llama_model & model, const llm_graph_params & params); }; +struct llm_build_vaetki : public llm_graph_context { + llm_build_vaetki(const llama_model & model, const llm_graph_params & params); +}; + struct llm_build_xverse : public llm_graph_context { llm_build_xverse(const llama_model & model, const llm_graph_params & params); }; diff --git a/src/models/vaetki.cpp b/src/models/vaetki.cpp new file mode 100644 index 0000000000..9621220462 --- /dev/null +++ b/src/models/vaetki.cpp @@ -0,0 +1,194 @@ +#include "models.h" + +llm_build_vaetki::llm_build_vaetki(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla; + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla; + + const int64_t n_embd_head_qk_rope = hparams.n_rot; + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + + const uint32_t kv_lora_rank = hparams.n_lora_kv; + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head_qk_nope + n_embd_head_qk_rope)); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + llm_graph_input_attn_kv_iswa * inp_attn = build_attn_inp_kv_iswa(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self_attention + { + ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur); + cb(q, "q_a", il); + + q = build_norm(q, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q, "q_a_norm", il); + + q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q); + cb(q, "q", il); + + ggml_tensor * q_nope = + ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k_mla), + ggml_row_size(q->type, n_embd_head_k_mla) * n_head, 0); + cb(q_nope, "q_nope", il); + + ggml_tensor * q_pe = ggml_view_3d( + ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k_mla), + ggml_row_size(q->type, n_embd_head_k_mla) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope)); + cb(q_pe, "q_pe", il); + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur); + cb(kv_cmpr_pe, "kv_cmpr_pe", il); + + ggml_tensor * kv_cmpr = + ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + cb(kv_cmpr, "kv_cmpr", il); + + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + cb(k_pe, "k_pe", il); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q_pe, "q_pe_rope", il); + + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(k_pe, "k_pe_rope", il); + + // convert interleaved RoPE to split format + q_pe = ggml_reshape_4d(ctx0, q_pe, 2, n_embd_head_qk_rope/2, n_head, n_tokens); + q_pe = ggml_permute(ctx0, q_pe, 1, 0, 2, 3); + q_pe = ggml_cont(ctx0, q_pe); + q_pe = ggml_reshape_3d(ctx0, q_pe, n_embd_head_qk_rope, n_head, n_tokens); + cb(q_pe, "q_pe_split", il); + + k_pe = ggml_reshape_4d(ctx0, k_pe, 2, n_embd_head_qk_rope/2, 1, n_tokens); + k_pe = ggml_permute(ctx0, k_pe, 1, 0, 2, 3); + k_pe = ggml_cont(ctx0, k_pe); + k_pe = ggml_reshape_3d(ctx0, k_pe, n_embd_head_qk_rope, 1, n_tokens); + cb(k_pe, "k_pe_split", il); + + kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv_cmpr, "kv_cmpr_norm", il); + + ggml_tensor * kv = ggml_mul_mat(ctx0, model.layers[il].wkv_b, kv_cmpr); + cb(kv, "kv", il); + + ggml_tensor * k_nope = + ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v_mla), + ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v_mla) * n_head, 0); + cb(k_nope, "k_nope_view", il); + + ggml_tensor * Vcur = ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v_mla), + ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v_mla) * n_head, + ggml_row_size(kv->type, n_embd_head_qk_nope)); + cb(Vcur, "Vcur_view", il); + + Vcur = ggml_cont(ctx0, Vcur); + cb(Vcur, "Vcur_cont", il); + + ggml_tensor * Qcur = ggml_concat(ctx0, q_nope, q_pe, 0); + cb(Qcur, "Qcur", il); + + ggml_tensor * Kcur = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_pe, q_pe), 0); + cb(Kcur, "Kcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, NULL, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + } + + cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + ggml_tensor * moe_out = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = + build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, NULL, + model.layers[il].ffn_gate_shexp, NULL, NULL, + model.layers[il].ffn_down_shexp, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } + + cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_post_norm", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/unicode.cpp b/src/unicode.cpp index b47dcbe619..6a0c970335 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -956,7 +956,7 @@ bool unicode_cpt_is_han(uint32_t cpt) { return false; } -std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs) { +std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs, bool skip_byte_encoding) { // unicode categories static const std::map k_ucat_enum = { { "\\p{N}", unicode_cpt_flags::NUMBER }, @@ -1143,5 +1143,8 @@ std::vector unicode_regex_split(const std::string & text, const std start += offset; } + if (skip_byte_encoding) { + return bpe_words; + } return unicode_byte_encoding_process(bpe_words); } diff --git a/src/unicode.h b/src/unicode.h index 5bd1362ff4..7856dbc98e 100644 --- a/src/unicode.h +++ b/src/unicode.h @@ -108,4 +108,4 @@ uint32_t unicode_tolower(uint32_t cpt); bool unicode_cpt_is_han(uint32_t cpt); -std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs); +std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs, bool skip_byte_encoding = false);