「最佳实践」腾讯云 Elasticsearch 8 向量化语义检索：使用自建 GPU 机器学习节点进行高效推理

原创

岳涛

修改于 2024-03-11 19:27:24

2.4K280

代码可运行

文章被收录于专栏：大数据生态大数据生态

运行总次数：0

代码可运行

说明

本文描述问题及解决方法同样适用于 腾讯云 Elasticsearch Service（ES）。

另外使用到：腾讯云云服务器（Cloud Virtual Machine，CVM）

声明

本文使用的文本样本数据系淘宝、京东等电商平台首页随机爬取的商品标题。

环境配置

客户端环境

版本

CVM 镜像：CentOS 7.9 64位 | img-l8og963d | 20GiB

Linux环境：Centos 7.9

Python：3.9.12

配置

GN7.2XLARGE32（8核32G）

Elasticsearch 服务端环境

版本

Elasticsearch版本：8.8.1（腾讯云 Elasticsearch Service 基础版）

配置

规格：ES.SA2.8XLARGE64（32核64G）

节点数量：3

硬盘：200G * 3

背景

不同于「最佳实践」腾讯云 Elasticsearch 8：预训练模型与一站式向量化语义检索的完美结合，上一篇全篇围绕一站式体验ES 向量检索 —— Elasticsearch Relevance Engine™（ESRE™），本文主要介绍的是使用已有的机器学习服务器进行推理，效率相较使用 CPU 进行模型推理要高很多。接下来，开始我们这次的向量检索之旅。

1. 部署客户端环境

点击购买机器

注意：

GPU 驱动一定要选最新版，否则会出现兼容问题
磁盘建议选择 增强型SSD
VPC 需和 ES 在同一个网络环境

安装成功后验证是否成功：

nvidia-smi

返回显卡驱动版本信息则说明驱动正常。

2. 创建 ES 集群

点击创建集群

版本这里我们选择 8.8.1 基础版，也可以选择白金版，白金版有更多的 X-PACK 高级特性，可以根据实际需求选择基础版还是白金版：

配置建议至少选择（32核64G 200G*3）3个节点，也可以根据实际需求进行调整：

提交集群构建之后，大概需要20分钟左右可以完成。

集群创建完成之后，为了方便测试，需要移步ES实例 > 访问快照 > 可视化访问控制 > 公网访问策略，将白名单修改为0.0.0.0/0

注意：此操作是为了方便测试，生产环境还需谨慎操作。

3. 客户端准备工作

Python 环境部署

一键安装环境：

yum install conda -y; conda init; source ~/.bashrc; echo y | conda create -n es_vector python=3.9.12; conda activate es_vector

也可以逐步安装：

安装python环境管理工具——conda

yum install conda -y

初始化conda环境变量

conda init

说明：结尾报错Operation failed.无需理会

加载环境变量

source ~/.bashrc

创建虚拟环境es_vector

conda create -n es_vector python=3.9.12

激活虚拟环境

conda activate es_vector

安装 Python 依赖包

pip install altair==5.2.0 attrs==23.2.0 blinker==1.7.0 cachetools==5.3.2 certifi==2023.11.17 charset-normalizer==3.3.2 click==8.1.7 elastic-transport==8.11.0 elasticsearch==8.8.1 filelock==3.13.1 fsspec==2023.12.2 gitdb==4.0.11 GitPython==3.1.41 huggingface-hub==0.20.2 idna==3.6 importlib-metadata==7.0.1 Jinja2==3.1.3 joblib==1.3.2 jsonschema==4.20.0 jsonschema-specifications==2023.12.1 markdown-it-py==3.0.0 MarkupSafe==2.1.3 mdurl==0.1.2 mpmath==1.3.0 networkx==3.2.1 nltk==3.8.1 numpy==1.26.3 nvidia-cublas-cu12==12.1.3.1 nvidia-cuda-cupti-cu12==12.1.105 nvidia-cuda-nvrtc-cu12==12.1.105 nvidia-cuda-runtime-cu12==12.1.105 nvidia-cudnn-cu12==8.9.2.26 nvidia-cufft-cu12==11.0.2.54 nvidia-curand-cu12==10.3.2.106 nvidia-cusolver-cu12==11.4.5.107 nvidia-cusparse-cu12==12.1.0.106 nvidia-nccl-cu12==2.18.1 nvidia-nvjitlink-cu12==12.3.101 nvidia-nvtx-cu12==12.1.105 packaging==23.2 pandas==2.1.4 pillow==10.2.0 protobuf==4.25.2 pyarrow==14.0.2 pydeck==0.8.1b0 Pygments==2.17.2 python-dateutil==2.8.2 pytz==2023.3.post1 PyYAML==6.0.1 referencing==0.32.1 regex==2023.12.25 requests==2.31.0 rich==13.7.0 rpds-py==0.17.1 safetensors==0.4.1 scikit-learn==1.3.2 scipy==1.11.4 sentence-transformers==2.2.2 sentencepiece==0.1.99 six==1.16.0 smmap==5.0.1 streamlit==1.30.0 sympy==1.12 tenacity==8.2.3 threadpoolctl==3.2.0 tokenizers==0.15.0 toml==0.10.2 toolz==0.12.0 torch==2.1.2 torchvision==0.16.2 tornado==6.4 tqdm==4.66.1 transformers==4.36.2 triton==2.1.0 typing_extensions==4.9.0 tzdata==2023.4 tzlocal==5.2 urllib3==2.1.0 validators==0.22.0 watchdog==3.0.0 zipp==3.17.0

这一步需要下载很多依赖，安装时间会比较久，需要耐心等待。

4. 下载整合包

博主已将依赖模型及脚本打包成 整合包，可下载后上传至客户端服务器家目录：/root

解压整合包

博主将整合包压缩成了 ZSTD 格式，该格式的好处是压缩/解压缩性能极高，所以解压也需要使用 ZSTD 算法解压。

安装 zstd 命令：

yum install -y zstd

执行解压：

zstd -T0 -d tencent-es_vector.tzst && tar -xvf tencent-es_vector.tar && rm -rf tencent-es_vector.tar

一键复制命令进行：解压 -> 反归档 -> 删除归档包 ，然后我们可以得到一个整合包目录：

一共1个目录，3个文件：

bge-base-zh：预训练 Embedding 中文推理模型
goods_sentence.txt：商品标题库
insert_vector.py：通过预训练模型进行推理，生成向量化数据写入 ES
vector_search.py：将查询的关键词推理转化成向量，然后对 ES 发起向量检索

5. Elasticsearch 准备工作

定义模板

执行模板创建：

PUT _template/goods_vector
{
  "index_patterns": [
    "goods_vector*"
  ],
  "settings": {
    "number_of_shards": 6
  },
  "mappings": {
    "properties": {
      "id": {
        "type": "long"
      },
      "text_field": {
        "type": "text",
        "analyzer": "ik_max_word",
        "search_analyzer": "ik_smart"
      },
      "vector": {
        "type": "dense_vector",
        "similarity": "cosine",
        "index": true,
        "dims": 768,
        "element_type": "float",
        "index_options": {
          "type": "hnsw",
          "m": 16,
          "ef_construction": 128
        }
      }
    }
  }
}

返回"acknowledged": true即提交成功。

6. 数据生成

这里我们使用预训练模型进行推理，并将数据写入到ES

cd /root/tencent-es_vector/
vim insert_vector.py

修改配置信息：es_password es_host

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from sentence_transformers import SentenceTransformer
import torch

es_username = 'elastic'
es_password = '******' # 修改ES密码
es_host = '10.0.0.1'     # 修改ES HOST
es_port = 9200

# 加载模型
model = SentenceTransformer(r'/root/tencent-es_vector/bge-base-zh')  # 修改为模型路径
if torch.cuda.is_available():
    model = model.to(torch.device("cuda"))

es = Elasticsearch(
    hosts=[{'host': es_host, 'port': es_port, 'scheme': 'http'}],
    basic_auth=(es_username, es_password),
)

# 读取文本
file_path = 'goods_sentence.txt'
index_name = 'goods_vector'

def gen_vector(sentences):
    embeddings = model.encode(sentences, normalize_embeddings=True)
    return embeddings

def read_data(file_path):
    with open(file_path, 'r') as file:
        counter = 1
        for line in file:
            yield {
                '_index': index_name,
                '_id': counter,  # 添加自增 ID 字段
                '_source': {
                    'id': counter,
                    'text_field': line.strip(),
                    'vector': gen_vector(line.strip())
                }
            }
            counter += 1

# 执行批量插入
def bulk_insert(file_path, chunk_size=5000):
    data = read_data(file_path)
    bulk(es, data, chunk_size=chunk_size)

bulk_insert(file_path)

执行文本导入：

cd /root/tencent-es_vector/
python insert_sentence.py

导入过程中可以使用英伟达显卡监控命令观测GPU使用率情况：

# 5秒打印1次
nvidia-smi -l 5

4万条数据进行推理运算转换成向量，博主测试耗时 8 分钟，导入完成后可以在 kibana 中检索到数据：

GET goods_vector/_search
{
  "_source": "text_field"
}

查询有返回则说明数据已经开始导入。

7. 语义检索

所有准备工作就绪，下面将演示向量检索，我们分别用向量检索和分词检索测试两者的检索效果：

cd /root/tencent-es_vector/
vim vector_search.py

修改配置信息：password host

import torch, streamlit as st
from sentence_transformers import SentenceTransformer
from elasticsearch import Elasticsearch

username = 'elastic'
password = '******' # 修改为ES密码
host = 'http://10.0.0.1:9200' # 修改为ES URL
index = 'goods_vector'

# 加载模型
model = SentenceTransformer(r'/root/tencent-es_vector/bge-base-zh')  # 修改为模型路径
if torch.cuda.is_available():
    # 使用GPU推理
    model = model.to(torch.device("cuda"))

# 使用IK分词器
def ik_search(es, index_name, query_text):
    source_fields = ["text_field"]
    query = {
        "match": {
            "text_field": query_text
        }
    }

    resp = es.search(
        index=index_name,
        fields=source_fields,
        query=query,
        size=10,
        source=False)
    return resp

# 使用k-NN搜索
def knn_search(es, index_name, vector):
    source_fields = ["text_field"]
    knn = [{
        "field": "vector",
        "query_vector": vector,
        "k": 10,
        "num_candidates": 100
    }]

    resp = es.search(
        index=index_name,
        fields=source_fields,
        knn=knn,
        source=False)
    return resp

# 生成句子向量
def gen_vector(sentences):
    embeddings = model.encode(sentences, normalize_embeddings=True)
    return embeddings

# 创建界面
st.title("腾讯云Elaticsearch 8.8.1 向量检索")

with st.form("chat_form"):
    query = st.text_input("请输入文本：")
    submit_button = st.form_submit_button("查询")

# 连接ES
es = Elasticsearch(hosts=[host],
                   basic_auth=(username, password))

# 当点击查询按钮时
if submit_button:
    # 调用ik检索
    ik_resp = ik_search(es, index, query)

    # 调用knn检索
    knn_resp = knn_search(es, index, gen_vector(query))

    # 创建两列
    col1, col2 = st.columns(2)

    counter = 1
    with col1:
        st.write("### 分词检索结果")
        for hit in ik_resp['hits']['hits']:
            text_field = hit['fields']['text_field'][0]
            with st.container():
                st.text(f"{counter}. {text_field}")
            counter += 1

    counter = 1
    with col2:
        st.write("### 向量检索结果")
        for hit in knn_resp['hits']['hits']:
            text_field = hit['fields']['text_field'][0]
            with st.container():
                st.text(f"{counter}. {text_field}")
            counter += 1

cd /root/tencent-es_vector/
streamlit run vector_search.py

访问返回的公网地址，进行向量测试：

成功访问则表示部署成功。

8. 效果展示

关键词：iphone 12 pro max

关键词：iphone 11 pro max

可以看到传统分词检索和向量检索的差异还是比较明显的。

总结

语义搜索之所以如此重要，是因为它能够进行更广泛的搜索范围。得益于向量搜索的支持，语义搜索能够提供更加直观的搜索体验，并根据查询的上下文和搜索意图生成相匹配的结果。

相较于关键字，语义搜索更具优势，因为它通过匹配概念而非关键字来生成更精确的搜索结果。通过维度嵌入，一个向量能够代表一个词的概念。

因此，基于向量搜索的语义搜索超越了简单匹配由词元表示的关键字概念的局限，从而实现了更高效准确的搜索体验。

我正在参与2024腾讯技术创作特训营第五期有奖征文，快来和我瓜分大奖！

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

elasticsearch

2024腾讯·技术创作特训营第五期

大数据解决方案

Elasticsearch Service