近期在研究开源的 rust 实现的向量数据库 qdrant。顾名思义，向量数据是用于存储和查询向量的数据库，而向量本质上是一个多维空间中的点。如果要用向量数据库处理文本数据，就需要将文本转换为向量表示，机器学习术语叫做文本嵌入（Text Embedding）。
传统的文本嵌入方法是基于统计的，比如 TF-IDF，Word2Vec 等。随着 transformer 架构的出现和发展，基于 transformer 的文本嵌入方法也越来越流行，并且在很多任务上取得了很好的效果。sentence-transformers 就是一个基于 transformer 的文本嵌入工具包，可以用于生成句子的向量表示。
安装 sentence-transformers
安装 pytorch
可以前往 pytorch 官网 根据自己的环境选择合适的安装方式，我的设备是 RTX 2060s 的 Windows PC，安装了 windows 的 cuda 版本。
代码语言：javascript复制pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118安装好 pytorch 后可以验证一下。
代码语言：javascript复制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 安装。
代码语言：javascript复制pip3 install sentence-transformers使用 sentence-transformers
sentence-transformers 提供了很多预训练模型，可以直接使用。我们这里使用的是 paraphrase-multilingual-MiniLM-L12-v2 模型，支持多语言，模型尺寸也比较大（480M）。只处理英文文本的话，可以使用 all-MiniLM-L6-v2 模型（80M）。
代码语言：javascript复制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]，值越大表示两个向量越相似。其他的相似度度量方法还有欧氏距离、曼哈顿距离等。
我们还可以使用中文文本进行测试。
代码语言：javascript复制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}")输出结果如下：
代码语言：javascript复制湖南的省会是长沙 长沙是湖南的首府 tensor([[0.9138]])
长沙的特色小吃有臭豆腐 绍兴的臭豆腐也很有名 tensor([[0.8868]])
湖南的省会是长沙 长沙的特色小吃有臭豆腐 tensor([[0.3102]])
湖南的省会是长沙 绍兴的臭豆腐也很有名 tensor([[0.2694]])
长沙是湖南的首府 长沙的特色小吃有臭豆腐 tensor([[0.2276]])
长沙是湖南的首府 绍兴的臭豆腐也很有名 tensor([[0.1900]])可以看到相似度最高的一组是“湖南的省会是长沙”和“长沙是湖南的首府”，这两句本质是同一个意思。第二高的一组是“长沙的特色小吃有臭豆腐”和“绍兴的臭豆腐也很有名”，分别说明了两个以臭豆腐闻名的城市。剩下的几组相关性就很低了，尤其是最后一组“长沙是湖南的首府”和“绍兴的臭豆腐也很有名”，这两句话也确实没什么关联。
使用 sentence-transformers 进行语义搜索
通过比较不同向量间的余弦相似度，我们可以找到最相似的向量，这就是语义搜索的基本原理。下面是一个来自 sentence-transformers 官方文档的例子。
代码语言：javascript复制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("nn======================nn")
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))这里贴出第一个查询的输出。
代码语言：javascript复制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 函数，简化了语义搜索的过程。可以使用一些中文文本来测试一下。
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"]])输出结果如下：
代码语言：javascript复制Query: 张三今年几岁？ Answer: 张三今年二十岁。
Query: 李四的女朋友是谁？ Answer: 李四和娜娜是一对情侣。
Query: 谁是医生？ Answer: 王五是一名医生。
Query: 小明喜欢吃什么？ Answer: 小明喜欢吃水果。
Query: 谁会弹钢琴？ Answer: 小红会弹钢琴。
Query: 谁教数学？ Answer: 刘老师教数学。
Query: 这个城市有什么？ Answer: 这个城市有很多高楼大厦。可以看到，语义搜索的结果还是比较准确的，并且模型正确识别出了在“情侣”、“兄妹”两个关系中，“女朋友”与前者更接近。
总结
sentence-transformers 是一个非常好用的文本嵌入工具包，可以用于生成句子的向量表示，也可以用于语义搜索。sentence-transformers 还提供了很多预训练模型，可以根据自己的需求选择合适的模型。
本文代码中的所有向量数据都是存在内存中的，可以使用多种方式持久化向量数据，比如存储到 JSON 文件中，或者存储到关系型数据库中。不过为了更好的查询性能，我们可以使用专门的向量数据库，比如 qdrant，这也是后续文章的主题。
