es-将知识库中的数据转换为向量存储到es并进行相似性检索

目录

为什么要将数据转为向量存入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 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()

得到的结果(找到的相似度前十的数据):

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

相关推荐
bxlj_jcj1 小时前
深入Flink核心概念:解锁大数据流处理的奥秘
大数据·flink
云资源服务商1 小时前
阿里云Flink:开启大数据实时处理新时代
大数据·阿里云·云计算
Aurora_NeAr2 小时前
Spark SQL架构及高级用法
大数据·后端·spark
王小王-1232 小时前
基于Hadoop的公共自行车数据分布式存储和计算平台的设计与实现
大数据·hive·hadoop·分布式·hadoop公共自行车·共享单车大数据分析·hadoop共享单车
数据与人工智能律师2 小时前
数字资产革命中的信任之锚:RWA法律架构的隐形密码
大数据·网络·人工智能·云计算·区块链
Edingbrugh.南空3 小时前
Flink OceanBase CDC 环境配置与验证
大数据·flink·oceanbase
全星0073 小时前
解锁研发高效密码:全星研发项目管理APQP软件的多维助力
大数据·汽车
时序数据说5 小时前
为什么时序数据库IoTDB选择Java作为开发语言
java·大数据·开发语言·数据库·物联网·时序数据库·iotdb
Codebee5 小时前
OneCode图表配置速查手册
大数据·前端·数据可视化
HGW6896 小时前
基于 Elasticsearch 实现地图点聚合
java·elasticsearch·高德地图