在 ik_llama.cpp 中用 llama-gritlm 运行 GRIT 模型：一条命令同时完成向量嵌入与文本生成

在 ik_llama.cpp 中用 llama-gritlm 运行 GRIT 模型：一条命令同时完成向量嵌入与文本生成

【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp

本文是一份基于 ik_llama.cpp 仓库中 examples/gritlm/README.md 及其示例实现 gritlm.cpp 的实战指南。GRIT（Generative Representational Instruction Tuning）是一类通过指令区分任务、既能产出文本嵌入（embedding）又能进行常规文本生成（generation）的模型，本指南将带你完成模型下载、示例编译与运行，并从源码层面拆解指令模板构造、均值池化、归一化与余弦相似度计算等核心环节，让你能在自己的 RAG 管道中直接复用这套方案。

GRIT 是什么：一个模型，两种能力

GRIT 是 "Generative Representational Instruction Tuning"（生成式表示指令微调）的缩写。它基于"指令即任务开关"的思路训练而成：同一个模型，根据 prompt 中的指令不同，既可以输出可用于向量检索的文本嵌入，也可以像普通 LLM 一样进行文本生成。

以 ik_llama.cpp 中的示例程序llama-gritlm为例，它演示了两类行为：

• 给定检索指令（如"Given a scientific paper title, retrieve the paper's abstract"），模型对查询与文档分别输出归一化嵌入向量，随后计算两两之间的余弦相似度，用于判断文档与查询的相关性；
• 给定生成指令（如"Please write me a poem about my recent hike of Mt. Fuji at midnight in the style of Shakespeare"），模型切换回生成模式，逐 token 输出一首莎士比亚风格的诗歌。

这种"一模型两用"的能力正是 GRIT 论文（arXiv:2402.09906）提出的核心观点：当前大多数模型只擅长生成或只擅长嵌入二者之一，而 GRIT 通过指令区分两者，在 Massive Text Embedding Benchmark（MTEB）上取得了当时同规模模型中的领先嵌入表现，同时保持了一流的生成能力，且两种任务统一训练互不损失性能。

为什么 GRIT 特别适合 RAG

原始文档中专门讨论了 GRIT 在Retrieval-Augmented Generation（RAG）场景下的价值。RAG 的常规流程是：

1. 把需要作为上下文使用的文档切分并做token 化，生成向量嵌入后存入向量数据库；
2. 用户发起查询时，先为查询生成嵌入，再在向量数据库中检索最相似的向量，取回对应文档；
3. 将"查询 + 取回的上下文"一起交给 LLM，由 LLM 基于上下文生成回答。

问题出在第 3 步：传统方案中，查询会先后被 token 化并计算两次——一次在检索阶段做嵌入，一次在生成阶段作为 prompt 输入。而 GRIT 方案的关键收益在于：

• 同一个模型既能做嵌入又能做生成，无需部署两套模型（检索模型 + 生成模型），从而简化架构、减少显存/内存占用；
• 检索阶段计算出的查询嵌入可以直接缓存复用，生成阶段无需再次对查询进行 token 化与嵌入计算，对长文档场景可显著提速。

这正是llama-gritlm示例存在的意义：在同一个推理上下文里，先用嵌入模式完成检索打分，再切换回生成模式产出回答，验证"一次查询、两种用途"的完整闭环。

运行前准备：下载 GRIT 模型与构建示例

1. 下载 GGUF 格式的 GRIT 模型

原始文档给出了使用仓库自带下载脚本 scripts/hf.sh 获取模型的命令：

$ scripts/hf.sh --repo cohesionet/GritLM-7B_gguf --file gritlm-7b_q4_1.gguf --outdir models

该脚本支持--url、--repo、--file、--outdir等参数（参见 scripts/hf.sh 中的 usage 说明），这里--repo指定 Hugging Face 上的 GGUF 仓库，--file指定要下载的文件（示例使用 Q4_1 量化），--outdir指定输出目录为models。你也可以用其他 GRIT 系 GGUF 模型替换，只需保证模型遵循 GritLM 的指令模板约定（见下文"指令构造"一节）。

2. 构建 llama-gritlm 示例

示例程序通过 CMake 构建。在 examples/CMakeLists.txt 中，gritlm子目录被add_subdirectory(gritlm)默认纳入构建；其自身的构建规则定义在 examples/gritlm/CMakeLists.txt：

set(TARGET llama-gritlm) add_executable(${TARGET} gritlm.cpp) install(TARGETS ${TARGET} RUNTIME) target_link_libraries(${TARGET} PRIVATE common llama ${CMAKE_THREAD_LIBS_INIT}) target_compile_features(${TARGET} PRIVATE cxx_std_11)

可见目标二进制名为llama-gritlm，静态链接common与llama两个库，采用 C++11 标准。在仓库根目录按常规流程配置并构建即可（例如cmake -B build && cmake --build build --target llama-gritlm），产物中即包含llama-gritlm可执行文件。

运行示例与预期输出

使用上一步下载的模型直接运行：

$ ./llama-gritlm -m models/gritlm-7b_q4_1.gguf

程序会依次执行两部分演示，预期输出（与原始文档一致）如下。

第一部分：嵌入与相似度打分。程序先用检索指令嵌入两个查询，再对两个文档做嵌入，输出 2×2 的余弦相似度矩阵：

Cosine similarity between "Bitcoin: A Peer-to-Peer Electronic Cash System" and "A purely peer-to-peer version of electronic cash w" is: 0.605 Cosine similarity between "Bitcoin: A Peer-to-Peer Electronic Cash System" and "All text-based language problems can be reduced to" is: 0.103 Cosine similarity between "Generative Representational Instruction Tuning" and "A purely peer-to-peer version of electronic cash w" is: 0.112 Cosine similarity between "Generative Representational Instruction Tuning" and "All text-based language problems can be reduced to" is: 0.547

注意观察结果：查询 1（比特币论文标题）与文档 1（比特币论文摘要）相似度 0.605 远高于与文档 2（GRIT 论文摘要）的 0.103；查询 2（GRIT 论文标题）与文档 2 的相似度 0.547 也远高于与文档 1 的 0.112。这证明嵌入模式正确捕捉了语义相关性，正是 RAG 检索打分所需的信号。

第二部分：文本生成。程序随后切换回生成模式，输出一首模拟莎士比亚风格的富士山夜攀诗歌：

Oh, brave adventurer, who dared to climb The lofty peak of Mt. Fuji in the night, When shadows lurk and ghosts do roam, And darkness reigns, a fearsome sight. Thou didst set out, with heart aglow, To conquer this mountain, so high, And reach the summit, where the stars do glow, And the moon shines bright, up in the sky. Through the mist and fog, thou didst press on, With steadfast courage, and a steadfast will, Through the darkness, thou didst not be gone, But didst climb on, with a steadfast skill. At last, thou didst reach the summit's crest, And gazed upon the world below, And saw the beauty of the night's best, And felt the peace, that only nature knows. Oh, brave adventurer, who dared to climb The lofty peak of Mt. Fuji in the night, Thou art a hero, in the eyes of all, For thou didst conquer this mountain, so bright.

两段输出共同验证了 GRIT 模型"同一个权重、按指令切换嵌入/生成两种模式"的核心能力。需要说明的是，具体相似度数值与生成的诗歌内容会随模型版本、量化方式与采样细节略有浮动，以上为参考输出。

源码拆解：指令模板与嵌入编码流程

示例程序的主体逻辑在 examples/gritlm/gritlm.cpp 中，其中main()依次演示嵌入与生成两条路径。下面按关键函数逐层拆解。

1. 指令模板构造：gritlm_instruction()

static std::string gritlm_instruction(const std::string & instruction) { return !instruction.empty() ? "<|user|>\n" + instruction + "\n<|embed|>\n" : "<|embed|>\n"; }

这是整个示例最关键的约定之一：GritLM 通过特殊 token<|embed|>触发嵌入模式。构造规则为：

• 有检索指令时，格式为<|user|>\n{instruction}\n<|embed|>\n；
• 无指令（如对入库文档做嵌入）时，仅使用<|embed|>\n。

示例中正是这样使用的——文档无需附加检索指令（gritlm_instruction("")），只有查询需要附带任务指令（gritlm_instruction(instruction)），这与 GritLM 原版的使用方式保持一致。

2. 嵌入编码：encode()

encode()函数对一组句子批量生成归一化嵌入，核心步骤（对应 gritlm.cpp 中encode实现）如下：

const std::string input_string = instruction + sentences[i]; std::vector<llama_token> inputs = common_tokenize(mdl, input_string, true, false);

• 将指令与句子拼接后，通过 common.h 中声明的common_tokenize完成 token 化；
• 单独对指令再做一次 token 化得到n_inst，用于记录指令占用的 token 数——后续均值池化时要跳过这些 token，避免指令污染句子表示；
• 源码中注释提到 GritLM 的 EOS 为空串，因此没有像常规做法那样在末尾追加 EOS token。

随后进入批处理与解码阶段：

// add input to batch (this increments n_tokens) for (int32_t j = 0; j < n_toks; j++) { common_batch_add(batch, inputs[j], j, { 0 }, j >= n_inst); } // clear previous kv_cache values (irrelevant for embeddings) llama_kv_cache_clear(ctx); llama_set_embeddings(ctx, true); llama_set_causal_attn(ctx, false); // run model llama_decode(ctx, batch);

这里有几个值得注意的底层 API（均声明于 include/llama.h）：

• common_batch_add(batch, token, pos, seq_ids, logits)：把 token 加入llama_batch，其中最后一个布尔参数控制是否对该位置输出 logits。示例中j >= n_inst意味着只有句子部分（非指令部分）才请求输出；
• llama_kv_cache_clear(ctx)：清空 KV 缓存（见 llama.h 中 "Clear the KV cache - both cell info is erased and KV data is zeroed"），因为嵌入计算不需要跨句子复用缓存；
• llama_set_embeddings(ctx, true)：切换到嵌入模式——"If true, embeddings will be returned but logits will not"（见 llama.h 对应注释）；
• llama_set_causal_attn(ctx, false)：关闭因果注意力。这是嵌入模式的关键：双向注意力让每个 token 都能看到上下文两侧的信息，从而产生更好的句子表示；而生成模式必须开启因果注意力（只看到过去 token），这两个模式正是通过该开关在同一个上下文中切换。

解码完成后读取嵌入并做均值池化（mean pooling）：

uint64_t n_embd = llama_model_n_embd(mdl); std::vector<float> emb_unorm(n_embd, 0.0f); // sum up all token embeddings for (int32_t k = n_inst; k < n_toks; k++) { float * emb = llama_get_embeddings_ith(ctx, k); for (uint64_t j = 0; j < n_embd; j++) { emb_unorm[j] += emb[j]; } } // divide by number of tokens (mean pooling) const uint64_t n_sent = n_toks - n_inst; for (uint64_t j = 0; j < n_embd; j++) { emb_unorm[j] /= n_sent; }

• llama_model_n_embd(mdl)获取模型嵌入维度；
• llama_get_embeddings_ith(ctx, k)取第 k 个 token 的嵌入向量（shape 为[n_embd]，见 llama.h 中对该 API 的说明）；
• 只累加k >= n_inst的 token 嵌入（即跳过指令 token），再除以句子 token 数n_toks - n_inst，得到句子的平均嵌入。

最后调用归一化工具：

std::vector<float> emb_norm(emb_unorm.size()); common_embd_normalize(emb_unorm.data(), emb_norm.data(), n_embd); result.push_back(emb_norm);

3. 归一化与相似度：common_embd_normalize与common_embd_similarity_cos

这两个工具函数实现在 common/common.cpp 中（声明见 common/common.h）。

common_embd_normalize(inp, out, n, embd_norm)的归一化方式由第四个参数控制：

• embd_norm = -1：不归一化；
• embd_norm = 0：按最大绝对值归一化，并缩放到 int16 范围（除以 32760.0）；
• embd_norm = 2：欧几里得范数归一化（L2）；
• 其他默认分支：p-范数归一化。

示例采用默认值2（L2 归一化），使嵌入向量落在单位超球面上，从而让后续的余弦相似度计算退化为简单的向量点积。源码中norm = sum > 0.0 ? 1.0 / sum : 0.0f还对零向量做了保护。

common_embd_similarity_cos(embd1, embd2, n)则直接计算两个向量的余弦相似度，其实现为：

sum += embd1[i] * embd2[i]; sum1 += embd1[i] * embd1[i]; sum2 += embd2[i] * embd2[i]; // Handle the case where one or both vectors are zero vectors if (sum1 == 0.0 || sum2 == 0.0) { if (sum1 == 0.0 && sum2 == 0.0) { return 1.0f; // two zero vectors are similar } return 0.0f; } return sum / (sqrt(sum1) * sqrt(sum2));

即cos(θ) = a·b / (|a|·|b|)，并显式处理了零向量边界情况。由于示例中查询与文档均已做 L2 归一化，此时分母恒为 1，相似度即点积。示例输出中的 0.605 / 0.103 / 0.112 / 0.547 正是由该函数计算得出（见 gritlm.cpp 中common_embd_similarity_cos的四处调用）。

源码拆解：文本生成路径

生成逻辑封装在generate()函数中，与嵌入模式形成鲜明的对照：

llama_kv_cache_clear(ctx); llama_set_embeddings(ctx, false); llama_set_causal_attn(ctx, true);

• llama_set_embeddings(ctx, false)：退出嵌入模式；
• llama_set_causal_attn(ctx, true)：重新开启因果注意力，使模型只能看到历史 token——这正是自回归生成的要求。

随后是典型的自回归循环：对 prompt 逐 token 构建 batch 并llama_decode，用llama_get_logits_ith(ctx, bat.n_tokens - 1)取最后一个位置的 logits，构造llama_token_data_array候选后调用llama_sample_token_greedy做贪心采样，遇到 EOS（llama_token_eos）即停止；common_token_to_piece负责把 token 还原为文本片段并实时流式打印（stream=true时）。

示例的生成 prompt 为：

<|user|> Please write me a poem about my recent hike of Mt. Fuji at midnight in the style of Shakespeare. <|assistant|>

源码注释特别指出：GritLM 模型没有使用 system prompt 微调，因此系统级指令应直接并入 user 指令中，这正是上述 prompt 把"以莎士比亚风格写诗"作为 user 指令一部分的原因，也解释了示例开头<|user|>与结尾<|assistant|>的模板结构。

关键 API 一览与 RAG 落地要点

嵌入/生成双模式的核心 API（均见 include/llama.h）

落地为自有 RAG 管道的要点

1. 指令模板不可省：查询必须包裹为<|user|>\n{指令}\n<|embed|>\n，入库文档仅用<|embed|>\n；特殊 token 的拼写必须与模型训练约定完全一致；
2. 池化策略需对齐：示例采用"跳过指令 token 的均值池化"，若改用其他池化（如 CLS、加权池化），需确保与模型训练时的做法一致，否则检索质量会明显下降；
3. 归一化与度量：统一使用 L2 归一化（common_embd_normalize默认参数），相似度用common_embd_similarity_cos；向量入库后可进一步配合向量数据库做 ANN 检索；
4. 长文档提速：文档嵌入可在离线阶段一次性完成入库，查询嵌入在线计算并缓存，生成阶段直接用缓存结果，避免二次 token 化——这正是原始文档强调的 RAG 加速收益；
5. 调试辅助：源码中保留了#define GRIT_DEBUG宏，取消注释可打印每个 token 的[id:piece]对照与完整嵌入向量，便于核对 token 化结果是否与原版 GritLM 参考实现一致。

小结

llama-gritlm示例用不足 220 行的代码，完整演示了 GRIT 模型在 ik_llama.cpp 上的双模式用法：以<|user|>/<|embed|>指令模板触发嵌入、以llama_set_embeddings与llama_set_causal_attn双开关切换模式、以均值池化 + L2 归一化 + 余弦相似度完成检索打分、再切回因果注意力完成自回归生成。对于希望构建"检索与生成一体化"RAG 管道的开发者，本文档与示例代码（examples/gritlm/README.md、gritlm.cpp）以及底层 API（include/llama.h、common/common.cpp）构成了从"跑通示例"到"自建流程"的完整参考路径。

【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp