如何使用Rust实现语义搜索引擎?

22-11-08 banq

语义搜索引擎是一种推荐系统,它依靠词语的含义来提供更好的搜索结果。它与传统的全文搜索引擎不同,后者依靠关键词匹配来提供结果。
语义搜索引擎允许你搜索概念,而不仅仅是关键词。它了解意义和不同概念之间的关系,并能根据这些关系提供更相关的结果。

在这篇文章中,我们将讨论如何在Rust中建立一个语义搜索引擎。我们将解释什么是嵌入、转化器和最近邻搜索,以及如何使用KD树来执行这些搜索。

什么是嵌入?
嵌入是机器学习和统计学中的一个术语。嵌入是数据的数学表示,可用于降维、相似性测量和许多其他任务。
在自然语言处理(NLP)中,文本嵌入是将一个词或短语映射到一个实数的向量中。该向量在一个高维空间中代表该词或短语。
词嵌入的想法在21世纪初首次提出,但直到2013年,随着word2vec算法的发布和2018年BERT的发布,词嵌入才被广泛使用。

什么是Transformer?
嵌入器被用于NLP任务,如机器翻译、问题回答和文本分类。它们也被用于推荐系统和知识图谱中。
Transformer变换器是一种神经网络,可用于NLP任务,如机器翻译和问题回答。变换器于2017年在论文《Attention Is All You Need》中首次提出。
Transformer变换器与传统的递归神经网络(RNN)不同,它们不需要连续的数据。这使得它们很适合像机器翻译这样的任务,在这种情况下,输入和输出不一定是同一顺序的。
变换器也比RNN更容易并行,这使得它们的训练速度更快。
句子变换器Sentence Transformers是文本和图像嵌入的最先进的模型。

如何使用句子变换器进行文本嵌入?

rust-bert crate 在 Rust 中提供了现成的 NLP 管道和基于转换器的模型(BERT、DistilBERT、GPT2……)。
crate 还提供了一个 Sentence Embeddings 模型。

在本文中,我们将使用 rust-bert 库来生成文本嵌入。我们将使用全 MiniLM-L12-v2模型,这是一个使用 MiniLM[url=https://arxiv.org/abs/2002.10957]架构[/url]的预训练模型。
例如,要获取句子“这是一个例句”和“每个句子都被转换”的文本嵌入:

let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2).create_model()?;

let sentences = [
    "this is an example sentence", 
    "each sentence is converted"
];
    
let output = model.predict(&sentences);


最近邻居搜索
最近邻居搜索是一种寻找离给定查询点最近的数据点的方法。查询点可以是任何数据点,而不仅仅是数据集中的一个点。
近邻搜索的目标是找到与查询点最相似的数据点。相似性可以用任何距离指标来衡量,如欧氏距离或余弦相似度。

什么是KD树?
KD树是一种数据结构,可用于高效的近邻搜索。KD树是一种二进制树,每个节点根据分割标准被分成两个子节点。
分割准则可以是任何将数据点映射为实数的函数。最常见的分割标准是欧氏距离。
KD树对于近邻搜索是非常有效的,因为它们在树的每一层都将搜索空间减少一半。
例如,如果我们想找到与点(3.1, 0.9, 2.1)最接近的数据点,我们可以使用KD树将搜索空间从整个数据集减少到仅与(3.1, 0.9, 2.1)最接近的数据点。

在Rust中,我们可以使用 kd-tree  crate来构建一个KD树。
我们可以使用build_by_ordered_float函数来从一个点的列表中创建一个KD树。

let kdtree = kd_tree::KdTree::build_by_ordered_float(vec![
    [1.0, 2.0, 3.0],
    [3.0, 1.0, 2.0],
    [2.0, 3.0, 1.0],
]);

// search the nearest neighbor
let found = kdtree.nearest(&[3.1, 0.9, 2.1]).unwrap();



实现语义搜索引擎
在本节中,我们将讨论如何在Rust中构建一个语义搜索引擎。
我们将使用Rust NLP crates rust-bert和kd-tree来进行书籍的相似性搜索。在应用中,它是一个通常被称为 "相似书籍 "的功能。
现在我们可以为我们的语义搜索引擎编写代码了。我们将首先加载存储在json中的数据,对嵌入进行编码,并将其输入我们的kdtree。

开始
创建项目和依赖:

cargo new semantic-search

cd semantic-search 
cargo add serde --features derive 
cargo add rust-bert 
cargo add anyhow 
cargo add kd-tree 
cargo add typenum


创建一个books.json存储json数据:

{
    "books": [
        {
            "title": "The Great Gatsby",
            "author": "F. Scott Fitzgerald",
            "summary": "The story primarily concerns the young and mysterious millionaire Jay Gatsby and his quixotic passion and obsession with the beautiful former debutante Daisy Buchanan."
        },
        {
            "title": "The Catcher in the Rye",
            "author": "J. D. Salinger",
            "summary": "The story is told in the first person by Holden Caulfield, a cynical teenager who recently has been expelled from prep school."
        },
        {
            "title": "The Grapes of Wrath",
            "author": "John Steinbeck",
            "summary": "The novel tells the story of the Joad family, who are driven from their Oklahoma homestead and forced to travel west to the promised land of California."
        }
    ]
}


我们的程序将把json数据加载到一个结构中。我们可以使用serde crate将json数据反序列化到我们的结构中。

use serde::{Deserialize};

#[derive(Debug, Deserialize)]
pub struct Library {
    pub books: Vec<Book>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct Book {
    pub title: String,

    pub author: String,

    pub summary: String,
}


现在我们可以使用std::fs模块从文件中读取json,并将json反序列化为我们的Book结构。

fn main() -> anyhow::Result<()> {

let json = fs::read_to_string("data/books.json")?;
let library: Library = serde_json::from_str(&json)?;
for book in library.books.clone() {
    println!("Embedding book: {}", book.title);
}

    Ok(())
}


将书籍摘要编码为嵌入
我们现在可以使用我们的模型将每本图书摘要编码为嵌入。我们将使用rust-bert crate来进行编码。我们还需要一个方便的方法来将我们的图书模型转换为嵌入式图书模型。

#[derive(Debug)]
pub struct EmbeddedBook {
    pub title: String,

    pub author: String,

    pub summary: String,

    pub embeddings: [f32; 384],
}

impl Book {
    fn to_embedded(self, embeddings: [f32; 384]) -> EmbeddedBook {
        EmbeddedBook {
            title: self.title,
            author: self.author,
            summary: self.summary,
            embeddings: embeddings,
        }
    }
}

// convenient to convert a slice to a fixed size array
fn to_array(barry: &[f32]) -> [f32; 384] {
    barry.try_into().expect("slice with incorrect length")
}


最后,为每本书调用编码,并将其添加到嵌入式书籍向量中:

fn main() -> anyhow::Result<()> {

+ let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2).create_model()?;
let json = fs::read_to_string("data/books.json")?;
let library: Library = serde_json::from_str(&json)?;
+ let mut embeddedbooks = Vec::new();
for book in library.books.clone() {
-    println!("Embedding book: {}", book.title);
+    let embeddings = model.encode(&[book.clone().summary])?;
+    let embedding = to_array(embeddings[0].as_slice());
+    embeddedbooks.push(book.to_embedded(embedding));
}

    Ok(())
}


创建KD树
在执行搜索之前,首先,我们需要从我们的嵌入式books向量创建一个KD树。我们可以使用kd-tree工具箱中的sort_by函数来做这件事。

我们还需要为我们的嵌入式图书结构实现KdPoint特性。这个特性允许我们在kd-tree工具箱中使用我们的EmbeddedBook结构。

impl KdPoint for EmbeddedBook {
    type Scalar = f32;
    type Dim = typenum::U2; // 2 dimensional tree.
    fn at(&self, k: usize) -> f32 {
        self.embeddings[k]
    }
}


现在我们可以使用sort_by函数从我们的嵌入式books向量中创建一个KD树。

fn main() -> anyhow::Result<()> {

let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2).create_model()?;
let json = fs::read_to_string("data/books.json")?;
let library: Library = serde_json::from_str(&json)?;
 let mut embeddedbooks = Vec::new();
for book in library.books.clone() {
    println!("Embedding book: {}", book.title);
    let embeddings = model.encode(&[book.clone().summary])?;
    let embedding = to_array(embeddings[0].as_slice());
    embeddedbooks.push(book.to_embedded(embedding));
}
+  let kdtree = kd_tree::KdSlice::sort_by(&mut embeddedbooks, |item1, item2, k| {
+        item1.embeddings[k]
+            .partial_cmp(&item2.embeddings[k])
+            .unwrap()
+    });

Ok(())
}



执行最近的邻居搜索
现在我们已经将书籍编码为嵌入,我们可以使用最近的邻居搜索来找到与给定查询最相似的书籍。

我们需要在EmbeddedBook上有一个方便的主题方法来创建一个带有作者、标题等的假书,以便进行主题嵌入。

impl EmbeddedBook {
    fn topic(embeddings: [f32; 384]) -> Self {
        Self {
            title: None,
            author: None,
            summary: None,
            embeddings: embeddings,
        }
    }
}


我们可以使用kd-tree crate的nearests函数来进行搜索。这个函数需要一个查询点和一些要返回的结果。

在本例中,我们将使用查询点 "rich "并返回10个结果。

fn main() -> anyhow::Result<()> {

let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2).create_model()?;
let json = fs::read_to_string("data/books.json")?;
let library: Library = serde_json::from_str(&json)?;
 let mut embeddedbooks = Vec::new();
for book in library.books.clone() {
    println!("Embedding book: {}", book.title);
    let embeddings = model.encode(&[book.clone().summary])?;
    let embedding = to_array(embeddings[0].as_slice());
    embeddedbooks.push(book.to_embedded(embedding));
}
let kdtree = kd_tree::KdSlice::sort_by(&mut embeddedbooks, |item1, item2, k| {
    item1.embeddings[k]
    .partial_cmp(&item2.embeddings[k]).unwrap()
});
+ let query = "rich";
+ println!("Querying: {}", query);
+ let rich_embeddings = model.encode(&[query])?;
+ let rich_embedding = to_array(rich_embeddings[0].as_slice());
+ let rich_topic = EmbeddedBook::topic(rich_embedding);
+ let nearests = kdtree.nearests(&rich_topic, 10);
+ for nearest in nearests {
+     println!("nearest: {:?}", nearest.item.title);
+     println!("distance: {:?}", nearest.squared_distance);
+}
Ok(())
}



结论
在本文中,我们讨论了如何在 Rust 中构建一个基本的语义搜索引擎。我们已经解释了什么是嵌入、变换器和最近邻搜索,以及如何使用 KD 树来执行它们。
我们还讨论了如何使用 rust-bert crate 生成文本嵌入,以及如何使用 kd-tree crate 执行最近邻搜索。

您可以在此存储库中找到源代码。

在现实世界的场景中,将这些嵌入保存在内存中并在每次运行程序时运行模型并不是一个非常可扩展的解决方案。此外,这个 kd-tree crate 不支持删除或更新我们树中的节点。
你可以试试这个rs-annoy项目,Annoy 是 Spotify 用来存储歌曲嵌入的东西。
还有一个完全用 Rust 编写的hora