目录
为什么要将数据转为向量存入es?
数据准备
创建索引库
向量存储
验证
为什么要将数据转为向量存入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 pddef 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_statecls_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 resultdef main():# 模型下载的地址model_name = 'D:\\model\\chinese-roberta-wwm-ext-large'# ES 信息es_host = "http://your_ip"es_port = 9200es_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))# 读取数据写入ESdata = 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 torchdef 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_statecls_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_documentsdef main():# 模型下载的地址model_name = 'D:\\model\\chinese-roberta-wwm-ext-large'# ES 信息es_host = "http://your_ip"es_port = 9200es_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()得到的结果（找到的相似度前十的数据）：
可以看到，第一个数据确实最为相似。
