在今天的文章里,我是继上一篇文章 "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).zip
inflating: 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_dotenv
load_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 BadRequestError
try:
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 ConnectionTimeout
def 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.")
continue
for 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 started
if model_id not in models or model_status_map[model_id]:
continue
if "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] = True
if 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")
下载。