元旦将近，窗外寒风阵阵，你独倚窗边，思绪万千，于是拿出手机，想发一条朋友圈抒发情怀，顺便展示一下文采。
好不容易按出几个字，"今天的风好大……"，但这展示不出你的文采，于是全部删除。
于是，你灵机一动，如果有一个搜索引擎，能搜索出和"今天的风好大"意思相近的古诗词，岂不妙哉！
使用向量数据库就可以实现，代码还不到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"# 根据设备选择是否使用fp16use_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，现在可以搜索了吗？别急，为了提高搜索效率，我们还需要创建索引。什么是索引？一些大部头图书的最后，一般都会有索引，它列出了书中出现的关键术语以及对应的页码，帮助你快速找到它们的位置。如果没有索引，那就只能用笨方法，从第一页开始一页一页往后找了。
# 创建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']# 查看索引详情index_details = client.describe_index(collection_name=collection_name,# 指定索引名称，这里假设使用第一个索引index_name="vector_index")print(f"索引vector_index详情：{index_details}")
索引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'}print(f"正在加载集合：{collection_name}")client.load_collection(collection_name=collection_name)
print(client.get_load_state(collection_name=collection_name)){'state': <LoadState: Loaded>}# 获取查询向量text = "今天的雨好大"query_vectors = [vectorize_text([text])['dense'][0].tolist()]
# 设置搜索参数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': '张昭'}}]"]
# 打印向量搜索结果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}}
res2 = client.search(collection_name=collection_name,# 指定查询向量data=query_vectors,# 指定搜索的字段anns_field="vector",# 设置搜索参数search_params=search_params,# 删除limit参数# 指定返回的字段output_fields=output_fields)print(res2)
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': '雨零感节，云飞应序。'}}]"]
# 通过表达式过滤字段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: ['[]']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': '去年战桑干源，今年战葱河道。'}}]"]
# paragraphs字段包含"雨"字filter = f"paragraphs like '%雨%'"res4 = client.query(collection_name=collection_name,filter=filter,output_fields=output_fields,limit=limit)print(res4)
data: ["{"author": "太宗皇帝","title": "咏雨","paragraphs": "罩云飘远岫，喷雨泛长河。","id": 20000305},{"author": "太宗皇帝","title": "咏雨","paragraphs": "和气吹绿野，梅雨洒芳田。","id": 20000402},{"author": "太宗皇帝","title": "赋得花庭雾","paragraphs": "还当杂行雨，髣髴隐遥空。","id": 20000421}"]
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': '杜甫'}"]
推荐阅读
没有评论:
发表评论