如何使用Rust+Candle框架基于NV-Embed-v2实现输入任意内容输出Embedding的CLI程序？

三、编写CLI交互逻辑 #

用clap库让你的CLI支持参数输入，比如：

• 支持通过--input参数直接传入要处理的文本
• 支持从标准输入读取文本（适合批量处理场景）
• 可选加--output-format参数，让用户选择输出是纯文本向量还是JSON格式

四、核心代码示例 #

这里给你一个简化的可运行代码片段，你可以基于这个扩展：

use clap::Parser;
use candle_core::{Device, Tensor};
use candle_transformers::models::mistral::{MistralConfig, MistralForCausalLM};
use tokenizers::Tokenizer;

#[derive(Parser, Debug)]
#[command(about = "CLI to generate embeddings using NV-Embed-v2 with Candle")]
struct Args {
    /// Input text to generate embedding for
    #[arg(short, long)]
    input: String,

    /// Optional: Output format (text/json), default is text
    #[arg(short, long, default_value = "text")]
    output_format: String,
}

fn main() -> anyhow::Result<()> {
    // 解析CLI参数
    let args = Args::parse();
    // 选择运行设备，有GPU可以换成Device::Cuda(0)
    let device = Device::Cpu;

    // 加载NV-Embed-v2的tokenizer
    let tokenizer = Tokenizer::from_pretrained("nvidia/NV-Embed-v2", None)?;
    let encoding = tokenizer.encode(args.input, true)?;
    let tokens = encoding.get_ids().to_vec();
    // 转成模型需要的张量格式
    let tokens = Tensor::new(&tokens, &device)?.unsqueeze(0)?;

    // 加载模型配置和权重
    let config = MistralConfig::from_pretrained("nvidia/NV-Embed-v2")?;
    let model = MistralForCausalLM::from_pretrained(&config, "nvidia/NV-Embed-v2", &device)?;

    // 运行模型得到最后一层隐藏状态
    let logits = model.forward(&tokens)?;
    let last_hidden_state = logits.squeeze(0)?; // 形状为[序列长度, 隐藏层维度]

    // 均值池化得到最终Embedding
    let embedding = last_hidden_state.mean(0)?; // 形状为[隐藏层维度]
    let embedding_vec: Vec<f32> = embedding.to_vec1()?;

    // 根据输出格式打印结果
    match args.output_format.as_str() {
        "json" => {
            println!("{}", serde_json::to_string(&embedding_vec)?);
        }
        _ => {
            println!("Embedding vector: {}", embedding_vec.iter()
                .map(|x| format!("{:.6}", x))
                .collect::<Vec<_>>()
                .join(", "));
        }
    }

    Ok(())
}

五、常见问题调试 #

• 权重加载失败：如果遇到维度不匹配的问题，可能是原Mistral示例的配置是针对生成任务的，你可以尝试手动加载模型权重，禁用生成相关的头部层
• Tokenizer加载失败：如果自动加载有问题，可以手动下载tokenizer的相关文件（比如tokenizer.json），然后从本地路径加载
• 性能优化：CPU运行时可以启用MKL或OpenBLAS加速，GPU运行时记得切换到Cuda或Metal设备

备注：内容来源于stack exchange，提问作者Zomagk