mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-29 07:45:55 +02:00
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8190848bb3 | |||
| 7e1e28cae3 | |||
| 6e2bc65fb2 | |||
| ad77bd31a6 | |||
| ee3d1b54c1 |
+29
-3
@@ -1518,23 +1518,49 @@ done:
|
||||
return res;
|
||||
}
|
||||
|
||||
void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
static void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) {
|
||||
GGML_ABORT("%s", string_format("failed to remove sequence %d with p0=%d, p1=%d\n", seq_id, p0, p1).c_str());
|
||||
}
|
||||
}
|
||||
|
||||
void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
||||
static void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1);
|
||||
}
|
||||
|
||||
void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
|
||||
static void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
llama_memory_seq_add(mem, seq_id, p0, p1, delta);
|
||||
}
|
||||
|
||||
void common_memory::init(llama_context * ctx_tgt, llama_context * ctx_dft) {
|
||||
this->ctx_tgt = ctx_tgt;
|
||||
this->ctx_dft = ctx_dft;
|
||||
}
|
||||
|
||||
void common_memory::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) const {
|
||||
common_context_seq_rm(ctx_tgt, seq_id, p0, p1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm(ctx_dft, seq_id, p0, p1);
|
||||
}
|
||||
}
|
||||
|
||||
void common_memory::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const {
|
||||
common_context_seq_cp(ctx_tgt, seq_id_src, seq_id_dst, p0, p1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_cp(ctx_dft, seq_id_src, seq_id_dst, p0, p1);
|
||||
}
|
||||
}
|
||||
|
||||
void common_memory::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const {
|
||||
common_context_seq_add(ctx_tgt, seq_id, p0, p1, delta);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_add(ctx_dft, seq_id, p0, p1, delta);
|
||||
}
|
||||
}
|
||||
|
||||
void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora) {
|
||||
std::vector<llama_adapter_lora *> loras;
|
||||
std::vector<float> scales;
|
||||
|
||||
+11
-4
@@ -949,10 +949,17 @@ enum common_context_seq_rm_type {
|
||||
// note: clears the memory of the context
|
||||
common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx);
|
||||
|
||||
// aborts execution on failure
|
||||
void common_context_seq_rm (llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1);
|
||||
void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta);
|
||||
void common_context_seq_cp (llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1);
|
||||
struct common_memory {
|
||||
llama_context * ctx_tgt = nullptr;
|
||||
llama_context * ctx_dft = nullptr;
|
||||
|
||||
void init(llama_context * ctx_tgt, llama_context * ctx_dft = nullptr);
|
||||
|
||||
// aborts execution on failure
|
||||
void seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) const;
|
||||
void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const;
|
||||
void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const;
|
||||
};
|
||||
|
||||
//
|
||||
// Batch utils
|
||||
|
||||
+50
-7
@@ -39,28 +39,48 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
}
|
||||
return vision_config
|
||||
|
||||
def get_audio_config(self) -> dict[str, Any] | None:
|
||||
return self.global_config.get("sound_config")
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
if "image_mean" not in self.preprocessor_config:
|
||||
self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406]
|
||||
if "image_std" not in self.preprocessor_config:
|
||||
self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225]
|
||||
|
||||
if self.hparams_audio is not None:
|
||||
self.has_vision_encoder = True
|
||||
self.has_audio_encoder = True
|
||||
self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])
|
||||
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
|
||||
self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"])
|
||||
self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"])
|
||||
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET)
|
||||
self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
else:
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
|
||||
super().set_gguf_parameters()
|
||||
hparams = self.global_config
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
|
||||
self.gguf_writer.add_vision_use_gelu(True)
|
||||
downsample_ratio = hparams.get("downsample_ratio", 0.5)
|
||||
self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
if ".position_embd." in new_name or "pos_embed" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if "sound_encoder" in name or new_name.startswith("mm.a."):
|
||||
if "bias" in new_name or "norm" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if "conv" in new_name and "weight" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
if (titem := super().filter_tensors(item)) is None:
|
||||
return None
|
||||
name, gen = titem
|
||||
|
||||
if "input_conditioner" in name:
|
||||
return None
|
||||
@@ -69,14 +89,18 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
if "radio_model.model.patch_generator.video_embedder" in name:
|
||||
return None
|
||||
|
||||
if not name.startswith("vision_model.radio_model.model.") and not name.startswith("mlp1."):
|
||||
if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")):
|
||||
return None
|
||||
|
||||
if "patch_generator.pos_embed" in name:
|
||||
if not name.endswith(".weight"):
|
||||
name += ".weight"
|
||||
|
||||
return super().filter_tensors((name, gen))
|
||||
# num_batches is only used for training not inference.
|
||||
if "conv.norm" in name and "num_batches" in name:
|
||||
return None
|
||||
|
||||
return name, gen
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
# RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it
|
||||
@@ -104,7 +128,26 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
n_embd = self.hparams["hidden_size"]
|
||||
data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size)
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
if "depthwise_conv.weight" in name:
|
||||
data_torch = data_torch.unsqueeze(-1)
|
||||
data_torch = data_torch.permute(3, 1, 0, 2).contiguous()
|
||||
|
||||
if "pointwise_conv" in name and name.endswith(".weight"):
|
||||
if len(data_torch.shape) == 3 and data_torch.shape[2] == 1:
|
||||
data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1])
|
||||
|
||||
if "subsampling.layers" in name and name.endswith(".bias"):
|
||||
if len(data_torch.shape) == 1:
|
||||
data_torch = data_torch.reshape(1, -1, 1, 1)
|
||||
|
||||
if "pointwise_conv" in name and name.endswith(".bias"):
|
||||
if len(data_torch.shape) == 1:
|
||||
data_torch = data_torch.reshape(1, -1, 1, 1)
|
||||
|
||||
for mapped_name, tensor in super().modify_tensors(data_torch, name, bid):
|
||||
if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."):
|
||||
mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.")
|
||||
yield mapped_name, tensor
|
||||
|
||||
|
||||
@ModelBase.register("NemotronForCausalLM")
|
||||
|
||||
+5
-5
@@ -16,22 +16,22 @@ conda-forge provides builds for:
|
||||
- Apple Metal (macOS)
|
||||
|
||||
```sh
|
||||
conda install -c conda-forge llama-cpp
|
||||
conda install -c conda-forge llama.cpp
|
||||
```
|
||||
|
||||
```sh
|
||||
mamba install -c conda-forge llama-cpp
|
||||
mamba install -c conda-forge llama.cpp
|
||||
```
|
||||
|
||||
```sh
|
||||
# Project-local installation
|
||||
pixi add llama-cpp
|
||||
pixi add llama.cpp
|
||||
|
||||
# Global installation
|
||||
pixi global install llama-cpp
|
||||
pixi global install llama.cpp
|
||||
```
|
||||
|
||||
This distribution is managed on [`conda-forge/llama-cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/).
|
||||
This distribution is managed on [`conda-forge/llama.cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/).
|
||||
|
||||
Shall you have any problems, please open an issue on [its issue tracker](https://github.com/conda-forge/llama.cpp-feedstock/issues).
|
||||
|
||||
|
||||
@@ -15675,7 +15675,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten
|
||||
// <--------------------------------------------> //
|
||||
extra0 = src0->view_src ? (ggml_tensor_extra_cl *)src0->view_src->extra : (ggml_tensor_extra_cl *)src0->extra;
|
||||
|
||||
region.origin = (extra0->offset);
|
||||
region.origin = (extra0->offset + src0->view_offs);
|
||||
if (nb01 > nb02) {
|
||||
// KQ
|
||||
region.size = nb01 * ne01;
|
||||
@@ -15691,7 +15691,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten
|
||||
|
||||
// create sub-buffer for B
|
||||
// <--------------------------------------------> //
|
||||
region.origin = (extra1->offset);
|
||||
region.origin = (extra1->offset + src1->view_offs);
|
||||
region.size = nb10 * ne10 * ne11 * ne12;
|
||||
B_sub_buffer = clCreateSubBuffer((extra1->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
|
||||
CL_CHECK(status);
|
||||
@@ -15712,7 +15712,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten
|
||||
|
||||
// create sub-buffer for output C
|
||||
// <--------------------------------------------> //
|
||||
region.origin = (extrad->offset);
|
||||
region.origin = (extrad->offset + dst->view_offs);
|
||||
region.size = ne0 * ne1 * dst->ne[2] * dst->nb[0]; // size of C in bytes
|
||||
D_sub_buffer = clCreateSubBuffer((extrad->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
|
||||
CL_CHECK(status);
|
||||
@@ -18591,6 +18591,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
|
||||
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
|
||||
if(src0t == GGML_TYPE_F16 && src1t == GGML_TYPE_F32){
|
||||
if (ne01 >= 64 && ne1 >= 32 && ne00 >= 16 && (ne12 % ne02) == 0 &&
|
||||
// the KQ/KQV image kernels do not handle dim 3 (multi-stream batches)
|
||||
ne03 == 1 && ne13 == 1 &&
|
||||
// dst is wrapped with image1d_buffer, the size limit applies, also src0
|
||||
(ne0 * ne1 * dst->ne[2] * dst->nb[0] / 4 <= backend_ctx->image_max_buffer_size)) {
|
||||
// For KQ
|
||||
|
||||
@@ -373,6 +373,7 @@ class Keys:
|
||||
FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length"
|
||||
PROJECTION_DIM = "clip.audio.projection_dim"
|
||||
BLOCK_COUNT = "clip.audio.block_count"
|
||||
SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor"
|
||||
CHUNK_SIZE = "clip.audio.chunk_size"
|
||||
CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size"
|
||||
MAX_POS_EMB = "clip.audio.max_pos_emb"
|
||||
@@ -1002,6 +1003,10 @@ class MODEL_TENSOR(IntEnum):
|
||||
A_ENC_CONV_NORM = auto() # SSM conv
|
||||
A_ENC_CONV_PW1 = auto()
|
||||
A_ENC_CONV_PW2 = auto()
|
||||
A_ENC_CONV_NORM_MEAN = auto() # parakeet
|
||||
A_ENC_CONV_NORM_VAR = auto() # parakeet
|
||||
A_ENC_MEL_FILTERS = auto() # parakeet
|
||||
A_ENC_WINDOW = auto() # parakeet
|
||||
A_CTC_OUT = auto()
|
||||
A_CTC_OUT_MID = auto()
|
||||
A_ENC_ATTN_REL_POS_EMB = auto()
|
||||
@@ -1591,6 +1596,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM: "a.blk.{bid}.conv_norm",
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1: "a.blk.{bid}.conv_pw1",
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2: "a.blk.{bid}.conv_pw2",
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: "a.blk.{bid}.conv_norm_mean",
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR: "a.blk.{bid}.conv_norm_var",
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS: "a.mel_filters",
|
||||
MODEL_TENSOR.A_ENC_WINDOW: "a.window",
|
||||
MODEL_TENSOR.A_CTC_OUT: "a.enc_ctc_out",
|
||||
MODEL_TENSOR.A_CTC_OUT_MID: "a.enc_ctc_out_mid",
|
||||
MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB: "a.blk.{bid}.attn_rel_pos_emb",
|
||||
@@ -1810,6 +1819,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM,
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1,
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS,
|
||||
MODEL_TENSOR.A_ENC_WINDOW,
|
||||
MODEL_TENSOR.A_MM_INP_PROJ,
|
||||
MODEL_TENSOR.A_MM_SOFT_EMB_NORM,
|
||||
MODEL_TENSOR.A_MM_EMBEDDING,
|
||||
@@ -4861,6 +4874,7 @@ class VisionProjectorType:
|
||||
YOUTUVL = "youtuvl"
|
||||
NEMOTRON_V2_VL = "nemotron_v2_vl"
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
PARAKEET = "parakeet" # audio
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
MINICPMV4_6 = "minicpmv4_6"
|
||||
GRANITE_SPEECH = "granite_speech" # audio
|
||||
|
||||
@@ -1374,6 +1374,9 @@ class GGUFWriter:
|
||||
def add_audio_stack_factor(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value)
|
||||
|
||||
def add_audio_subsampling_factor(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.SUBSAMPLING_FACTOR, value)
|
||||
|
||||
def add_audio_chunk_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.CHUNK_SIZE, value)
|
||||
|
||||
|
||||
@@ -2107,6 +2107,7 @@ class TensorNameMap:
|
||||
"conformer.pre_encode.conv.{bid}", # lfm2
|
||||
"model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n
|
||||
"conformer.subsample_conv_projection.layer{bid}.conv", # gemma4
|
||||
"sound_encoder.encoder.subsampling.layers.{bid}", # parakeet
|
||||
"encoder.conv{bid}", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
@@ -2140,6 +2141,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_q", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.q_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.q_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.q_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_q", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2149,6 +2151,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_k", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.k_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.k_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.k_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv)
|
||||
"encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2158,6 +2161,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_v", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.v_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.v_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.v_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv)
|
||||
"encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2187,6 +2191,7 @@ class TensorNameMap:
|
||||
"audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
|
||||
"conformer.layers.{bid}.norm_self_att", # lfm2
|
||||
"conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
|
||||
"encoder.layers.{bid}.attn.pre_norm", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2196,6 +2201,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_out", # lfm2
|
||||
"conformer.layers.{bid}.attention.post", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.post", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_out", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2204,6 +2210,7 @@ class TensorNameMap:
|
||||
"audio_tower.layers.{bid}.final_layer_norm", # ultravox
|
||||
"conformer.layers.{bid}.norm_out", # lfm2
|
||||
"conformer.layers.{bid}.attention.post_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
|
||||
"encoder.layers.{bid}.post_norm", # granite_speech
|
||||
"encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2212,6 +2219,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.norm_feed_forward1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.pre_layer_norm", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.pre_layer_norm", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.norm_feed_forward1", # parakeet
|
||||
"encoder.layers.{bid}.ff1.pre_norm", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2229,6 +2237,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward1.linear1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet
|
||||
"encoder.layers.{bid}.ff1.up_proj", # granite_speech
|
||||
"encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2240,6 +2249,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward1.linear2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet
|
||||
"encoder.layers.{bid}.ff1.down_proj", # granite_speech
|
||||
"encoder.layers.{bid}.fc2", # mimo-audio-tokenizer
|
||||
),
|
||||
@@ -2248,6 +2258,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward2.linear1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.ffw_layer_1", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.ffw_layer_1", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward2.linear1", # parakeet
|
||||
"encoder.layers.{bid}.ff2.up_proj", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2255,6 +2266,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward2.linear2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.ffw_layer_2", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.ffw_layer_2", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward2.linear2", # parakeet
|
||||
"encoder.layers.{bid}.ff2.down_proj", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2262,6 +2274,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.norm_feed_forward2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.pre_layer_norm", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.pre_layer_norm", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.norm_feed_forward2", # parakeet
|
||||
"encoder.layers.{bid}.ff2.pre_norm", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2290,20 +2303,24 @@ class TensorNameMap:
|
||||
MODEL_TENSOR.A_ENC_LINEAR_POS: (
|
||||
"conformer.layers.{bid}.self_attn.linear_pos", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.relative_position_embedding.pos_proj", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.relative_k_proj", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_POS_BIAS_U: (
|
||||
"conformer.layers.{bid}.self_attn.pos_bias_u", # lfm2
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.bias_u", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_POS_BIAS_V: (
|
||||
"conformer.layers.{bid}.self_attn.pos_bias_v", # lfm2
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.bias_v", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_OUT: (
|
||||
"conformer.pre_encode.out", # lfm2
|
||||
"model.audio_tower.subsample_conv_projection.input_proj_linear", # gemma3n (note: it should be A_ENC_INP_PROJ, this is a mistake; it should be corrected in C++ code when it's supported)
|
||||
"conformer.output_proj", # gemma4
|
||||
"sound_encoder.encoder.subsampling.linear", # parakeet
|
||||
),
|
||||
|
||||
# note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors
|
||||
@@ -2313,6 +2330,7 @@ class TensorNameMap:
|
||||
"audio.multi_modal_projector.linear_{bid}", # ultravox, meralion
|
||||
"audio_adapter.model.{bid}", # lfm2
|
||||
"audio_tower.proj{bid}", # qwen3omni
|
||||
"sound_projection.linear{bid}", # parakeet (linear1, linear2)
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MMPROJ_FC: (
|
||||
@@ -2323,6 +2341,7 @@ class TensorNameMap:
|
||||
|
||||
MODEL_TENSOR.A_MM_NORM_PRE: (
|
||||
"audio.multi_modal_projector.ln_pre", # ultravox
|
||||
"sound_projection.norm", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MM_NORM_MID: (
|
||||
@@ -2368,30 +2387,43 @@ class TensorNameMap:
|
||||
MODEL_TENSOR.A_ENC_CONV_DW: (
|
||||
"conformer.layers.{bid}.conv.depthwise_conv", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.depthwise_conv", # parakeet
|
||||
"encoder.layers.{bid}.conv.depth_conv.conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM: (
|
||||
"conformer.layers.{bid}.conv.batch_norm", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.pre_layer_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: (
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm.running_mean", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR: (
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm.running_var", # parakeet
|
||||
"encoder.layers.{bid}.conv.batch_norm", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1: (
|
||||
"conformer.layers.{bid}.conv.pointwise_conv1", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.linear_start", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet
|
||||
"encoder.layers.{bid}.conv.up_conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2: (
|
||||
"conformer.layers.{bid}.conv.pointwise_conv2", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.linear_end", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet
|
||||
"encoder.layers.{bid}.conv.down_conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_NORM_CONV: (
|
||||
"conformer.layers.{bid}.norm_conv", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.conv_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_conv", # parakeet
|
||||
"encoder.layers.{bid}.conv.norm", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2403,6 +2435,14 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.attention.attn.per_dim_scale", # gemma4
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS: (
|
||||
"sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_WINDOW: (
|
||||
"sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MM_EMBEDDING: (
|
||||
"model.embed_audio.embedding", # gemma3n
|
||||
),
|
||||
|
||||
@@ -60,6 +60,7 @@ add_library(mtmd
|
||||
models/mobilenetv5.cpp
|
||||
models/youtuvl.cpp
|
||||
models/yasa2.cpp
|
||||
models/parakeet.cpp
|
||||
)
|
||||
|
||||
set_target_properties(mtmd PROPERTIES
|
||||
|
||||
@@ -88,6 +88,7 @@
|
||||
#define KEY_A_ATTN_WINDOW_SIZE "clip.audio.window_size" // mimo-audio-tokenizer: sliding-window radius
|
||||
#define KEY_A_LOCAL_BLOCK_COUNT "clip.audio.local_block_count" // mimo-v2.5: input_local_transformer layer count
|
||||
#define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size
|
||||
#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor"
|
||||
|
||||
//
|
||||
// tensor name constants
|
||||
@@ -338,6 +339,12 @@
|
||||
#define TN_YASA_STAGE_DOWN_CONV "v.stage.%d.down.conv.%s"
|
||||
#define TN_YASA_STAGE_BLK "v.stage.%d.blk.%d.%s.%s"
|
||||
|
||||
// parakeet
|
||||
#define TN_MEL_FILTERS "a.mel_filters"
|
||||
#define TN_WINDOW "a.window"
|
||||
#define TN_CONV_NORM_MEAN "%s.blk.%d.conv_norm_mean"
|
||||
#define TN_CONV_NORM_VAR "%s.blk.%d.conv_norm_var"
|
||||
|
||||
// align x to upper multiple of n
|
||||
#define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n))
|
||||
|
||||
@@ -392,6 +399,7 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_KIMIK25,
|
||||
PROJECTOR_TYPE_NEMOTRON_V2_VL,
|
||||
PROJECTOR_TYPE_HUNYUANVL,
|
||||
PROJECTOR_TYPE_PARAKEET,
|
||||
PROJECTOR_TYPE_EXAONE4_5,
|
||||
PROJECTOR_TYPE_MINICPMV4_6,
|
||||
PROJECTOR_TYPE_GRANITE_SPEECH,
|
||||
@@ -455,6 +463,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"},
|
||||
{ PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"},
|
||||
{ PROJECTOR_TYPE_MIMO_AUDIO, "mimo_audio"},
|
||||
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
|
||||
};
|
||||
|
||||
static projector_type clip_projector_type_from_string(const std::string & str) {
|
||||
|
||||
+16
-8
@@ -110,6 +110,8 @@ struct clip_hparams {
|
||||
// audio
|
||||
int32_t n_mel_bins = 0; // whisper preprocessor
|
||||
int32_t proj_stack_factor = 0; // ultravox
|
||||
int32_t subsampling_factor = 0; // parakeet
|
||||
|
||||
int32_t audio_chunk_size = 0;
|
||||
int32_t audio_conv_kernel_size = 0;
|
||||
int32_t audio_max_pos_emb = 0;
|
||||
@@ -124,6 +126,10 @@ struct clip_hparams {
|
||||
int32_t audio_window_len = -1;
|
||||
int32_t audio_hop_len = -1;
|
||||
|
||||
// parakeet
|
||||
std::vector<float> mel_filters;
|
||||
std::vector<float> window;
|
||||
|
||||
// mimo-audio-tokenizer: residual vector quantizer
|
||||
int32_t rvq_num_quantizers = 0;
|
||||
std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
|
||||
@@ -245,14 +251,16 @@ struct clip_layer {
|
||||
ggml_tensor * norm_conv_b = nullptr;
|
||||
ggml_tensor * linear_pos_w = nullptr;
|
||||
|
||||
ggml_tensor * conv_norm_w = nullptr;
|
||||
ggml_tensor * conv_norm_b = nullptr;
|
||||
ggml_tensor * conv_dw_w = nullptr;
|
||||
ggml_tensor * conv_dw_b = nullptr;
|
||||
ggml_tensor * conv_pw1_w = nullptr;
|
||||
ggml_tensor * conv_pw1_b = nullptr;
|
||||
ggml_tensor * conv_pw2_w = nullptr;
|
||||
ggml_tensor * conv_pw2_b = nullptr;
|
||||
ggml_tensor * conv_norm_w = nullptr;
|
||||
ggml_tensor * conv_norm_b = nullptr;
|
||||
ggml_tensor * conv_norm_mean = nullptr; // parakeet
|
||||
ggml_tensor * conv_norm_var = nullptr; // parakeet
|
||||
ggml_tensor * conv_dw_w = nullptr;
|
||||
ggml_tensor * conv_dw_b = nullptr;
|
||||
ggml_tensor * conv_pw1_w = nullptr;
|
||||
ggml_tensor * conv_pw1_b = nullptr;
|
||||
ggml_tensor * conv_pw2_w = nullptr;
|
||||
ggml_tensor * conv_pw2_b = nullptr;
|
||||
|
||||
// gemma4 audio conformer per-layer
|
||||
ggml_tensor * attn_pre_norm_w = nullptr;
|
||||
|
||||
+204
-6
@@ -1033,6 +1033,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
{
|
||||
builder = std::make_unique<clip_graph_yasa2>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_parakeet>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GRANITE4_VISION:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_granite4_vision>(ctx, img);
|
||||
@@ -1356,6 +1360,20 @@ struct clip_model_loader {
|
||||
{
|
||||
get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor);
|
||||
GGML_ASSERT(hparams.subsampling_factor == 8 &&
|
||||
"subsampling_factor must match the conv strides in clip_graph_parakeet::build()");
|
||||
get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size);
|
||||
GGML_ASSERT(hparams.audio_conv_kernel_size > 0 && hparams.audio_conv_kernel_size % 2 == 1 &&
|
||||
"audio_conv_kernel_size must be a positive odd integer");
|
||||
hparams.audio_chunk_len = 0;
|
||||
hparams.audio_sample_rate = 16000;
|
||||
hparams.audio_n_fft = 512;
|
||||
hparams.audio_window_len = 400;
|
||||
hparams.audio_hop_len = 160;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_IDEFICS3:
|
||||
{
|
||||
// use default llava-uhd preprocessing params
|
||||
@@ -1893,16 +1911,46 @@ struct clip_model_loader {
|
||||
return cur;
|
||||
};
|
||||
|
||||
auto get_scalar = [&](const std::string & name, float default_val) {
|
||||
auto get_vector = [&](const std::string & name) {
|
||||
std::vector<float> result;
|
||||
auto it = tensor_offset.find(name);
|
||||
if (it == tensor_offset.end()) {
|
||||
return result;
|
||||
}
|
||||
|
||||
const int64_t idx = gguf_find_tensor(ctx_gguf.get(), name.c_str());
|
||||
if (idx < 0) {
|
||||
throw std::runtime_error(string_format("%s: failed to find tensor %s\n", __func__, name.c_str()));
|
||||
}
|
||||
|
||||
if (const auto type = gguf_get_tensor_type(ctx_gguf.get(), idx); type != GGML_TYPE_F32) {
|
||||
throw std::runtime_error(string_format("%s: %s must be %s, was %s\n", __func__,
|
||||
name.c_str(), ggml_type_name(GGML_TYPE_F32), ggml_type_name(type)));
|
||||
}
|
||||
|
||||
const size_t n_bytes = gguf_get_tensor_size(ctx_gguf.get(), idx);
|
||||
if (n_bytes == 0) {
|
||||
throw std::runtime_error(string_format("%s: tensor %s is empty\n", __func__, name.c_str()));
|
||||
}
|
||||
|
||||
const size_t n_elems = n_bytes / sizeof(float);
|
||||
result.resize(n_elems);
|
||||
fin.seekg(it->second, std::ios::beg);
|
||||
fin.read(reinterpret_cast<char*>(result.data()), n_bytes);
|
||||
return result;
|
||||
};
|
||||
|
||||
auto get_scalar = [&](const std::string & name, float default_val) {
|
||||
auto v = get_vector(name);
|
||||
if (v.empty()) {
|
||||
return default_val;
|
||||
}
|
||||
size_t offset = it->second;
|
||||
fin.seekg(offset, std::ios::beg);
|
||||
float value;
|
||||
fin.read(reinterpret_cast<char*>(&value), sizeof(float));
|
||||
return value;
|
||||
if (v.size() != 1) {
|
||||
throw std::runtime_error(string_format("%s: expected scalar tensor '%s' but got %d elements\n",
|
||||
__func__, name.c_str(), (int) v.size()));
|
||||
}
|
||||
|
||||
return v[0];
|
||||
};
|
||||
|
||||
model.class_embedding = get_tensor(TN_CLASS_EMBD, false);
|
||||
@@ -2800,6 +2848,68 @@ struct clip_model_loader {
|
||||
layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias"));
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
|
||||
hparams.mel_filters = get_vector(TN_MEL_FILTERS);
|
||||
hparams.window = get_vector(TN_WINDOW);
|
||||
|
||||
// Subsampling layers (conv1d)
|
||||
for (int i : {0, 2, 3, 5, 6}) {
|
||||
model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight"));
|
||||
model.pre_encode_conv_X_b[i] = get_tensor(string_format(TN_CONV1D, i, "bias"));
|
||||
}
|
||||
model.pre_encode_out_w = get_tensor(string_format(TN_PRE_ENCODE_OUT, "weight"));
|
||||
model.pre_encode_out_b = get_tensor(string_format(TN_PRE_ENCODE_OUT, "bias"));
|
||||
|
||||
// Projection layers
|
||||
model.mm_norm_pre_w = get_tensor(string_format(TN_MM_NORM_PRE, "weight"), false);
|
||||
model.mm_0_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"), false);
|
||||
model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"), false);
|
||||
|
||||
// Encoder layers
|
||||
for (int il = 0; il < hparams.n_layer; ++il) {
|
||||
auto & layer = model.layers[il];
|
||||
|
||||
// Attention (from shared above)
|
||||
|
||||
// Relative position encoding
|
||||
layer.linear_pos_w = get_tensor(string_format(TN_LINEAR_POS, prefix, il, "weight"));
|
||||
layer.pos_bias_u = get_tensor(string_format(TN_POS_BIAS_U, prefix, il));
|
||||
layer.pos_bias_v = get_tensor(string_format(TN_POS_BIAS_V, prefix, il));
|
||||
|
||||
// Convolution module
|
||||
layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, prefix, il, "weight"));
|
||||
layer.conv_pw1_b = get_tensor(string_format(TN_CONV_PW1, prefix, il, "bias"), false);
|
||||
layer.conv_dw_w = get_tensor(string_format(TN_CONV_DW, prefix, il, "weight"));
|
||||
layer.conv_dw_b = get_tensor(string_format(TN_CONV_DW, prefix, il, "bias"), false);
|
||||
layer.conv_norm_w = get_tensor(string_format(TN_CONV_NORM, prefix, il, "weight"));
|
||||
layer.conv_norm_b = get_tensor(string_format(TN_CONV_NORM, prefix, il, "bias"));
|
||||
layer.conv_norm_mean = get_tensor(string_format(TN_CONV_NORM_MEAN, prefix, il));
|
||||
layer.conv_norm_var = get_tensor(string_format(TN_CONV_NORM_VAR, prefix, il));
|
||||
layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, prefix, il, "weight"));
|
||||
layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias"), false);
|
||||
|
||||
// Feed-forward networks
|
||||
layer.ff_norm_w = get_tensor(string_format(TN_FFN_NORM, prefix, il, "weight"));
|
||||
layer.ff_norm_b = get_tensor(string_format(TN_FFN_NORM, prefix, il, "bias"));
|
||||
|
||||
layer.ff_norm_1_w = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "weight"));
|
||||
layer.ff_norm_1_b = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "bias"));
|
||||
layer.ff_up_1_w = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "weight"));
|
||||
layer.ff_up_1_b = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "bias"), false);
|
||||
layer.ff_down_1_w = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "weight"));
|
||||
layer.ff_down_1_b = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "bias"), false);
|
||||
|
||||
// Layer norms
|
||||
layer.norm_conv_w = get_tensor(string_format(TN_NORM_CONV, prefix, il, "weight"));
|
||||
layer.norm_conv_b = get_tensor(string_format(TN_NORM_CONV, prefix, il, "bias"));
|
||||
}
|
||||
|
||||
model.mm_model_mlp_1_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 0, "weight"));
|
||||
model.mm_model_mlp_2_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 1, "weight"));
|
||||
model.mm_model_mlp_3_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 3, "weight"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GRANITE_SPEECH:
|
||||
{
|
||||
model.inp_proj_w = get_tensor(string_format(TN_INP_PROJ, "weight"));
|
||||
@@ -3645,6 +3755,10 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
}
|
||||
n_patches = n;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
n_patches = (img->nx() + (params.subsampling_factor - 1)) / params.subsampling_factor;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GEMMA4UA:
|
||||
{
|
||||
n_patches = img->nx(); // no downsampling: one token per raw waveform frame
|
||||
@@ -4558,6 +4672,88 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
}
|
||||
set_input_f32("pos_emb", pos_emb);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
GGML_ASSERT(imgs.entries.size() == 1);
|
||||
struct ggml_tensor * attn_mask = ggml_graph_get_tensor(gf, "attn_mask");
|
||||
const int n_q = attn_mask->ne[1];
|
||||
const int n_k = attn_mask->ne[0];
|
||||
const int n_frames = imgs.entries.front().nx();
|
||||
const int n_tokens_real = (n_frames + hparams.subsampling_factor-1) / hparams.subsampling_factor;
|
||||
const float mask_value = -1e30f;
|
||||
|
||||
std::vector<float> mask_data(n_q * n_k);
|
||||
if (n_k == n_q) {
|
||||
// full attention: mask keys that are padding
|
||||
for (int q = 0; q < n_q; ++q) {
|
||||
for (int k = 0; k < n_k; ++k) {
|
||||
mask_data[q * n_k + k] = (k >= n_tokens_real) ? mask_value : 0.0f;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// local attention: mask keys outside the valid window
|
||||
const int att_left = n_k / 2;
|
||||
for (int q = 0; q < n_q; ++q) {
|
||||
for (int k = 0; k < n_k; ++k) {
|
||||
const int key = q - att_left + k;
|
||||
mask_data[q * n_k + k] = (key >= 0 && key < n_tokens_real) ? 0.0f : mask_value;
|
||||
}
|
||||
}
|
||||
}
|
||||
set_input_f32(attn_mask->name, mask_data);
|
||||
|
||||
// local attention skew mask: zeroes out the probs that were
|
||||
// computed for keys outside the valid sliding window.
|
||||
if (struct ggml_tensor * local_mask = ggml_graph_get_tensor(gf, "local_mask")) {
|
||||
const int lm_k = local_mask->ne[0];
|
||||
const int lm_q = local_mask->ne[1];
|
||||
const int window_size = lm_k - lm_q + 1;
|
||||
std::vector<float> lm_data(lm_q * lm_k);
|
||||
for (int q = 0; q < lm_q; ++q) {
|
||||
for (int k = 0; k < lm_k; ++k) {
|
||||
const int rel = k - q;
|
||||
lm_data[q * lm_k + k] = (rel >= 0 && rel < window_size) ? 1.0f : 0.0f;
|
||||
}
|
||||
}
|
||||
set_input_f32(local_mask->name, lm_data);
|
||||
}
|
||||
|
||||
// Generate rotation frequencies for relative positional encoding.
|
||||
{
|
||||
const int n_state = hparams.n_embd;
|
||||
const int d_half = n_state / 2;
|
||||
const float log_10000 = logf(10000.0f);
|
||||
std::vector<float> freqs(d_half);
|
||||
for (int k = 0; k < d_half; ++k) {
|
||||
freqs[k] = expf(-(float(k * 2) * log_10000 / float(n_state)));
|
||||
}
|
||||
set_input_f32("pos_freqs", freqs);
|
||||
}
|
||||
|
||||
// Generate relative positional distance values which scaled by
|
||||
// the frequency to produce the angles for sin/cos.
|
||||
{
|
||||
// window_size is only known after graph construction since it depends on
|
||||
// n_time from the conv output, so we read it back from the graph tensor.
|
||||
struct ggml_tensor * rel_pos = ggml_graph_get_tensor(gf, "rel_positions");
|
||||
const int window_size = rel_pos->ne[1];
|
||||
std::vector<float> pos(window_size);
|
||||
// local attention: window is fixed at [att_left, att_right]
|
||||
// full attention: window covers the full sequence, centered
|
||||
if (ggml_graph_get_tensor(gf, "local_mask")) {
|
||||
const int att_left = window_size / 2;
|
||||
for (int t = 0; t < window_size; ++t) {
|
||||
pos[t] = float(att_left - t);
|
||||
}
|
||||
} else {
|
||||
const int n_time = (window_size + 1) / 2;
|
||||
for (int t = 0; t < window_size; ++t) {
|
||||
pos[t] = float(n_time - 1 - t);
|
||||
}
|
||||
}
|
||||
set_input_f32(rel_pos->name, pos);
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GRANITE_SPEECH:
|
||||
{
|
||||
const int context_size = ctx->model.hparams.audio_chunk_size;
|
||||
@@ -4841,6 +5037,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.mm_ffn_down_w->ne[1];
|
||||
case PROJECTOR_TYPE_MIMO_AUDIO:
|
||||
return ctx->model.mm_2_w->ne[1];
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
return ctx->model.mm_1_w->ne[1];
|
||||
default:
|
||||
GGML_ABORT("Unknown projector type");
|
||||
}
|
||||
|
||||
@@ -222,6 +222,11 @@ struct clip_graph_kimik25 : clip_graph {
|
||||
ggml_tensor * resize_position_embeddings_3d(uint32_t interpolation_mode);
|
||||
};
|
||||
|
||||
struct clip_graph_parakeet : clip_graph {
|
||||
clip_graph_parakeet(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
};
|
||||
|
||||
struct clip_graph_exaone4_5 : clip_graph {
|
||||
clip_graph_exaone4_5(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
#include "models.h"
|
||||
|
||||
static constexpr int PARAKEET_LOCAL_ATTN_THRESHOLD = 8192;
|
||||
static constexpr int PARAKEET_LOCAL_ATTN_WINDOW = 128;
|
||||
|
||||
// conv subsampling + conformer encoder
|
||||
ggml_cgraph * clip_graph_parakeet::build() {
|
||||
|
||||
// Conv subsampling
|
||||
ggml_tensor * inp = build_inp_raw(1);
|
||||
inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp));
|
||||
|
||||
// [freq, time, channels, batch]
|
||||
ggml_tensor * cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[0], inp, 2, 2, 1, 1, 1, 1);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[0]);
|
||||
cb(cur, "pre_conv_0", -1);
|
||||
|
||||
cur = ggml_relu(ctx0, cur);
|
||||
cb(cur, "pre_conv_0_relu", -1);
|
||||
|
||||
// [freq, time, channels, batch]
|
||||
cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[2], cur, 2, 2, 1, 1, 1, 1);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[2]);
|
||||
cb(cur, "pre_conv_2", -1);
|
||||
|
||||
// [freq, time, channels, batch]
|
||||
cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]);
|
||||
cb(cur, "pre_conv_3", -1);
|
||||
|
||||
cur = ggml_relu(ctx0, cur);
|
||||
cb(cur, "pre_conv_3_relu", -1);
|
||||
|
||||
// [freq, time, channels, batch]
|
||||
cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[5], cur, 2, 2, 1, 1, 1, 1);
|
||||
cb(cur, "pre_conv_5_direct", -1);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[5]);
|
||||
cb(cur, "pre_conv_5", -1);
|
||||
|
||||
// [freq, time, channels, batch]
|
||||
cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]);
|
||||
cb(cur, "pre_conv_6", -1);
|
||||
|
||||
cur = ggml_relu(ctx0, cur);
|
||||
cb(cur, "pre_conv_6_relu", -1);
|
||||
|
||||
// [freq, time, chan]
|
||||
cur = ggml_permute(ctx0, cur, 0, 2, 1, 3);
|
||||
// [freq, chan, time]
|
||||
cur = ggml_cont(ctx0, cur);
|
||||
|
||||
const int n_freq = cur->ne[0];
|
||||
const int n_chan = cur->ne[1];
|
||||
const int n_frames = cur->ne[2];
|
||||
|
||||
// [freq, time, chan, batch] -> [(freq * chan), time]
|
||||
cur = ggml_reshape_2d(ctx0, cur, n_freq * n_chan, n_frames);
|
||||
|
||||
cur = build_mm(model.pre_encode_out_w, cur);
|
||||
cur = ggml_add(ctx0, cur, model.pre_encode_out_b);
|
||||
|
||||
ggml_set_name(cur, "pre_enc_out");
|
||||
|
||||
// Encoder
|
||||
|
||||
const auto & hparams = model.hparams;
|
||||
const int n_layer = hparams.n_layer;
|
||||
const int n_state = hparams.n_embd;
|
||||
const float fc_factor = 0.5f;
|
||||
|
||||
const int n_time = cur->ne[1];
|
||||
const bool local_attn = n_time > PARAKEET_LOCAL_ATTN_THRESHOLD;
|
||||
const int att_left = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
|
||||
const int att_right = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
|
||||
const int window_size = local_attn ? att_left + att_right + 1 : 2 * n_time - 1;
|
||||
const int d_half = n_state / 2;
|
||||
const int mask_dim = local_attn ? window_size : n_time;
|
||||
|
||||
// mask [key, n_time]
|
||||
struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, mask_dim, n_time);
|
||||
ggml_set_name(attn_mask, "attn_mask");
|
||||
ggml_set_input(attn_mask);
|
||||
|
||||
struct ggml_tensor * local_mask = nullptr;
|
||||
if (local_attn) {
|
||||
const int chunk = att_left + att_right;
|
||||
local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, chunk + window_size - 1, chunk);
|
||||
ggml_set_name(local_mask, "local_mask");
|
||||
ggml_set_input(local_mask);
|
||||
}
|
||||
|
||||
struct ggml_tensor * pos_freqs = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, d_half);
|
||||
ggml_set_name(pos_freqs, "pos_freqs");
|
||||
ggml_set_input(pos_freqs);
|
||||
|
||||
struct ggml_tensor * rel_positions = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, window_size);
|
||||
ggml_set_name(rel_positions, "rel_positions");
|
||||
ggml_set_input(rel_positions);
|
||||
|
||||
struct ggml_tensor * freqs = ggml_repeat_4d(ctx0, pos_freqs, d_half, window_size, 1, 1);
|
||||
struct ggml_tensor * theta = ggml_mul(ctx0, freqs, rel_positions);
|
||||
|
||||
struct ggml_tensor * sin = ggml_reshape_3d(ctx0, ggml_sin(ctx0, theta), 1, d_half, window_size);
|
||||
struct ggml_tensor * cos = ggml_reshape_3d(ctx0, ggml_cos(ctx0, theta), 1, d_half, window_size);
|
||||
struct ggml_tensor * pos_emb = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_concat(ctx0, sin, cos, 0)), n_state, window_size);
|
||||
ggml_set_name(pos_emb, "pos_emb");
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const auto & layer = model.layers[il];
|
||||
// FFN1
|
||||
{
|
||||
struct ggml_tensor * residual = cur;
|
||||
ggml_format_name(cur, "enc_%d_res", il);
|
||||
|
||||
// norm
|
||||
cur = ggml_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_w), layer.ff_norm_b);
|
||||
ggml_format_name(cur, "enc_%d_ffn_norm_1", il);
|
||||
|
||||
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il);
|
||||
ggml_format_name(cur, "enc_%d_ffn_1", il);
|
||||
|
||||
cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, fc_factor));
|
||||
ggml_format_name(cur, "enc_%d_res_ffn", il);
|
||||
}
|
||||
|
||||
// self attention block using relative positional encoding from model.position_embedding.
|
||||
{
|
||||
// [feat, time_frames, 1, 1]
|
||||
struct ggml_tensor * residual = cur;
|
||||
|
||||
cur = ggml_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_1_w), layer.ln_1_b);
|
||||
ggml_format_name(cur, "enc_%d_attn_norm", il);
|
||||
|
||||
const int n_head = hparams.n_head;
|
||||
const int d_head = n_state / n_head;
|
||||
|
||||
// [feat, time_frames, 1, 1]
|
||||
struct ggml_tensor * Q_cur = build_mm(layer.q_w, cur);
|
||||
struct ggml_tensor * K_cur = build_mm(layer.k_w, cur);
|
||||
struct ggml_tensor * V_cur = build_mm(layer.v_w, cur);
|
||||
|
||||
// [d_head, n_heads, n_time, 1]
|
||||
Q_cur = ggml_reshape_3d(ctx0, Q_cur, d_head, n_head, n_time);
|
||||
K_cur = ggml_reshape_3d(ctx0, K_cur, d_head, n_head, n_time);
|
||||
V_cur = ggml_reshape_3d(ctx0, V_cur, d_head, n_head, n_time);
|
||||
|
||||
// [n_state, window_size]
|
||||
struct ggml_tensor * pos = build_mm(layer.linear_pos_w, pos_emb);
|
||||
// [feat, head, window_size, 1]
|
||||
pos = ggml_reshape_3d(ctx0, pos, d_head, n_head, pos_emb->ne[1]);
|
||||
// [feat, window_size, head, 1]
|
||||
pos = ggml_cont(ctx0, ggml_permute(ctx0, pos, 0, 2, 1, 3));
|
||||
ggml_format_name(pos, "enc_%d_attn_pos", il);
|
||||
|
||||
if (local_attn) {
|
||||
const int chunk = att_left + att_right;
|
||||
const int n_group = (n_time + chunk - 1) / chunk;
|
||||
const int n_time_padded = n_group * chunk;
|
||||
const int n_kv_chunk = chunk + window_size - 1;
|
||||
const int n_kv_dense = n_kv_chunk * n_group;
|
||||
const bool need_padding = n_time_padded > n_time;
|
||||
|
||||
Q_cur = ggml_cont(ctx0, ggml_permute(ctx0, Q_cur, 0, 2, 1, 3));
|
||||
K_cur = ggml_cont(ctx0, ggml_permute(ctx0, K_cur, 0, 2, 1, 3));
|
||||
V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 0, 2, 1, 3));
|
||||
|
||||
// content bias
|
||||
struct ggml_tensor * bias_u = ggml_reshape_3d(ctx0, layer.pos_bias_u, d_head, 1, n_head);
|
||||
struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, bias_u);
|
||||
|
||||
// position bias
|
||||
struct ggml_tensor * bias_v = ggml_reshape_3d(ctx0, layer.pos_bias_v, d_head, 1, n_head);
|
||||
struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, bias_v);
|
||||
|
||||
// right pad the time dimension
|
||||
struct ggml_tensor * Q_u_padded = need_padding ?
|
||||
ggml_pad_ext(ctx0, Q_u, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : Q_u;
|
||||
Q_u_padded = ggml_reshape_4d(ctx0, Q_u_padded, d_head, chunk, n_group, n_head);
|
||||
|
||||
// pad front and back for the first and last time frames
|
||||
struct ggml_tensor * K_padded = ggml_pad_ext(ctx0, K_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);
|
||||
if (n_kv_dense > K_padded->ne[1]) {
|
||||
K_padded = ggml_pad_ext(ctx0, K_padded, 0, 0, 0, n_kv_dense - K_padded->ne[1], 0, 0, 0, 0);
|
||||
}
|
||||
|
||||
// sliding window view: each group spans n_kv_chunk keys but steps by chunk
|
||||
struct ggml_tensor * K_chunk = ggml_view_4d(ctx0, K_padded,
|
||||
d_head, n_kv_chunk, n_group, n_head,
|
||||
K_padded->nb[1],
|
||||
(size_t) chunk * K_padded->nb[1],
|
||||
K_padded->nb[2],
|
||||
0);
|
||||
K_chunk = ggml_cont(ctx0, K_chunk);
|
||||
|
||||
struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_chunk, Q_u_padded);
|
||||
|
||||
// trim the dense output down to window_size scores per query
|
||||
content_scores = ggml_view_4d(ctx0, content_scores,
|
||||
window_size, chunk, n_group, n_head,
|
||||
(size_t) (chunk + window_size) * content_scores->nb[0],
|
||||
content_scores->nb[2],
|
||||
content_scores->nb[3],
|
||||
0);
|
||||
content_scores = ggml_cont(ctx0, content_scores);
|
||||
|
||||
// ungroup: [window_size, n_time_padded, n_head]
|
||||
content_scores = ggml_reshape_3d(ctx0, content_scores, window_size, n_time_padded, n_head);
|
||||
if (need_padding) {
|
||||
content_scores = ggml_view_3d(ctx0, content_scores,
|
||||
window_size, n_time, n_head,
|
||||
content_scores->nb[1],
|
||||
content_scores->nb[2],
|
||||
0);
|
||||
}
|
||||
|
||||
// Q_v: [d_head, time, head]
|
||||
Q_v = ggml_cont(ctx0, ggml_permute(ctx0, Q_v, 0, 2, 1, 3));
|
||||
struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);
|
||||
|
||||
struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);
|
||||
attn_scores = ggml_soft_max_ext(ctx0, attn_scores, attn_mask, 1.0f / std::sqrt(d_head), 0.0f);
|
||||
ggml_format_name(attn_scores, "enc_%d_attn_probs", il);
|
||||
|
||||
// expand probs back to n_kv_chunk width for the V matmul
|
||||
struct ggml_tensor * probs_padded = need_padding ?
|
||||
ggml_pad_ext(ctx0, attn_scores, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : attn_scores;
|
||||
|
||||
probs_padded = ggml_reshape_4d(ctx0, probs_padded, window_size, chunk, n_group, n_head);
|
||||
probs_padded = ggml_pad_ext(ctx0, probs_padded, 0, chunk, 0, 0, 0, 0, 0, 0);
|
||||
probs_padded = ggml_view_4d(ctx0, probs_padded,
|
||||
n_kv_chunk, chunk, n_group, n_head,
|
||||
(size_t) n_kv_chunk * probs_padded->nb[0],
|
||||
probs_padded->nb[2],
|
||||
probs_padded->nb[3],
|
||||
0);
|
||||
probs_padded = ggml_cont(ctx0, probs_padded);
|
||||
probs_padded = ggml_mul(ctx0, probs_padded, local_mask);
|
||||
|
||||
struct ggml_tensor * V_padded = ggml_pad_ext(ctx0, V_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);
|
||||
if (n_kv_dense > V_padded->ne[1]) {
|
||||
V_padded = ggml_pad_ext(ctx0, V_padded, 0, 0, 0, n_kv_dense - V_padded->ne[1], 0, 0, 0, 0);
|
||||
}
|
||||
V_padded = ggml_cont(ctx0, ggml_transpose(ctx0, V_padded));
|
||||
|
||||
struct ggml_tensor * V_chunk = ggml_view_4d(ctx0, V_padded,
|
||||
n_kv_chunk, d_head, n_group, n_head,
|
||||
V_padded->nb[1],
|
||||
(size_t) chunk * V_padded->nb[0],
|
||||
V_padded->nb[2],
|
||||
0);
|
||||
V_chunk = ggml_cont(ctx0, V_chunk);
|
||||
|
||||
cur = ggml_mul_mat(ctx0, V_chunk, probs_padded);
|
||||
cur = ggml_reshape_3d(ctx0, cur, d_head, n_time_padded, n_head);
|
||||
if (need_padding) {
|
||||
cur = ggml_view_3d(ctx0, cur, d_head, n_time, n_head, cur->nb[1], cur->nb[2], 0);
|
||||
}
|
||||
cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3));
|
||||
cur = ggml_reshape_2d(ctx0, cur, n_state, n_time);
|
||||
cur = build_mm(layer.o_w, cur);
|
||||
} else {
|
||||
// full attention
|
||||
struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, layer.pos_bias_u);
|
||||
ggml_format_name(Q_u, "enc_%d_attn_q_u", il);
|
||||
|
||||
struct ggml_tensor * K_prep = ggml_permute(ctx0, K_cur, 0, 2, 1, 3);
|
||||
struct ggml_tensor * Q_prep = ggml_permute(ctx0, Q_u, 0, 2, 1, 3);
|
||||
struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_prep, Q_prep);
|
||||
ggml_format_name(content_scores, "enc_%d_attn_content_scores", il);
|
||||
|
||||
struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, layer.pos_bias_v);
|
||||
ggml_format_name(Q_v, "enc_%d_attn_q_v", il);
|
||||
|
||||
Q_v = ggml_permute(ctx0, Q_v, 0, 2, 1, 3);
|
||||
Q_v = ggml_cont(ctx0, Q_v);
|
||||
ggml_format_name(Q_v, "enc_%d_attn_q_v_perm", il);
|
||||
|
||||
struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);
|
||||
ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos", il);
|
||||
|
||||
// Relative positional shift
|
||||
{
|
||||
const auto pos_window = rel_pos_scores->ne[0];
|
||||
const auto n_frame = rel_pos_scores->ne[1];
|
||||
const auto n_head = rel_pos_scores->ne[2];
|
||||
|
||||
rel_pos_scores = ggml_pad(ctx0, rel_pos_scores, 1, 0, 0, 0);
|
||||
rel_pos_scores = ggml_roll(ctx0, rel_pos_scores, 1, 0, 0, 0);
|
||||
|
||||
rel_pos_scores = ggml_reshape_3d(ctx0, rel_pos_scores, n_frame, pos_window + 1, n_head);
|
||||
rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
|
||||
ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_reshaped", il);
|
||||
|
||||
int center = pos_window / 2;
|
||||
size_t offset = rel_pos_scores->nb[0] * (center+1);
|
||||
|
||||
rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
|
||||
n_frame, pos_window, n_head,
|
||||
(pos_window) * 4,
|
||||
rel_pos_scores->nb[2],
|
||||
offset);
|
||||
rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
|
||||
ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted", il);
|
||||
|
||||
rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
|
||||
content_scores->ne[0],
|
||||
content_scores->ne[1],
|
||||
rel_pos_scores->ne[2],
|
||||
rel_pos_scores->nb[1],
|
||||
rel_pos_scores->nb[2],
|
||||
0);
|
||||
rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
|
||||
ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted_view", il);
|
||||
}
|
||||
|
||||
struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);
|
||||
ggml_format_name(attn_scores, "enc_%d_attn_scores", il);
|
||||
attn_scores = ggml_scale(ctx0, attn_scores, 1.0f / std::sqrt(d_head));
|
||||
attn_scores = ggml_add(ctx0, attn_scores, attn_mask);
|
||||
ggml_format_name(attn_scores, "enc_%d_attn_scores_scaled", il);
|
||||
|
||||
struct ggml_tensor * probs = ggml_soft_max(ctx0, attn_scores);
|
||||
ggml_format_name(probs, "enc_%d_attn_probs", il);
|
||||
|
||||
V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 1, 2, 0, 3));
|
||||
ggml_format_name(V_cur, "enc_%d_attn_v_cur", il);
|
||||
cur = ggml_mul_mat(ctx0, probs, V_cur);
|
||||
ggml_format_name(cur, "enc_%d_attn_inp", il);
|
||||
|
||||
cur = ggml_permute(ctx0, cur, 2, 0, 1, 3);
|
||||
cur = ggml_cont_2d(ctx0, cur, n_state, n_time);
|
||||
cur = build_mm(layer.o_w, cur);
|
||||
}
|
||||
ggml_format_name(cur, "enc_%d_attn_out", il);
|
||||
|
||||
cur = ggml_add(ctx0, residual, cur);
|
||||
ggml_format_name(cur, "enc_%d_attn_res", il);
|
||||
}
|
||||
|
||||
// Convolution
|
||||
{
|
||||
struct ggml_tensor * residual = cur;
|
||||
ggml_format_name(cur, "enc_%d_residual_conv", il);
|
||||
|
||||
cur = ggml_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_conv_w), layer.norm_conv_b);
|
||||
ggml_format_name(cur, "enc_%d_norm_conv", il);
|
||||
|
||||
// pointwise 1d convolution:
|
||||
cur = build_mm(layer.conv_pw1_w, cur);
|
||||
ggml_format_name(cur, "enc_%d_conv_pw1", il);
|
||||
|
||||
{
|
||||
int64_t d = cur->ne[0] / 2;
|
||||
struct ggml_tensor * signal = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], 0);
|
||||
struct ggml_tensor * gate = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], d * cur->nb[0]);
|
||||
|
||||
cur = ggml_mul(ctx0, signal, ggml_sigmoid(ctx0, gate));
|
||||
ggml_format_name(cur, "enc_%d_conv_glu", il);
|
||||
}
|
||||
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
|
||||
// use ggml_ssm_conv for f32 precision
|
||||
const int dw_pad = (hparams.audio_conv_kernel_size - 1) / 2;
|
||||
cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
|
||||
cur = ggml_roll(ctx0, cur, dw_pad, 0, 0, 0);
|
||||
cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
|
||||
ggml_format_name(cur, "enc_%d_conv_dw_pad", il);
|
||||
|
||||
cur = ggml_ssm_conv(ctx0, cur, layer.conv_dw_w);
|
||||
ggml_format_name(cur, "enc_%d_conv_1d_dw", il);
|
||||
|
||||
cur = ggml_sub(ctx0, cur, layer.conv_norm_mean);
|
||||
struct ggml_tensor * std = ggml_sqrt(ctx0, layer.conv_norm_var);
|
||||
cur = ggml_div(ctx0, cur, std);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.conv_norm_w), layer.conv_norm_b);
|
||||
ggml_format_name(cur, "enc_%d_conv_bn", il);
|
||||
|
||||
cur = ggml_silu(ctx0, cur);
|
||||
ggml_format_name(cur, "enc_%d_conv_silu", il);
|
||||
|
||||
cur = build_mm(layer.conv_pw2_w, cur);
|
||||
ggml_format_name(cur, "enc_%d_conv_pw2", il);
|
||||
|
||||
cur = ggml_add(ctx0, residual, cur);
|
||||
ggml_format_name(cur, "enc_%d_conv_res", il);
|
||||
}
|
||||
|
||||
// FFN2
|
||||
{
|
||||
struct ggml_tensor * residual = cur;
|
||||
cur = ggml_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_1_w), layer.ff_norm_1_b);
|
||||
ggml_format_name(cur, "enc_%d_ffn_norm_2", il);
|
||||
|
||||
cur = build_ffn(cur, layer.ff_up_1_w, nullptr, nullptr, nullptr, layer.ff_down_1_w, nullptr, FFN_SILU, il);
|
||||
cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, 0.5));
|
||||
ggml_format_name(cur, "enc_%d_ffn_res", il);
|
||||
}
|
||||
|
||||
cur = ggml_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_2_w), layer.ln_2_b);
|
||||
}
|
||||
|
||||
cb(cur, "encoder_out", -1);
|
||||
|
||||
cur = ggml_rms_norm(ctx0, cur, 1e-6);
|
||||
cur = ggml_mul(ctx0, cur, model.mm_norm_pre_w);
|
||||
cb(cur, "sound_projection.norm", -1);
|
||||
|
||||
cur = build_ffn(cur, model.mm_0_w, model.mm_0_b, nullptr, nullptr, model.mm_1_w, model.mm_1_b, FFN_RELU_SQR, -1);
|
||||
cb(cur, "projected", -1);
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
return gf;
|
||||
}
|
||||
@@ -1022,6 +1022,209 @@ bool mtmd_audio_preprocessor_gemma4a::preprocess(const float * s
|
||||
}
|
||||
|
||||
//
|
||||
// mtmd_audio_preprocessor_parakeet implementation
|
||||
//
|
||||
|
||||
void mtmd_audio_preprocessor_parakeet::worker_thread(
|
||||
int ith,
|
||||
const float * window_func,
|
||||
int window_size,
|
||||
const std::vector<float> & samples,
|
||||
int n_samples,
|
||||
int frame_size,
|
||||
int frame_step,
|
||||
int n_threads,
|
||||
int n_fft_bins,
|
||||
const mtmd_audio_cache & cache,
|
||||
mtmd_audio_mel & mel) {
|
||||
std::vector<float> fft_in(frame_size * 2, 0.0);
|
||||
std::vector<float> fft_out(frame_size * 2 * 2 * 2);
|
||||
|
||||
int n_fb = n_fft_bins;
|
||||
int i = ith;
|
||||
|
||||
GGML_ASSERT(n_fb == 1 + (frame_size / 2));
|
||||
|
||||
const double eps = 5.960464477539063e-08;
|
||||
|
||||
for (; i < std::min(n_samples / frame_step + 1, (int) mel.n_len); i += n_threads) {
|
||||
const int offset = i * frame_step;
|
||||
const int window_pad_left = (frame_size - window_size) / 2;
|
||||
|
||||
// Zero-pad left.
|
||||
std::fill(fft_in.begin(), fft_in.begin() + window_pad_left, 0.0f);
|
||||
|
||||
// Apply windowed samples in the center.
|
||||
const int n_to_process = std::min({window_size, n_samples - offset});
|
||||
for (int j = 0; j < n_to_process; j++) {
|
||||
fft_in[window_pad_left + j] = window_func[j] * samples[offset + window_pad_left + j];
|
||||
}
|
||||
|
||||
// Zero-pad right.
|
||||
std::fill(fft_in.begin() + window_pad_left + n_to_process, fft_in.begin() + frame_size, 0.0f);
|
||||
|
||||
// FFT.
|
||||
fft(cache, fft_in.data(), frame_size, fft_out.data());
|
||||
|
||||
// Calculate modulus^2 of complex numbers.
|
||||
for (int j = 0; j < n_fb; j++) {
|
||||
fft_out[j] = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]);
|
||||
}
|
||||
|
||||
// mel spectrogram.
|
||||
for (int j = 0; j < mel.n_mel; j++) {
|
||||
double sum = 0.0;
|
||||
int k = 0;
|
||||
for (k = 0; k < n_fb - 3; k += 4) {
|
||||
sum +=
|
||||
fft_out[k + 0] * cache.filters.data[j * n_fb + k + 0] +
|
||||
fft_out[k + 1] * cache.filters.data[j * n_fb + k + 1] +
|
||||
fft_out[k + 2] * cache.filters.data[j * n_fb + k + 2] +
|
||||
fft_out[k + 3] * cache.filters.data[j * n_fb + k + 3];
|
||||
}
|
||||
for (; k < n_fb; k++) {
|
||||
sum += fft_out[k] * cache.filters.data[j * n_fb + k];
|
||||
}
|
||||
mel.data[j * mel.n_len + i] = std::log(sum + eps);
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise fft_out are all zero.
|
||||
const double empty_sum = std::log(eps);
|
||||
for (; i < mel.n_len; i += n_threads) {
|
||||
for (int j = 0; j < mel.n_mel; j++) {
|
||||
mel.data[j * mel.n_len + i] = empty_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void mtmd_audio_preprocessor_parakeet::initialize() {
|
||||
cache.fill_sin_cos_table(hparams.audio_n_fft);
|
||||
|
||||
const size_t n_fft = hparams.audio_n_fft / 2 + 1;
|
||||
GGML_ASSERT(hparams.mel_filters.size() == (size_t)hparams.n_mel_bins * n_fft);
|
||||
cache.filters.n_mel = hparams.n_mel_bins;
|
||||
cache.filters.n_fft = n_fft;
|
||||
cache.filters.data = hparams.mel_filters;
|
||||
|
||||
GGML_ASSERT(hparams.window.size() == (size_t)hparams.audio_window_len);
|
||||
GGML_ASSERT(hparams.window.size() <= (size_t) hparams.audio_n_fft);
|
||||
cache.hann_window = hparams.window;
|
||||
}
|
||||
|
||||
bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
|
||||
size_t n_samples_in,
|
||||
std::vector<mtmd_audio_mel> & output) {
|
||||
if (n_samples_in == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
filter_params params;
|
||||
params.n_mel = hparams.n_mel_bins;
|
||||
params.n_fft_bins = 1 + (hparams.audio_n_fft / 2);
|
||||
params.hann_window_size = hparams.audio_window_len;
|
||||
params.hop_length = hparams.audio_hop_len;
|
||||
params.sample_rate = hparams.audio_sample_rate;
|
||||
|
||||
GGML_ASSERT(!cache.sin_vals.empty());
|
||||
GGML_ASSERT(!cache.cos_vals.empty());
|
||||
GGML_ASSERT(!cache.filters.data.empty());
|
||||
|
||||
const float * window_func = cache.hann_window.data();
|
||||
const int window_size = params.hann_window_size;
|
||||
const int frame_size = (params.n_fft_bins - 1) * 2;
|
||||
const int frame_step = params.hop_length;
|
||||
|
||||
// Apply preemphasis filter (high-pass): x[i] = x[i] - 0.97 * x[i-1]
|
||||
std::vector<float> samples_preprocessed(samples, samples + n_samples_in);
|
||||
{
|
||||
const float preemph = 0.97f;
|
||||
for (int i = n_samples_in - 1; i > 0; i--) {
|
||||
samples_preprocessed[i] = samples_preprocessed[i] - preemph * samples_preprocessed[i - 1];
|
||||
}
|
||||
}
|
||||
|
||||
// Parakeet uses centered constant padding
|
||||
const size_t pad = (size_t)(frame_size / 2);
|
||||
std::vector<float> samples_padded(n_samples_in + 2 * pad, 0.0f);
|
||||
std::copy(samples_preprocessed.begin(), samples_preprocessed.end(), samples_padded.begin() + pad);
|
||||
|
||||
mtmd_audio_mel out_full;
|
||||
out_full.n_mel = params.n_mel;
|
||||
out_full.n_len = (samples_padded.size() - frame_size) / frame_step + 1;
|
||||
out_full.n_len_org = out_full.n_len;
|
||||
out_full.data.resize(out_full.n_mel * out_full.n_len);
|
||||
|
||||
const int n_threads = 4;
|
||||
std::vector<std::thread> workers(n_threads - 1);
|
||||
for (int iw = 0; iw < n_threads - 1; ++iw) {
|
||||
workers[iw] = std::thread(
|
||||
worker_thread, iw + 1,
|
||||
window_func,
|
||||
window_size,
|
||||
std::cref(samples_padded),
|
||||
samples_padded.size(),
|
||||
frame_size,
|
||||
frame_step,
|
||||
n_threads,
|
||||
params.n_fft_bins,
|
||||
std::cref(cache),
|
||||
std::ref(out_full)
|
||||
);
|
||||
}
|
||||
|
||||
worker_thread(0,
|
||||
window_func,
|
||||
window_size,
|
||||
samples_padded,
|
||||
samples_padded.size(),
|
||||
frame_size,
|
||||
frame_step,
|
||||
n_threads,
|
||||
params.n_fft_bins,
|
||||
cache,
|
||||
out_full);
|
||||
|
||||
for (int iw = 0; iw < n_threads - 1; ++iw) {
|
||||
workers[iw].join();
|
||||
}
|
||||
|
||||
// Per-feature normalization (only on valid frames)
|
||||
{
|
||||
const double eps = 1e-5;
|
||||
int valid_frames = n_samples_in / frame_step;
|
||||
|
||||
for (int j = 0; j < out_full.n_mel; j++) {
|
||||
double sum = 0.0;
|
||||
double sq_diff_sum = 0.0;
|
||||
|
||||
// Calculate Mean ONLY on valid audio frames
|
||||
for (int i = 0; i < valid_frames; i++) {
|
||||
sum += (double)out_full.data[j * out_full.n_len + i];
|
||||
}
|
||||
double mean = sum / valid_frames;
|
||||
|
||||
// Calculate Variance ONLY on valid audio frames
|
||||
for (int i = 0; i < valid_frames; i++) {
|
||||
double diff = (double)out_full.data[j * out_full.n_len + i] - mean;
|
||||
sq_diff_sum += diff * diff;
|
||||
}
|
||||
|
||||
double std_dev = std::sqrt(sq_diff_sum / (valid_frames - 1.0));
|
||||
double denominator = std_dev + eps;
|
||||
|
||||
// Apply to ALL frames (including the padded ones)
|
||||
for (int i = 0; i < out_full.n_len; i++) {
|
||||
out_full.data[j * out_full.n_len + i] = (float)((out_full.data[j * out_full.n_len + i] - mean) / denominator);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output.push_back(std::move(out_full));
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
// mtmd_audio_preprocessor_gemma4ua
|
||||
//
|
||||
|
||||
|
||||
@@ -120,6 +120,21 @@ struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor {
|
||||
mtmd_audio_cache cache;
|
||||
};
|
||||
|
||||
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
|
||||
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
|
||||
void initialize() override;
|
||||
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
|
||||
|
||||
private:
|
||||
mtmd_audio_cache cache;
|
||||
|
||||
static void worker_thread(int ith, const float * window_func, int window_size,
|
||||
const std::vector<float> & samples, int n_samples,
|
||||
int frame_size, int frame_step, int n_threads,
|
||||
int n_fft_bins,
|
||||
const mtmd_audio_cache & cache, mtmd_audio_mel & mel);
|
||||
};
|
||||
|
||||
//
|
||||
// streaming ISTFT - converts spectrogram frames back to audio one frame at a time
|
||||
//
|
||||
|
||||
@@ -724,6 +724,10 @@ struct mtmd_context {
|
||||
aud_end = "<audio|>";
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_gemma4a>(ctx_a);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_parakeet>(ctx_a);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GEMMA4UA:
|
||||
{
|
||||
aud_beg = "<|audio>";
|
||||
|
||||
@@ -164,6 +164,8 @@ struct server_slot {
|
||||
llama_context * ctx_tgt = nullptr;
|
||||
llama_context * ctx_dft = nullptr;
|
||||
|
||||
common_memory mem;
|
||||
|
||||
// multimodal
|
||||
mtmd_context * mctx = nullptr;
|
||||
mtmd::batch_ptr mbatch = nullptr;
|
||||
@@ -253,10 +255,7 @@ struct server_slot {
|
||||
void prompt_clear() {
|
||||
SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size());
|
||||
|
||||
common_context_seq_rm(ctx_tgt, id, -1, -1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm(ctx_dft, id, -1, -1);
|
||||
}
|
||||
mem.seq_rm(id, -1, -1);
|
||||
|
||||
prompt.clear();
|
||||
}
|
||||
@@ -668,13 +667,8 @@ struct server_slot {
|
||||
void copy_state_to(server_slot & other) const {
|
||||
GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT);
|
||||
|
||||
common_context_seq_rm(ctx_tgt, other.id, -1, -1);
|
||||
common_context_seq_cp(ctx_tgt, id, other.id, -1, -1);
|
||||
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm(ctx_dft, other.id, -1, -1);
|
||||
common_context_seq_cp(ctx_dft, id, other.id, -1, -1);
|
||||
}
|
||||
mem.seq_rm(other.id, -1, -1);
|
||||
mem.seq_cp(id, other.id, -1, -1);
|
||||
|
||||
other.n_decoded = n_decoded;
|
||||
other.n_remaining = n_remaining;
|
||||
@@ -1302,6 +1296,7 @@ private:
|
||||
slot.id = i;
|
||||
slot.ctx_tgt = ctx_tgt;
|
||||
slot.ctx_dft = ctx_dft;
|
||||
slot.mem.init(ctx_tgt, ctx_dft);
|
||||
slot.spec = spec.get();
|
||||
slot.n_ctx = n_ctx_slot;
|
||||
|
||||
@@ -2881,13 +2876,8 @@ private:
|
||||
|
||||
SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard);
|
||||
|
||||
common_context_seq_rm (ctx_tgt, slot.id, n_keep , n_keep + n_discard);
|
||||
common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard);
|
||||
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard);
|
||||
common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
|
||||
}
|
||||
slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard);
|
||||
slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
|
||||
|
||||
// add generated tokens to cache
|
||||
// ref: https://github.com/ggml-org/llama.cpp/pull/16818#discussion_r2473269481
|
||||
@@ -2998,7 +2988,9 @@ private:
|
||||
ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
||||
}
|
||||
|
||||
common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1);
|
||||
if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) {
|
||||
GGML_ABORT("failed to remove sequence %d\n", slot.id);
|
||||
}
|
||||
}
|
||||
|
||||
if (!draft.empty()) {
|
||||
@@ -3201,13 +3193,8 @@ private:
|
||||
|
||||
const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c;
|
||||
|
||||
common_context_seq_rm (ctx_tgt, slot.id, head_p, head_c);
|
||||
common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift);
|
||||
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm (ctx_dft, slot.id, head_p, head_c);
|
||||
common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift);
|
||||
}
|
||||
slot.mem.seq_rm (slot.id, head_p, head_c);
|
||||
slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift);
|
||||
|
||||
for (size_t i = 0; i < n_match; i++) {
|
||||
slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]);
|
||||
@@ -3379,10 +3366,7 @@ private:
|
||||
|
||||
SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0);
|
||||
|
||||
common_context_seq_rm(ctx_tgt, slot.id, p0, -1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm(ctx_dft, slot.id, p0, -1);
|
||||
}
|
||||
slot.mem.seq_rm(slot.id, p0, -1);
|
||||
|
||||
// If using an alora, there may be uncached tokens that come
|
||||
// before the invocation sequence. When this happens, the
|
||||
@@ -3837,18 +3821,14 @@ private:
|
||||
|
||||
SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size());
|
||||
|
||||
{
|
||||
ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
||||
|
||||
common_context_seq_rm(slot.ctx_tgt, slot.id, ckpt.pos_max + 1, -1);
|
||||
}
|
||||
ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
||||
|
||||
if (slot.ctx_dft) {
|
||||
ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
||||
|
||||
common_context_seq_rm(slot.ctx_dft, slot.id, ckpt.pos_max + 1, -1);
|
||||
}
|
||||
|
||||
slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
|
||||
|
||||
slot.prompt.tokens.keep_first(ckpt.n_tokens);
|
||||
slot.smpl = std::move(smpl_save);
|
||||
|
||||
@@ -3889,10 +3869,7 @@ private:
|
||||
slot.sampled = ids.back(); // last accepted token
|
||||
SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft);
|
||||
|
||||
common_context_seq_rm(slot.ctx_tgt, slot.id, slot.prompt.tokens.pos_next(), -1);
|
||||
if (slot.ctx_dft) {
|
||||
common_context_seq_rm(slot.ctx_dft, slot.id, slot.prompt.tokens.pos_next(), -1);
|
||||
}
|
||||
slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1);
|
||||
|
||||
for (size_t i = 0; i < ids.size(); ++i) {
|
||||
completion_token_output result;
|
||||
|
||||
+3
-5
@@ -11,8 +11,7 @@
|
||||
classifyToolResult,
|
||||
formatJsonPretty,
|
||||
parseToolResultWithImages,
|
||||
type AgenticSection,
|
||||
type ToolResultLine
|
||||
type AgenticSection
|
||||
} from '$lib/utils';
|
||||
import { getBuiltinToolUi } from '$lib/constants/built-in-tools';
|
||||
import type { DatabaseMessageExtra } from '$lib/types';
|
||||
@@ -29,11 +28,10 @@
|
||||
let { section, open, isStreaming, attachments, onToggle }: Props = $props();
|
||||
|
||||
const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? '');
|
||||
|
||||
const parsedLines: ToolResultLine[] = $derived(
|
||||
const outputKind = $derived(classifyToolResult(section.toolResult));
|
||||
const parsedLines = $derived(
|
||||
section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
|
||||
);
|
||||
const outputKind = $derived(classifyToolResult(section.toolResult));
|
||||
</script>
|
||||
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={null} {title} {onToggle}>
|
||||
|
||||
-1
@@ -15,7 +15,6 @@
|
||||
let { section, open, isStreaming, onToggle }: Props = $props();
|
||||
|
||||
const editFileMeta = $derived(parseEditFileMeta(section));
|
||||
|
||||
const editDiffs = $derived(
|
||||
(editFileMeta?.edits ?? []).map((edit) => computeLineDiff(edit.oldText, edit.newText))
|
||||
);
|
||||
|
||||
+1
-1
@@ -27,7 +27,7 @@
|
||||
const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING);
|
||||
const showSpinner = $derived(isPending || (isStreamingCall && isStreaming));
|
||||
|
||||
const results: SearchResult[] = $derived(extractSearchResults(section.toolResult));
|
||||
const results = $derived(extractSearchResults(section.toolResult));
|
||||
const query = $derived(extractSearchQuery(section.toolArgs));
|
||||
|
||||
// Same icon-resolution chain as ChatMessageToolCallBlockDefault so
|
||||
|
||||
@@ -28,6 +28,18 @@ export const LATEX_MATH_AND_CODE_PATTERN =
|
||||
/** Regex to capture the content of a $$...\\\\...$$ block (display-formula with line-break) */
|
||||
export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/;
|
||||
|
||||
/**
|
||||
* Matches the unescaped `\[...\]` display-math delimiter and surrounding
|
||||
* context so callers can insert line-breaks around the placeholder or convert
|
||||
* to inline when the formula has a non-empty trailing context (e.g. a table
|
||||
* cell that opens with `\[` and closes with content after `\]`).
|
||||
*
|
||||
* group 1: prefix before `\[`
|
||||
* group 2: formula body
|
||||
* group 3: trailing context after `\]`
|
||||
*/
|
||||
export const LATEX_DISPLAY_BLOCK_REGEXP = /([\S].*?)\\\[([\s\S]*?)\\\](.*)/g;
|
||||
|
||||
/**
|
||||
* Cheap gate for `preprocessLaTeX`. Every transformation it performs is triggered
|
||||
* by a `$` (inline/display math, currency escaping) or a backslash escape
|
||||
@@ -36,6 +48,76 @@ export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/;
|
||||
*/
|
||||
export const LATEX_TRIGGER_REGEXP = /[$\\]/;
|
||||
|
||||
/** Inline LaTeX math delimiter (the dollar sign). */
|
||||
export const LATEX_INLINE_DELIMITER = '$';
|
||||
|
||||
/** Display LaTeX math delimiter (paired dollar signs). */
|
||||
export const LATEX_DISPLAY_DELIMITER = '$$';
|
||||
|
||||
/** Matches a single non-whitespace character. */
|
||||
export const LATEX_NON_WHITESPACE_REGEXP = /\S/;
|
||||
|
||||
/** Matches a character that may appear adjacent to `$`, indicating a non-TeX
|
||||
* context such as an identifier (`var$`, `$var`), currency ($5), or code. */
|
||||
export const LATEX_NEIGHBOR_CHAR_REGEXP = /[A-Za-z0-9_$-]/;
|
||||
|
||||
/** Matches a single digit (used to detect currency-like `$5`). */
|
||||
export const LATEX_DIGIT_REGEXP = /[0-9]/;
|
||||
|
||||
/** Matches the leading blockquote prefix (`> ` or `>`) on a markdown line. */
|
||||
export const LATEX_BLOCKQUOTE_PREFIX_REGEXP = /^(>\s*)/;
|
||||
|
||||
/** Matches the placeholder inserted by the protect/restore pipeline for a
|
||||
* protected LaTeX expression. Group 1 is the index into `latexExpressions`. */
|
||||
export const LATEX_PLACEHOLDER_REGEXP = /<<LATEX_(\d+)>>/g;
|
||||
|
||||
/** Matches the placeholder inserted by the protect/restore pipeline for a
|
||||
* protected code block. Group 1 is the index into `codeBlocks`. */
|
||||
export const CODE_BLOCK_PLACEHOLDER_REGEXP = /<<CODE_BLOCK_(\d+)>>/g;
|
||||
|
||||
/** Matches a `$` immediately followed by a digit, which is treated as a
|
||||
* currency amount (e.g. `$5`) and escaped to `\$5` so it isn't parsed as math. */
|
||||
export const LATEX_CURRENCY_DOLLAR_REGEXP = /\$(?=\d)/g;
|
||||
|
||||
/** Captures remaining `$$...$$`, `\[...\]`, `\(...\)` (only unescaped via
|
||||
* `(?<!\\)`) after the display-block pass has run. Group 1 holds the
|
||||
* matched formula. */
|
||||
export const LATEX_PROTECT_REGEXP =
|
||||
/(\$\$[\s\S]*?\$\$|(?<!\\)\\\[[\s\S]*?\\\]|(?<!\\)\\\(.*?\\\))/g;
|
||||
|
||||
/** Matches unescaped inline `\(...\)` (at least one char inside) used to
|
||||
* convert `\(` → `$` after the protect pass. */
|
||||
export const LATEX_INLINE_CONVERT_REGEXP = /(?<!\\)\\\((.+?)\\\)/g;
|
||||
|
||||
/** Matches unescaped display `\[...\]` used to convert `\[` → `$$`
|
||||
* after the protect pass. */
|
||||
export const LATEX_DISPLAY_CONVERT_REGEXP = /(?<!\\)\\\[([\s\S]*?)\\\]/g;
|
||||
|
||||
/** `\(` — opens an inline LaTeX math block. */
|
||||
export const LATEX_INLINE_OPEN = '\\(';
|
||||
|
||||
/** `\)` — closes an inline LaTeX math block. */
|
||||
export const LATEX_INLINE_CLOSE = '\\)';
|
||||
|
||||
/** `\[` — opens a display LaTeX math block. */
|
||||
export const LATEX_DISPLAY_OPEN = '\\[';
|
||||
|
||||
/** `\]` — closes a display LaTeX math block. */
|
||||
export const LATEX_DISPLAY_CLOSE = '\\]';
|
||||
|
||||
/** `\` — the LaTeX escape character. */
|
||||
export const LATEX_BACKSLASH = '\\';
|
||||
|
||||
/** `\$` — dollar sign escaped so it isn't parsed as math (used to disambiguate
|
||||
* currency amounts like `$5`). */
|
||||
export const LATEX_CURRENCY_ESCAPE = '\\$';
|
||||
|
||||
/** `\ce{` — mhchem chemistry command prefix. */
|
||||
export const LATEX_MHCHEM_CE = '\\ce{';
|
||||
|
||||
/** `\pu{` — mhchem physics-unit command prefix. */
|
||||
export const LATEX_MHCHEM_PU = '\\pu{';
|
||||
|
||||
/** map from mchem-regexp to replacement */
|
||||
export const MHCHEM_PATTERN_MAP: readonly [RegExp, string][] = [
|
||||
[/(\s)\$\\ce{/g, '$1$\\\\ce{'],
|
||||
|
||||
@@ -92,8 +92,17 @@ function deriveSingleTurnSections(
|
||||
|
||||
// 3. Persisted tool calls (from message.toolCalls field)
|
||||
const toolCalls = parseToolCalls(message.toolCalls);
|
||||
|
||||
// Index tool messages by toolCallId for O(1) lookup instead of O(n) find()
|
||||
const toolMsgById = new Map<string, DatabaseMessage>();
|
||||
for (const tm of toolMessages) {
|
||||
if (tm.toolCallId && !toolMsgById.has(tm.toolCallId)) {
|
||||
toolMsgById.set(tm.toolCallId, tm);
|
||||
}
|
||||
}
|
||||
|
||||
for (const tc of toolCalls) {
|
||||
const resultMsg = toolMessages.find((m) => m.toolCallId === tc.id);
|
||||
const resultMsg = tc.id ? toolMsgById.get(tc.id) : undefined;
|
||||
// Only show as pending/loading if we're actively streaming; otherwise it's just a tool call without result
|
||||
const type = resultMsg
|
||||
? AgenticSectionType.TOOL_CALL
|
||||
@@ -112,9 +121,10 @@ function deriveSingleTurnSections(
|
||||
}
|
||||
|
||||
// 4. Streaming tool calls (not yet persisted - currently being received)
|
||||
const persistedIds = new Set(toolCalls.map((t) => t.id).filter(Boolean));
|
||||
for (const tc of streamingToolCalls) {
|
||||
// Skip if already in persisted tool calls
|
||||
if (tc.id && toolCalls.find((t) => t.id === tc.id)) continue;
|
||||
if (tc.id && persistedIds.has(tc.id)) continue;
|
||||
sections.push({
|
||||
type: AgenticSectionType.TOOL_CALL_STREAMING,
|
||||
content: '',
|
||||
@@ -281,15 +291,31 @@ export function splitSearchSummaryList(
|
||||
return { lines };
|
||||
}
|
||||
|
||||
/** Bounded cache for parseToolResultWithImages results. */
|
||||
const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32;
|
||||
const toolResultLinesCache = new Map<string, ToolResultLine[]>();
|
||||
|
||||
/**
|
||||
* Parse tool result text into lines, matching image attachments by name.
|
||||
* Memoized: called per render during streaming on unchanged tool result
|
||||
* strings with unchanged extras.
|
||||
*/
|
||||
export function parseToolResultWithImages(
|
||||
toolResult: string,
|
||||
extras?: DatabaseMessageExtra[]
|
||||
): ToolResultLine[] {
|
||||
// Cache key includes image attachment names so we recompute when
|
||||
// attachments change, even if the count stays the same.
|
||||
const imageNames = (extras ?? [])
|
||||
.filter((e): e is DatabaseMessageExtraImageFile => e.type === AttachmentType.IMAGE)
|
||||
.map((e) => e.name)
|
||||
.join(NEWLINE);
|
||||
const cacheKey = `${imageNames}:${toolResult}`;
|
||||
const cached = toolResultLinesCache.get(cacheKey);
|
||||
if (cached !== undefined) return cached;
|
||||
|
||||
const lines = toolResult.split(NEWLINE);
|
||||
return lines.map((line) => {
|
||||
const result = lines.map((line) => {
|
||||
const match = line.match(ATTACHMENT_SAVED_REGEX);
|
||||
if (!match || !extras) return { text: line };
|
||||
|
||||
@@ -301,8 +327,19 @@ export function parseToolResultWithImages(
|
||||
|
||||
return { text: line, image };
|
||||
});
|
||||
|
||||
if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) {
|
||||
toolResultLinesCache.delete(toolResultLinesCache.keys().next().value!);
|
||||
}
|
||||
toolResultLinesCache.set(cacheKey, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Bounded cache for classifyToolResult results. */
|
||||
const CLASSIFY_CACHE_MAX_SIZE = 32;
|
||||
const classifyCache = new Map<string, ToolResultKind>();
|
||||
|
||||
/**
|
||||
* Pick a renderer tier for a tool's result content.
|
||||
*
|
||||
@@ -312,25 +349,39 @@ export function parseToolResultWithImages(
|
||||
* through MarkdownContent for proper formatting.
|
||||
* text - everything else, rendered as plain text lines (with image
|
||||
* attachment resolution as a side effect).
|
||||
* Memoized: called per render during streaming on unchanged content.
|
||||
*/
|
||||
export function classifyToolResult(content: string | undefined): ToolResultKind {
|
||||
if (!content) return ToolResultKind.TEXT;
|
||||
|
||||
const cached = classifyCache.get(content);
|
||||
if (cached !== undefined) return cached;
|
||||
|
||||
const trimmed = content.trim();
|
||||
if (!trimmed) return ToolResultKind.TEXT;
|
||||
|
||||
let result: ToolResultKind = ToolResultKind.TEXT;
|
||||
|
||||
// Strongest signal: JSON object/array round-trips through JSON.parse.
|
||||
if (TOOL_RESULT_JSON_OPEN_REGEX.test(trimmed)) {
|
||||
try {
|
||||
JSON.parse(trimmed);
|
||||
return ToolResultKind.JSON;
|
||||
result = ToolResultKind.JSON;
|
||||
} catch (error) {
|
||||
console.error('[agentic] tool result looked like JSON but failed to parse:', error);
|
||||
}
|
||||
}
|
||||
|
||||
if (looksLikeMarkdown(trimmed)) return ToolResultKind.MARKDOWN;
|
||||
if (result === ToolResultKind.TEXT && looksLikeMarkdown(trimmed)) {
|
||||
result = ToolResultKind.MARKDOWN;
|
||||
}
|
||||
|
||||
return ToolResultKind.TEXT;
|
||||
if (classifyCache.size >= CLASSIFY_CACHE_MAX_SIZE) {
|
||||
classifyCache.delete(classifyCache.keys().next().value!);
|
||||
}
|
||||
classifyCache.set(content, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -370,19 +421,35 @@ function looksLikeMarkdown(content: string): boolean {
|
||||
return false;
|
||||
}
|
||||
|
||||
/** Bounded cache for parsed tool-call JSON blobs. */
|
||||
const TOOL_CALLS_CACHE_MAX_SIZE = 64;
|
||||
const toolCallsParseCache = new Map<string, ApiChatCompletionToolCall[]>();
|
||||
|
||||
/**
|
||||
* Safely parse the toolCalls JSON string from a DatabaseMessage.
|
||||
* Memoized: the same JSON string is re-parsed on every render during
|
||||
* streaming, which is wasted CPU since tool calls don't change mid-stream.
|
||||
*/
|
||||
function parseToolCalls(toolCallsJson?: string): ApiChatCompletionToolCall[] {
|
||||
if (!toolCallsJson) return [];
|
||||
|
||||
const cached = toolCallsParseCache.get(toolCallsJson);
|
||||
if (cached) return cached;
|
||||
|
||||
let result: ApiChatCompletionToolCall[];
|
||||
try {
|
||||
const parsed = JSON.parse(toolCallsJson);
|
||||
|
||||
return Array.isArray(parsed) ? parsed : [];
|
||||
result = Array.isArray(parsed) ? parsed : [];
|
||||
} catch {
|
||||
return [];
|
||||
result = [];
|
||||
}
|
||||
|
||||
if (toolCallsParseCache.size >= TOOL_CALLS_CACHE_MAX_SIZE) {
|
||||
toolCallsParseCache.delete(toolCallsParseCache.keys().next().value!);
|
||||
}
|
||||
toolCallsParseCache.set(toolCallsJson, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -34,6 +34,10 @@ function escapeCode(code: string): string {
|
||||
return code.replace(AMPERSAND_REGEX, '&').replace(LT_REGEX, '<').replace(GT_REGEX, '>');
|
||||
}
|
||||
|
||||
/** Bounded cache for highlightCode results. */
|
||||
const HIGHLIGHT_CACHE_MAX_SIZE = 64;
|
||||
const highlightCache = new Map<string, string>();
|
||||
|
||||
/**
|
||||
* Highlights code using highlight.js
|
||||
* @param code - The code to highlight
|
||||
@@ -47,23 +51,37 @@ function escapeCode(code: string): string {
|
||||
export function highlightCode(code: string, language: string, autoDetect = true): string {
|
||||
if (!code) return '';
|
||||
|
||||
// Cache key includes language and autoDetect flag since results differ.
|
||||
// During streaming, the same code string may be highlighted repeatedly
|
||||
// (e.g., when text after a code block changes but the code itself doesn't).
|
||||
const cacheKey = `${language}:${autoDetect}:${code}`;
|
||||
const cached = highlightCache.get(cacheKey);
|
||||
if (cached) return cached;
|
||||
|
||||
const trimmed = trimCodePadding(code);
|
||||
let result: string;
|
||||
|
||||
try {
|
||||
const lang = language.toLowerCase();
|
||||
const isSupported = hljs.getLanguage(lang);
|
||||
|
||||
if (isSupported) {
|
||||
return hljs.highlight(trimmed, { language: lang }).value;
|
||||
result = hljs.highlight(trimmed, { language: lang }).value;
|
||||
} else if (autoDetect) {
|
||||
return hljs.highlightAuto(trimmed).value;
|
||||
result = hljs.highlightAuto(trimmed).value;
|
||||
} else {
|
||||
return escapeCode(trimmed);
|
||||
result = escapeCode(trimmed);
|
||||
}
|
||||
} catch {
|
||||
// Fallback to escaped plain text
|
||||
return escapeCode(trimmed);
|
||||
result = escapeCode(trimmed);
|
||||
}
|
||||
|
||||
if (highlightCache.size >= HIGHLIGHT_CACHE_MAX_SIZE) {
|
||||
highlightCache.delete(highlightCache.keys().next().value!);
|
||||
}
|
||||
highlightCache.set(cacheKey, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
export { trimCodePadding };
|
||||
|
||||
@@ -1,9 +1,31 @@
|
||||
import {
|
||||
CODE_BLOCK_PLACEHOLDER_REGEXP,
|
||||
CODE_BLOCK_REGEXP,
|
||||
LATEX_BACKSLASH,
|
||||
LATEX_BLOCKQUOTE_PREFIX_REGEXP,
|
||||
LATEX_CURRENCY_DOLLAR_REGEXP,
|
||||
LATEX_CURRENCY_ESCAPE,
|
||||
LATEX_DIGIT_REGEXP,
|
||||
LATEX_DISPLAY_BLOCK_REGEXP,
|
||||
LATEX_DISPLAY_CLOSE,
|
||||
LATEX_DISPLAY_CONVERT_REGEXP,
|
||||
LATEX_DISPLAY_DELIMITER,
|
||||
LATEX_DISPLAY_OPEN,
|
||||
LATEX_INLINE_CLOSE,
|
||||
LATEX_INLINE_CONVERT_REGEXP,
|
||||
LATEX_INLINE_DELIMITER,
|
||||
LATEX_INLINE_OPEN,
|
||||
LATEX_MATH_AND_CODE_PATTERN,
|
||||
LATEX_MHCHEM_CE,
|
||||
LATEX_MHCHEM_PU,
|
||||
LATEX_LINEBREAK_REGEXP,
|
||||
LATEX_NEIGHBOR_CHAR_REGEXP,
|
||||
LATEX_NON_WHITESPACE_REGEXP,
|
||||
LATEX_PLACEHOLDER_REGEXP,
|
||||
LATEX_PROTECT_REGEXP,
|
||||
LATEX_TRIGGER_REGEXP,
|
||||
MHCHEM_PATTERN_MAP
|
||||
MHCHEM_PATTERN_MAP,
|
||||
NEWLINE
|
||||
} from '$lib/constants';
|
||||
|
||||
/**
|
||||
@@ -20,13 +42,13 @@ import {
|
||||
* @returns The processed string with LaTeX replaced by placeholders.
|
||||
*/
|
||||
export function maskInlineLaTeX(content: string, latexExpressions: string[]): string {
|
||||
if (!content.includes('$')) {
|
||||
if (!content.includes(LATEX_INLINE_DELIMITER)) {
|
||||
return content;
|
||||
}
|
||||
return content
|
||||
.split('\n')
|
||||
.split(NEWLINE)
|
||||
.map((line) => {
|
||||
if (line.indexOf('$') == -1) {
|
||||
if (line.indexOf(LATEX_INLINE_DELIMITER) == -1) {
|
||||
return line;
|
||||
}
|
||||
|
||||
@@ -34,7 +56,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st
|
||||
let currentPosition = 0;
|
||||
|
||||
while (currentPosition < line.length) {
|
||||
const openDollarIndex = line.indexOf('$', currentPosition);
|
||||
const openDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, currentPosition);
|
||||
|
||||
if (openDollarIndex == -1) {
|
||||
processedLine += line.slice(currentPosition);
|
||||
@@ -42,7 +64,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st
|
||||
}
|
||||
|
||||
// Is there a next $-sign?
|
||||
const closeDollarIndex = line.indexOf('$', openDollarIndex + 1);
|
||||
const closeDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, openDollarIndex + 1);
|
||||
|
||||
if (closeDollarIndex == -1) {
|
||||
processedLine += line.slice(currentPosition);
|
||||
@@ -62,14 +84,14 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st
|
||||
shouldSkipAsNonLatex = true;
|
||||
}
|
||||
|
||||
if (/[A-Za-z0-9_$-]/.test(charBeforeOpen)) {
|
||||
if (LATEX_NEIGHBOR_CHAR_REGEXP.test(charBeforeOpen)) {
|
||||
// Character, digit, $, _ or - before first '$', no TeX.
|
||||
shouldSkipAsNonLatex = true;
|
||||
}
|
||||
|
||||
if (
|
||||
/[0-9]/.test(charAfterOpen) &&
|
||||
(/[A-Za-z0-9_$-]/.test(charAfterClose) || ' ' == charBeforeClose)
|
||||
LATEX_DIGIT_REGEXP.test(charAfterOpen) &&
|
||||
(LATEX_NEIGHBOR_CHAR_REGEXP.test(charAfterClose) || ' ' == charBeforeClose)
|
||||
) {
|
||||
// First $ seems to belong to an amount.
|
||||
shouldSkipAsNonLatex = true;
|
||||
@@ -92,7 +114,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st
|
||||
|
||||
return processedLine;
|
||||
})
|
||||
.join('\n');
|
||||
.join(NEWLINE);
|
||||
}
|
||||
|
||||
function escapeBrackets(text: string): string {
|
||||
@@ -107,9 +129,9 @@ function escapeBrackets(text: string): string {
|
||||
if (codeBlock != null) {
|
||||
return codeBlock;
|
||||
} else if (squareBracket != null) {
|
||||
return `$$${squareBracket}$$`;
|
||||
return `${LATEX_DISPLAY_DELIMITER}${squareBracket}${LATEX_DISPLAY_DELIMITER}`;
|
||||
} else if (roundBracket != null) {
|
||||
return `$${roundBracket}$`;
|
||||
return `${LATEX_INLINE_DELIMITER}${roundBracket}${LATEX_INLINE_DELIMITER}`;
|
||||
}
|
||||
|
||||
return match;
|
||||
@@ -145,32 +167,49 @@ const doEscapeMhchem = false;
|
||||
* preprocessLaTeX("Price: $10. The equation is \\(x^2\\).")
|
||||
* // → "Price: $10. The equation is $x^2$."
|
||||
*/
|
||||
/** Bounded cache for preprocessLaTeX results. */
|
||||
const LATEX_CACHE_MAX_SIZE = 64;
|
||||
const latexCache = new Map<string, string>();
|
||||
|
||||
export function preprocessLaTeX(content: string): string {
|
||||
// See also:
|
||||
// https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts
|
||||
|
||||
// Memoize on the input string. During streaming the prefix before an
|
||||
// incomplete code block stays the same across multiple tokens, so the
|
||||
// full protect/restore pipeline would re-run unnecessarily.
|
||||
const cached = latexCache.get(content);
|
||||
if (cached !== undefined) return cached;
|
||||
|
||||
// Save original before the function mutates `content` through steps 0-8
|
||||
const originalContent = content;
|
||||
|
||||
// Every step below keys off a `$` or a backslash escape (\[ \] \( \) \ce{ \pu{).
|
||||
// With neither present the protect/restore passes round-trip the input
|
||||
// unchanged, so skip them: the step 2 scan is O(n^2) in line length and costs
|
||||
// ~90ms on a 26KB single-line message that contains no math at all. This
|
||||
// matters during streaming, where the whole message is reprocessed per frame.
|
||||
if (!LATEX_TRIGGER_REGEXP.test(content)) {
|
||||
if (latexCache.size >= LATEX_CACHE_MAX_SIZE) {
|
||||
latexCache.delete(latexCache.keys().next().value!);
|
||||
}
|
||||
latexCache.set(originalContent, content);
|
||||
return content;
|
||||
}
|
||||
|
||||
// Step 0: Temporarily remove blockquote markers (>) to process LaTeX correctly
|
||||
// Store the structure so we can restore it later
|
||||
const blockquoteMarkers: Map<number, string> = new Map();
|
||||
const lines = content.split('\n');
|
||||
const lines = content.split(NEWLINE);
|
||||
const processedLines = lines.map((line, index) => {
|
||||
const match = line.match(/^(>\s*)/);
|
||||
const match = line.match(LATEX_BLOCKQUOTE_PREFIX_REGEXP);
|
||||
if (match) {
|
||||
blockquoteMarkers.set(index, match[1]);
|
||||
return line.slice(match[1].length);
|
||||
}
|
||||
return line;
|
||||
});
|
||||
content = processedLines.join('\n');
|
||||
content = processedLines.join(NEWLINE);
|
||||
|
||||
// Step 1: Protect code blocks
|
||||
const codeBlocks: string[] = [];
|
||||
@@ -187,58 +226,52 @@ export function preprocessLaTeX(content: string): string {
|
||||
// Match \S...\[...\] and protect them and insert a line-break.
|
||||
// Guarded: with no `\[` present this pattern still probes every start offset,
|
||||
// expanding `.*?` to the end of each line before failing - O(n^2) for nothing.
|
||||
if (content.includes('\\[')) {
|
||||
content = content.replace(
|
||||
/([\S].*?)\\\[([\s\S]*?)\\\](.*)/g,
|
||||
(match, group1, group2, group3) => {
|
||||
// Check if there are characters following the formula (display-formula in a table-cell?)
|
||||
if (group1.endsWith('\\')) {
|
||||
return match; // Backslash before \[, do nothing.
|
||||
}
|
||||
const hasSuffix = /\S/.test(group3);
|
||||
let optBreak;
|
||||
|
||||
if (hasSuffix) {
|
||||
latexExpressions.push(`\\(${group2.trim()}\\)`); // Convert into inline.
|
||||
optBreak = '';
|
||||
} else {
|
||||
latexExpressions.push(`\\[${group2}\\]`);
|
||||
optBreak = '\n';
|
||||
}
|
||||
|
||||
return `${group1}${optBreak}<<LATEX_${latexExpressions.length - 1}>>${optBreak}${group3}`;
|
||||
if (content.includes(LATEX_DISPLAY_OPEN)) {
|
||||
content = content.replace(LATEX_DISPLAY_BLOCK_REGEXP, (match, group1, group2, group3) => {
|
||||
// Check if there are characters following the formula (display-formula in a table-cell?)
|
||||
if (group1.endsWith(LATEX_BACKSLASH)) {
|
||||
return match; // Backslash before \[, do nothing.
|
||||
}
|
||||
);
|
||||
const hasSuffix = LATEX_NON_WHITESPACE_REGEXP.test(group3);
|
||||
let optBreak;
|
||||
|
||||
if (hasSuffix) {
|
||||
latexExpressions.push(`${LATEX_INLINE_OPEN}${group2.trim()}${LATEX_INLINE_CLOSE}`); // Convert into inline.
|
||||
optBreak = '';
|
||||
} else {
|
||||
latexExpressions.push(`${LATEX_DISPLAY_OPEN}${group2}${LATEX_DISPLAY_CLOSE}`);
|
||||
optBreak = NEWLINE;
|
||||
}
|
||||
|
||||
return `${group1}${optBreak}<<LATEX_${latexExpressions.length - 1}>>${optBreak}${group3}`;
|
||||
});
|
||||
}
|
||||
|
||||
// Match \(...\), \[...\], $$...$$ and protect them
|
||||
content = content.replace(
|
||||
/(\$\$[\s\S]*?\$\$|(?<!\\)\\\[[\s\S]*?\\\]|(?<!\\)\\\(.*?\\\))/g,
|
||||
(match) => {
|
||||
latexExpressions.push(match);
|
||||
content = content.replace(LATEX_PROTECT_REGEXP, (match) => {
|
||||
latexExpressions.push(match);
|
||||
|
||||
return `<<LATEX_${latexExpressions.length - 1}>>`;
|
||||
}
|
||||
);
|
||||
return `<<LATEX_${latexExpressions.length - 1}>>`;
|
||||
});
|
||||
|
||||
// Protect inline $...$ but NOT if it looks like money (e.g., $10, $3.99)
|
||||
content = maskInlineLaTeX(content, latexExpressions);
|
||||
|
||||
// Step 3: Escape standalone $ before digits (currency like $5 → \$5)
|
||||
// (Now that inline math is protected, this will only escape dollars not already protected)
|
||||
content = content.replace(/\$(?=\d)/g, '\\$');
|
||||
content = content.replace(LATEX_CURRENCY_DOLLAR_REGEXP, LATEX_CURRENCY_ESCAPE);
|
||||
|
||||
// Step 4: Restore protected LaTeX expressions (they are valid)
|
||||
content = content.replace(/<<LATEX_(\d+)>>/g, (_, index) => {
|
||||
content = content.replace(LATEX_PLACEHOLDER_REGEXP, (_, index) => {
|
||||
let expr = latexExpressions[parseInt(index)];
|
||||
const match = expr.match(LATEX_LINEBREAK_REGEXP);
|
||||
if (match) {
|
||||
// Katex: The $$-delimiters should be in their own line
|
||||
// if there are \\-line-breaks.
|
||||
const formula = match[1];
|
||||
const prefix = formula.startsWith('\n') ? '' : '\n';
|
||||
const suffix = formula.endsWith('\n') ? '' : '\n';
|
||||
expr = '$$' + prefix + formula + suffix + '$$';
|
||||
const prefix = formula.startsWith(NEWLINE) ? '' : NEWLINE;
|
||||
const suffix = formula.endsWith(NEWLINE) ? '' : NEWLINE;
|
||||
expr = LATEX_DISPLAY_DELIMITER + prefix + formula + suffix + LATEX_DISPLAY_DELIMITER;
|
||||
}
|
||||
return expr;
|
||||
});
|
||||
@@ -247,7 +280,7 @@ export function preprocessLaTeX(content: string): string {
|
||||
// This must happen BEFORE restoring code blocks to avoid affecting code content
|
||||
content = escapeBrackets(content);
|
||||
|
||||
if (doEscapeMhchem && (content.includes('\\ce{') || content.includes('\\pu{'))) {
|
||||
if (doEscapeMhchem && (content.includes(LATEX_MHCHEM_CE) || content.includes(LATEX_MHCHEM_PU))) {
|
||||
content = escapeMhchem(content);
|
||||
}
|
||||
|
||||
@@ -257,31 +290,38 @@ export function preprocessLaTeX(content: string): string {
|
||||
// Using the look‑behind pattern `(?<!\\)` we skip matches
|
||||
// that are preceded by a backslash, e.g.
|
||||
// `Definitions\\(also called macros)` (title of chapter 20 in The TeXbook).
|
||||
.replace(/(?<!\\)\\\((.+?)\\\)/g, '$$$1$') // inline
|
||||
.replace(LATEX_INLINE_CONVERT_REGEXP, (_, formula: string) => {
|
||||
return `${LATEX_INLINE_DELIMITER}${formula}${LATEX_INLINE_DELIMITER}`;
|
||||
}) // inline
|
||||
.replace(
|
||||
// Using the look‑behind pattern `(?<!\\)` we skip matches
|
||||
// that are preceded by a backslash, e.g. `\\[4pt]`.
|
||||
/(?<!\\)\\\[([\s\S]*?)\\\]/g, // display, see also PR #16599
|
||||
(_, content: string) => {
|
||||
return `$$${content}$$`;
|
||||
LATEX_DISPLAY_CONVERT_REGEXP, // display, see also PR #16599
|
||||
(_, formula: string) => {
|
||||
return `${LATEX_DISPLAY_DELIMITER}${formula}${LATEX_DISPLAY_DELIMITER}`;
|
||||
}
|
||||
);
|
||||
|
||||
// Step 7: Restore code blocks
|
||||
// This happens AFTER all LaTeX conversions to preserve code content
|
||||
content = content.replace(/<<CODE_BLOCK_(\d+)>>/g, (_, index) => {
|
||||
content = content.replace(CODE_BLOCK_PLACEHOLDER_REGEXP, (_, index) => {
|
||||
return codeBlocks[parseInt(index)];
|
||||
});
|
||||
|
||||
// Step 8: Restore blockquote markers
|
||||
if (blockquoteMarkers.size > 0) {
|
||||
const finalLines = content.split('\n');
|
||||
const finalLines = content.split(NEWLINE);
|
||||
const restoredLines = finalLines.map((line, index) => {
|
||||
const marker = blockquoteMarkers.get(index);
|
||||
return marker ? marker + line : line;
|
||||
});
|
||||
content = restoredLines.join('\n');
|
||||
content = restoredLines.join(NEWLINE);
|
||||
}
|
||||
|
||||
if (latexCache.size >= LATEX_CACHE_MAX_SIZE) {
|
||||
latexCache.delete(latexCache.keys().next().value!);
|
||||
}
|
||||
latexCache.set(originalContent, content);
|
||||
|
||||
return content;
|
||||
}
|
||||
|
||||
@@ -14,70 +14,94 @@ const JSON_ARRAY_CLOSE = ']';
|
||||
// comma when the model cut off mid-key.
|
||||
const TRAILING_JSON_PUNCTUATION_REGEX = /,?\s*$/;
|
||||
|
||||
/** Bounded cache for parsePartialJsonArgs results. */
|
||||
const PARTIAL_JSON_CACHE_MAX_SIZE = 32;
|
||||
const partialJsonCache = new Map<string, Record<string, unknown> | null>();
|
||||
|
||||
function cacheResult(input: string, result: Record<string, unknown> | null): void {
|
||||
if (partialJsonCache.size >= PARTIAL_JSON_CACHE_MAX_SIZE) {
|
||||
partialJsonCache.delete(partialJsonCache.keys().next().value!);
|
||||
}
|
||||
partialJsonCache.set(input, result);
|
||||
}
|
||||
|
||||
// Parse partial tool-arg JSON streamed token-by-token. Closes any
|
||||
// unterminated string and dangling open containers (in reverse order),
|
||||
// so parsers can still surface keys already received while the call
|
||||
// is still in flight.
|
||||
// is still in flight. Memoized: the char-by-char scanner runs on every
|
||||
// render during streaming even when toolArgs hasn't changed.
|
||||
export function parsePartialJsonArgs(toolArgsString: string): Record<string, unknown> | null {
|
||||
const cached = partialJsonCache.get(toolArgsString);
|
||||
if (cached !== undefined) return cached;
|
||||
|
||||
let result: Record<string, unknown> | null;
|
||||
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(toolArgsString);
|
||||
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
|
||||
return parsed as Record<string, unknown>;
|
||||
}
|
||||
return null;
|
||||
result =
|
||||
parsed && typeof parsed === 'object' && !Array.isArray(parsed)
|
||||
? (parsed as Record<string, unknown>)
|
||||
: null;
|
||||
} catch {
|
||||
let inString = false;
|
||||
let escape = false;
|
||||
const stack: ('{' | '[')[] = [];
|
||||
result = scanPartialJson(toolArgsString);
|
||||
}
|
||||
|
||||
for (let i = 0; i < toolArgsString.length; i++) {
|
||||
const ch = toolArgsString[i];
|
||||
if (escape) {
|
||||
escape = false;
|
||||
continue;
|
||||
}
|
||||
if (ch === JSON_BACKSLASH && inString) {
|
||||
escape = true;
|
||||
continue;
|
||||
}
|
||||
if (ch === JSON_QUOTE) {
|
||||
inString = !inString;
|
||||
continue;
|
||||
}
|
||||
if (inString) continue;
|
||||
if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN);
|
||||
else if (ch === JSON_OBJECT_CLOSE) {
|
||||
if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null;
|
||||
stack.pop();
|
||||
} else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN);
|
||||
else if (ch === JSON_ARRAY_CLOSE) {
|
||||
if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null;
|
||||
stack.pop();
|
||||
}
|
||||
}
|
||||
cacheResult(toolArgsString, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
let completed = toolArgsString;
|
||||
/** Char-by-char scanner for unterminated partial JSON. */
|
||||
function scanPartialJson(toolArgsString: string): Record<string, unknown> | null {
|
||||
let inString = false;
|
||||
let escape = false;
|
||||
const stack: ('{' | '[')[] = [];
|
||||
|
||||
for (let i = 0; i < toolArgsString.length; i++) {
|
||||
const ch = toolArgsString[i];
|
||||
if (escape) {
|
||||
// Dangling escape at end of partial JSON: escape the trailing
|
||||
// backslash as a literal so we can close the string cleanly.
|
||||
completed += JSON_BACKSLASH;
|
||||
escape = false;
|
||||
continue;
|
||||
}
|
||||
if (inString) completed += JSON_QUOTE;
|
||||
if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, '');
|
||||
|
||||
// Close in reverse nesting order: innermost container first.
|
||||
for (let i = stack.length - 1; i >= 0; i--) {
|
||||
completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE;
|
||||
if (ch === JSON_BACKSLASH && inString) {
|
||||
escape = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(completed);
|
||||
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
|
||||
return parsed as Record<string, unknown>;
|
||||
}
|
||||
return null;
|
||||
} catch {
|
||||
return null;
|
||||
if (ch === JSON_QUOTE) {
|
||||
inString = !inString;
|
||||
continue;
|
||||
}
|
||||
if (inString) continue;
|
||||
if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN);
|
||||
else if (ch === JSON_OBJECT_CLOSE) {
|
||||
if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null;
|
||||
stack.pop();
|
||||
} else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN);
|
||||
else if (ch === JSON_ARRAY_CLOSE) {
|
||||
if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null;
|
||||
stack.pop();
|
||||
}
|
||||
}
|
||||
|
||||
let completed = toolArgsString;
|
||||
if (escape) {
|
||||
// Dangling escape at end of partial JSON: escape the trailing
|
||||
// backslash as a literal so we can close the string cleanly.
|
||||
completed += JSON_BACKSLASH;
|
||||
}
|
||||
if (inString) completed += JSON_QUOTE;
|
||||
if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, '');
|
||||
|
||||
// Close in reverse nesting order: innermost container first.
|
||||
for (let i = stack.length - 1; i >= 0; i--) {
|
||||
completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE;
|
||||
}
|
||||
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(completed);
|
||||
return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
|
||||
? (parsed as Record<string, unknown>)
|
||||
: null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -155,41 +155,71 @@ function parseChunk(chunk: string): SearchResult | null {
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Bounded cache for extractSearchResults results. */
|
||||
const SEARCH_RESULTS_CACHE_MAX_SIZE = 32;
|
||||
const searchResultsCache = new Map<string, SearchResult[]>();
|
||||
|
||||
/**
|
||||
* Extract a SearchResult[] from a tool-result string. Returns `[]` when
|
||||
* the input does not match the expected shape — useful for branching
|
||||
* between dedicated search-results rendering and the generic tool-call
|
||||
* block.
|
||||
* block. Memoized: called per render during streaming on unchanged
|
||||
* tool result strings.
|
||||
*/
|
||||
export function extractSearchResults(text: string | undefined | null): SearchResult[] {
|
||||
if (!text) return [];
|
||||
|
||||
const cached = searchResultsCache.get(text);
|
||||
if (cached) return cached;
|
||||
|
||||
const results: SearchResult[] = [];
|
||||
for (const chunk of splitChunks(text)) {
|
||||
const parsed = parseChunk(chunk);
|
||||
if (parsed) results.push(parsed);
|
||||
}
|
||||
|
||||
if (searchResultsCache.size >= SEARCH_RESULTS_CACHE_MAX_SIZE) {
|
||||
searchResultsCache.delete(searchResultsCache.keys().next().value!);
|
||||
}
|
||||
searchResultsCache.set(text, results);
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/** Bounded cache for extractSearchQuery results. */
|
||||
const SEARCH_QUERY_CACHE_MAX_SIZE = 32;
|
||||
const searchQueryCache = new Map<string, string>();
|
||||
|
||||
/**
|
||||
* Best-effort extraction of the search query out of a tool call's JSON
|
||||
* argument blob. Currently looks for a `query` field (the convention
|
||||
* used by Exa and most web-search MCP servers); returns an empty string
|
||||
* if it cannot be located.
|
||||
* if it cannot be located. Memoized: called per render during streaming
|
||||
* on unchanged tool args strings.
|
||||
*/
|
||||
export function extractSearchQuery(toolArgs: string | undefined | null): string {
|
||||
if (!toolArgs) return '';
|
||||
|
||||
const cached = searchQueryCache.get(toolArgs);
|
||||
if (cached !== undefined) return cached;
|
||||
|
||||
let result = '';
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(toolArgs);
|
||||
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
|
||||
const candidate = (parsed as Record<string, unknown>)[SEARCH_TOOL_QUERY_FIELD];
|
||||
if (typeof candidate === 'string') return candidate.trim();
|
||||
if (typeof candidate === 'string') result = candidate.trim();
|
||||
}
|
||||
} catch {
|
||||
return '';
|
||||
result = '';
|
||||
}
|
||||
return '';
|
||||
|
||||
if (searchQueryCache.size >= SEARCH_QUERY_CACHE_MAX_SIZE) {
|
||||
searchQueryCache.delete(searchQueryCache.keys().next().value!);
|
||||
}
|
||||
searchQueryCache.set(toolArgs, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
// Tests for the memoized parseToolCalls and O(1) tool message lookup in
|
||||
// deriveAgenticSections. These were added to prevent regressions where
|
||||
// streaming text tokens trigger redundant JSON.parse calls on unchanged
|
||||
// tool call data.
|
||||
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { deriveAgenticSections } from '$lib/utils/agentic';
|
||||
import type { ApiChatCompletionToolCall } from '$lib/types/api';
|
||||
import type { DatabaseMessage } from '$lib/types/database';
|
||||
import { MessageRole, AgenticSectionType } from '$lib/enums';
|
||||
|
||||
function makeMessage(overrides: Partial<DatabaseMessage>): DatabaseMessage {
|
||||
return {
|
||||
id: 'm1',
|
||||
convId: 'c1',
|
||||
type: 'text',
|
||||
timestamp: 0,
|
||||
role: MessageRole.ASSISTANT,
|
||||
content: '',
|
||||
parent: null,
|
||||
children: [],
|
||||
...overrides
|
||||
} as DatabaseMessage;
|
||||
}
|
||||
|
||||
describe('parseToolCalls memoization', () => {
|
||||
it('returns the same array reference for the same JSON string', () => {
|
||||
// parseToolCalls is not exported, but deriveAgenticSections uses it
|
||||
// internally. We verify memoization through behavior: calling
|
||||
// deriveAgenticSections twice with the same toolCalls should not
|
||||
// re-parse (which we verify by checking the returned sections
|
||||
// are equivalent).
|
||||
const toolCallsJson = JSON.stringify([
|
||||
{ id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } }
|
||||
]);
|
||||
|
||||
const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson });
|
||||
const sections1 = deriveAgenticSections(msg, [], [], false);
|
||||
const sections2 = deriveAgenticSections(msg, [], [], false);
|
||||
|
||||
expect(sections1).toHaveLength(sections2.length);
|
||||
expect(sections1[0].type).toBe(sections2[0].type);
|
||||
});
|
||||
|
||||
it('does not re-parse JSON on cache hit', () => {
|
||||
const toolCallsJson = JSON.stringify([
|
||||
{ id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } }
|
||||
]);
|
||||
|
||||
const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson });
|
||||
const spy = vi.spyOn(JSON, 'parse');
|
||||
|
||||
deriveAgenticSections(msg, [], [], false);
|
||||
const callsAfterFirst = spy.mock.calls.length;
|
||||
|
||||
deriveAgenticSections(msg, [], [], false);
|
||||
expect(spy.mock.calls.length).toBe(callsAfterFirst);
|
||||
|
||||
spy.mockRestore();
|
||||
});
|
||||
|
||||
it('handles empty/undefined toolCalls without error', () => {
|
||||
const msg = makeMessage({ content: 'hello' });
|
||||
const sections = deriveAgenticSections(msg, [], [], false);
|
||||
|
||||
expect(sections).toHaveLength(1);
|
||||
expect(sections[0].type).toBe(AgenticSectionType.TEXT);
|
||||
});
|
||||
|
||||
it('handles invalid JSON gracefully', () => {
|
||||
const msg = makeMessage({ content: 'hello', toolCalls: '{invalid' });
|
||||
const sections = deriveAgenticSections(msg, [], [], false);
|
||||
|
||||
// Should return just the text section, no tool call sections
|
||||
expect(sections).toHaveLength(1);
|
||||
expect(sections[0].type).toBe(AgenticSectionType.TEXT);
|
||||
});
|
||||
});
|
||||
|
||||
describe('deriveAgenticSections O(1) tool message lookup', () => {
|
||||
it('matches tool messages to tool calls by toolCallId', () => {
|
||||
const toolCallsJson = JSON.stringify([
|
||||
{ id: 'call_1', type: 'function', function: { name: 'test_1', arguments: '{}' } },
|
||||
{ id: 'call_2', type: 'function', function: { name: 'test_2', arguments: '{}' } }
|
||||
]);
|
||||
|
||||
const toolMessages = [
|
||||
makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_1', content: 'result_1' }),
|
||||
makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_2', content: 'result_2' })
|
||||
];
|
||||
|
||||
const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson });
|
||||
const sections = deriveAgenticSections(msg, toolMessages, [], false);
|
||||
|
||||
// Expect: TEXT + 2 TOOL_CALL sections
|
||||
const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL);
|
||||
expect(toolCallSections).toHaveLength(2);
|
||||
expect(toolCallSections[0].toolResult).toBe('result_1');
|
||||
expect(toolCallSections[1].toolResult).toBe('result_2');
|
||||
});
|
||||
|
||||
it('handles missing tool messages (pending calls during streaming)', () => {
|
||||
const toolCallsJson = JSON.stringify([
|
||||
{ id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } }
|
||||
]);
|
||||
|
||||
const msg = makeMessage({ content: '', toolCalls: toolCallsJson });
|
||||
const sections = deriveAgenticSections(msg, [], [], true);
|
||||
|
||||
const toolCallSection = sections.find((s) => s.type === AgenticSectionType.TOOL_CALL_PENDING);
|
||||
expect(toolCallSection).toBeDefined();
|
||||
expect(toolCallSection?.content).toBe('');
|
||||
});
|
||||
|
||||
it('scales with many tool calls (no O(n^2) blowup)', () => {
|
||||
const N = 100;
|
||||
const toolCalls = Array.from(
|
||||
{ length: N },
|
||||
(_, i): ApiChatCompletionToolCall => ({
|
||||
id: `call_${i}`,
|
||||
type: 'function',
|
||||
function: { name: `tool_${i}`, arguments: '{}' }
|
||||
})
|
||||
);
|
||||
const toolCallsJson = JSON.stringify(toolCalls);
|
||||
|
||||
const toolMessages = Array.from({ length: N }, (_, i) =>
|
||||
makeMessage({
|
||||
role: MessageRole.TOOL,
|
||||
toolCallId: `call_${i}`,
|
||||
content: `result_${i}`
|
||||
})
|
||||
);
|
||||
|
||||
const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson });
|
||||
|
||||
// If the lookup were still O(n^2), this would be noticeably slow
|
||||
const start = Date.now();
|
||||
const sections = deriveAgenticSections(msg, toolMessages, [], false);
|
||||
const elapsed = Date.now() - start;
|
||||
|
||||
const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL);
|
||||
expect(toolCallSections).toHaveLength(N);
|
||||
expect(elapsed).toBeLessThan(100); // Should be fast with O(1) lookup
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user