当前位置：首页 > article >正文

elasticsearch+sentencetransformer检索样例

article 2025/3/26 15:53:10

elasticsearch+sentencetransformer结合实现搜索确实效率很高，但是跟着官方版本进行操作会有很多的误区，小编踩了很多坑才得到一版不报错的代码，感兴趣的小伙伴可以运行一下玩玩。
在开始前需要做一些准备工作：

一、准备工作

1、脚本中用到了elasticsearch、sentence_transformers，需要先行安装

这里我用的elasticsearch版本是8.8.1，不同版本可能会有所差异

pip install elasticsearch
pip install sentence_transformers

2、在运行代码前需要先行启动elasticsearch，出现以下类似的页面

在这里插入图片描述

3、在https://hf-mirror.com/中找到对应的模型

我在这里使用到的是all-MiniLM-L6-v2

4、准备待搜索的文档

官方文档提供的数据集引入方式是load_dataset，这个是下载现成的datasets来使用，但是我发现很难找到对应的数据集，所以我自己做了一个，其实很简单的csv文档，这样你可以方便的使用你的数据来替换搜索的内容
在这里插入图片描述

二、具体实现代码

from elasticsearch import Elasticsearch
from sentence_transformers import SentenceTransformer
import csv

# 配置参数
ES_HOST = "http://localhost:9200"
MODEL_PATH = "E:/all/pythonProject/2025srt/models/sentence-transformers/all-MiniLM-L6-v2" #将模型替换为你的模型对应的位置
CSV_PATH = "E:/all/pythonProject/2025srt/ess/quora/quornew.csv" #替换为你需要的路径和名字
INDEX_NAME = "quora_questions"  # 确保与后续查询一致

# 初始化连接和模型
es = Elasticsearch(ES_HOST)
model = SentenceTransformer(MODEL_PATH)

# 步骤1：创建包含 dense_vector 的索引
if not es.indices.exists(index=INDEX_NAME):
    index_settings = {
        "mappings": {
            "properties": {
                "text_vector": {
                    "type": "dense_vector",  # 关键修正点
                    "dims": 384,              # all-MiniLM-L6-v2 的向量维度是384
                    "index": True,           # 启用kNN索引
                    "similarity": "cosine"   # 根据模型选择相似度算法
                },
                "id": {"type": "integer"},
                "text": {"type": "text"}
            }
        }
    }
    es.indices.create(index=INDEX_NAME, body=index_settings)

# 步骤2：索引数据（处理BOM字符问题）
with open(CSV_PATH, newline='', encoding='utf-8-sig') as csvfile:  # 使用utf-8-sig处理BOM
    reader = csv.DictReader(csvfile)
    for row in reader:
        doc_id = int(row['id'])  # 修正BOM导致的字段名问题
        text = row['text']
        
        # 生成向量
        text_vector = model.encode(text).tolist()
        
        # 构建文档
        document = {
            "id": doc_id,
            "text": text,
            "text_vector": text_vector
        }
        
        # 写入ES
        es.index(index=INDEX_NAME, id=doc_id, document=document)

print(f"数据已索引到 {INDEX_NAME}")

# 步骤3：执行kNN查询（修正查询结构）
query_text = "What is the meaning of democracy?"
query_vector = model.encode(query_text).tolist()

# 修正后的查询体
search_body = {
    "knn": {
        "field": "text_vector",
        "query_vector": query_vector,
        "k": 5,
        "num_candidates": 100
    },
    "_source": ["text", "id"]  # 返回字段
}
# knn：从海量数据中‌快速找出与目标向量最接近的 k 个结果
try:
    response = es.search(index=INDEX_NAME, body=search_body)  # 关键修正：使用body参数
    print("搜索结果:")
    for hit in response['hits']['hits']:
        print("查找对象：" + query_text)
        print(f"ID: {hit['_id']}, Score: {hit['_score']:.4f}")
        print(f"Text: {hit['_source']['text']}\n")
except Exception as e:
    print(f"查询失败: {str(e)}")