Compare commits

..

5 Commits

Author SHA1 Message Date
Hongqiang Wang 8190848bb3 opencl: skip the Adreno KQ/KQV image kernels for multi-stream batches (#26189)
The Adreno KQ/KQV image1d kernels (ggml_cl_mul_mat_kq_kqv_adreno) ignore
dim 3 entirely: the sub-buffer covers only nb02*ne02 bytes and the kernel
receives no ne03/ne13/nb03/nb13 arguments. With the unified KV cache,
multi-sequence batches (e.g. llama-perplexity with its default -b 2048,
n_seq=4, or a multi-slot llama-server) present KQ/KQV as 4D tensors with
ne3 = n_stream, so every stream past the first reads the first stream's
K/V and produces garbage. Flash attention masks the bug where it is
enabled; devices where FA is declined (e.g. Adreno 740) hit it with
default settings.

Route ne03/ne13 > 1 to the general path, which handles dim 3, and honor
view_offs when creating the sub-buffers (currently always 0 for tensors
reaching this function, but the function would silently misread any
future view).

Llama-3.2-1B-Instruct Q4_0, wiki.test.raw, 8 chunks, -ngl 99:
- Adreno 740, default:            PPL 1817.64 -> 15.61
- Adreno 740, -fa 0:              PPL 1941.64 -> 15.61
- Adreno 840, -fa 0:              PPL 1943.90 -> 15.50
- single-stream (-b 512) results unchanged (15.6090)
- test-backend-ops -o MUL_MAT on 740: identical before/after (909 OK,
  12 pre-existing q6_K failures)
2026-07-28 11:04:42 -07:00
Daniel Bevenius 7e1e28cae3 mtmd : add Nemotron 3 Nano Omni support (parakeet) (#22520)
* mtmd : add Nemotron 3 Nano Omni support (parakeet)

This commit adds support for the subsampling and encoder part of
Nemotron Nemo 3 omni model.

The Parakeet subsampling/encoder were taken from parakeet.cpp which
is currently a pull request against whisper.cpp. I've tried to copy the
code a close as possible to hopefully enable easy patching between the
these two project later.

Refs: https://github.com/ggml-org/whisper.cpp/pull/3735

* mtmd : generate rel pos tensor in graph instead of in conversion [no ci]

This commit removes the generation of the relative positional tensor in
the model conversion script and instead computes it in the encoder
graph. This is only done for the window of positions required for the
current audio sample.

* mtmd : add clip_get_model to clip API [no ci]

This commit adds a function to get access to the clip_model. It also
removes the two functions clip_get_mel_filter_tensor, and
clip_get_window_tensor(const struct clip_ctx * ctx) which can now use
clip_get_model to access the model tensors that it needs.

* mtmd : read mel_filters and window into hparams

* mtmd : use set_input_f32 lambda [no ci]

* mtmd : add better asserts for mel_filters and hann window [no ci]

* mtmd : add missing size_t cast

* mtmd : change type of pad to size_t

* mtmd : zero initialize samples_padded

* mtmd : remove unsued ctx member from parakeet preprocessor

* mtmd : make log_mel_spectrogram_parakeet_worker_thread private static

* mtmd : sync/update parakeeet impl with latest whisper.cpp

This commit updates the parakeet code in mtmd to reflect the latest
updates to parakeet.cpp in whisper.cpp.

A follow up commit will address the currently hardcoded dw_pad and see
if we can add n_conv_kernel as a model metadata field.

* mtmd : add audio_conv_kernel_size to model conversion

This commit updates the model conversion to read the conv_kernel_size
field from the sound_config section of the models config.json file.
It then uses this field instead of the hardcoded values in parakeet.cpp.

* mtmd : cleanup [no ci]

* conversion : call super().filter_tensors [no ci]

* do not discard result of super filter_tensors

* mtmd : use build_mm instead of ggml_mul_mat

* mtmd : use build_ffn

* mtmd : move and reuse get_vector lambda

* mtmd : use build_inp_raw for parakeet

* mtmd : throw exception in get_scalar instead of assert

* mtmd : fix std::min call

* mtmt : use .c_str in throw clause in get_vector

* mtmd : check for F32 type and non-empty tensor in get_vector

The get_vector lambda is used by get_scalar but also standalone to read
in the mel_filters and the window data. Therefor we are not checking
for 1D tensors but allowing multiple dimensions. We do have a check in
get_scalar to verify the size of the vector.

* mtmd : replace hardcoded 1101 for n_tokens_real

* mtmd : assert subsampling_factor is 8

This commit adds an assert of the parakeet subsampling factor to check
that it is 8.

The motivation for this is that this model currently has three
convolutions with a stride of 2. If the underlying model updates the
subsampling factor these convolution operations will need to be updated
and this will produce and error if this occurs.

* mtmd : remove unused ggml_tensors attn_pos_w and mm_norm_w

* mtmd : remove single thread path

This commit removes the single thread path which was a left over from
the original parakeet.cpp where n_threads is configurable.

* fix some security issues

---------

Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
2026-07-28 17:20:25 +02:00
Aleksander Grygier 6e2bc65fb2 ui: rendering performance follow-up (#26097) 2026-07-28 17:13:25 +02:00
Julien Jerphanion ad77bd31a6 docs: Adapt conda-forge package name (#26229)
Co-authored-by: dev-tinker <dev-tinker@users.noreply.github.com>
2026-07-28 16:51:20 +02:00
Xuan-Son Nguyen ee3d1b54c1 server: abstract llama_memory calls to common_memory (#26221) 2026-07-28 16:35:20 +02:00
28 changed files with 1591 additions and 211 deletions
+29 -3
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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).
+5 -3
View File
@@ -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, &region, &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, &region, &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
+14
View File
@@ -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
+3
View File
@@ -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)
+40
View File
@@ -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
),
+1
View File
@@ -60,6 +60,7 @@ add_library(mtmd
models/mobilenetv5.cpp
models/youtuvl.cpp
models/yasa2.cpp
models/parakeet.cpp
)
set_target_properties(mtmd PROPERTIES
+9
View File
@@ -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
View File
@@ -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
View File
@@ -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");
}
+5
View File
@@ -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;
+421
View File
@@ -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;
}
+203
View File
@@ -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
//
+15
View File
@@ -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
//
+4
View File
@@ -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>";
+18 -41
View File
@@ -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;
@@ -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}>
@@ -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))
);
@@ -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{'],
+76 -9
View File
@@ -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;
}
/**
+23 -5
View File
@@ -34,6 +34,10 @@ function escapeCode(code: string): string {
return code.replace(AMPERSAND_REGEX, '&amp;').replace(LT_REGEX, '&lt;').replace(GT_REGEX, '&gt;');
}
/** 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 };
+96 -56
View File
@@ -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 lookbehind 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 lookbehind 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;
}
}
+35 -5
View File
@@ -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
});
});