CandleEmbed 🧨

在Hugging Face上使用任何模型进行文本嵌入。嵌入您的应用程序中。用GPU替换您的API账单。

特性

• 常用嵌入模型的枚举或指定来自HF（查看排行榜）的自定义模型

• CUDA支持的GPU

• 带有易于访问的配置设置的构建器

安装

将以下内容添加到您的Cargo.toml文件中

[dependencies]
candle_embed = "*"

[dependencies]
candle_embed = { version = "*", features = ["cuda"] } // For CUDA support

基础 🫡

use candle_embed::{CandleEmbedBuilder, WithModel};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create a builder with default settings
    //
    let candle_embed = CandleEmbedBuilder::new().build()?;
    
    // Embed a single text
    //
    let text = "This is a test sentence.";
    let embeddings = candle_embed.embed_one(text)?;
    
    // Embed a batch of texts
    //
    let texts = vec![
        "This is the first sentence.",
        "This is the second sentence.",
        "This is the third sentence.",
    ];
    let batch_embeddings = candle_embed.embed_batch(texts)?;
    

    Ok(())
}

自定义 🤓

    // Custom settings
    //
    builder
        .approximate_gelu(false)
        .mean_pooling(true)
        .normalize_embeddings(false)
        .truncate_text_len_overflow(true)

    // Set model from preset
    //
    builder
        .set_model_from_presets(WithModel::UaeLargeV1);

    // Or use a custom model and revision
    //
    builder
        .custom_embedding_model("avsolatorio/GIST-small-Embedding-v0")
        .custom_model_revision("d6c4190")

    // Will use the first available CUDA device (Default)
    //
    builder.with_device_any_cuda()

    // Use a specific CUDA device
    //
    builder.with_device_specific_cuda(ordinal: usize);

    // Use CPU (CUDA options fail over to this)
    //
    builder.with_device_cpu()

    // Unload the model and tokenizer, dropping them from memory
    //
    candle_embed.unload();

    // ---

    // These are automatically loaded from the model's `config.json` after builder init

    // model_dimensions
    // This is the same as "hidden_size"
    //
    let dimensions = candle_embed.model_dimensions;

    // model_max_input 
    // This is the same as "max_position_embeddings"
    // If `truncate_text_len_overflow == false`, and your input exceeds this a panic will result
    // If you don't want to worry about this, the default truncation strategy will just chop the end off the input
    // However, you lose accuracy by mindlessly truncating your inputs
    //
    let max_position_embeddings = candle_embed.model_max_input;
    

    // ---

标记化 🧮

    // Generate tokens using the model

    let texts = vec![
            "This is the first sentence.",
            "This is the second sentence.",
            "This is the third sentence.",
        ];
    let text = "This is the first sentence.";

    // Get tokens from a batch of texts
    //
    let batch_tokens = candle_embed.tokenize_batch(&texts)?;
    assert_eq!(batch_tokens.len(), 3);
    // Get tokens from a single text
    //
    let tokens = candle_embed.tokenize_one(text)?;
    assert!(!tokens.is_empty());

    // Get a count of tokens
    // This is important to use if you are using your own chunking strategy
    // For example, using a text splitter on any text string whose token count exceeds candle_embed.model_max_input

    // Get token counts from a batch of texts
    //
    let batch_tokens = candle_embed.token_count_batch(&texts)?;
   
    // Get token count from a single text
    //
    let tokens = candle_embed.token_count(text)?;
    assert!(tokens > 0);

这与text-embeddings-inference 🤗有何不同

• TEI实现了一个客户端-服务器模型。这需要作为它自己的外部服务器、docker容器或本地服务器运行。
• CandleEmbed是为了嵌入而设计的：它可以作为一个crate安装并在进程中运行。
• TEI非常优化，并且可扩展性非常好。
• CandleEmbed运行速度快（使用GPU），但它并不是为像HuggingFace的文本嵌入API那样大规模提供服务而创建的。

text-embeddings-inference是一个更成熟的项目，并且备受尊敬。我建议您查看它！

这与fastembed.rs 🦀有何不同

• 两者都可以作为crate使用！
• 通过输入自定义模型的repo_name/model_id从hf-hub下载并运行。
• CandleEmbed的设计使得项目可以实现自己的截断和/或分块策略。
• CandleEmbed使用CUDA。FastEmbed使用ONNX。
• 最后，CandleEmbed使用Candle。

fastembed.rs是一个更成熟的项目，并且备受尊敬。我建议您查看它！

路线图

• 多GPU支持
• 基准测试系统

许可证

本项目受MIT许可证许可。

贡献

发布我的动机是希望有人指出我可能做错了什么！