目录
为什么要将数据转为向量存入es?
我之前把数据作为文档存入 ES,主要用于全文检索(BM25 算法),但是它不适合语义匹配,比如如果用户输入的是"求解 x² - 5x + 6 = 0",但我的文档是"二次方程求解方法",ES 可能不会返回这个文档,因为它不包含完全匹配的关键词。
而如果将数据作为向量存入ES(语义搜索),则可以查找语义相似的内容。
数据准备
我这里准备包含250条数学文档的csv文件:

需要检查该文件的编码格式,用文本的方式打开该文件,如果不是utf-8,则另存为,并指定编码格式为utf-8,在改了编码格式后,再用excel打开,可能会变为乱码,不用管,只要保证编码为utf-8,且用文本打开能正常显示就行。
创建索引库
我使用的es版本为7.12.1
打开devtools工具,创建一个名为math_index的索引库
PUT /math_index
{
"settings": {
"number_of_shards": 3,
"number_of_replicas": 1
},
"mappings": {
"properties": {
"ask_vector": {
"type": "dense_vector",
"dims": 1024
},
"ask": {
"type": "text",
"analyzer": "ik_max_word",
"search_analyzer": "ik_smart"
},
"answer": {
"type": "text",
"analyzer": "ik_max_word",
"search_analyzer": "ik_smart"
}
}
}
}
1. 索引设定 (settings
)
number_of_shards: 3
:该索引被分成 3 个分片,提高查询和索引的吞吐量。number_of_replicas: 1
:每个主分片有 1 个副本,提高数据的可用性和容错性。
2. 字段映射 (mappings
)
ask_vector
:- 类型为
dense_vector
,维度为1024
,用于存储文本的向量表示(通常用于语义搜索,如基于向量相似度的检索)。
- 类型为
ask
:- 类型为
text
,用于存储用户问题文本。 analyzer: "ik_max_word"
:索引时使用ik_max_word
(细粒度分词)。search_analyzer: "ik_smart"
:搜索时使用ik_smart
(较粗粒度分词,提升搜索效率)。
- 类型为
answer
:- 类型为
text
,用于存储回答文本。 - 同样采用
ik_max_word
进行索引,ik_smart
进行搜索。
- 类型为
向量存储
from elasticsearch import Elasticsearch
from transformers import BertTokenizer, BertModel
import torch
import pandas as pd
def embeddings_doc(doc, tokenizer, model, max_length=300):
encoded_dict = tokenizer.encode_plus(
doc,
add_special_tokens=True,
max_length=max_length,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt'
)
input_id = encoded_dict['input_ids']
attention_mask = encoded_dict['attention_mask']
# 前向传播
with torch.no_grad():
outputs = model(input_id, attention_mask=attention_mask)
# 提取最后一层的CLS向量作为文本表示
last_hidden_state = outputs.last_hidden_state
cls_embeddings = last_hidden_state[:, 0, :]
return cls_embeddings[0]
def add_doc(index_name, id, embedding_ask, ask, answer, es):
body = {
"ask_vector": embedding_ask.tolist(),
"ask": ask,
"answer": answer
}
result = es.create(index=index_name, id=id, doc_type="_doc", body=body)
return result
def main():
# 模型下载的地址
model_name = 'D:\\model\\chinese-roberta-wwm-ext-large'
# ES 信息
es_host = "http://your_ip"
es_port = 9200
es_user = ""
es_password = ""
index_name = "math_index"
# 数据地址
path = "D:\\Downloads\\知识库1.4.csv"
# 分词器和模型
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)
# ES 连接
es = Elasticsearch(
[es_host],
port=es_port,
http_auth=(es_user, es_password)
)
# 读取数据写入ES
data = pd.read_csv(path, encoding='utf-8')
for index, row in data.iterrows():
ask = row["题目"]
answer = row["答案"]
# 文本转向量
embedding_ask = embeddings_doc(ask, tokenizer, model)
result = add_doc(index_name, index, embedding_ask, ask, answer, es)
print(result)
if __name__ == '__main__':
main()
里面的模型文件在这里下载:
https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/
把这几个文件下载下来,放到一个文件夹中:

然后运行脚本就可以了。(这里的es依赖用到7.12.1版本:pip install elasticsearch==7.12.1 -i https://pypi.tuna.tsinghua.edu.cn/simple)
运行结束后,es就存储了知识库数据以及生成的向量:

验证
这里使用余弦相似度进行相似性检索
from elasticsearch import Elasticsearch
from transformers import BertTokenizer, BertModel
import torch
def embeddings_doc(doc, tokenizer, model, max_length=300):
encoded_dict = tokenizer.encode_plus(
doc,
add_special_tokens=True,
max_length=max_length,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt'
)
input_id = encoded_dict['input_ids']
attention_mask = encoded_dict['attention_mask']
# 前向传播
with torch.no_grad():
outputs = model(input_id, attention_mask=attention_mask)
# 提取最后一层的CLS向量作为文本表示
last_hidden_state = outputs.last_hidden_state
cls_embeddings = last_hidden_state[:, 0, :]
return cls_embeddings[0]
def search_similar(index_name, query_text, tokenizer, model, es, top_k=3):
query_embedding = embeddings_doc(query_text, tokenizer, model)
print(query_embedding.tolist())
query = {
"query": {
"script_score": {
"query": {"match_all": {}},
"script": {
"source": "cosineSimilarity(params.queryVector, 'ask_vector') + 1.0",
"lang": "painless",
"params": {
"queryVector": query_embedding.tolist()
}
}
}
},
"size": top_k
}
res = es.search(index=index_name, body=query)
hits = res['hits']['hits']
similar_documents = []
for hit in hits:
similar_documents.append(hit['_source'])
return similar_documents
def main():
# 模型下载的地址
model_name = 'D:\\model\\chinese-roberta-wwm-ext-large'
# ES 信息
es_host = "http://your_ip"
es_port = 9200
es_user = ""
es_password = ""
index_name = "math _index"
# 分词器和模型
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)
# ES 连接
es = Elasticsearch(
[es_host],
port=es_port,
http_auth=(es_user, es_password)
)
query_text = "在复平面内,(1+3i)(3-i) 对应的点位于哪个象限"
similar_documents = search_similar(index_name, query_text, tokenizer, model, es)
for item in similar_documents:
print("================================")
print('ask:', item['ask'])
print('answer:', item['answer'])
if __name__ == '__main__':
main()
得到的结果(找到的相似度前十的数据):

可以看到,第一个数据确实最为相似。