中秋将近,午后小雨绵绵(江浙沪的朋友们,可能是台风猛吹),你独倚窗边,思绪万千,于是拿脱手机,想发一条朋友圈抒发情怀,顺便展示一下文采。

好不容易按出几个字,“本日的雨好大……”,但这展示不出你的文采,于是全部删除。

于是,你灵机一动,如果有一个搜索引擎,能搜索出和“本日的雨好大”意思附近的古诗词,岂不妙哉!

利用向量数据库就可以实现,代码还不到100行,一起来试试吧。
我们会一起从零开始安装向量数据库 Milvus,向量化古诗词数据集,然后创建凑集,导入数据,创建索引,末了实现语义搜索功能。

同伙圈装腔指南若何用向量数据库把大年夜白话变成古诗词

准备事情

首先安装向量数据库 Milvus。
Milvus 支持本地,Docker 和 K8s 支配。
本文利用 Docker 运行 Milvus,以是须要先安装 Docker Desktop。
MacOS 系统安装方法:Install Docker Desktop on Mac (https://docs.docker.com/desktop/install/mac-install/),Windows 系统安装方法:Install Docker Desktop on Windows(https://docs.docker.com/desktop/install/windows-install/)

然后安装 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": "唐诗" }, ...]

准备就绪,正式开始啦。

01 向量化文本

为了实现语义搜索,我们须要有一个办法表示语义。

传统的全文检索是通过提取关键词的办法,但是非常机器,比如语义附近的内容不一定关键词完备重合,例如“多大了?”和“年事是?”关键词重合度并不高但意思险些一样。

反之也不尽然,关键词重合但多一个“不”字,意思大相径庭。

目前盛行的办法叫做嵌入模型(embedding),通过深度神经网络模型把各种非构造化数据(如笔墨、图像、声音)提取成一堆数字。
这一堆数字叫做向量,它可以表示语义,在空间中两个向量之间的间隔可以表示他们所代表的数据语义的近似度,例如两个诗句的语义是否干系,或者两张图片是否很像。
向量可以存储在 Milvus 向量数据库中,以用来高效的查询。

理解了这个事理之后,让我们先写一个把文本向量化的函数 vectorize_text,方便后面调用。

import torchimport jsonfrom 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, ... ] }, ...]
02 创建凑集

接下来我们要把向量数据导入向量数据库。
当然,我们得先在向量数据库中创建一个凑集,用来容纳向量数据。

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")# 添加字段到schemaschema.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}")03 入库

接下来把文件导入到 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']}")

返回插入实体的数量,看来是成功了。

插入的实体数量: 430704 创建索引

向量已经导入 Milvus,现在可以搜索了吗?别急,为了提高搜索效率,我们还须要创建索引。
什么是索引?一些大部头图书的末了,一样平常都会有索引,它列出了书中涌现的关键术语以及对应的页码,帮助你快速找到它们的位置。
如果没有索引,那就只能用笨方法,从第一页开始一页一页今后找了。

图片来源:《英国皇家园艺学会植物繁育手册:用已有植物打造完美新植物》

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'}05 加载索引

索引创建成功了,现在可以搜索了吗?等等,我们还须要把凑集中的数据和索引,从硬盘加载到内存中。
由于在内存中搜索更快。

print(f"正在加载凑集:{collection_name}")client.load_collection(collection_name=collection_name)

加载完成了,仍旧验证下。

print(client.get_load_state(collection_name=collection_name))

返回加载状态 Loaded,没问题,加载完成。

{'state': <LoadState: Loaded>}06 搜索

经由前面的一系列准备,现在我们终于可以回到开头的问题了,用当代口语文搜索语义相似的古诗词。

首先,把我们要搜索的当代口语文提取向量。

# 获取查询向量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

(https://milvus.io/docs/get-and-scalar-query.md#Use-Basic-Operators)

总结

恭喜,到这里你已经完成了一个大口语搜古诗词的 demo!

完全的数据集和代码拜会文末的下载链接,推举直接去Colab 里运行体验。

可能前面的搜索结果并没有让你完备满意,这里面有多个缘故原由:首先,数据集太小了。
只有4000多个句子,语义更靠近的句子可能没有包含个中。
其次,嵌入模型虽然支持中文,但是古诗词并不是它的专长。
这就彷佛你找了个翻译帮你和老外交流,翻译虽然懂普通话,但是你满嘴四川方言,翻译也只能也蒙带猜,翻译质量可想而知。

如果你希望优化搜索功能,可以在 chinese-poetry (https://github.com/BushJiang/chinese-poetry)下载完全的古诗词数据集,再把古诗词翻译成口语文,搜索口语文,返回对应的古诗词。

参考:

[1]古诗词数据集来自 chinese-poetry,数据构造做了调度。

复当代码及数据集:

searchPoems.ipynb(https://zilliverse.feishu.cn/docx/Tot4d3p8kogir8xbUGtcGH65nyg)

TangShi.json(https://zilliverse.feishu.cn/docx/Tot4d3p8kogir8xbUGtcGH65nyg)