朋友圈装腔指南:如何用向量数据库把大白话变成古诗词
午后细雨绵绵,你独倚窗边,思绪万千,于是拿出手机,想发一条朋友圈抒发情怀,顺便展示一下文采。奈何好不容易按出几个字,又全部删除。“今天的雨好大”展示不出你的文采。你灵机一动,如果有一个搜索引擎,能搜索出和“今天的雨好大”意思相近的古诗词,岂不妙哉!
使用向量数据库就可以实现,代码还不到100行,一起来试试吧。我们会一起从零开始安装向量数据库 Milvus,向量化古诗词数据集,然后创建集合,导入数据,创建索引,最后实现语义搜索功能。
0 准备工作
首先安装向量数据库 Milvus。Milvus 支持本地,Docker 和 K8s 部署。本文使用 Docker 运行 Milvus,所以需要先安装 Docker Desktop。MacOS 系统安装方法:Install Docker Desktop on Mac ,Windows 系统安装方法:Install Docker Desktop on Windows
然后安装 Milvus。
下载安装脚本:
curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh
运行 Milvus:
standalone_embed.sh start
安装依赖。
pip install pymilvus "pymilvus[model]" torch
下载古诗词数据集[1] TangShi.json
。它的格式是这样的:
[
{
"author": "太宗皇帝",
"paragraphs": [
"秦川雄帝宅,函谷壮皇居。"
],
"title": "帝京篇十首 一",
"id": 20000001,
"type": "唐诗"
},
...
]
准备就绪,正式开始啦。
1 向量化文本
为了实现语义搜索,我们需要有一个办法表示语义。传统的全文检索是通过提取关键词的方式,但是非常机械,比如语义相近的内容不一定关键词完全重合,例如“多大了?”和“年龄是?”关键词重合度并不高但意思几乎一样。反之也不尽然,关键词重合但多一个“不”字,意思大相径庭。目前流行的方式叫做嵌入模型(embedding),通过深度神经网络模型把各种非结构化数据(如文字、图像、声音)提取成一堆数字。这一堆数字叫做向量,它可以表示语义,在空间中两个向量之间的距离可以表示他们所代表的数据语义的近似度,例如两个诗句的语义是否相关,或者两张图片是否很像。向量可以存储在 Milvus 向量数据库中,以用来高效的查询。
理解了这个原理之后,让我们先写一个把文本向量化的函数 vectorize_text
,方便后面调用。
import torch
import json
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
# 1 向量化文本数据
def vectorize_text(text, model_name="BAAI/bge-small-zh-v1.5"):
# 检查是否有可用的CUDA设备
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# 根据设备选择是否使用fp16
use_fp16 = device.startswith("cuda")
# 创建嵌入模型实例
bge_m3_ef = BGEM3EmbeddingFunction(
model_name=model_name,
device=device,
use_fp16=use_fp16
)
# 把输入的文本向量化
vectors = bge_m3_ef.encode_documents(text)
return vectors
函数 vectorize_text
中使用了嵌入模型 BGEM3EmbeddingFunction
,就是它把文本提取成向量。
准备好后,我们就可以对整个数据集进行向量化了。 下面我们读取TangShi.json
中的数据,把其中的 paragraphs
字段转成向量,然后写入 TangShi_vector.json
文件。如果你是第一次使用 Milvus,运行下面的代码时还会安装必要的依赖。
# 读取 json 文件,把paragraphs字段向量化
with open("TangShi.json", 'r', encoding='utf-8') as file:
data_list = json.load(file)
# 提取该json文件中的所有paragraphs字段的值
text = [data['paragraphs'][0] for data in data_list]
# 向量化文本数据
vectors = vectorize_text(text)
# 将向量添加到原始文本中
for data, vector in zip(data_list, vectors['dense']):
data['vector'] = vector.tolist()
# 将更新后的文本内容写入新的json文件
with open("TangShi_vector.json", 'w', encoding='utf-8') as outfile:
json.dump(data_list, outfile, ensure_ascii=False, indent=4)
如果一切顺利,你会得到 TangShi_vector.json
文件,它增加了 vector
字段,它的值是一个浮点数列表,也就是“向量”。
[
{
"author": "太宗皇帝",
"paragraphs": [
"秦川雄帝宅,函谷壮皇居。"
],
"title": "帝京篇十首 一",
"id": 20000001,
"type": "唐诗",
"vector": [
0.005114779807627201,
0.033538609743118286,
0.020395483821630478,
...
]
},
{
"author": "太宗皇帝",
"paragraphs": [
"绮殿千寻起,离宫百雉余。"
],
"title": "帝京篇十首 一",
"id": 20000002,
"type": "唐诗",
"vector": [
-0.06334448605775833,
0.0017451602034270763,
-0.0010646708542481065,
...
]
},
...
]
2 创建集合
接下来我们要把向量数据导入向量数据库。当然,我们得先在向量数据库中创建一个集合,用来容纳向量数据。
from pymilvus import MilvusClient
# 连接向量数据库,创建client实例
client = MilvusClient(uri="http://localhost:19530")
# 指定集合名称
collection_name = "TangShi"
注意,为了避免向量数据库中存在同名集合,产生干扰,创建集合前先删除同名集合。
# 检查同名集合是否存在,如果存在则删除
if client.has_collection(collection_name):
print(f"Collection {collection_name} already exists")
try:
# 删除同名集合
client.drop_collection(collection_name)
print(f"Deleted the collection {collection_name}")
except Exception as e:
print(f"Error occurred while dropping collection: {e}")
就像我们把数据填入 excel 表格前,需要先设计好表头,规定有哪些字段,各个字段的数据类型是怎样的,向量数据库也需要定义表结构,它的“表头”就是 schema
。
from pymilvus import DataType
# 创建集合模式
schema = MilvusClient.create_schema(
auto_id=False,
enable_dynamic_field=True,
description="TangShi"
)
# 添加字段到schema
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=512)
schema.add_field(field_name="title", datatype=DataType.VARCHAR, max_length=1024)
schema.add_field(field_name="author", datatype=DataType.VARCHAR, max_length=256)
schema.add_field(field_name="paragraphs", datatype=DataType.VARCHAR, max_length=10240)
schema.add_field(field_name="type", datatype=DataType.VARCHAR, max_length=128)
schema 创建好了,接下来就可以创建集合了。
# 创建集合
try:
client.create_collection(
collection_name=collection_name,
schema=schema,
shards_num=2
)
print(f"Created collection {collection_name}")
except Exception as e:
print(f"Error occurred while creating collection: {e}")
3 入库
接下来把文件导入到 Milvus。
# 读取和处理文件
with open("TangShi_vector.json", 'r') as file:
data = json.load(file)
# paragraphs的值是列表,需要从列表中取出字符串,取代列表,以符合Milvus插入数据的要求
for item in data:
item["paragraphs"] = item["paragraphs"][0]
# 将数据插入集合
print(f"正在将数据插入集合:{collection_name}")
res = client.insert(
collection_name=collection_name,
data=data
)
导入成功了吗?我们来验证下。
print(f"插入的实体数量: {res['insert_count']}")
返回插入实体的数量,看来是成功了。
插入的实体数量: 4307
4 创建索引
向量已经导入 Milvus,现在可以搜索了吗?别急,为了提高搜索效率,我们还需要创建索引。什么是索引?一些大部头图书的最后,一般都会有索引,它列出了书中出现的关键术语以及对应的页码,帮助你快速找到它们的位置。如果没有索引,那就只能用笨方法,从第一页开始一页一页往后找了。
8eae27f1-a473-4a67-acaf-a6ab41fca664.jpg
图片来源:《英国皇家园艺学会植物繁育手册:用已有植物打造完美新植物》
Milvus 的索引也是如此。如果不创建索引,虽然也可以搜索,但是速度很慢,它会逐一比较查询向量与数据库中每一个向量,通过指定方法计算出两个向量之间的 距离,找出距离最近的几个向量。而创建索引之后,搜索速度会大大提升。
索引有不同的类型,适合不同的场景使用,我们以后会详细讨论这个问题。这里我们使用较为简单的索引类型 IVF_FLAT
。另外,计算距离的方法也有多种,我们使用 IP
,也就是计算两个向量的内积。这些都是索引的参数,我们先创建这些参数。
# 创建IndexParams对象,用于存储索引的各种参数
index_params = client.prepare_index_params()
# 设置索引名称
vector_index_name = "vector_index"
# 设置索引的各种参数
index_params.add_index(
# 指定为"vector"字段创建索引
field_name="vector",
# 设置索引类型
index_type="IVF_FLAT",
# 设置度量类型
metric_type="IP",
# 设置索引聚类中心的数量
params={"nlist": 128},
# 指定索引名称
index_name=vector_index_name
)
索引参数创建好了,现在终于可以创建索引了。
print(f"开始创建索引:{vector_index_name}")
# 创建索引
client.create_index(
# 指定为哪个集合创建索引
collection_name=collection_name,
# 使用前面创建的索引参数创建索引
index_params=index_params
)
我们来验证下索引是否创建成功了。
indexes = client.list_indexes(
collection_name=collection_name
)
print(f"列出创建的索引:{indexes}")
返回了包含索引名称的列表,索引名称 vector_index
正是我们之前创建的。
列出创建的索引:['vector_index']
再来查看下索引的详情。
# 查看索引详情
index_details = client.describe_index(
collection_name=collection_name,
# 指定索引名称,这里假设使用第一个索引
index_name="vector_index"
)
print(f"索引vector_index详情:{index_details}")
返回了一个包含索引详细信息的字典,可以看到我们之前设置的索引参数,比如 nlist
,index_type
和 metric_type
等等。nlist 是索引聚类中心的数量,简单来说就是把一个数据分片里的向量数据分成 nlist 个聚类,方便后续取出最相关的若干个聚类加速检索。
索引vector_index详情:{'nlist': '128', 'index_type': 'IVF_FLAT', 'metric_type': 'IP', 'field_name': 'vector', 'index_name': 'vector_index', 'total_rows': 0, 'indexed_rows': 0, 'pending_index_rows': 0, 'state': 'Finished'}
5 加载索引
索引创建成功了,现在可以搜索了吗?等等,我们还需要把集合中的数据和索引,从硬盘加载到内存中。因为在内存中搜索更快。
print(f"正在加载集合:{collection_name}")
client.load_collection(collection_name=collection_name)
加载完成了,仍然验证下。
print(client.get_load_state(collection_name=collection_name))
返回加载状态 Loaded
,没问题,加载完成。
{'state': <LoadState: Loaded>}
6 搜索
经过前面的一系列准备,现在我们终于可以回到开头的问题了,用现代白话文搜索语义相似的古诗词。
首先,把我们要搜索的现代白话文提取向量。
# 获取查询向量
text = "今天的雨好大"
query_vectors = [vectorize_text([text])['dense'][0].tolist()]
然后,设置搜索参数,告诉 Milvus 怎么搜索。
# 设置搜索参数
search_params = {
# 设置度量类型
"metric_type": "IP",
# 指定在搜索过程中要查询的聚类单元数量,增加nprobe值可以提高搜索精度,但会降低搜索速度
"params": {"nprobe": 16}
}
最后,我们还得告诉它要返回什么结果。
# 指定搜索结果的数量,“limit=3”表示返回最相近的前3个搜索结果
limit = 3
# 指定返回的字段
output_fields = ["author", "title", "paragraphs"]
一切就绪,让我们开始搜索吧!
res1 = client.search(
collection_name=collection_name,
# 指定查询向量
data=query_vectors,
# 指定搜索的字段
anns_field="vector",
# 设置搜索参数
search_params=search_params,
# 指定返回搜索结果的数量
limit=limit,
# 指定返回的字段
output_fields=output_fields
)
print(res1)
得到下面的结果:
data: [
"[
{
'id': 20002740,
'distance': 0.6542239189147949,
'entity': {
'title': '郊庙歌辞 享太庙乐章 大明舞',
'paragraphs': '旱望春雨,云披大风。',
'author': '张说'
}
},
{
'id': 20001658,
'distance': 0.6228379011154175,
'entity': {
'title': '三学山夜看圣灯',
'paragraphs': '细雨湿不暗,好风吹更明。',
'author': '蜀太妃徐氏'
}
},
{
'id': 20003360,
'distance': 0.6123768091201782,
'entity': {
'title': '郊庙歌辞 汉宗庙乐舞辞 积善舞',
'paragraphs': '云行雨施,天成地平。',
'author': '张昭'
}
}
]"
]
在搜索结果中,id
、title
等字段我们都了解了,只有 distance
是新出现的。它指的是搜索结果与查询向量之间的“距离”,具体含义和度量类型有关。我们使用的度量类型是 IP
内积,数字越大表示搜索结果和查询向量越接近。
为了增加可读性,我们写一个输出函数:
# 打印向量搜索结果
def print_vector_results(res):
# hit是搜索结果中的每一个匹配的实体
res = [hit["entity"] for hit in res[0]]
for item in res:
print(f"title: {item['title']}")
print(f"author: {item['author']}")
print(f"paragraphs: {item['paragraphs']}")
print("-"*50)
print(f"数量:{len(res)}")
重新输出结果:
print_vector_results(res1)
这下搜索结果容易阅读了。
title: 郊庙歌辞 享太庙乐章 大明舞
author: 张说
paragraphs: 旱望春雨,云披大风。
--------------------------------------------------
title: 三学山夜看圣灯
author: 蜀太妃徐氏
paragraphs: 细雨湿不暗,好风吹更明。
--------------------------------------------------
title: 郊庙歌辞 汉宗庙乐舞辞 积善舞
author: 张昭
paragraphs: 云行雨施,天成地平。
--------------------------------------------------
数量:3
如果你不想限制搜索结果的数量,而是返回所有质量符合要求的搜索结果,可以修改搜索参数:
# 修改搜索参数,设置距离的范围
search_params = {
"metric_type": "IP",
"params": {
"nprobe": 16,
"radius": 0.55,
"range_filter": 1.0
}
}
这里比较特殊的是在搜索参数中增加 radius
和 range_filter
参数,它们限制了距离 distance
的范围在0.55到1之间。然后调整搜索代码,删除 limit
参数。
res2 = client.search(
collection_name=collection_name,
# 指定查询向量
data=query_vectors,
# 指定搜索的字段
anns_field="vector",
# 设置搜索参数
search_params=search_params,
# 删除limit参数
# 指定返回的字段
output_fields=output_fields
)
print(res2)
可以看到,输出结果的 distance
都大于0.55。
data: [
"[
{
'id': 20002740,
'distance': 0.6542239189147949,
'entity': {
'author': '张说',
'title': '郊庙歌辞 享太庙乐章 大明舞',
'paragraphs': '旱望春雨,云披大风。'
}
},
{
'id': 20001658,
'distance': 0.6228379011154175,
'entity': {
'author': '蜀太妃徐氏',
'title': '三学山夜看圣灯',
'paragraphs': '细雨湿不暗,好风吹更明。'
}
},
{
'id': 20003360,
'distance': 0.6123768091201782,
'entity': {
'author': '张昭',
'title': '郊庙歌辞 汉宗庙乐舞辞 积善舞',
'paragraphs': '云行雨施,天成地平。'
}
},
{
'id': 20003608,
'distance': 0.5755923986434937,
'entity': {
'author': '李端',
'title': '鼓吹曲辞 巫山高',
'paragraphs': '回合云藏日,霏微雨带风。'
}
},
{
'id': 20000992,
'distance': 0.5700664520263672,
'entity': {
'author': '德宗皇帝',
'title': '九月十八赐百僚追赏因书所怀',
'paragraphs': '雨霁霜气肃,天高云日明。'
}
},
{
'id': 20002246,
'distance': 0.5583387613296509,
'entity': {
'author': '不详',
'title': '郊庙歌辞 祭方丘乐章 顺和',
'paragraphs': '雨零感节,云飞应序。'
}
}
]"
]
也许你还想知道你最喜欢的李白,有没有和你一样感慨今天的雨真大,没问题,我们增加filter
参数就可以指定只搜索'author'为李白的内容。
# 通过表达式过滤字段author,筛选出字段“author”的值为“李白”的结果
filter = f"author == '李白'"
res3 = client.search(
collection_name=collection_name,
# 指定查询向量
data=query_vectors,
# 指定搜索的字段
anns_field="vector",
# 设置搜索参数
search_params=search_params,
# 通过表达式实现标量过滤,筛选结果
filter=filter,
# 指定返回搜索结果的数量
limit=limit,
# 指定返回的字段
output_fields=output_fields
)
print(res3)
返回的结果为空值。
data: ['[]']
这是因为我们前面设置了 distance
的范围在0.55到1之间,放大范围可以获得更多结果。把 "radius"
的值修改为0.2,再次运行命令,让我们看看李白是怎么感慨的。
data: [
"[
{
'id': 20004246,
'distance': 0.46472394466400146,
'entity': {
'author': '李白',
'title': '横吹曲辞 关山月',
'paragraphs': '明月出天山,苍茫云海间。'
}
},
{
'id': 20003707,
'distance': 0.4347272515296936,
'entity': {
'author': '李白',
'title': '鼓吹曲辞 有所思',
'paragraphs': '海寒多天风,白波连山倒蓬壶。'
}
},
{
'id': 20003556,
'distance': 0.40778297185897827,
'entity': {
'author': '李白',
'title': '鼓吹曲辞 战城南',
'paragraphs': '去年战桑干源,今年战葱河道。'
}
}
]"
]
我们观察搜索结果发现, distance
在0.4左右,小于之前设置的0.55,所以被排除了。另外,distance
数值较小,说明搜索结果并不是特别接近查询向量,而这几句诗词的确和“雨”的关系比较远。
如果你希望搜索结果中直接包含“雨”字,可以使用 query
方法做标量搜索。
# paragraphs字段包含“雨”字
filter = f"paragraphs like '%雨%'"
res4 = client.query(
collection_name=collection_name,
filter=filter,
output_fields=output_fields,
limit=limit
)
print(res4)
标量查询的代码更简单,因为它免去了和向量搜索相关的参数,比如查询向量 data
,指定搜索字段的 anns_field
和搜索参数 search_params
,查询参数只有 filter
。
观察搜索结果发现,标量查询结果的数据结构少了一个 []
,我们在提取具体字段时需要注意这一点。
data: [
"{
"author": "太宗皇帝",
"title": "咏雨",
"paragraphs": "罩云飘远岫,喷雨泛长河。",
"id": 20000305
},
{
"author": "太宗皇帝",
"title": "咏雨",
"paragraphs": "和气吹绿野,梅雨洒芳田。",
"id": 20000402
},
{
"author": "太宗皇帝",
"title": "赋得花庭雾",
"paragraphs": "还当杂行雨,髣髴隐遥空。",
"id": 20000421
}"
]
filter
表达式还有丰富的用法,比如同时查询两个字段,author
字段指定为 “杜甫”,同时 paragraphs
字段仍然要求包含“雨”字:
filter = f"author == '杜甫' && paragraphs like '%雨%'"
res5 = client.query(
collection_name=collection_name,
filter=filter,
output_fields=output_fields,
limit=limit
)
print(res5)
返回杜甫含有“雨”字的诗句:
data: [
"{
'title': '横吹曲辞 前出塞九首 七',
'paragraphs': '驱马天雨雪,军行入高山。',
'id': 20004039,
'author': '杜甫'
}"
]
更多标量搜索的表达式可以参考Get & Scalar Query。
总结
恭喜,到这里你已经完成了一个大白话搜古诗词的 demo!完整的数据集和代码参见文末的下载链接,推荐直接去Colab 里运行体验。可能前面的搜索结果并没有让你完全满意,这里面有多个原因。首先,数据集太小了。只有4000多个句子,语义更接近的句子可能没有包含其中。其次,嵌入模型虽然支持中文,但是古诗词并不是它的专长。这就好像你找了个翻译帮你和老外交流,翻译虽然懂普通话,但是你满嘴四川方言,翻译也只能也蒙带猜,翻译质量可想而知。
如果你希望优化搜索功能,可以在 chinese-poetry 下载完整的古诗词数据集,再把古诗词翻译成白话文,搜索白话文,返回对应的古诗词。
参考
[1]古诗词数据集来自 chinese-poetry,数据结构做了调整。
代码及数据集可通过链接获取:https://pan.baidu.com/s/1RPycmocva_Bivszyir28lA?pwd=1234 提取码: 1234
技术干货
重磅版本发布|三大关键特性带你认识 Milvus 2.2.9 :JSON、PartitionKey、Dynamic Schema
随着 LLM 的持续火爆,众多应用开发者将目光投向了向量数据库领域,而作为开源向量数据库的领先者,Milvus 也充分吸收了大量来自社区、用户、AI 从业者的建议,把重心投入到了开发者使用体验上,以简化开发者的使用门槛。
2023-6-5技术干货
我决定给 ChatGPT 做个缓存层 >>> Hello GPTCache
我们从自己的开源项目 Milvus 和一顿没有任何目的午饭中分别获得了灵感,做出了 OSSChat、GPTCache。在这个过程中,我们也在不断接受「从 0 到 1」的考验。作为茫茫 AI 领域开发者和探索者中的一员,我很愿意与诸位分享这背后的故事、逻辑和设计思考,希望大家能避坑避雷、有所收获。
2023-4-14技术干货
LangChain 查询使用指「北」
LangChain 是一种 AI 代理工具,可以为以 ChatGPT 为代表的额大语言模型(LLM)增添更多功能。此外,LangChain 还具备 token 和上下文管理功能。本文主要通过查询 GPT 和查询文档两个示例介绍如何使用 LangChain。
2023-5-30