BGE-M3联合嵌入在FastEmbed-rs中的应用: dense、sparse与ColBERT三合一

BGE-M3联合嵌入在FastEmbed-rs中的应用: dense、sparse与ColBERT三合一

【免费下载链接】fastembed-rsRust library for generating vector embeddings, reranking locally!项目地址: https://gitcode.com/gh_mirrors/fa/fastembed-rs

FastEmbed-rs是一个功能强大的Rust库,专门用于在本地生成向量嵌入和重排序。它支持多种模型,包括文本嵌入、稀疏文本嵌入、图像嵌入和重排序模型。其中,BGE-M3模型是一个非常重要的模型,它能够在单次前向传递中同时生成密集、稀疏和ColBERT嵌入,为各种自然语言处理任务提供了全面的向量表示。

BGE-M3模型简介

BGE-M3模型是由BAAI开发的一款先进的嵌入模型,它具有以下特点:

  • 多模态输出:能够同时生成dense、sparse和ColBERT三种类型的嵌入,满足不同场景的需求。
  • 高效性能:在单次前向传递中完成三种嵌入的生成,大大提高了处理效率。
  • 广泛语言支持:支持100多种语言,具有很强的通用性。
  • 量化版本:提供INT8量化版本,优化了CPU上的推理性能。

BGE-M3模型的这些特点使得它在信息检索、文本匹配、语义理解等任务中具有广泛的应用前景。

FastEmbed-rs中BGE-M3的实现

在FastEmbed-rs中,BGE-M3的实现主要集中在src/bgem3_embedding目录下,包括impl.rsinit.rsmod.rs等文件。这些文件定义了BGE-M3模型的初始化、嵌入生成等核心功能。

初始化BGE-M3模型

要使用BGE-M3模型,首先需要进行初始化。FastEmbed-rs提供了多种初始化方式,以满足不同的使用场景。

默认初始化

使用默认选项初始化BGE-M3模型非常简单,代码如下:

use fastembed::{Bgem3Embedding, Bgem3InitOptions, Bgem3Model}; // With default options let mut model = Bgem3Embedding::try_new(Default::default())?;

这种方式将使用默认的模型参数,包括默认的量化模型BGEM3Q,它经过优化,适合在CPU上运行。

自定义初始化

如果需要自定义模型参数,可以使用Bgem3InitOptions结构体来设置,例如:

// With custom options (supporting custom max length up to 8192 tokens) let mut model = Bgem3Embedding::try_new( Bgem3InitOptions::new(Bgem3Model::BGEM3Q) .with_max_length(1024) .with_show_download_progress(true), )?;

在这个例子中,我们设置了最大序列长度为1024,并启用了下载进度显示。

生成BGE-M3联合嵌入

初始化模型后,就可以使用embed方法生成联合嵌入了。该方法接受一个文本列表,并返回包含dense、sparse和ColBERT嵌入的Bgem3EmbeddingOutput结构体。

let documents = vec![ "Hello, World!", "This is an example passage.", "fastembed-rs is licensed under Apache 2.0", "i dont know" ]; // Generate all three representations in a single forward pass let output = model.embed(documents, None)?; println!("Dense dimension: {}", output.dense[0].len()); // -> Dense dimension: 1024 let sparse_emb = &output.sparse[0]; println!("Sparse non-zero tokens: {}", sparse_emb.indices.len()); println!("ColBERT token count: {}", output.colbert[0].len());

从输出结果可以看出,dense嵌入的维度为1024,sparse嵌入包含一定数量的非零 tokens,而ColBERT嵌入则包含多个 token 向量。

BGE-M3嵌入的应用

BGE-M3生成的三种嵌入各有特点,可以在不同的场景中发挥作用:

  • Dense嵌入:适用于传统的向量检索任务,如相似性搜索。
  • Sparse嵌入:适合关键词匹配和稀疏检索,能够捕捉文本中的重要关键词。
  • ColBERT嵌入:提供了细粒度的 token 级向量表示,有助于更精确的语义匹配。

通过结合这三种嵌入,可以构建更强大的检索系统,提高检索的准确性和召回率。

使用注意事项

在使用BGE-M3模型时,需要注意以下几点:

模型缓存

模型在首次使用时会下载,并缓存到本地。缓存位置可以通过FASTEMBED_CACHE_DIR环境变量设置,默认位置为.fastembed_cache。如果设置了HF_HOME环境变量,则会优先使用该位置。

自定义模型

如果需要使用自定义的BGE-M3模型,可以通过try_new_from_user_definedtry_new_from_path方法加载本地模型文件。例如:

// Create a Bgem3Embedding instance from model files provided by the user. pub fn try_new_from_user_defined( model: UserDefinedBgem3Model, options: InitOptionsUserDefined, ) -> Result<Self> { // ... } // Create a Bgem3Embedding instance from a model directory on disk. pub fn try_new_from_path( model_path: impl AsRef<std::path::Path>, tokenizer_files: TokenizerFiles, options: InitOptionsUserDefined, ) -> Result<Self> { // ... }

量化模型限制

默认的量化模型BGEM3Q是为CPU优化的,如果尝试在GPU上运行,可能会失败。如果需要在GPU上进行推理,或者有其他自定义需求,可以使用ONNX导出脚本导出自己的模型(FP32、FP16或INT8),并通过try_new_from_path加载。

总结

BGE-M3联合嵌入在FastEmbed-rs中的应用为自然语言处理任务提供了强大的工具。通过单次前向传递生成dense、sparse和ColBERT三种嵌入,不仅提高了处理效率,还为各种检索和匹配任务提供了全面的向量表示。无论是构建信息检索系统、文本匹配工具,还是进行语义理解研究,BGE-M3模型都能发挥重要作用。

如果你还没有尝试过FastEmbed-rs和BGE-M3模型,不妨通过以下命令克隆仓库,开始你的探索之旅:

git clone https://gitcode.com/gh_mirrors/fa/fastembed-rs

相信BGE-M3联合嵌入将为你的项目带来新的可能性,帮助你构建更高效、更准确的自然语言处理应用。

【免费下载链接】fastembed-rsRust library for generating vector embeddings, reranking locally!项目地址: https://gitcode.com/gh_mirrors/fa/fastembed-rs

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考