如何假装文艺青年，怎么把大白话“变成”古诗词？

午后细雨绵绵，你独倚窗边，思绪万千，于是拿出手机，想发一条朋友圈抒发情怀，随便展示一下文采。奈何好不容易按出几个字，又全部删除。“今天的雨好大”展示不出你的文采。你灵机一动，如果有一个搜索引擎，能搜索出和“今天的雨好大”意思相近的古诗词，岂不妙哉！

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

本文首发于 Zilliz 公众号。文中代码的 Notebook 在这里下载。

0 准备工作

首先安装向量数据库 Milvus。因为 Milvus 是运行在 docker 中的，所以需要先安装 Docker Desktop。MacOS 系统安装方法：Install Docker Desktop on Mac ，Windows 系统安装方法：Install Docker Desktop on Windows

然后安装 Milvus。Milvus 版本：>=2.5.0
下载安装脚本：

1
curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh

运行 Milvus：

1
bash standalone_embed.sh start

安装依赖。pymilvus >= 2.5.0

1
pip install pymilvus==2.5.0 "pymilvus[model]" torch

下载古诗词数据集[1] TangShi.json。它的格式是这样的：

1
2
3
4
5
6
7
8
9
10
11
12
[
	{
		"author": "太宗皇帝",
		"paragraphs": [
			"秦川雄帝宅，函谷壮皇居。"
		],
		"title": "帝京篇十首 一",
		"id": 20000001,
		"type": "唐诗"
	},
	...
]

准备就绪，正式开始啦。

1 向量化文本

为了实现语义搜索，我们需要先把文本向量化。你可以理解为把不同类型的信息（如文字、图像、声音）翻译成计算机可以理解的数字语言。计算机理解了，才能帮你找到语义相近的诗句。

先定义两个函数，一个用来初始化嵌入模型（也就是用来向量化的模型）的实例，另一个是调用嵌入模型的实例，把输入的文档向量化。

初始化嵌入模型的实例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from tqdm import tqdm
import torch
from pymilvus.model.hybrid import BGEM3EmbeddingFunction

# 初始化嵌入模型的实例
def init_embedding_model():
    # 检查是否有可用的CUDA设备
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    # 根据设备选择是否使用fp16
    use_fp16 = device.startswith("cuda")
    # 创建嵌入模型实例
    bge_m3_ef = BGEM3EmbeddingFunction(
        model_name="BAAI/bge-m3",
        device=device,
        use_fp16=use_fp16
    )
    return bge_m3_ef

bge_m3_ef = init_embedding_model()

向量化文档：

1
2
3
4
5
6
7
8
# 把文档向量化
def vectorize_docs(docs, encoder):
    # 验证参数是否符合要求
    if encoder is None:
        raise ValueError("嵌入模型未初始化。")
    if not (isinstance(docs, list) and all(isinstance(text, str) for text in docs)):
        raise ValueError("docs必须为字符串列表。")
    return encoder.encode_documents(docs)

准备好后，我们就可以向量化整个数据集了。首先读取数据集 TangShi.json 中的数据，把其中的 paragraphs 字段向量化，然后写入 TangShi_vector.json 文件。如果你是第一次使用 Milvus，运行下面的代码时还会安装必要的依赖。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import json
# 读取 json 文件，把paragraphs字段向量化
with open("TangShi.json", 'r', encoding='utf-8') as file:
	data_list = json.load(file)
	docs = [data['paragraphs'][0] for data in data_list]

# 向量化文本数据
vectors = vectorize_docs(docs, bge_m3_ef)

# 将向量添加到原始文本中
for data, vector in zip(data_list, vectors['dense']):
    # 将 NumPy 数组转换为 Python 的普通列表
	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 字段，它的值是一个字符串列表，也就是“向量”。

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
[
	{
		"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 创建集合

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

1
2
3
4
5
6
# 创建milvus_client实例
from pymilvus import MilvusClient, DataType

milvus_client = MilvusClient(uri="http://localhost:19530")
# 指定集合名称
collection_name = "TangShi"

为了避免向量数据库中存在同名集合，产生干扰，创建集合前先删除同名集合。

1
2
3
4
5
6
7
8
9
# 检查同名集合是否存在，如果存在则删除
if milvus_client.has_collection(collection_name):
    print(f"集合 {collection_name} 已经存在")
    try:
        # 删除同名集合
        milvus_client.drop_collection(collection_name)
        print(f"删除集合：{collection_name}")
    except Exception as e:
        print(f"删除集合时出现错误: {e}")

我们把数据填入 excel 表格前，需要先设计好表头，规定有哪些字段，各个字段的数据类型是怎样的。向量数据库也是一样，它的“表头”就是 schema，模式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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=1024)
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)

模式创建好了，接下来就可以创建集合了。

1
2
3
4
5
6
7
8
9
10
# 创建集合
try:
    milvus_client.create_collection(
        collection_name=collection_name,
        schema=schema,
        shards_num=2
    )
    print(f"开始创建集合：{collection_name}")
except Exception as e:
    print(f"创建集合的过程中出现了错误: {e}")

3 入库

接下来把文件导入到 Milvus。

1
2
3
4
5
6
7
8
9
10
11
12
13
# 读取和处理文件
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 = milvus_client.insert(
    collection_name=collection_name,
    data=data
)

导入成功了吗？我们来验证下。

1
print(f"插入的实体数量: {res['insert_count']}")

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

1
插入的实体数量: 4307

4 创建索引

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

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

Milvus 的索引也是如此。如果不创建索引，虽然也可以搜索，但是速度很慢，它会逐一比较查询向量与数据库中每一个向量，通过指定方法计算出两个向量之间的 距离，找出距离最近的几个向量。而创建索引之后，搜索速度会大大提升。

索引有不同的类型，适合不同的场景使用，我们以后会详细讨论这个问题。这里我们使用 IVF_FLAT。另外，计算距离的方法也有多种，我们使用 IP，也就是计算两个向量的内积。这些都是索引的参数，我们先创建这些参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 创建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
)

索引参数创建好了，现在终于可以创建索引了。

1
2
3
4
5
6
7
8
9
print(f"开始创建索引：{vector_index_name}")

# 创建索引
client.create_index(
	# 指定为哪个集合创建索引
	collection_name=collection_name,
	# 使用前面创建的索引参数创建索引
	index_params=index_params
)

我们来验证下索引是否创建成功了。

1
2
3
4
indexes = client.list_indexes(
	collection_name=collection_name
)
print(f"列出创建的索引：{indexes}")

返回了包含索引名称的列表，索引名称 vector_index 正是我们之前创建的。

1
列出创建的索引：['vector_index']

再来查看下索引的详情。

1
2
3
4
5
6
7
# 查看索引详情
index_details = client.describe_index(
	collection_name=collection_name,
	# 指定索引名称，这里假设使用第一个索引
	index_name="vector_index"
)
print(f"索引vector_index详情：{index_details}")

返回了一个包含索引详细信息的字典，可以我们之前设置的索引参数，比如 nlist，index_type 和 metric_type 等等。

1
索引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 加载索引

索引创建成功了，现在可以搜索了吗？等等，我们还需要把集合中的数据和索引，从硬盘加载到内存中。因为在内存中搜索更快。

1
2
print(f"正在加载集合：{collection_name}")
client.load_collection(collection_name=collection_name)

加载完成了，仍然验证下。

1
print(client.get_load_state(collection_name=collection_name))

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

1
{'state': <LoadState: Loaded>}

6 搜索

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

首先，把我们要搜索的现代白话文“翻译”成向量。

1
2
3
# 获取查询向量
query = "今天的雨好大"
query_vectors = [vectorize_docs([query], bge_m3_ef)['dense'][0].tolist()]

然后，设置搜索参数，告诉 Milvus 怎么搜索。

1
2
3
4
5
6
7
# 设置搜索参数
search_params = {
	# 设置度量类型
	"metric_type": "IP",
	# 指定在搜索过程中要查询的聚类单元数量，增加nprobe值可以提高搜索精度，但会降低搜索速度
	"params": {"nprobe": 16}
}

最后，我们还得告诉它怎么输出结果。

1
2
3
4
# 指定返回搜索结果的数量
limit = 3
# 指定返回的字段
output_fields = ["author", "title", "paragraphs"]

一切就绪，让我们开始搜索吧！

1
2
3
4
5
6
7
8
9
10
11
12
res1 = milvus_client.search(
    collection_name=collection_name,
    # 指定查询向量
    data=query_vectors,
    # 指定搜索的字段
    anns_field="vector",
    # 设置搜索参数
    search_params=search_params,
    limit=limit,
    output_fields=output_fields
)
print(res1)

得到下面的结果：

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
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 内积，数字越大表示搜索结果和查询向量越接近。

为了增加可读性，我们写一个输出函数：

1
2
3
4
5
6
7
8
9
10
# 打印向量搜索结果
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)}")

重新输出结果：

1
print_vector_results(res1)

这下搜索结果容易阅读了。

1
2
3
4
5
6
7
8
9
10
11
12
13
title: 郊庙歌辞 享太庙乐章 大明舞
author: 张说
paragraphs: 旱望春雨，云披大风。
--------------------------------------------------
title: 三学山夜看圣灯
author: 蜀太妃徐氏
paragraphs: 细雨湿不暗，好风吹更明。
--------------------------------------------------
title: 郊庙歌辞 汉宗庙乐舞辞 积善舞
author: 张昭
paragraphs: 云行雨施，天成地平。
--------------------------------------------------
数量：3

如果你不想限制搜索结果的数量，而是返回所有质量符合要求的搜索结果，可以修改搜索参数。比如，在搜索参数中增加 radius 和 range_filter 参数，它们限制了距离 distance 的范围在0.55到1之间。

1
2
3
4
5
6
7
8
9
# 修改搜索参数，设置距离的范围
search_params = {
    "metric_type": "IP",
    "params": {
        "nprobe": 16,
        "radius": 0.55,
        "range_filter": 1.0
    }
}

然后调整搜索代码，删除 limit 参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14

res2 = milvus_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。

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
data: [
    "[
        {
            'id': 20002740,
            'distance': 0.6542239189147949,
            'entity': {
                'author': '张说',
                'title': '郊庙歌辞 享太庙乐章 大明舞',
                'paragraphs': '旱望春雨，云披大风。'
            }
        },
        {
            'id': 20001658,
            'distance': 0.6228379011154175,
            'entity': {
                'author': '蜀太妃徐氏',
                'title': '三学山夜看圣灯',
                'paragraphs': '细雨湿不暗，好风吹更明。'
            }
        },
        ...
        {
	        'id': 20001375, 
	        'distance': 0.5891480445861816, 
	        'entity': {
		        'author': '上官昭容', 
		        'title': '游长宁公主流杯池二十五首 二十', 
		        'paragraphs': '瀑溜晴疑雨，丛篁昼似昏。'
		    }
		}
    ]"
]

也许你还想知道你最喜欢的李白，有没有和你一样感慨今天的雨真大，没问题，我们增加filter参数就可以了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 通过表达式过滤字段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)

返回的结果为空值。

1
data: ['[]']

这是因为我们前面设置了 distance 的范围在0.55到1之间，放大范围可以获得更多结果。把 “radius” 的值修改为0.2，再次运行命令，让我们看看李白是怎么感慨的。

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
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 方法做标量搜索。

1
2
3
4
5
6
7
8
9
10
# 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 。

观察搜索结果发现，标量搜索结果的数据结构少了一个 “[]”，我们在提取具体字段时需要注意这一点。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
data: [
    "{
        "author": "太宗皇帝",
        "title": "咏雨",
        "paragraphs": "罩云飘远岫，喷雨泛长河。",
        "id": 20000305
    },
    {
        "author": "太宗皇帝",
        "title": "咏雨",
        "paragraphs": "和气吹绿野，梅雨洒芳田。",
        "id": 20000402
    },
    {
        "author": "太宗皇帝",
        "title": "赋得花庭雾",
        "paragraphs": "还当杂行雨，髣髴隐遥空。",
        "id": 20000421
    }"
]

filter 表达式还有丰富的用法，比如同时搜索两个字段，author 字段指定为 “杜甫”，同时 paragraphs 字段仍然要求包含“雨”字：

1
2
3
4
5
6
7
8
9
filter = f"author == '杜甫' && paragraphs like '%雨%'"

res5 = client.query(
	collection_name=collection_name,
	filter=filter,
	output_fields=output_fields,
	limit=limit
)
print(res5)

返回杜甫含有“雨”字的诗句：

1
2
3
4
5
6
7
8
data: [
	"{
		'title': '横吹曲辞 前出塞九首 七', 
		'paragraphs': '驱马天雨雪，军行入高山。', 
		'id': 20004039, 
		'author': '杜甫'
	}"
]

更多标量搜索的表达式可以参考Get & Scalar Query。

总结

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

如果你希望优化搜索功能，可以在 chinese-poetry 下载完整的古诗词数据集，再找找专门用于古诗词的嵌入模型，相信搜索效果会有较大提升。

另外，我在以上代码的基础上，开发了一个命令行应用，有兴趣可以玩玩：语义搜索古诗词

注释

  1. 古诗词数据集来自 chinese-poetry，数据结构做了调整。 ↩
向量数据库 > 趣味应用
如何假装文艺青年，怎么把大白话“变成”古诗词？
http://example.com/2024/09/16/如何假装文艺青年，怎么把大白话“变成”古诗词？/
作者
江浩
发布于
2024年9月16日
许可协议