Show HN: Model2vec-Rs – 基于 Rust 的快速静态文本嵌入

内容跳转

这是一个轻量级的、使用 Rust 实现的，用于加载和推断 Model2Vec 静态嵌入模型的库。对于模型蒸馏和训练，可以使用 Python Model2Vec package。

快速开始

添加依赖库:

cargo add model2vec-rs

生成嵌入向量:

use anyhow::Result;
use model2vec_rs::model::StaticModel;
fn main() -> Result<()> {
  // 从 Hugging Face Hub 或本地路径加载模型
  // args = (repo_or_path, token, normalize, subfolder)
  let model = StaticModel::from_pretrained("minishlab/potion-base-8M", None, None, None)?;
  // 准备句子列表
  let sentences = vec![
    "Hello world".to_string(),
    "Rust is awesome".to_string(),
  ];
  // 创建嵌入向量
  let embeddings = model.encode(&sentences);
  println!("Embeddings: {:?}", embeddings);
  Ok(())
}

使用 CLI 生成嵌入向量：

# 单个句子
cargo run -- encode "Hello world" minishlab/potion-base-8M
# 从文件中读取多行
echo -e "Hello world\nRust is awesome" > input.txt
cargo run -- encode input.txt minishlab/potion-base-8M --output embeds.json

使用自定义参数生成嵌入向量:

let embeddings = model.encode_with_args(
  &sentences,   // input texts
  Some(512), // max length
  1024,    // batch size
);

模型

我们提供了一些可以直接使用的模型。这些模型可以在 HuggingFace hub 上找到，并可以使用 from_pretrained 方法加载。模型列表如下：

模型 | 语言 | Sentence Transformer | 参数量 | 任务 ---|---|---|---|--- potion-base-32M | 英语 | bge-base-en-v1.5 | 32.3M | 通用 potion-base-8M | 英语 | bge-base-en-v1.5 | 7.5M | 通用 potion-base-4M | 英语 | bge-base-en-v1.5 | 3.7M | 通用 potion-base-2M | 英语 | bge-base-en-v1.5 | 1.8M | 通用 potion-retrieval-32M | 英语 | bge-base-en-v1.5 | 32.3M | 检索 M2V_multilingual_output | 多语言 | LaBSE | 471M | 通用

性能

我们将 Rust 实现的性能与 Python 版本的 Model2Vec 进行了比较。该基准测试在 CPU 上以单线程运行。

实现 | 吞吐量 ---|--- Rust | 8000 samples/秒 Python | 4650 samples/秒

Rust 版本大约比 Python 版本快 1.7 倍。

许可

MIT

关于

Model2Vec 的官方 Rust 实现。

minishlab.github.io/