From 36744cd4b592ffb1ca3349b08d9a98ee486a1f19 Mon Sep 17 00:00:00 2001 From: Kwa Jie Hao Date: Tue, 31 Mar 2026 17:07:30 +0800 Subject: [PATCH 01/23] feat: (vocab) fix stray text appended in llama_decode_text Remove accidental concatenation of the full `text` string when formatting UNK_BYTE hex escapes. Only the closing "]" should be appended. --- src/llama-vocab.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 163f222ef612..66385c6c2ffc 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -3056,7 +3056,7 @@ static std::string llama_decode_text(const std::string & text) { for (const auto c : utf8) { decoded_text += format("%02x", (uint8_t) c); } - decoded_text += text + "]"; + decoded_text += "]"; } } From 9eaf3abe81588c3f481095f77710184502fb6366 Mon Sep 17 00:00:00 2001 From: Kwa Jie Hao Date: Tue, 31 Mar 2026 17:16:48 +0800 Subject: [PATCH 02/23] feat(mtmd): add Yasa2 vision encoder support Add a Yasa2 (ConvNeXtV2-based) vision encoder for reka-edge: - Register PROJECTOR_TYPE_YASA2 and tensor name definitions - Add yasa2_block/yasa2_stage model structs - Implement graph builder with ConvNeXt stages, GRN, adaptive pooling - Wire into clip.cpp switch statements and mtmd.cpp init_vision - Use mtmd_image_preprocessor_fixed_size for image preprocessing --- tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 13 +++ tools/mtmd/clip-model.h | 30 +++++ tools/mtmd/clip.cpp | 66 +++++++++++ tools/mtmd/models/models.h | 8 ++ tools/mtmd/models/yasa2.cpp | 217 ++++++++++++++++++++++++++++++++++++ tools/mtmd/mtmd.cpp | 6 + 7 files changed, 341 insertions(+) create mode 100644 tools/mtmd/models/yasa2.cpp diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 3bafde178de2..9031d063c14a 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -40,6 +40,7 @@ add_library(mtmd models/deepseekocr.cpp models/mobilenetv5.cpp models/youtuvl.cpp + models/yasa2.cpp ) set_target_properties(mtmd PROPERTIES diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 17cb703f7fbb..ef1ef7fa9508 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -242,6 +242,17 @@ #define TN_STD_BIAS "v.std_bias" #define TN_STD_SCALE "v.std_scale" +// yasa2 +#define TN_YASA_PATCH_W "v.patch_embd.weight" +#define TN_YASA_PATCH_B "v.patch_embd.bias" +#define TN_YASA_PATCH_LN_W "v.patch_ln.weight" +#define TN_YASA_PATCH_LN_B "v.patch_ln.bias" +#define TN_YASA_BACKBONE_LN_W "v.backbone_ln.weight" +#define TN_YASA_BACKBONE_LN_B "v.backbone_ln.bias" +#define TN_YASA_POS_EMBD "v.vision_pos_embed" +#define TN_YASA_STAGE_DOWN_LN "v.stage.%d.down.ln.%s" +#define TN_YASA_STAGE_DOWN_CONV "v.stage.%d.down.conv.%s" +#define TN_YASA_STAGE_BLK "v.stage.%d.blk.%d.%s.%s" // align x to upper multiple of n #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n)) @@ -290,6 +301,7 @@ enum projector_type { PROJECTOR_TYPE_LFM2A, PROJECTOR_TYPE_GLM4V, PROJECTOR_TYPE_YOUTUVL, + PROJECTOR_TYPE_YASA2, PROJECTOR_TYPE_KIMIK25, PROJECTOR_TYPE_NEMOTRON_V2_VL, PROJECTOR_TYPE_HUNYUANOCR, @@ -335,6 +347,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_LFM2A, "lfm2a"}, { PROJECTOR_TYPE_GLM4V, "glm4v"}, { PROJECTOR_TYPE_YOUTUVL, "youtuvl"}, + { PROJECTOR_TYPE_YASA2, "yasa2"}, { PROJECTOR_TYPE_KIMIK25, "kimik25"}, { PROJECTOR_TYPE_NEMOTRON_V2_VL, "nemotron_v2_vl"}, { PROJECTOR_TYPE_HUNYUANOCR, "hunyuanocr"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 9a93584d9beb..bf8031b55b28 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -268,6 +268,27 @@ struct mobilenetv5_block { ggml_tensor * attn_norm_w = nullptr; }; +struct yasa2_block { + ggml_tensor * dw_w = nullptr; + ggml_tensor * dw_b = nullptr; + ggml_tensor * ln_w = nullptr; + ggml_tensor * ln_b = nullptr; + ggml_tensor * pw1_w = nullptr; + ggml_tensor * pw1_b = nullptr; + ggml_tensor * grn_w = nullptr; + ggml_tensor * grn_b = nullptr; + ggml_tensor * pw2_w = nullptr; + ggml_tensor * pw2_b = nullptr; +}; + +struct yasa2_stage { + ggml_tensor * down_ln_w = nullptr; + ggml_tensor * down_ln_b = nullptr; + ggml_tensor * down_conv_w = nullptr; + ggml_tensor * down_conv_b = nullptr; + std::vector blocks; +}; + struct clip_model { clip_modality modality = CLIP_MODALITY_VISION; projector_type proj_type = PROJECTOR_TYPE_MLP; @@ -402,6 +423,15 @@ struct clip_model { ggml_tensor * msfa_ffn_expand_bn = nullptr; ggml_tensor * msfa_ffn_project_bn = nullptr; + // yasa2 + ggml_tensor * yasa_patch_w = nullptr; + ggml_tensor * yasa_patch_b = nullptr; + ggml_tensor * yasa_patch_ln_w = nullptr; + ggml_tensor * yasa_patch_ln_b = nullptr; + ggml_tensor * yasa_backbone_ln_w = nullptr; + ggml_tensor * yasa_backbone_ln_b = nullptr; + ggml_tensor * yasa_vision_pos_embed = nullptr; + std::vector yasa_stages; // pixtral, glm4v ggml_tensor * token_embd_img_break = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index f0e8786b6601..c3028f3d8006 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -947,6 +947,10 @@ static ggml_cgraph * clip_image_build_graph(clip_ctx * ctx, const clip_image_f32 { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_YASA2: + { + builder = std::make_unique(ctx, img); + } break; default: GGML_ABORT("missing cgraph builder"); } @@ -1389,6 +1393,12 @@ struct clip_model_loader { hparams.set_limit_image_tokens(1, 62500); hparams.set_warmup_n_tokens(16*16); // avoid OOM on warmup } break; + case PROJECTOR_TYPE_YASA2: + { + hparams.ffn_op = FFN_GELU_ERF; + log_ffn_op = "gelu_erf"; + hparams.set_warmup_n_tokens(64); + } break; case PROJECTOR_TYPE_GLM4V: { hparams.rope_theta = 10000.0f; @@ -1839,6 +1849,52 @@ struct clip_model_loader { model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); // merger.mlp.2 model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias")); } break; + case PROJECTOR_TYPE_YASA2: + { + model.yasa_patch_w = get_tensor(TN_YASA_PATCH_W); + model.yasa_patch_b = get_tensor(TN_YASA_PATCH_B, false); + model.yasa_patch_ln_w = get_tensor(TN_YASA_PATCH_LN_W, false); + model.yasa_patch_ln_b = get_tensor(TN_YASA_PATCH_LN_B, false); + model.yasa_backbone_ln_w = get_tensor(TN_YASA_BACKBONE_LN_W, false); + model.yasa_backbone_ln_b = get_tensor(TN_YASA_BACKBONE_LN_B, false); + model.yasa_vision_pos_embed = get_tensor(TN_YASA_POS_EMBD, false); + model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); + model.mm_0_b = get_tensor(string_format(TN_LLAVA_PROJ, 0, "bias"), false); + model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); + model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"), false); + + model.yasa_stages.clear(); + for (int s = 0; ; ++s) { + yasa2_stage stage; + stage.down_ln_w = get_tensor(string_format(TN_YASA_STAGE_DOWN_LN, s, "weight"), false); + stage.down_ln_b = get_tensor(string_format(TN_YASA_STAGE_DOWN_LN, s, "bias"), false); + stage.down_conv_w = get_tensor(string_format(TN_YASA_STAGE_DOWN_CONV, s, "weight"), false); + stage.down_conv_b = get_tensor(string_format(TN_YASA_STAGE_DOWN_CONV, s, "bias"), false); + + for (int bi = 0; ; ++bi) { + yasa2_block blk; + blk.dw_w = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "dw", "weight"), false); + if (!blk.dw_w) { + break; + } + blk.dw_b = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "dw", "bias"), false); + blk.ln_w = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "ln", "weight"), false); + blk.ln_b = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "ln", "bias"), false); + blk.pw1_w = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "pw1", "weight"), false); + blk.pw1_b = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "pw1", "bias"), false); + blk.grn_w = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "grn", "weight"), false); + blk.grn_b = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "grn", "bias"), false); + blk.pw2_w = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "pw2", "weight"), false); + blk.pw2_b = get_tensor(string_format(TN_YASA_STAGE_BLK, s, bi, "pw2", "bias"), false); + stage.blocks.push_back(blk); + } + + if (!stage.down_conv_w && stage.blocks.empty()) { + break; + } + model.yasa_stages.push_back(std::move(stage)); + } + } break; case PROJECTOR_TYPE_GLM4V: { model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight")); @@ -2801,6 +2857,8 @@ int clip_n_output_tokens_x(const struct clip_ctx * ctx, struct clip_image_f32 * return (img->nx / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: return img->nx / (params.patch_size * params.n_merge); + case PROJECTOR_TYPE_YASA2: + return 8; default: break; } @@ -2820,6 +2878,8 @@ int clip_n_output_tokens_y(const struct clip_ctx * ctx, struct clip_image_f32 * return (img->ny / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: return img->ny / (params.patch_size * params.n_merge); + case PROJECTOR_TYPE_YASA2: + return 8; default: break; } @@ -2843,6 +2903,10 @@ int clip_n_output_tokens(const struct clip_ctx * ctx, struct clip_image_f32 * im { // do nothing } break; + case PROJECTOR_TYPE_YASA2: + { + n_patches = 64; // adaptive average pooling to 8x8 tokens + } break; case PROJECTOR_TYPE_LDP: case PROJECTOR_TYPE_LDPV2: case PROJECTOR_TYPE_GLM_EDGE: @@ -3463,6 +3527,7 @@ bool clip_image_batch_encode(clip_ctx * ctx, const int n_threads, const clip_ima case PROJECTOR_TYPE_PHI4: case PROJECTOR_TYPE_COGVLM: case PROJECTOR_TYPE_HUNYUANOCR: + case PROJECTOR_TYPE_YASA2: { // do nothing } break; @@ -3689,6 +3754,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_KIMIVL: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_KIMIK25: + case PROJECTOR_TYPE_YASA2: return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_HUNYUANOCR: return ctx->model.mm_model_proj->ne[1]; diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 03d99e15b054..c30d79133efe 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -43,6 +43,14 @@ struct clip_graph_youtuvl : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_yasa2 : clip_graph { + clip_graph_yasa2(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + + ggml_tensor * layer_norm_channels(ggml_tensor * inp, ggml_tensor * w, ggml_tensor * b, float eps = 1e-6f); + ggml_tensor * convnext_grn(ggml_tensor * inp, ggml_tensor * w, ggml_tensor * b); +}; + struct clip_graph_minicpmv : clip_graph { clip_graph_minicpmv(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/yasa2.cpp b/tools/mtmd/models/yasa2.cpp new file mode 100644 index 000000000000..277ad3cba6c5 --- /dev/null +++ b/tools/mtmd/models/yasa2.cpp @@ -0,0 +1,217 @@ +// ABOUTME: Yasa2 vision encoder graph builder for ConvNeXt-based architecture. +// ABOUTME: Implements patch embedding, ConvNeXt stages with GRN, and adaptive pooling. + +#include "models.h" + +static ggml_tensor * add_channel_bias( + ggml_context * ctx0, + ggml_tensor * x_whcb, + ggml_tensor * b_c) { + if (!b_c) { + return x_whcb; + } + ggml_tensor * bc = ggml_cast(ctx0, b_c, x_whcb->type); + ggml_tensor * b4 = ggml_reshape_4d(ctx0, bc, 1, 1, bc->ne[0], 1); + b4 = ggml_repeat(ctx0, b4, x_whcb); + return ggml_add(ctx0, x_whcb, b4); +} + +static ggml_tensor * mul_channel_weight( + ggml_context * ctx0, + ggml_tensor * x_whcb, + ggml_tensor * w_c) { + if (!w_c) { + return x_whcb; + } + ggml_tensor * wc = ggml_cast(ctx0, w_c, x_whcb->type); + ggml_tensor * w4 = ggml_reshape_4d(ctx0, wc, 1, 1, wc->ne[0], 1); + w4 = ggml_repeat(ctx0, w4, x_whcb); + return ggml_mul(ctx0, x_whcb, w4); +} + +ggml_tensor * clip_graph_yasa2::layer_norm_channels(ggml_tensor * inp, ggml_tensor * w, ggml_tensor * b, float eps) { + // Match HF ConvNextLayerNorm(channels_first): + // u = mean_c(x), s = mean_c((x-u)^2), x = (x-u)/sqrt(s+eps) + // cast back to input dtype before affine. + ggml_tensor * cur = ggml_cast(ctx0, inp, GGML_TYPE_F32); + cur = ggml_permute(ctx0, cur, 2, 1, 0, 3); // [W,H,C,B] -> [C,H,W,B] + cur = ggml_cont(ctx0, cur); + + ggml_tensor * u = ggml_mean(ctx0, cur); // [1,H,W,B] + u = ggml_repeat(ctx0, u, cur); // [C,H,W,B] + ggml_tensor * xm = ggml_sub(ctx0, cur, u); // [C,H,W,B] + + ggml_tensor * s = ggml_mul(ctx0, xm, xm); // [C,H,W,B] + s = ggml_mean(ctx0, s); // [1,H,W,B] + s = ggml_clamp(ctx0, s, eps, 1e30f); // avoid div-by-zero in no-alloc warmup + s = ggml_sqrt(ctx0, s); // [1,H,W,B] + s = ggml_repeat(ctx0, s, xm); // [C,H,W,B] + + ggml_tensor * xhat = ggml_div(ctx0, xm, s); // [C,H,W,B] + xhat = ggml_permute(ctx0, xhat, 2, 1, 0, 3); // [W,H,C,B] + xhat = ggml_cont(ctx0, xhat); + xhat = ggml_cast(ctx0, xhat, inp->type); // HF casts back before affine + + xhat = mul_channel_weight(ctx0, xhat, w); + xhat = add_channel_bias(ctx0, xhat, b); + return xhat; +} + +ggml_tensor * clip_graph_yasa2::convnext_grn(ggml_tensor * inp, ggml_tensor * w, ggml_tensor * b) { + // Exact ConvNeXtV2 GRN: + // Gx = ||x||_2 over spatial dims (W,H), Nx = Gx / (mean_c(Gx) + eps) + // y = w * (x * Nx) + b + x + const int64_t wdim = inp->ne[0]; + const int64_t hdim = inp->ne[1]; + const int64_t cdim = inp->ne[2]; + const int64_t bdim = inp->ne[3]; + + // Keep GRN math in fp32 for stability; fp16/bf16 accumulation can drift. + ggml_tensor * inp_f = ggml_cast(ctx0, inp, GGML_TYPE_F32); + ggml_tensor * sq = ggml_mul(ctx0, inp_f, inp_f); + ggml_tensor * sq_flat = ggml_reshape_4d(ctx0, sq, wdim * hdim, cdim, 1, bdim); // [WH,C,1,B] + ggml_tensor * gx = ggml_sum_rows(ctx0, sq_flat); // [1,C,1,B] + gx = ggml_sqrt(ctx0, gx); // [1,C,1,B] + + ggml_tensor * gx_ch_first = ggml_permute(ctx0, gx, 1, 0, 2, 3); // [C,1,1,B] + gx_ch_first = ggml_cont(ctx0, gx_ch_first); + ggml_tensor * gx_mean = ggml_mean(ctx0, gx_ch_first); // [1,1,1,B] + + ggml_tensor * gx_mean_rep = ggml_repeat(ctx0, gx_mean, gx); // [1,C,1,B] + gx_mean_rep = ggml_clamp(ctx0, gx_mean_rep, 1e-6f, 1e30f); // approx +eps, warmup-safe + ggml_tensor * nx = ggml_div(ctx0, gx, gx_mean_rep); // [1,C,1,B] + nx = ggml_permute(ctx0, nx, 0, 2, 1, 3); // [1,1,C,B] + nx = ggml_cont(ctx0, nx); + nx = ggml_repeat(ctx0, nx, inp_f); // [W,H,C,B] + + ggml_tensor * xnx = ggml_mul(ctx0, inp_f, nx); + xnx = mul_channel_weight(ctx0, xnx, w); + xnx = add_channel_bias(ctx0, xnx, b); + ggml_tensor * out = ggml_add(ctx0, inp_f, xnx); + return ggml_cast(ctx0, out, inp->type); +} + +ggml_cgraph * clip_graph_yasa2::build() { + ggml_tensor * cur = build_inp_raw(); + + // Patch embedding Conv2d(kernel=4, stride=4) + ggml_tensor * patch_w = ggml_cast(ctx0, model.yasa_patch_w, GGML_TYPE_F32); + ggml_tensor * patch_inp = ggml_cast(ctx0, cur, GGML_TYPE_F32); + cur = ggml_conv_2d(ctx0, patch_w, patch_inp, patch_size, patch_size, 0, 0, 1, 1); + cur = add_channel_bias(ctx0, cur, model.yasa_patch_b); + cur = ggml_cont(ctx0, cur); + ggml_set_name(cur, "yasa2_patch_conv_out"); + cb(cur, "yasa2_patch_conv_out", -1); + cur = layer_norm_channels(cur, model.yasa_patch_ln_w, model.yasa_patch_ln_b, eps); + ggml_set_name(cur, "yasa2_patch_ln_out"); + cb(cur, "yasa2_patch_ln_out", -1); + + // ConvNeXt stages + for (size_t s = 0; s < model.yasa_stages.size(); ++s) { + const auto & stage = model.yasa_stages[s]; + + if (stage.down_conv_w) { + cur = layer_norm_channels(cur, stage.down_ln_w, stage.down_ln_b, eps); + cur = ggml_cont(ctx0, cur); + cur = ggml_conv_2d(ctx0, stage.down_conv_w, cur, 2, 2, 0, 0, 1, 1); + cur = add_channel_bias(ctx0, cur, stage.down_conv_b); + ggml_format_name(cur, "yasa2_stage%zu_down_out", s); + } + + for (size_t bi = 0; bi < stage.blocks.size(); ++bi) { + const auto & blk = stage.blocks[bi]; + ggml_tensor * res = cur; + + cur = ggml_cont(ctx0, cur); + ggml_tensor * x = ggml_conv_2d_dw(ctx0, blk.dw_w, cur, 1, 1, 3, 3, 1, 1); + x = add_channel_bias(ctx0, x, blk.dw_b); + x = layer_norm_channels(x, blk.ln_w, blk.ln_b, eps); + + // pwconv1/pwconv2 are HF Linear layers over channels; implement via matmul on tokens. + const int64_t w = x->ne[0]; + const int64_t h = x->ne[1]; + const int64_t b = x->ne[3]; + + ggml_tensor * tok = ggml_reshape_3d(ctx0, x, w * h, x->ne[2], b); // [T,C,B] + tok = ggml_permute(ctx0, tok, 1, 0, 2, 3); // [C,T,B] + tok = ggml_cont(ctx0, tok); + + tok = ggml_mul_mat(ctx0, blk.pw1_w, tok); // [4C,T,B] + if (blk.pw1_b) { + ggml_tensor * b1 = ggml_cast(ctx0, blk.pw1_b, tok->type); + b1 = ggml_reshape_3d(ctx0, b1, b1->ne[0], 1, 1); // [4C,1,1] + b1 = ggml_repeat(ctx0, b1, tok); + tok = ggml_add(ctx0, tok, b1); + } + x = ggml_permute(ctx0, tok, 1, 0, 2, 3); // [T,4C,B] + x = ggml_cont(ctx0, x); + x = ggml_reshape_4d(ctx0, x, w, h, tok->ne[0], b); // [W,H,4C,B] + x = ggml_gelu_erf(ctx0, x); + x = convnext_grn(x, blk.grn_w, blk.grn_b); + + tok = ggml_reshape_3d(ctx0, x, w * h, x->ne[2], b); // [T,4C,B] + tok = ggml_permute(ctx0, tok, 1, 0, 2, 3); // [4C,T,B] + tok = ggml_cont(ctx0, tok); + + tok = ggml_mul_mat(ctx0, blk.pw2_w, tok); // [C,T,B] + if (blk.pw2_b) { + ggml_tensor * b2 = ggml_cast(ctx0, blk.pw2_b, tok->type); + b2 = ggml_reshape_3d(ctx0, b2, b2->ne[0], 1, 1); // [C,1,1] + b2 = ggml_repeat(ctx0, b2, tok); + tok = ggml_add(ctx0, tok, b2); + } + x = ggml_permute(ctx0, tok, 1, 0, 2, 3); // [T,C,B] + x = ggml_cont(ctx0, x); + x = ggml_reshape_4d(ctx0, x, w, h, tok->ne[0], b); // [W,H,C,B] + + cur = ggml_add(ctx0, res, x); + ggml_format_name(cur, "yasa2_stage%zu_blk%zu_out", s, bi); + } + } + + // HF path adds vision position embeddings BEFORE adaptive pooling. + const int64_t pre_w = cur->ne[0]; + const int64_t pre_h = cur->ne[1]; + ggml_tensor * tokens_pre = ggml_reshape_3d(ctx0, cur, pre_w * pre_h, cur->ne[2], cur->ne[3]); // [T,C,B] + tokens_pre = ggml_permute(ctx0, tokens_pre, 1, 0, 2, 3); // [C,T,B] + tokens_pre = ggml_cont(ctx0, tokens_pre); + if (model.yasa_vision_pos_embed && tokens_pre->ne[1] == model.yasa_vision_pos_embed->ne[1]) { + const int64_t n_ch = model.yasa_vision_pos_embed->ne[0]; + const int64_t n_tokens = model.yasa_vision_pos_embed->ne[1]; + ggml_tensor * pos = ggml_reshape_3d(ctx0, model.yasa_vision_pos_embed, (int) n_ch, (int) n_tokens, 1); + pos = ggml_cont(ctx0, pos); + pos = ggml_cast(ctx0, pos, tokens_pre->type); + pos = ggml_repeat(ctx0, pos, tokens_pre); + tokens_pre = ggml_add(ctx0, tokens_pre, pos); + } + cur = ggml_permute(ctx0, tokens_pre, 1, 0, 2, 3); // [T,C,B] + cur = ggml_cont(ctx0, cur); + cur = ggml_reshape_4d(ctx0, cur, pre_w, pre_h, cur->ne[1], cur->ne[2]); // [W,H,C,B] + + // AdaptiveAvgPool2d target is 8x8 for real inputs, but warmup can use tiny images. + const int pooled_w = std::min(8, (int) cur->ne[0]); + const int pooled_h = std::min(8, (int) cur->ne[1]); + const int kw = std::max(1, (int) cur->ne[0] / pooled_w); + const int kh = std::max(1, (int) cur->ne[1] / pooled_h); + cur = ggml_pool_2d(ctx0, cur, GGML_OP_POOL_AVG, kw, kh, kw, kh, 0, 0); + + // [W,H,C,B] -> [C,T,B] + ggml_tensor * tokens = ggml_reshape_3d(ctx0, cur, cur->ne[0] * cur->ne[1], cur->ne[2], cur->ne[3]); + tokens = ggml_permute(ctx0, tokens, 1, 0, 2, 3); + tokens = ggml_cont(ctx0, tokens); + cb(tokens, "yasa2_tokens", -1); + cb(tokens, "yasa2_tokens_pos", -1); + + GGML_ASSERT(model.mm_0_w && model.mm_2_w); + ggml_tensor * embeddings = build_ffn( + tokens, + model.mm_0_w, model.mm_0_b, + nullptr, nullptr, + model.mm_2_w, model.mm_2_b, + FFN_GELU_ERF, + -1); + cb(embeddings, "yasa2_emb", -1); + + ggml_build_forward_expand(gf, embeddings); + return gf; +} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index d0a0a4865ef8..a21930a8d05c 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -293,6 +293,12 @@ struct mtmd_context { img_end = "<|vision_end|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_YASA2: + { + img_beg = ""; + img_end = ""; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_GEMMA3NV: { From b6f0af2682eb93a059183a6d792ae0d911a0845d Mon Sep 17 00:00:00 2001 From: Kwa Jie Hao Date: Tue, 31 Mar 2026 17:34:58 +0800 Subject: [PATCH 03/23] feat(chat): add reka-edge template handler (tools, thinking) - Add chat-reka.cpp/h implementing PEG-based parser for reka-edge format - Add Reka-Edge.jinja chat template - Detect reka-edge template in try_specialized_template() - Add LLAMA_EXAMPLE_MTMD to chat-template-file arg --- common/CMakeLists.txt | 2 + common/arg.cpp | 2 +- common/chat-reka.cpp | 132 ++++++++++++++++++++++++++++++++++++++++++ common/chat-reka.h | 14 +++++ common/chat.cpp | 10 ++++ 5 files changed, 159 insertions(+), 1 deletion(-) create mode 100644 common/chat-reka.cpp create mode 100644 common/chat-reka.h diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index b313a7320e56..7a3f40ce1794 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -53,6 +53,8 @@ add_library(${TARGET} STATIC chat-diff-analyzer.cpp chat-peg-parser.cpp chat-peg-parser.h + chat-reka.cpp + chat-reka.h chat.cpp chat.h common.cpp diff --git a/common/arg.cpp b/common/arg.cpp index 3d0183ed7026..8743901647f5 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3140,7 +3140,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params, const std::string & value) { params.chat_template = read_file(value); } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_TEMPLATE_FILE")); + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD}).set_env("LLAMA_ARG_CHAT_TEMPLATE_FILE")); add_opt(common_arg( {"--skip-chat-parsing"}, {"--no-skip-chat-parsing"}, diff --git a/common/chat-reka.cpp b/common/chat-reka.cpp new file mode 100644 index 000000000000..de175c71501f --- /dev/null +++ b/common/chat-reka.cpp @@ -0,0 +1,132 @@ +// ABOUTME: Reka Edge chat template handler with tool calling and thinking support. +// ABOUTME: Implements PEG-based parsing for / XML format. + +#include "chat-reka.h" + +#include "chat-auto-parser.h" +#include "chat-peg-parser.h" +#include "json-schema-to-grammar.h" + +static void reka_foreach_function(const json & tools, const std::function & fn) { + for (const auto & tool : tools) { + if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) { + continue; + } + fn(tool); + } +} + +static json reka_sort_tools_for_parsing(const json & tools) { + if (!tools.is_array()) { + return tools; + } + + json sorted = tools; + std::stable_sort(sorted.begin(), sorted.end(), [](const json & a, const json & b) { + const std::string a_name = a.contains("function") ? a.at("function").value("name", "") : ""; + const std::string b_name = b.contains("function") ? b.at("function").value("name", "") : ""; + return a_name.size() > b_name.size(); + }); + return sorted; +} + +common_chat_params common_chat_params_init_reka_edge(const common_chat_template & tmpl, + const autoparser::templates_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.preserved_tokens = { + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + }; + + const bool has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + const bool extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + const bool thinking_forced_open = + extract_reasoning && + inputs.enable_thinking && + inputs.add_generation_prompt && + (inputs.messages.empty() || inputs.messages.back().value("role", "") != "assistant"); + const bool include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + const json parser_tools = has_tools ? reka_sort_tools_for_parsing(inputs.tools) : inputs.tools; + + if (extract_reasoning) { + data.thinking_start_tag = ""; + data.thinking_end_tag = ""; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + const std::string THINK_START = ""; + const std::string THINK_END = ""; + const std::string TOOL_CALL = ""; + + auto end = p.end(); + auto reasoning = p.eps(); + + if (extract_reasoning) { + auto reasoning_body = p.reasoning(p.until_one_of({ THINK_END, TOOL_CALL })); + if (thinking_forced_open) { + reasoning = reasoning_body + p.optional(p.literal(THINK_END)) + p.space(); + } else { + reasoning = p.optional(p.literal(THINK_START) + reasoning_body + p.optional(p.literal(THINK_END)) + p.space()); + } + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return reasoning + p.content(p.rest()) + end; + } + + auto single_tool = p.standard_json_tools( + "", + "", + parser_tools, + /* parallel_tool_calls = */ false, + /* force_tool_calls = */ true); + + auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; + auto max_calls = inputs.parallel_tool_calls ? -1 : 1; + auto tool_calls = p.rule("tool-calls", p.repeat(single_tool + p.space(), min_calls, max_calls)); + auto content = p.content(p.until(TOOL_CALL)); + + return reasoning + content + tool_calls + end; + }); + + data.additional_stops = { "" }; + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + reka_foreach_function(parser_tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + if (inputs.parallel_tool_calls) { + string_replace_all( + data.grammar, + "root ::= tool-call\n", + "root ::= tool-call (space tool-call)*\n"); + } + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "" } + }; + } + + return data; +} diff --git a/common/chat-reka.h b/common/chat-reka.h new file mode 100644 index 000000000000..8ae11afc6bf0 --- /dev/null +++ b/common/chat-reka.h @@ -0,0 +1,14 @@ +// ABOUTME: Declares the Reka Edge chat template handler. +// ABOUTME: Provides tool calling, thinking, and vision token support for Reka Edge. + +#pragma once + +#include "chat.h" + +namespace autoparser { +struct templates_params; +} + +common_chat_params common_chat_params_init_reka_edge( + const common_chat_template & tmpl, + const autoparser::templates_params & inputs); diff --git a/common/chat.cpp b/common/chat.cpp index e27b6c3413c9..58284ec57dd4 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -3,6 +3,7 @@ #include "chat-auto-parser-helpers.h" #include "chat-auto-parser.h" #include "chat-peg-parser.h" +#include "chat-reka.h" #include "common.h" #include "ggml.h" #include "json-schema-to-grammar.h" @@ -2127,6 +2128,15 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_gemma4(tmpl, params); } + // Reka Edge - uses / with delimiters and vision tokens + if (src.find("") != std::string::npos && + src.find("continue_final_message") != std::string::npos && + src.find("") != std::string::npos && + src.find("num_video_frames") != std::string::npos) { + LOG_DBG("Using specialized template: Reka Edge\n"); + return common_chat_params_init_reka_edge(tmpl, params); + } + return std::nullopt; } From e5852faa2a6d0e2aae11105f4e504fc4316abc9d Mon Sep 17 00:00:00 2001 From: Kwa Jie Hao Date: Tue, 31 Mar 2026 17:38:38 +0800 Subject: [PATCH 04/23] feat: add reka vlm to gguf conversion script Converts Reka Yasa2 hf checkpoints to GGUF format: - Text decoder: Llama-arch with tiktoken/BPE vocab - Mmproj (--mmproj): ConvNeXt vision backbone + language_projection - Generates 2D sincos positional embeddings for vision encoder --- convert_reka_vlm_to_gguf.py | 311 ++++++++++++++++++++++++++++++++++++ 1 file changed, 311 insertions(+) create mode 100755 convert_reka_vlm_to_gguf.py diff --git a/convert_reka_vlm_to_gguf.py b/convert_reka_vlm_to_gguf.py new file mode 100755 index 000000000000..c6376415a8a6 --- /dev/null +++ b/convert_reka_vlm_to_gguf.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +""" +Convert Reka Yasa2 checkpoints to GGUF (text decoder and optional mmproj for MTMD). + +- Text: Llama-arch decoder + tiktoken/BPE vocab (bytes keys normalized for GGUF). +- Mmproj (--mmproj): ConvNeXt vision + language_projection for Yasa2 MTMD path. +""" + +from __future__ import annotations + +import re +from typing import Iterable + +import numpy as np +import torch +from transformers import AutoTokenizer + +import convert_hf_to_gguf as base + + +def _get_2d_sincos_pos_embed_yasa2(embed_dim: int, image_size: int = 50, seq_len: int = 256) -> np.ndarray: + """Match HF get_2d_sincos_pos_embed(hidden, image_size=50) then slice [:seq_len].""" + assert embed_dim % 2 == 0 + + def _get_1d(embed_dim_1d: int, pos_2d: np.ndarray) -> np.ndarray: + assert embed_dim_1d % 2 == 0 + omega = np.arange(embed_dim_1d // 2, dtype=np.float32) + omega /= embed_dim_1d / 2.0 + omega = 1.0 / (10000.0**omega) + out = np.einsum("hw,d->hwd", pos_2d, omega) + return np.concatenate([np.sin(out), np.cos(out)], axis=-1) + + grid_h = np.arange(image_size, dtype=np.float32) + grid_w = np.arange(image_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) + grid = np.stack(grid, axis=0) + emb_h = _get_1d(embed_dim // 2, grid[0]) + emb_w = _get_1d(embed_dim // 2, grid[1]) + pos = np.concatenate([emb_h, emb_w], axis=-1).reshape(image_size * image_size, embed_dim) + return pos[:seq_len].astype(np.float32, copy=False) + + +class RekaYasa2TextDecoderModel(base.LlamaModel): + model_arch = base.gguf.MODEL_ARCH.LLAMA + + def get_vocab_base(self) -> tuple[list[str], list[int], str]: + tokens: list[str] = [] + toktypes: list[int] = [] + + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + vocab_dict = tokenizer.get_vocab() + vocab_size = self.hparams.get("vocab_size", len(vocab_dict)) + assert max(vocab_dict.values()) < vocab_size + + tokpre = self.get_vocab_base_pre(tokenizer) + reverse_vocab = {idx: tok for tok, idx in vocab_dict.items()} + added_vocab = tokenizer.get_added_vocab() + added_tokens_decoder = getattr(tokenizer, "added_tokens_decoder", {}) + tiktoken_special = set(getattr(tokenizer, "tiktoken_special_tokens", {}).keys()) + + def token_to_str(tok: str | bytes) -> str: + if isinstance(tok, bytes): + return base.QwenModel.token_bytes_to_string(tok) + return tok + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(base.gguf.TokenType.UNUSED) + continue + + token = reverse_vocab[i] + token_str = token_to_str(token) + if token_str in tiktoken_special: + toktypes.append(base.gguf.TokenType.CONTROL) + elif token in added_vocab or token_str in added_vocab: + if i in added_tokens_decoder and getattr(added_tokens_decoder[i], "special", False): + toktypes.append(base.gguf.TokenType.CONTROL) + else: + toktypes.append(base.gguf.TokenType.USER_DEFINED) + else: + toktypes.append(base.gguf.TokenType.NORMAL) + tokens.append(token_str) + + return tokens, toktypes, tokpre + + def set_vocab(self): + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + mergeable_ranks = getattr(tokenizer, "mergeable_ranks", None) + if mergeable_ranks is None and hasattr(tokenizer, "tiktoken"): + mergeable_ranks = getattr(tokenizer.tiktoken, "_mergeable_ranks", None) + + if mergeable_ranks: + merges: list[str] = [] + for token, rank in mergeable_ranks.items(): + if len(token) == 1: + continue + merged = base.QwenModel.bpe(mergeable_ranks, token, max_rank=rank) + if len(merged) == 2: + merges.append(" ".join(map(base.QwenModel.token_bytes_to_string, merged))) + self.gguf_writer.add_token_merges(merges) + + special_vocab = base.gguf.SpecialVocab(self.dir_model, load_merges=False) + special_vocab.add_to_gguf(self.gguf_writer) + + # SpecialVocab finds no eos_token in the tokenizer config and writes nothing, + # causing llama.cpp to fall back to the gpt2 default of token 11 (comma) as EOS. + # Explicitly write the correct values: <|endoftext|> as EOS, as EOT. + self.gguf_writer.add_eos_token_id(tokenizer.tiktoken.encode_single_token("<|endoftext|>")) + self.gguf_writer.add_eot_token_id(tokenizer.tiktoken.encode_single_token("")) + + # Tiktoken/BPE typically does not add BOS; avoid shifting logits. + self.gguf_writer.add_add_bos_token(False) + + def modify_tensors( + self, data_torch, name: str, bid: int | None + ) -> Iterable[tuple[str, object]]: + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model.") :] + elif name.startswith("language_model."): + name = name[len("language_model.") :] + elif name.startswith("model.vision_model.") or name.startswith("model.connector."): + return + elif name.startswith("vision_model.") or name.startswith("connector."): + return + + if not ( + name == "lm_head.weight" + or name == "model.embed_tokens.weight" + or name == "model.norm.weight" + or name.startswith("model.layers.") + ): + return + + yield from super().modify_tensors(data_torch, name, bid) + + +class RekaYasa2VisionMmprojModel(base.MmprojModel): + """Vision backbone + language_projection tensors for MTMD Yasa2 clip graph.""" + + model_arch = base.gguf.MODEL_ARCH.MMPROJ + has_vision_encoder = True + has_audio_encoder = False + + def get_vision_config(self) -> dict[str, object] | None: + cfg = self.global_config.get("vision_config") + if not isinstance(cfg, dict): + return cfg + + out = dict(cfg) + depths = out.get("depths") + if isinstance(depths, list) and "num_hidden_layers" not in out: + out["num_hidden_layers"] = int(sum(int(x) for x in depths)) + return out + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_clip_has_vision_encoder(True) + self.gguf_writer.add_clip_projector_type("yasa2") + self.gguf_writer.add_vision_projection_dim(self.n_embd_text) + + vcfg = self.global_config.get("vision_config", {}) + self.gguf_writer.add_vision_image_size(int(vcfg.get("image_size", 512))) + self.gguf_writer.add_vision_patch_size(int(vcfg.get("patch_size", 4))) + self.gguf_writer.add_vision_embedding_length(int(vcfg.get("hidden_size", 2816))) + self.gguf_writer.add_vision_feed_forward_length(int(vcfg.get("hidden_size", 2816))) + self.gguf_writer.add_vision_block_count(0) + self.gguf_writer.add_vision_head_count(int(vcfg.get("num_attention_heads", 1) or 1)) + self.gguf_writer.add_vision_attention_layernorm_eps(float(vcfg.get("layer_norm_eps", 1e-6))) + self.gguf_writer.add_vision_use_gelu(True) + + mean = self.preprocessor_config.get("image_mean", [0.485, 0.456, 0.406]) + std = self.preprocessor_config.get("image_std", [0.229, 0.224, 0.225]) + self.gguf_writer.add_vision_image_mean(mean) + self.gguf_writer.add_vision_image_std(std) + use_vision_pos = self.global_config.get("use_vision_pos_embed", True) + if use_vision_pos: + hidden = int(vcfg.get("hidden_size", 2816)) + self._yasa_vision_pos_embed = _get_2d_sincos_pos_embed_yasa2(hidden, 50, 256) + else: + self._yasa_vision_pos_embed = None + self._yasa_pos_embed_yielded = False + + def _ensure_yasa_pos_embed(self) -> None: + if hasattr(self, "_yasa_vision_pos_embed"): + return + self._yasa_pos_embed_yielded = False + vcfg = self.global_config.get("vision_config", {}) + use_vision_pos = self.global_config.get("use_vision_pos_embed", True) + if use_vision_pos: + hidden = int(vcfg.get("hidden_size", 2816)) + self._yasa_vision_pos_embed = _get_2d_sincos_pos_embed_yasa2(hidden, 50, 256) + else: + self._yasa_vision_pos_embed = None + + def modify_tensors( + self, data_torch, name: str, bid: int | None + ) -> Iterable[tuple[str, object]]: + del bid + + if name.startswith("model.vision_model.") or name.startswith("model.language_projection."): + self._ensure_yasa_pos_embed() + short = name[len("model.") :] + if not self._yasa_pos_embed_yielded and self._yasa_vision_pos_embed is not None and "vision_model." in short: + self._yasa_pos_embed_yielded = True + yield "v.vision_pos_embed", torch.from_numpy(self._yasa_vision_pos_embed.copy()) + out_name = self._map_mmproj_name(short) + if out_name is not None: + yield out_name, data_torch + return + + if name.startswith("vision_model.") or name.startswith("language_projection."): + self._ensure_yasa_pos_embed() + if not self._yasa_pos_embed_yielded and self._yasa_vision_pos_embed is not None and "vision_model." in name: + self._yasa_pos_embed_yielded = True + yield "v.vision_pos_embed", torch.from_numpy(self._yasa_vision_pos_embed.copy()) + out_name = self._map_mmproj_name(name) + if out_name is not None: + yield out_name, data_torch + return + + return + + @staticmethod + def _map_mmproj_name(name: str) -> str | None: + if name.startswith("language_projection."): + mapping = { + "language_projection.0.weight": "mm.0.weight", + "language_projection.0.bias": "mm.0.bias", + "language_projection.2.weight": "mm.2.weight", + "language_projection.2.bias": "mm.2.bias", + } + return mapping.get(name) + + simple = { + "vision_model.backbone.embeddings.patch_embeddings.weight": "v.patch_embd.weight", + "vision_model.backbone.embeddings.patch_embeddings.bias": "v.patch_embd.bias", + "vision_model.backbone.embeddings.layernorm.weight": "v.patch_ln.weight", + "vision_model.backbone.embeddings.layernorm.bias": "v.patch_ln.bias", + "vision_model.backbone.layernorm.weight": "v.backbone_ln.weight", + "vision_model.backbone.layernorm.bias": "v.backbone_ln.bias", + } + if name in simple: + return simple[name] + + m = re.match( + r"vision_model\.backbone\.encoder\.stages\.(\d+)\.downsampling_layer\.(0|1)\.(weight|bias)$", + name, + ) + if m: + stage = int(m.group(1)) + layer_idx = int(m.group(2)) + wb = m.group(3) + if layer_idx == 0: + return f"v.stage.{stage}.down.ln.{wb}" + return f"v.stage.{stage}.down.conv.{wb}" + + m = re.match( + r"vision_model\.backbone\.encoder\.stages\.(\d+)\.layers\.(\d+)\.(dwconv|layernorm|pwconv1|grn|pwconv2)\.(weight|bias)$", + name, + ) + if m: + stage = int(m.group(1)) + blk = int(m.group(2)) + part = m.group(3) + wb = m.group(4) + part_map = { + "dwconv": "dw", + "layernorm": "ln", + "pwconv1": "pw1", + "grn": "grn", + "pwconv2": "pw2", + } + return f"v.stage.{stage}.blk.{blk}.{part_map[part]}.{wb}" + + return None + + def tensor_force_quant( + self, name: str, new_name: str, bid: int | None, n_dims: int + ) -> base.gguf.GGMLQuantizationType | bool: + del name, new_name, bid + if n_dims > 1: + return base.gguf.GGMLQuantizationType.F16 + return False + + +def register_reka_architectures() -> None: + model_classes = base.ModelBase._model_classes[base.ModelType.TEXT] + mmproj_classes = base.ModelBase._model_classes[base.ModelType.MMPROJ] + for arch_name in ( + "Yasa2ForConditionalGeneration", + "Yasa2Model", + "YasaCausalLM", + ): + model_classes[arch_name] = RekaYasa2TextDecoderModel + mmproj_classes["Yasa2ForConditionalGeneration"] = RekaYasa2VisionMmprojModel + + +def main() -> None: + register_reka_architectures() + base.main() + + +if __name__ == "__main__": + main() From 51e6f02abe7b5669da5b951ef11e61ab8674b5d8 Mon Sep 17 00:00:00 2001 From: Kwa Jie Hao Date: Tue, 31 Mar 2026 17:44:24 +0800 Subject: [PATCH 05/23] test: add Reka Edge chat template and parser tests - test-chat-template: oracle tests comparing Jinja engine output vs common_chat_templates_apply for text, tools, thinking, images, video - test-chat: PEG parser tests for Reka Edge format, round-trip tests for image/video content parts, common path integration tests --- tests/test-chat-template.cpp | 212 +++++++++++++++++++++++++++++++++++ tests/test-chat.cpp | 150 +++++++++++++++++++++++++ 2 files changed, 362 insertions(+) diff --git a/tests/test-chat-template.cpp b/tests/test-chat-template.cpp index bf45d737c832..a009e1aea8da 100644 --- a/tests/test-chat-template.cpp +++ b/tests/test-chat-template.cpp @@ -225,6 +225,58 @@ static std::string normalize_newlines(const std::string & s) { #endif } +static std::string read_file(const std::string & path) { + std::ifstream infile(path, std::ios::binary); + if (!infile) { + infile = std::ifstream("../" + path, std::ios::binary); + } + if (!infile) { + throw std::runtime_error("Could not open file: " + path); + } + return std::string(std::istreambuf_iterator(infile), std::istreambuf_iterator()); +} + +static void apply_common_template_overrides(const json & input, common_chat_templates_inputs & inputs) { + if (input.contains("add_generation_prompt")) { + inputs.add_generation_prompt = input.at("add_generation_prompt").get(); + } + if (input.contains("enable_thinking")) { + inputs.enable_thinking = input.at("enable_thinking").get(); + } + if (input.contains("continue_final_message")) { + inputs.chat_template_kwargs["continue_final_message"] = input.at("continue_final_message").dump(); + } + if (input.contains("num_img_tokens")) { + inputs.chat_template_kwargs["num_img_tokens"] = input.at("num_img_tokens").dump(); + } + if (input.contains("num_video_frames")) { + inputs.chat_template_kwargs["num_video_frames"] = input.at("num_video_frames").dump(); + } + if (input.contains("chat_template_kwargs")) { + for (const auto & item : input.at("chat_template_kwargs").items()) { + inputs.chat_template_kwargs[item.key()] = item.value().dump(); + } + } +} + +static std::string format_using_common_json( + const std::string & template_str, + json input) { + auto tmpls = common_chat_templates_init( + /* model= */ nullptr, + template_str, + input.value("bos_token", ""), + input.value("eos_token", "")); + + common_chat_templates_inputs inputs; + inputs.use_jinja = true; + inputs.messages = common_chat_msgs_parse_oaicompat(input.value("messages", json::array())); + inputs.tools = common_chat_tools_parse_oaicompat(input.value("tools", json::array())); + apply_common_template_overrides(input, inputs); + + return normalize_newlines(common_chat_templates_apply(tmpls.get(), inputs).prompt); +} + static std::string format_using_common( const std::string & template_str, @@ -333,6 +385,164 @@ static common_chat_msg simple_msg(const std::string & role, const std::string & return msg; } +static void assert_string_equals(const std::string & label, const std::string & expected, const std::string & actual) { + if (expected != actual) { + std::cerr << "Mismatch in " << label << "\n"; + std::cerr << "Expected:\n" << expected << "\n"; + std::cerr << "Actual:\n" << actual << "\n"; + assert(expected == actual); + } +} + +static void run_reka_edge_oracle_tests() { + const std::string template_str = read_file("models/templates/Reka-Edge.jinja"); + + const json tool = { + {"type", "function"}, + {"function", { + {"name", "lookup_weather"}, + {"description", "Look up the weather."}, + {"parameters", { + {"type", "object"}, + {"properties", { + {"city", {{"type", "string"}}} + }}, + {"required", json::array({"city"})} + }} + }} + }; + + const std::vector> cases = { + { + "text_only", + { + {"messages", json::array({ + {{"role", "system"}, {"content", "Be brief."}}, + {{"role", "user"}, {"content", "Hello there"}}, + })}, + {"add_generation_prompt", true}, + {"enable_thinking", false}, + } + }, + { + "typed_text_spacing", + { + {"messages", json::array({ + {{"role", "user"}, {"content", json::array({ + {{"type", "text"}, {"text", "Hello"}}, + {{"type", "text"}, {"text", "world"}}, + {{"type", "text"}, {"text", "again"}}, + })}}, + })}, + {"add_generation_prompt", true}, + {"enable_thinking", false}, + } + }, + { + "tools_reasoning", + { + {"messages", json::array({ + {{"role", "system"}, {"content", "Use tools when needed."}}, + {{"role", "user"}, {"content", "What is the weather in Singapore?"}}, + })}, + {"tools", json::array({tool})}, + {"add_generation_prompt", true}, + {"enable_thinking", true}, + } + }, + { + "assistant_continue_final_message", + { + {"messages", json::array({ + {{"role", "user"}, {"content", "Continue this answer"}}, + {{"role", "assistant"}, {"content", "The first point is"}}, + })}, + {"add_generation_prompt", false}, + {"continue_final_message", true}, + {"enable_thinking", false}, + } + }, + { + "tool_response_history", + { + {"messages", json::array({ + {{"role", "user"}, {"content", "Run the lookup"}}, + {{"role", "assistant"}, {"content", ""}, + {"tool_calls", json::array({ + {{"type", "function"}, + {"function", {{"name", "lookup_weather"}, {"arguments", {{"city", "Singapore"}}}}}} + })}}, + {{"role", "tool"}, {"name", "lookup_weather"}, {"tool_call_id", "call0"}, {"content", "Sunny, 31C"}}, + })}, + {"tools", json::array({tool})}, + {"add_generation_prompt", true}, + {"enable_thinking", false}, + } + }, + { + "image_placeholder", + { + {"messages", json::array({ + {{"role", "user"}, {"content", json::array({ + {{"type", "text"}, {"text", "Describe this"}}, + {{"type", "image_url"}}, + })}}, + })}, + {"add_generation_prompt", true}, + {"enable_thinking", false}, + } + }, + { + "video_placeholder", + { + {"messages", json::array({ + {{"role", "user"}, {"content", json::array({ + {{"type", "video_url"}}, + {{"type", "text"}, {"text", "Summarize it"}}, + })}}, + })}, + {"add_generation_prompt", true}, + {"enable_thinking", false}, + } + }, + { + "mixed_history", + { + {"messages", json::array({ + {{"role", "system"}, {"content", "Use tools and inspect media carefully."}}, + {{"role", "user"}, {"content", json::array({ + {{"type", "image_url"}}, + {{"type", "text"}, {"text", "What is shown?"}}, + })}}, + {{"role", "assistant"}, {"content", "I will inspect it."}, {"reasoning_content", "Need to inspect the image first."}}, + {{"role", "assistant"}, {"content", ""}, + {"tool_calls", json::array({ + {{"type", "function"}, + {"function", {{"name", "lookup_weather"}, {"arguments", {{"city", "Singapore"}}}}}} + })}}, + {{"role", "tool"}, {"name", "lookup_weather"}, {"tool_call_id", "call1"}, {"content", "Cloudy"}}, + {{"role", "user"}, {"content", json::array({ + {{"type", "video_url"}}, + {{"type", "text"}, {"text", "And now summarize the clip"}}, + })}}, + })}, + {"tools", json::array({tool})}, + {"add_generation_prompt", true}, + {"enable_thinking", true}, + {"num_img_tokens", 4}, + {"num_video_frames", 3}, + } + }, + }; + + for (const auto & [name, input_case] : cases) { + json input = input_case; + auto expected = format_using_direct_engine(template_str, input)->as_string().str(); + auto actual = format_using_common_json(template_str, input_case); + assert_string_equals("Reka oracle case: " + name, normalize_newlines(expected), actual); + } +} + int main_automated_tests(void) { // jinja::enable_debug(true); @@ -706,6 +916,8 @@ int main_automated_tests(void) { } } + run_reka_edge_oracle_tests(); + std::cout << "\nOK: All tests passed successfully.\n"; return 0; diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 3b8de5ce02e5..e483d10081c4 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -816,6 +816,34 @@ const common_chat_msg message_user_parts{ /* .tool_call_id = */ "", }; +const common_chat_msg message_user_image_parts{ + "user", + /* .content = */ "", + /* .content_parts = */ + { + { "text", "Describe this" }, + { "image_url", "" }, + }, + /* .tool_calls = */ {}, + /* .reasoning_content = */ "", + /* .tool_name = */ "", + /* .tool_call_id = */ "", +}; + +const common_chat_msg message_user_video_parts{ + "user", + /* .content = */ "", + /* .content_parts = */ + { + { "video_url", "" }, + { "text", "Summarize it" }, + }, + /* .tool_calls = */ {}, + /* .reasoning_content = */ "", + /* .tool_name = */ "", + /* .tool_call_id = */ "", +}; + static common_chat_msg simple_assist_msg(const std::string & content, const std::string & reasoning_content = "", const std::string & tool_name = "", @@ -1407,6 +1435,8 @@ static void test_msgs_oaicompat_json_conversion() { std::vector msgs{ message_user, message_user_parts, + message_user_image_parts, + message_user_video_parts, message_assist_call, message_assist_call_thoughts, message_assist_call_thoughts_unparsed, @@ -1439,6 +1469,23 @@ static void test_msgs_oaicompat_json_conversion() { "]"), common_chat_msgs_to_json_oaicompat({ message_user_parts }).dump(2)); + assert_equals(std::string("[\n" + " {\n" + " \"role\": \"user\",\n" + " \"content\": [\n" + " {\n" + " \"type\": \"text\",\n" + " \"text\": \"Describe this\"\n" + " },\n" + " {\n" + " \"type\": \"image_url\",\n" + " \"text\": \"\"\n" + " }\n" + " ]\n" + " }\n" + "]"), + common_chat_msgs_to_json_oaicompat({ message_user_image_parts }).dump(2)); + // Note: content is "" instead of null due to workaround for templates that render null as "None" assert_equals(std::string("[\n" " {\n" @@ -3595,6 +3642,51 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // Reka Edge + { + auto tst = peg_tester("models/templates/Reka-Edge.jinja", detailed_debug); + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(false) + .expect(message_assist) + .run(); + tst.test("I'm\nthinking\n\nHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist_thoughts) + .run(); + tst.test("\n{\"name\": \"special_function\", \"arguments\": {\"arg1\": 1}}\n") + .enable_thinking(false) + .tools({ special_function_tool }) + .expect(message_assist_call) + .run(); + tst.test("Hello, world!\nWhat's up?\n\n{\"name\": \"special_function\", \"arguments\": {\"arg1\": 1}}\n") + .enable_thinking(false) + .tools({ special_function_tool }) + .expect(message_assist_call_content) + .run(); + tst.test("I'm\nthinking\n\n\n{\"name\": \"special_function\", \"arguments\": {\"arg1\": 1}}\n") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_assist_call_thoughts) + .run(); + tst.test("\n{\"name\": \"special_function\", \"arguments\": {\"arg1\": 1}}\n\n\n{\"name\": \"special_function_with_opt\", \"arguments\": {\"arg1\": 1, \"arg2\": 2}}\n") + .enable_thinking(false) + .parallel_tool_calls(true) + .tools({ special_function_tool, special_function_tool_with_optional_param }) + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + { "special_function_with_opt", R"({"arg1": 1, "arg2": 2})", {} }, + }) + .run(); + tst.test("\n{\"name\": \"special_function\", \"arguments\": {\"arg") + .enable_thinking(false) + .tools({ special_function_tool }) + .is_partial(true) + .expect(message_assist_call_cutoff_args) + .run(); + } + // Apriel 1.5 { auto tst = peg_tester("models/templates/unsloth-Apriel-1.5.jinja", detailed_debug); @@ -4077,6 +4169,63 @@ static void test_template_output_peg_parsers(bool detailed_debug) { } } +static void test_reka_edge_common_path() { + auto tmpls = read_templates("models/templates/Reka-Edge.jinja"); + + { + common_chat_templates_inputs inputs; + common_chat_msg system_msg; + system_msg.role = "system"; + system_msg.content = "Use tools and inspect media carefully."; + + common_chat_msg tool_call_msg = simple_assist_msg("", "", "special_function", "{\"arg1\": 1}"); + + common_chat_msg tool_msg; + tool_msg.role = "tool"; + tool_msg.tool_name = "special_function"; + tool_msg.tool_call_id = "call0"; + tool_msg.content = "Sunny"; + + inputs.messages = { system_msg, message_user_image_parts, tool_call_msg, tool_msg, message_user_video_parts }; + inputs.tools = { special_function_tool }; + inputs.enable_thinking = true; + inputs.add_generation_prompt = true; + inputs.chat_template_kwargs["num_img_tokens"] = "4"; + inputs.chat_template_kwargs["num_video_frames"] = "3"; + + auto params = common_chat_templates_apply(tmpls.get(), inputs); + + if (params.prompt.find("") == std::string::npos) { + throw std::runtime_error("Reka Edge prompt did not render image placeholder"); + } + if (params.prompt.find("