Elasticsearch：Retrievers 介绍 - Python Jupyter notebook

在今天的文章里，我是继上一篇文章 “Elasticsearch：介绍 retrievers - 搜索一切事物” 来使用一个可以在本地设置的 Elasticsearch 集群来展示 Retrievers 的使用。在本篇文章中，你将学到如下的内容：

• 从 Kaggle 下载 IMDB 数据集
• 创建两个推理服务
• 部署 ELSER
• 部署 e5-small
• 创建摄取管道
• 创建映射
• 摄取 IMDB 数据，在摄取过程中创建嵌入
• 缩小查询负载模型
• 运行示例检索器

安装

 Elasticsearch 及 Kibana

 如果你还没有安装好自己的 Elasticsearch 及 Kibana，请参考如下的链接来进行安装：

• 如何在 Linux，MacOS 及 Windows 上进行安装 Elasticsearch
• Kibana：如何在 Linux，MacOS 及 Windows上安装 Elastic 栈中的 Kibana

在安装的时候，我们选择 Elastic Stack 8.x 来进行安装。在首次启动 Elasticsearch 的时候，我们可以看到如下的输出：

​

在上面，我们可以看到 elastic 超级用户的密码。我们记下它，并将在下面的代码中进行使用。

我们还可以在安装 Elasticsearch 目录中找到 Elasticsearch 的访问证书：

$ pwd
/Users/liuxg/elastic/elasticsearch-8.14.1/config/certs
$ ls
http.p12      http_ca.crt   transport.p12

在上面，http_ca.crt 是我们需要用来访问 Elasticsearch 的证书。

 我们首先克隆已经写好的代码：

git clone https://github.com/liu-xiao-guo/elasticsearch-labs

我们然后进入到该项目的根目录下：

$ pwd
/Users/liuxg/python/elasticsearch-labs/supporting-blog-content/introducing-retrievers
$ ls
retrievers_intro_notebook.ipynb

如上所示，retrievers_intro_notebook.ipynb 就是我们今天想要工作的 notebook。

我们通过如下的命令来拷贝所需要的证书：

$ pwd
/Users/liuxg/python/elasticsearch-labs/supporting-blog-content/introducing-retrievers
$ cp ~/elastic/elasticsearch-8.14.1/config/certs/http_ca.crt .
$ ls
http_ca.crt                     retrievers_intro_notebook.ipynb

安装所需要的 python 依赖包

pip3 install -qqq pandas elasticsearch python-dotenv

我们可以使用如下的方法来查看 elasticsearch 的版本：

$ pip3 list | grep elasticsearch
elasticsearch                           8.14.0

 创建环境变量

为了能够使得下面的应用顺利执行，在项目当前的目录下运行如下的命令：

export ES_ENDPOINT="localhost"
export ES_USER="elastic"
export ES_PASSWORD="uK+7WbkeXMzwk9YvP-H3"

你需要根据自己的 Elasticsearch 设置进行相应的修改。

下载数据集

我们去到地址 IMDB movies dataset | Kaggle 下载数据集并解压缩。

$ pwd
/Users/liuxg/python/elasticsearch-labs/supporting-blog-content/introducing-retrievers
$ ls
archive (13).zip                http_ca.crt                     retrievers_intro_notebook.ipynb
$ unzip archive\ \(13\).zip 
Archive:  archive (13).zipinflating: imdb_movies.csv         
$ mkdir -p content
$ mv imdb_movies.csv content/

$ tree -L 2
.
├── archive\ (13).zip
├── content
│   └── imdb_movies.csv
├── http_ca.crt
└── retrievers_intro_notebook.ipynb

如上所示，我们吧 imdb_movies.csv 文件置于当前工作目录下的 content 目录下。

代码展示

我们在当前项目的根目录下打入如下的命令：

设置

import os
import zipfile
import pandas as pd
from elasticsearch import Elasticsearch, helpers
from elasticsearch.exceptions import ConnectionTimeout
from elastic_transport import ConnectionError
from time import sleep
import time
import logging# Get the logger for 'elastic_transport.node_pool'
logger = logging.getLogger("elastic_transport.node_pool")# Set its level to ERROR
logger.setLevel(logging.ERROR)# Suppress warnings from the elastic_transport module
logging.getLogger("elastic_transport").setLevel(logging.ERROR)

连接到 Elasticsearch

from dotenv import load_dotenvload_dotenv()ES_USER = os.getenv("ES_USER")
ES_PASSWORD = os.getenv("ES_PASSWORD")
ES_ENDPOINT = os.getenv("ES_ENDPOINT")
COHERE_API_KEY = os.getenv("COHERE_API_KEY")url = f"https://{ES_USER}:{ES_PASSWORD}@{ES_ENDPOINT}:9200"
print(url)es = Elasticsearch(url, ca_certs = "./http_ca.crt", verify_certs = True)
print(es.info())

如上所示，我们的客户端连接到 Elasticsearch 是成功的。

部署 ELSER 及 E5

下面的两个代码块将部署嵌入模型并自动扩展 ML 容量。

部署及启动 ELSER

from elasticsearch.exceptions import BadRequestErrortry:resp = es.options(request_timeout=5).inference.put_model(task_type="sparse_embedding",inference_id="my-elser-model",body={"service": "elser","service_settings": {"num_allocations": 64, "num_threads": 1},},)
except ConnectionTimeout:pass
except BadRequestError as e:print(e)

如果你之前已经部署过 ELSER，你可能会得到一个 resource already exists 这样的错误。你可以使用如下的命令来删除之前的 inference_id。

DELETE /_inference/my-elser-model

在运行完上面的命令后，需要经过一定的时间下载 ELSER 模型。这个依赖于你的网络速度。我们可以在 Kibana 中进行查看：

部署及启动 es-small

try:resp = es.inference.put_model(task_type="text_embedding",inference_id="my-e5-model",body={"service": "elasticsearch","service_settings": {"num_allocations": 8,"num_threads": 1,"model_id": ".multilingual-e5-small",},},)
except ConnectionTimeout:pass
except BadRequestError as e:print(e)

在运行完上面的代码后，我们可以在 Kibana 界面中：

点击上面的 "Add trained model" 来安装 .multilingual-e5-small 模型。

我们到最后能看到这个：

整个下载及部署需要很长的时间，需要你耐心等待！

提示：如果你的机器是在 x86 架构的机器上运行的话，那么你在上面可以选择 .multilingual-e5-small_linux-x86_64 作为其 model_id。

检查模型部署状态

这将循环检查，直到 ELSER 和 e5 都已完全部署。如果你在上面已经等了足够久的话，那么下面的代码讲很快地执行。

如果需要分配额外容量来运行模型，这可能需要几分钟

from time import sleep
from elasticsearch.exceptions import ConnectionTimeoutdef wait_for_models_to_start(es, models):model_status_map = {model: False for model in models}while not all(model_status_map.values()):try:model_status = es.ml.get_trained_models_stats()except ConnectionTimeout:print("A connection timeout error occurred.")continuefor x in model_status["trained_model_stats"]:model_id = x["model_id"]# Skip this model if it's not in our list or it has already startedif model_id not in models or model_status_map[model_id]:continueif "deployment_stats" in x:if ("nodes" in x["deployment_stats"]and len(x["deployment_stats"]["nodes"]) > 0):if (x["deployment_stats"]["nodes"][0]["routing_state"]["routing_state"]== "started"):print(f"{model_id} model deployed and started")model_status_map[model_id] = Trueif not all(model_status_map.values()):sleep(0.5)models = [".elser_model_2", ".multilingual-e5-small"]
wait_for_models_to_start(es, models)

.elser_model_2 model deployed and started
.multilingual-e5-small model deployed and started

创建索引模板并链接到摄取管道

template_body = {"index_patterns": ["imdb_movies*"],"template": {"settings": {"index": {"default_pipeline": "elser_e5_embed"}},"mappings": {"properties": {"budget_x": {"type": "double"},"country": {"type": "keyword"},"crew": {"type": "text"},"date_x": {"type": "date", "format": "MM/dd/yyyy||MM/dd/yyyy[ ]"},"genre": {"type": "keyword"},"names": {"type": "text"},"names_sparse": {"type": "sparse_vector"},"names_dense": {"type": "dense_vector"},"orig_lang": {"type": "keyword"},"orig_title": {"type": "text"},"overview": {"type": "text"},"overview_sparse": {"type": "sparse_vector"},"overview_dense": {"type": "dense_vector"},"revenue": {"type": "double"},"score": {"type": "double"},"status": {"type": "keyword"},}},},
}# Create the template
es.indices.put_index_template(name="imdb_movies", body=template_body)

创建采集管道

# Define the pipeline configuration
pipeline_body = {"processors": [{"inference": {"model_id": ".multilingual-e5-small","description": "embed names with e5 to names_dense nested field","input_output": [{"input_field": "names", "output_field": "names_dense"}],}},{"inference": {"model_id": ".multilingual-e5-small","description": "embed overview with e5 to names_dense nested field","input_output": [{"input_field": "overview", "output_field": "overview_dense"}],}},{"inference": {"model_id": ".elser_model_2","description": "embed overview with .elser_model_2 to overview_sparse nested field","input_output": [{"input_field": "overview", "output_field": "overview_sparse"}],}},{"inference": {"model_id": ".elser_model_2","description": "embed names with .elser_model_2 to names_sparse nested field","input_output": [{"input_field": "names", "output_field": "names_sparse"}],}},],"on_failure": [{"append": {"field": "_source._ingest.inference_errors","value": [{"message": "{{ _ingest.on_failure_message }}","pipeline": "{{_ingest.pipeline}}","timestamp": "{{{ _ingest.timestamp }}}",}],}}],
}# Create the pipeline
es.ingest.put_pipeline(id="elser_e5_embed", body=pipeline_body)

提取文档

这将

• 进行一些预处理
• 批量提取 10,178 条 IMDB 记录
• 使用 ELSER 模型为 overview 和 name 字段生成稀疏向量嵌入
• 使用 ELSER 模型为 overview 和 name 字段生成密集向量嵌入

使用上述分配设置通常需要一定的时间才能完成。这个依赖于你自己电脑的配置。

# Load CSV data into a pandas DataFrame
df = pd.read_csv("./content/imdb_movies.csv")# Replace all NaN values in DataFrame with None
df = df.where(pd.notnull(df), None)# Convert DataFrame into a list of dictionaries
# Each dictionary represents a document to be indexed
documents = df.to_dict(orient="records")# Define a function to generate actions for bulk API
def generate_bulk_actions(documents):for doc in documents:yield {"_index": "imdb_movies","_source": doc,}# Use the bulk helper to insert documents, 200 at a time
start_time = time.time()
helpers.bulk(es, generate_bulk_actions(documents), chunk_size=200)
end_time = time.time()print(f"The function took {end_time - start_time} seconds to run")

我们可以在 Kibana 中进行查看：

    

    

  

  

我们需要等一定的时间来完成上面的摄取工作。值得注意的是：在上面的代码中我把 chunk_size 设置为 20。这个是为了避免 "Connection timeout" 错误。如果我们把这个值设置很大，那么摄取的时间可能过长，那么就会发生 "Connection timeout" 这样的错误。我们在批量处理时，选择比较少的文档来完成摄取工作。有关如何设置这个 timeout 的时间，我们可以参考文章 “在 Elasticsearch 中扩展 ML 推理管道：如何避免问题并解决瓶颈”。

针对我的电脑，它花费了如下的时间来完成 10,178 个文档的摄取：

The function took 1292.8102316856384 seconds to run

这个将近20分钟。

缩小 ELSER 和 e5 模型

我们不需要大量的模型分配来进行测试查询，因此我们将每个模型分配缩小到 1 个

for model_id in [".elser_model_2","my-e5-model"]:result = es.perform_request("POST",f"/_ml/trained_models/{model_id}/deployment/_update",headers={"content-type": "application/json", "accept": "application/json"},body={"number_of_allocations": 1},)

Retriever 测试

我们将使用搜索输入 clueless slackers 在数据集中的 overview 字段（文本或嵌入）中搜索电影

请随意将下面的 movie_search 变量更改为其他内容

movie_search = "clueless slackers"

Standard - 搜索所有文本！ - bm25

response = es.search(index="imdb_movies",body={"query": {"match": {"overview": movie_search}},"size": 3,"fields": ["names", "overview"],"_source": False,},
)for hit in response["hits"]["hits"]:print(f"{hit['fields']['names'][0]}\n- {hit['fields']['overview'][0]}\n")

    

kNN-搜索所有密集向量！

response = es.search(index="imdb_movies",body={"retriever": {"knn": {"field": "overview_dense","query_vector_builder": {"text_embedding": {"model_id": "my-e5-model","model_text": movie_search,}},"k": 5,"num_candidates": 5,}},"size": 3,"fields": ["names", "overview"],"_source": False,},
)for hit in response["hits"]["hits"]:print(f"{hit['fields']['names'][0]}\n- {hit['fields']['overview'][0]}\n")

  

text_expansion - 搜索所有稀疏向量！ - elser

response = es.search(index="imdb_movies",body={"retriever": {"standard": {"query": {"text_expansion": {"overview_sparse": {"model_id": ".elser_model_2","model_text": movie_search,}}}}},"size": 3,"fields": ["names", "overview"],"_source": False,},
)for hit in response["hits"]["hits"]:print(f"{hit['fields']['names'][0]}\n- {hit['fields']['overview'][0]}\n")

  

rrf — 将所有事物结合起来！

response = es.search(index="imdb_movies",body={"retriever": {"rrf": {"retrievers": [{"standard": {"query": {"term": {"overview": movie_search}}}},{"knn": {"field": "overview_dense","query_vector_builder": {"text_embedding": {"model_id": "my-e5-model","model_text": movie_search,}},"k": 5,"num_candidates": 5,}},{"standard": {"query": {"text_expansion": {"overview_sparse": {"model_id": ".elser_model_2","model_text": movie_search,}}}}},],"window_size": 5,"rank_constant": 1,}},"size": 3,"fields": ["names", "overview"],"_source": False,},
)for hit in response["hits"]["hits"]:print(f"{hit['fields']['names'][0]}\n- {hit['fields']['overview'][0]}\n")

  

所有的源码可以在地址 elasticsearch-labs/supporting-blog-content/introducing-retrievers/retrievers_intro_notebook.ipynb at main · liu-xiao-guo/elasticsearch-labs · GitHub

下载。