是否可以使用 Candle 从 NV-Embed 获取嵌入？

Question

我想做的是一个输出任意输入嵌入的 CLI 程序。 为此，我想使用嵌入模型进行推理，我选择了

NV-Embed-v2

。我选择的框架是 Candle，但我也看了 Mistral-RS。

基本上，我想做的是这个代码片段： https://huggingface.co/nvidia/NV-Embed-v2 但有铁锈和蜡烛。

我尝试的是从 Mistral Candle 的示例开始，因为 NV-Embed 的 HF 页面显示：

Model Details / Base Decoder-only LLM: Mistral-7B-v0.1

。

我用

nvidia/NV-Embed-v2

替换了原始代码中的模型 id，并且能够从 Hugging Face 下载权重，但是在加载配置后，我得到了这个：

Error: missing field `vocab_size` at line 101 column 1

然后我将从 HF 加载的 JSON 配置中的值硬编码到新创建的

candle_transformers::models::mistral::Config

实例。之后，

Mistral::new(&config, vb)

失败并显示：

Error: cannot find tensor model.embed_tokens.weight

有没有办法解决这个问题——也许还有一些其他基于 Candle 的开源作品可以作为我的灵感？或者，也许这是一个很容易诊断的常见错误？

Answer 1

蜡烛寻找

model.embed_tokens.weight

，而原始张量名称是

embedding_model.embed_tokens.weight

。你只需要更改candle_transformers中的这行

mistral.rs

即可。

// from
let vb_m = vb.pp("model");
//to
let vb_m = vb.pp("embedding_model");