Rust SBert

Rust 端口的 sentence-transformers，使用 rust-bert 和 tch-rs。

支持 rust-tokenizers 和 Hugging Face 的 tokenizers。

支持的模型

• distiluse-base-multilingual-cased：支持语言：阿拉伯语、中文、荷兰语、英语、法语、德语、意大利语、韩语、波兰语、葡萄牙语、俄语、西班牙语、土耳其语。在扩展 STS2017 上的性能：80.1

• 基于 DistilRoBERTa 的分类器

用法

示例

API 设计得非常易于使用，可以让你以简单直接的方式创建高质量的多语言句子嵌入。

通过指定模型目录加载带有权重的 SBert 模型

let mut home: PathBuf = env::current_dir().unwrap();
home.push("path-to-model");

你可以使用不同版本的模型，这些模型使用不同的标记器

// To use Hugging Face tokenizer
let sbert_model = SBertHF::new(home.to_str().unwrap());

// To use Rust-tokenizers
let sbert_model = SBertRT::new(home.to_str().unwrap());

现在，你可以编码你的句子

let texts = ["You can encode",
             "As many sentences",
             "As you want",
             "Enjoy ;)"];

let batch_size = 64;

let output = sbert_model.forward(texts.to_vec(), batch_size).unwrap();

参数 batch_size 可以留空，让模型使用其默认值。

然后你可以将 output 句子嵌入用于任何你想要的应用。

将模型从 Python 转换为 Rust

首先，获取 UKPLabs 提供的模型（所有模型都在 这里）

mkdir -p models/distiluse-base-multilingual-cased

wget -P models https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/v0.2/distiluse-base-multilingual-cased.zip

unzip models/distiluse-base-multilingual-cased.zip -d models/distiluse-base-multilingual-cased

然后，你需要将模型转换为合适的格式（需要 pytorch）

python utils/prepare_distilbert.py models/distiluse-base-multilingual-cased

还提供了用于运行转换脚本的 Docker 化环境

docker build -t tch-converter -f utils/Dockerfile .

docker run \
  -v $(pwd)/models/distiluse-base-multilingual-cased:/model \
  tch-converter:latest \
  python prepare_distilbert.py /model

最后，在 "output_attentions"": true 中设置 distiluse-base-multilingual-cased/0_distilbert/config.json。