那时难决
0%

文本嵌入，语义搜索与sentence-transformers库

Posted on Edited on In

近期在研究开源的 rust 实现的向量数据库 qdrant。顾名思义，向量数据是用于存储和查询向量的数据库，而向量本质上是一个多维空间中的点。如果要用向量数据库处理文本数据，就需要将文本转换为向量表示，机器学习术语叫做文本嵌入（Text Embedding）。

传统的文本嵌入方法是基于统计的，比如 TF-IDF，Word2Vec 等。随着 transformer 架构的出现和发展，基于 transformer 的文本嵌入方法也越来越流行，并且在很多任务上取得了很好的效果。sentence-transformers 就是一个基于 transformer 的文本嵌入工具包，可以用于生成句子的向量表示。

安装 sentence-transformers

安装 pytorch

可以前往 pytorch 官网 根据自己的环境选择合适的安装方式，我的设备是 RTX 2060s 的 Windows PC，安装了 windows 的 cuda 版本。

1
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

安装好 pytorch 后可以验证一下。

1
2
3
4
5
6
7
8
9
10
11
import torch

# Check if PyTorch is installed
print("PyTorch version:", torch.__version__)

# Check if CUDA is available
print("CUDA available:", torch.cuda.is_available())


x = torch.rand(5, 3)
print(x)

安装 sentence-transformers

sentence-transformers 可以直接使用 pip 安装。

1
pip3 install sentence-transformers

使用 sentence-transformers

sentence-transformers 提供了很多预训练模型，可以直接使用。我们这里使用的是 paraphrase-multilingual-MiniLM-L12-v2 模型，支持多语言，模型尺寸也比较大（480M）。只处理英文文本的话，可以使用 all-MiniLM-L6-v2 模型（80M）。

1
2
3
4
5
6
7
8
9
10
from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

#Sentences are encoded by calling model.encode()
emb1 = model.encode("This is a red cat with a hat.")
emb2 = model.encode("Have you seen my red cat?")

cos_sim = util.cos_sim(emb1, emb2)
print("Cosine-Similarity:", cos_sim) # Cosine-Similarity: tensor([[0.7097]])

上述代码中，我们使用 sentence-transformers 加载了 paraphrase-multilingual-MiniLM-L12-v2 模型，并使用该模型将两个句子转换为向量表示，然后计算了两个向量的余弦相似度。

余弦相似度是一个常用的相似度度量方法，其值域为 [-1, 1]，值越大表示两个向量越相似。其他的相似度度量方法还有欧氏距离、曼哈顿距离等。

我们还可以使用中文文本进行测试。

1
2
3
4
5
6
7
8
9
10
11
12
13
zh_sentences = ["湖南的省会是长沙", "长沙是湖南的首府", "长沙的特色小吃有臭豆腐", "绍兴的臭豆腐也很有名"]

results = []
for a, b in itertools.combinations(zh_sentences, 2):
    emb_a = model.encode(a)
    emb_b = model.encode(b)
    cos_sim = util.cos_sim(emb_a, emb_b)
    results.append((a, b, cos_sim))

results = sorted(results, key=lambda x: x[2], reverse=True)

for a, b, cos_sim in results:
    print(f"{a} {b} {cos_sim}")

输出结果如下：

1
2
3
4
5
6
湖南的省会是长沙 长沙是湖南的首府 tensor([[0.9138]])
长沙的特色小吃有臭豆腐 绍兴的臭豆腐也很有名 tensor([[0.8868]])
湖南的省会是长沙 长沙的特色小吃有臭豆腐 tensor([[0.3102]])
湖南的省会是长沙 绍兴的臭豆腐也很有名 tensor([[0.2694]])
长沙是湖南的首府 长沙的特色小吃有臭豆腐 tensor([[0.2276]])
长沙是湖南的首府 绍兴的臭豆腐也很有名 tensor([[0.1900]])

可以看到相似度最高的一组是“湖南的省会是长沙”和“长沙是湖南的首府”，这两句本质是同一个意思。第二高的一组是“长沙的特色小吃有臭豆腐”和“绍兴的臭豆腐也很有名”，分别说明了两个以臭豆腐闻名的城市。剩下的几组相关性就很低了，尤其是最后一组“长沙是湖南的首府”和“绍兴的臭豆腐也很有名”，这两句话也确实没什么关联。

使用 sentence-transformers 进行语义搜索

通过比较不同向量间的余弦相似度，我们可以找到最相似的向量，这就是语义搜索的基本原理。下面是一个来自 sentence-transformers 官方文档的例子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
embedder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

# Corpus with example sentences
corpus = ['A man is eating food.',
          'A man is eating a piece of bread.',
          'The girl is carrying a baby.',
          'A man is riding a horse.',
          'A woman is playing violin.',
          'Two men pushed carts through the woods.',
          'A man is riding a white horse on an enclosed ground.',
          'A monkey is playing drums.',
          'A cheetah is running behind its prey.'
          ]
corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)

# Query sentences:
queries = ['A man is eating pasta.', 'Someone in a gorilla costume is playing a set of drums.', 'A cheetah chases prey on across a field.']


# Find the closest 5 sentences of the corpus for each query sentence based on cosine similarity
top_k = min(5, len(corpus))
for query in queries:
    query_embedding = embedder.encode(query, convert_to_tensor=True)

    # We use cosine-similarity and torch.topk to find the highest 5 scores
    cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
    top_results = torch.topk(cos_scores, k=top_k)

    print("\n\n======================\n\n")
    print("Query:", query)
    print("\nTop 5 most similar sentences in corpus:")

    for score, idx in zip(top_results[0], top_results[1]):
        print(corpus[idx], "(Score: {:.4f})".format(score))

这里贴出第一个查询的输出。

1
2
3
4
5
6
7
8
Query: A man is eating pasta.

Top 5 most similar sentences in corpus:
A man is eating food. (Score: 0.6734)
A man is eating a piece of bread. (Score: 0.4269)
A man is riding a horse. (Score: 0.2086)
A man is riding a white horse on an enclosed ground. (Score: 0.1020)
A cheetah is running behind its prey. (Score: 0.0566)

排名最高的两句都是关于“吃”的，看来模型确实一定程度上识别到了句子的含义。把”eating food”与”eating pasta”看作从属关系的话，”eating a piece of bread”与原始查询为并列关系，得分低一点也是合理的。

事实上，sentence-transformers 还提供了 utils.semantic_search 函数，简化了语义搜索的过程。可以使用一些中文文本来测试一下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
facts = [
    "张三今年二十岁。",
    "张三今年一百斤。",
    "李四和娜娜是一对情侣。",
    "王五是一名医生。",
    "李四和兰兰是一对兄妹",
    "小明喜欢吃水果。",
    "小红会弹钢琴。",
    "刘老师教数学。",
    "这个城市有很多高楼大厦。",
]

queries = ["张三今年几岁？", "李四的女朋友是谁？", "谁是医生？", "小明喜欢吃什么？", "谁会弹钢琴？", "谁教数学？", "这个城市有什么？"]

embedder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

corpus_embeddings = embedder.encode(facts, convert_to_tensor=True)
query_embeddings = embedder.encode(queries, convert_to_tensor=True)

corpus_embeddings = util.normalize_embeddings(corpus_embeddings)
query_embeddings = util.normalize_embeddings(query_embeddings)

hits = util.semantic_search(
    query_embeddings, corpus_embeddings, score_function=util.dot_score
)

for i, (closest, *_) in enumerate(hits):
    print("Query:", queries[i], "Answer:", facts[closest["corpus_id"]])

输出结果如下：

1
2
3
4
5
6
7
Query: 张三今年几岁？ Answer: 张三今年二十岁。
Query: 李四的女朋友是谁？ Answer: 李四和娜娜是一对情侣。
Query: 谁是医生？ Answer: 王五是一名医生。
Query: 小明喜欢吃什么？ Answer: 小明喜欢吃水果。
Query: 谁会弹钢琴？ Answer: 小红会弹钢琴。
Query: 谁教数学？ Answer: 刘老师教数学。
Query: 这个城市有什么？ Answer: 这个城市有很多高楼大厦。

可以看到，语义搜索的结果还是比较准确的，并且模型正确识别出了在“情侣”、“兄妹”两个关系中，“女朋友”与前者更接近。

总结

sentence-transformers 是一个非常好用的文本嵌入工具包，可以用于生成句子的向量表示，也可以用于语义搜索。sentence-transformers 还提供了很多预训练模型，可以根据自己的需求选择合适的模型。

本文代码中的所有向量数据都是存在内存中的，可以使用多种方式持久化向量数据，比如存储到 JSON 文件中，或者存储到关系型数据库中。不过为了更好的查询性能，我们可以使用专门的向量数据库，比如 qdrant，这也是后续文章的主题。

扫码加入技术交流群🖱️
QR code