用 Candle 在 Rust 中运行 SPLADE 稀疏检索：查询词项权重与句对相似度实战

人工智能大模型机器学习深度学习本地部署模型推理服务【免费下载链接】candleMinimalist ML framework for Rust项目地址https://gitcode.com/GitHub_Trending/ca/candle点击查看免费下载SPLADESparse Lexical and Expansion model是一种基于 BERT MLM掩码语言模型头与稀疏正则化训练的神经检索模型它把文本编码成稀疏的词项-权重向量同时具备倒排索引高效检索、显式词汇匹配与结果可解释性等优势且对域外数据的泛化能力优于稠密向量方案。本文以 Candle 示例 examples/splade 为主体完整讲解如何在 Rust 中加载 SPLADE 模型计算查询的稀疏词项权重向量并基于稀疏向量计算句子之间的语义相似度读完即可复现文中全部命令并掌握其底层实现log(1ReLU)稀疏化、token 级 max-pooling、L2 归一化的原理。SPLADE 示例能做什么该示例main.rs提供两条能力计算给定查询query的稀疏词项权重向量即对一段查询文本做词项扩展term expansion输出一组 token 及其权重权重非零的 token 构成该查询的稀疏表示用稀疏向量计算一组句子两两之间的相似度对一批句子批量编码输出按相似度降序排列的句对及得分用于快速验证稀疏向量的语义匹配质量。模型权重在首次运行时从 Hub 下载默认仓库为prithivida/Splade_PP_en_v1内部复用BertForMaskedLM模型因此该示例也是理解 Candle 中 BERT 系列模型加载与推理的良好入口。快速开始运行示例任务一计算查询的稀疏词项权重在仓库根目录执行建议--release以获得更好的推理性能cargo run --example splade --release -- --prompt Here is a test sentence程序会输出两行结果第一行是权重非零的 token 序列即词项扩展结果第二行是对应的权重值 the out there still house inside position outside stay standing hotel sitting dog animal sit bird cat statue cats [0.10270107, 0.269471, 0.047469813, 0.0016636598, 0.05394874, 0.23105666, 0.037475716, 0.45949644, 0.009062732, 0.06790692, 0.0327835, 0.33122346, 0.16863061, 0.12688516, 0.340983, 0.044972017, 0.47724655, 0.01765311, 0.37331146]可以看到一段看似普通的长句被扩展出了大量语义相关的词汇如从句子内容扩展出dog、cat、bird、statue、hotel、outside等这正是 SPLADE 通过 BERT MLM 头实现的词项扩展能力检索时不再依赖精确字面匹配而是把语义相关的词也纳入索引候选。任务二句子间稀疏相似度不传--prompt时示例会使用内置的 8 个句子样本见 main.rs对它们两两计算余弦相似度并输出最高的前 5 组cargo run --example splade --release score: 0.47 The new movie is awesome The new movie is so great score: 0.43 The cat sits outside The cat plays in the garden score: 0.14 I love pasta Do you like pizza? score: 0.11 A man is playing guitar The cat plays in the garden score: 0.05 A man is playing guitar A woman watches TV从得分排序可以看出语义相近的句对如同义改写、同主题陈述得分显著更高而主题无关的句对得分很低说明稀疏向量在保留显式词汇匹配的同时具备语义判别能力。命令行参数详解示例通过clap解析命令行参数定义见 main.rs全部参数如下参数类型默认值说明--cpu布尔开关关强制在 CPU 上运行而不是自动选择 GPU--tracing布尔开关关启用 tracing生成trace-时间戳.json文件可用于性能分析--model-id字符串prithivida/Splade_PP_en_v1指定 Hub 上的模型仓库 ID用于替换默认 SPLADE 模型--revision字符串main模型仓库的 revision分支/commit/tag用于固定模型版本--tokenizer-file字符串无自动下载本地 tokenizer 文件路径tokenizer.json--weight-files字符串无自动下载本地权重文件路径支持model.safetensors或pytorch_model.bin--config-file字符串无自动下载本地模型配置文件路径config.json--prompt字符串无设置后进入查询词项权重模式不设置则进入句对相似度模式其中--model-id、--revision、--tokenizer-file、--weight-files、--config-file的组合可以让你完全脱离 Hub 离线使用例如先把模型文件下载到本地再通过文件路径参数加载适合无外网环境或 CI 场景。另外模型推理默认优先使用 CUDA GPU其次 Metal GPU最后回退 CPU设备选择逻辑见 candle-examples/src/lib.rs非 GPU 环境会自动打印提示并在 CPU 上运行。模型与权重的加载流程1. 从 Hub 获取三件套程序通过candle_examples::hub::Api基于hf-hub客户端见 candle-examples/src/hub.rs在首次运行时下载以下文件并缓存到本地tokenizer.jsonBERT WordPiece tokenizer 的序列化配置config.jsonBERT 架构配置词表大小、隐藏层维度、层数等被反序列化为candle_transformers::models::bert::Config权重文件优先尝试model.safetensors以VarBuilder::from_mmaped_safetensors内存映射加载若不存在则回退到pytorch_model.bin以VarBuilder::from_pth加载两者都缺失时报错见 main.rs。下载行为受hf-hub的环境变量控制如HF_TOKEN、HF_HUB_CACHE这与 Candle 其他示例一致。2. 构建 BERT MLM 模型权重加载后直接构造 MLM 模型见 main.rslet vb if weights_filename.ends_with(model.safetensors) { unsafe { VarBuilder::from_mmaped_safetensors([weights_filename], dtype, device).unwrap() } } else { VarBuilder::from_pth(weights_filename, dtype, device).unwrap() }; let model BertForMaskedLM::load(vb, config)?;BertForMaskedLM是 candle-transformers/src/models/bert.rs 中提供的 MLM 模型由BertModel编码器加BertOnlyMLMHead含BertLMPredictionHead即 transform 词表大小的线性解码器组成。其forward先让输入经过完整 BERT 编码器得到序列输出再经 MLM 头映射到词表维度的 logits见 bert.rs。SPLADE 正是复用这个词表维度 logits作为每个 token 对每个词典词的激活信号。稀疏向量生成原理从 logits 到词项权重以查询模式为例见 main.rs核心计算分为三步1. 前向得到 logits对输入 token 序列做编码并前向推理得到形状为(batch, seq_len, vocab_size)的 MLM logitsys。2. 稀疏化log(1 ReLU(·))SPLADE 的核心激活公式为let vec Tensor::log( Tensor::try_from(1.0)? .to_dtype(dtype)? .to_device(device)? .broadcast_add(ys.relu()?)?, )? .max(1)?;即先对 logits 做ReLU将负激活清零再取log(1 x)压缩正激活的数值范围最后沿 token 维做max-pooling.max(1)得到每个词典词的最终权重。这里的数学形式对应 SPLADE 论文中的稀疏正则项约束权重为 0 的词项不参与后续检索从而形成真正的稀疏向量。3. L2 归一化let vec normalize_l2(vec)?;normalize_l2定义见 main.rs实现为pub fn normalize_l2(v: Tensor) - ResultTensor { Ok(v.broadcast_div(v.sqr()?.sum_keepdim(1)?.sqrt()?)?) }即对每个向量除以自身 L2 范数把向量投影到单位球面上便于后续用点积/余弦相似度直接比较。4. 输出词项与权重归一化后程序筛选出所有非零权重位置把这些索引用 tokenizer 解码成可读 token并同步输出对应权重值见 main.rs。权重越大说明该词项对查询语义的贡献越高。句对相似度的批量计算流程相似度模式与查询模式共享同一套稀疏化逻辑但需要处理批量输入与注意力掩码见 main.rs批量编码与 padding将 8 个句子一次性编码tokenizer 的 padding 策略被设置为BatchLongest按批内最长句补齐并显式取出每个样本的attention_mask构造批张量token_ids、attention_mask通过Tensor::stack拼接成批token_type_ids置零前向 稀疏化调用model.forward(..., Some(attention_mask))得到 logits应用相同的log(1 ReLU)与 max-pooling唯一区别是这里会用注意力掩码把 padding 位置的激活乘零broadcast_mul(attention_mask.unsqueeze(2)?)避免 padding token 污染词项权重两两计算余弦相似度对归一化后的稀疏向量两两点积(e_i * e_j)?.sum_all()即cosine Σ w_i·w_j由于向量已 L2 归一化点积即余弦相似度排序输出 Top-5相似度按降序排序后打印前 5 组句对。这也解释了为什么输出里会出现score: 0.43 The cat sits outside The cat plays in the garden这类语义相关但字面不同的匹配——稀疏向量在倒排检索中可以直接用权重高的词项命中彼此无需稠密向量的 ANN 索引。在 Candle 项目中进一步探索如果你想了解 SPLADE 用到的 BERT 架构细节嵌入层、编码器层、MLM 头可直接阅读 candle-transformers/src/models/bert.rs其中Config的默认结构与 BERT-base 一致12 层、768 隐藏维度、12 头注意力见该文件impl Default for Config本示例依赖的 Hub 下载封装在 candle-examples/src/hub.rs其他示例如 BERT 系列、稠密向量检索示例也复用同一套 API如果你想替换成自己的 SPLADE 微调模型把--model-id指向你的 Hub 仓库或通过--weight-files/--tokenizer-file/--config-file指向本地文件即可无需修改任何代码。小结本文完整覆盖了 Candle 中 SPLADE 稀疏检索示例的两大任务与全部命令行参数并深入其实现从 Hub 加载 BERT MLM 权重、log(1 ReLU)稀疏化激活、token 级 max-pooling、L2 归一化到批量编码与余弦相似度排序。这套流程演示了如何把 Hugging Face 生态的稀疏检索模型无缝迁移到 Rust 运行时也为基于倒排索引的高效检索系统提供了可直接落地的向量生成环节。赞分享人工智能大模型机器学习深度学习本地部署模型推理服务【免费下载链接】candleMinimalist ML framework for Rust项目地址https://gitcode.com/GitHub_Trending/ca/candle点击查看免费下载相关推荐Sentence-Transformers 稀疏编码器语义文本相似度STS实战指南SPLADE 嵌入与相似度计算Sentence Transformers 稀疏编码器语义文本相似度STS实战指南SPLADE 嵌入与相似度计算 导读 本文基于 sentence tra人工智能NLPEmbedding微调机器学习用 Rust 运行 BERT 句向量推理candle 示例的句子嵌入与相似度计算实战指南用 Rust 运行 BERT 句向量推理candle 示例的句子嵌入与相似度计算实战指南 candle 是一个用 Rust 编写的极简机器学习框架本文以仓库人工智能大模型机器学习深度学习本地部署模型推理服务PyPTO 实现 torch.t() 转置基于 view assemble 的搬运语义 Kernel 参考pypto-api-explore / t.mdPyPTO 实现 torch.t 转置基于 view assemble 的搬运语义 Kernel 参考pypto api explore / t.md人工智能大模型机器学习深度学习本地部署模型推理服务上一篇FortuneSheet移动端适配与响应式设计的最佳实践下一篇Newton物理参数调优弹性、摩擦与阻尼的设置方法创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考