Elasticsearch：使用 Elasticsearch 和 BERT 构建搜索引擎 - TensorFlow

Posted 2023-03-08 Elastic 中国社区官方博客

在本文中，我们使用预训练的 BERT 模型和 Elasticsearch 来构建搜索引擎。 Elasticsearch 最近发布了带有向量场的文本相似性（text similarity search with vector field）搜索。 另一方面，你可以使用 BERT 将文本转换为固定长度的向量。 因此，一旦我们将文档通过 BERT 转换为向量并存储到 Elasticsearch 中，我们就可以使用 Elasticsearch 和 BERT 搜索相似的文档。

这篇文章通过以下架构实现了一个带有 Elasticsearch 和 BERT 的搜索引擎。 在这里，我们使用 Docker 将整个系统分为三个部分：应用程序、BERT 和 Elasticsearch。 目的是使扩展每个服务更容易。

 整个系统是在 docker-compose.yamlin 中编写的，位于以下 GitHub 存储库中。 请查看存储库：GitHub - liu-xiao-guo/bertsearch: Elasticsearch with BERT for advanced document search.

$ pwd
/Users/liuxg/python/bertsearch
$ tree -L 3
.
├── LICENSE
├── README.md
├── bertserving
│   ├── Dockerfile
│   └── entrypoint.sh
├── docker-compose.yaml
├── docs
│   ├── architecture.png
│   ├── diagram.key
│   └── example.png
├── example
│   ├── __init__.py
│   ├── create_documents.py
│   ├── create_index.py
│   ├── example.csv
│   ├── index.json
│   ├── index_documents.py
│   └── requirements.txt
└── web
    ├── Dockerfile
    ├── app.py
    ├── requirements.txt
    └── templates
        └── index.html

请注意：在本文的展示中，我使用 TensorFlow 来进行展示。更对关于 Pytorch 的展示，请参阅我之前的文章 “Elastic：开发者上手指南” 中的 “NLP - 自然语言处理” 部分。另外，由于 TensorFlow 里的指令限制，该展示不支持 Apple chipset 的电脑。你需要在 Intel 运行的机器上运行。

这篇文章的计划是：

• 下载预训练的 BERT 模型
• 设置环境变量
• 启动 Docker 容器
• 创建 Elasticsearch 索引
• 创建文件
• 索引文件

安装

你需要安装好自己的 Docker 环境。你需要安装自己的 Python。需要在版本 3.0 及以上。另外，为了能够让项目里的 Python 能够正常运行，你需要按照如下的命令来安装如下的库：

pip install bert_serving_server
pip install bert_serving_client

你需要确保这里安装的库和 docker-compose.yml 里的所定义的库的版本是一致的。

web/requirements.txt

bert-serving-client==1.10.0
elasticsearch==8.6.1
Flask==2.2.2

bertserving/Dockerfile

FROM tensorflow/tensorflow:1.12.0-py3
RUN pip install -U pip
RUN pip install --no-cache-dir bert-serving-server==1.10.0
COPY ./ /app
COPY ./entrypoint.sh /app
WORKDIR /app
ENTRYPOINT ["/app/entrypoint.sh"]
CMD []

example/requirements.txt

bert-serving-client==1.10.0
elasticsearch==7.0.4
pandas==0.25.1

如上所示，我们选择的 bert-serving-client 及 bert-serving-server 的版本都是 1.10.0。

在本展示中，我将使用最新的 Elastic Stack 8.6.1 来进行展示，但是在 Elasticsearch 的配置中，我不使用安全配置。

docker-compose.yml

version: '3.9'
services:
  web:
    build: ./web
    ports:
      - "5100:5100"
    environment:
      - INDEX_NAME
    depends_on:
      - elasticsearch
      - bertserving
    networks:
      - elastic   
 
  elasticsearch:
    container_name: elasticsearch
    image: elasticsearch:8.6.1
    environment:
      - discovery.type=single-node
      - ES_JAVA_OPTS=-Xms1g -Xmx1g
      - xpack.security.enabled=false
    volumes:
      - es_data:/usr/share/elasticsearch/data
    ports:
       - target: 9200
         published: 9200
    networks:
      - elastic   
 
  kibana:
    container_name: kibana
    image: kibana:8.6.1
    ports:
      - target: 5601
        published: 5601
    depends_on:
      - elasticsearch
    networks:
      - elastic   
  bertserving:
    container_name: bertserving
    build: ./bertserving
    ports:
      - "5555:5555"
      - "5556:5556"
    environment:
      - PATH_MODEL=$PATH_MODEL
    volumes:
      - "$PATH_MODEL:/model"
    networks:
      - elastic
volumes:
  es_data:
    driver: local

networks:
  elastic:
    name: elastic
    driver: bridge   

下载 pre-trained BERT 模型

首先，下载预训练的 BERT 模型。 以下命令是下载英文模型的示例：

wget https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip
unzip cased_L-12_H-768_A-12.zip

$ pwd
/Users/liuxg/python/bertsearch
$ wget https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip
--2023-02-27 09:40:16--  https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip
Resolving storage.googleapis.com (storage.googleapis.com)... 142.251.36.16
Connecting to storage.googleapis.com (storage.googleapis.com)|142.251.36.16|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 404261442 (386M) [application/zip]
Saving to: ‘cased_L-12_H-768_A-12.zip’

cased_L-12_H-768_A- 100%[===================>] 385.53M  16.0MB/s    in 24s     

2023-02-27 09:40:41 (16.0 MB/s) - ‘cased_L-12_H-768_A-12.zip’ saved [404261442/404261442]

$ unzip cased_L-12_H-768_A-12.zip
Archive:  cased_L-12_H-768_A-12.zip
   creating: cased_L-12_H-768_A-12/
  inflating: cased_L-12_H-768_A-12/bert_model.ckpt.meta  
  inflating: cased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001  
  inflating: cased_L-12_H-768_A-12/vocab.txt  
  inflating: cased_L-12_H-768_A-12/bert_model.ckpt.index  
  inflating: cased_L-12_H-768_A-12/bert_config.json  
$ ls
LICENSE                   cased_L-12_H-768_A-12     docs
README.md                 cased_L-12_H-768_A-12.zip example
bertserving               docker-compose.yaml       web

从上面，我们可以看出来在当前的目录里创建了一个叫做 cased_L-12_H-768_A-12 的子目录。它将在下面的 bertserving 容器中被使用到。

设置环境变量

你需要将预训练的 BERT 模型和 Elasticsearch 的索引名称设置为环境变量。 这些变量在 Docker 容器中使用。 下面是一个将 jobsearch 指定为索引名称并将 ./cased_L-12_H-768_A-12 指定为模型路径的示例： 

export PATH_MODEL=./cased_L-12_H-768_A-12
export INDEX_NAME=jobsearch

启动 docker

现在，让我们使用 Docker compose 启动 Docker 容器。 这里要启动三个容器：应用程序容器、BERT 容器和 Elasticsearch 容器。我们按照如下的命令来启动所有的容器：

docker-compose up

一旦成功启动完毕后，我们可以用在 http://localhost:9200 来访问 Elasticsearch，并在地址 http://localhost:5601 访问 Kibana。

$ curl http://localhost:9200

  "name" : "d996f0e69e91",
  "cluster_name" : "docker-cluster",
  "cluster_uuid" : "KiXF66HWSw2RSEXTiHKP2Q",
  "version" : 
    "number" : "8.6.1",
    "build_flavor" : "default",
    "build_type" : "docker",
    "build_hash" : "180c9830da956993e59e2cd70eb32b5e383ea42c",
    "build_date" : "2023-01-24T21:35:11.506992272Z",
    "build_snapshot" : false,
    "lucene_version" : "9.4.2",
    "minimum_wire_compatibility_version" : "7.17.0",
    "minimum_index_compatibility_version" : "7.0.0"
  ,
  "tagline" : "You Know, for Search"

我们可以通过如下的命令来查看所有的真正运行的容器：

docker ps

$ docker ps
CONTAINER ID   IMAGE                    COMMAND                  CREATED          STATUS          PORTS                                                  NAMES
a043a2b1cabc   bertsearch_web           "python app.py"          3 minutes ago    Up 3 minutes    0.0.0.0:5100->5100/tcp                                 bertsearch_web_1
2c1e20206eff   bertsearch_bertserving   "/app/entrypoint.sh"     3 minutes ago    Up 3 minutes    6006/tcp, 0.0.0.0:5555-5556->5555-5556/tcp, 8888/tcp   bertserving
7eb3a9422c50   kibana:8.6.1             "/bin/tini -- /usr/l…"   16 minutes ago   Up 16 minutes   0.0.0.0:5601->5601/tcp                                 kibana
d996f0e69e91   elasticsearch:8.6.1      "/bin/tini -- /usr/l…"   16 minutes ago   Up 16 minutes   0.0.0.0:9200->9200/tcp, 9300/tcp                       elasticsearch

从上面，我们可以看出来有四个正在运行的容器。请注意，我建议你为 Docker 分配更多内存（超过 8GB）。 因为 BERT 容器需要大内存。

我们的 bertserving 服务运行于 http://localhost:5555，而 web 服务运行于 http://localhost:5100。我们可以在 docker-compose.yml 里进行查看。

创建 Elasticsearch 索引

你可以使用创建索引 API 将新索引添加到 Elasticsearch 集群。 创建索引时，你可以指定以下内容：

• 索引的设置
• 索引中字段的映射
• 索引别名

例如，如果要创建包含 title、text 和 text_vector 字段的 jobsearch 索引，可以通过以下命令创建索引：

python example/create_index.py --index_file=example/index.json --index_name=jobsearch

上面书命令在 Elasticsearch 中创建了具有如下配置的一个叫做 jobsearch 的索引：

  "settings": 
    "number_of_shards": 2,
    "number_of_replicas": 1
  ,
  "mappings": 
    "dynamic": "true",
    "_source": 
      "enabled": "true"
    ,
    "properties": 
      "title": 
        "type": "text"
      ,
      "text": 
        "type": "text"
      ,
      "text_vector": 
        "type": "dense_vector",
        "dims": 768
      
    
  

我们可以通过如下的命令来查看：

GET jobsearch

注意：text_vector 的 dims 值必须与预训练的 BERT 模型的 dims 相匹配。

创建文档

创建索引后，你就可以为一些文档编制索引了。 这里的重点是使用 BERT 将你的文档转换为向量。 结果向量存储在 text_vector 字段中。 让我们将你的数据转换为 JSON 文档。在本示例中，我们是用了一个简单的 example.csv 文件：

example/example.csv

"Title","Description"
"Saleswoman","a woman whose job is to sell a product or service in a given territory, in a store, or by telephone"
"Software Developer","Hire Expert Software Engineers and Developers With Crowdbotics"
"Chief Financial Officer","a senior executive responsible for managing the financial actions of a company. "
"General Manager","esponsible for improving efficiency and increasing departmental profits while managing the company’s overall operations."
"Network Administrator","installing, monitoring, troubleshooting, and upgrading network infrastructure, including both hardware and software components"

如上所示，我们的 csv 文件中，含有两个字段：Title 及 Description。我们将把 Description 这个部分矢量化，以方便我们下面的搜索。

python example/create_documents.py --data=example/example.csv --index_name=jobsearch

完成脚本后，你可以得到如下的 JSON 文档：

$ pwd
/Users/liuxg/python/bertsearch
$ ls
LICENSE                   cased_L-12_H-768_A-12     docs
README.md                 cased_L-12_H-768_A-12.zip example
bertserving               docker-compose.yaml       web
$ python example/create_documents.py --data=example/example.csv --index_name=jobsearch
$ ls
LICENSE                   cased_L-12_H-768_A-12.zip example
README.md                 docker-compose.yaml       web
bertserving               docs
cased_L-12_H-768_A-12     documents.jsonl

从上面的输出中，我们可以看出来，运行命令后，当前目录下多了一个文件  documents.jsonl。它的文件格式如下：

上述格式显然是易于我们使用 bulk 命令来进行批写入的格式。 

写入文档到 Elasticsearch

将数据转换为 JSON 后，你可以将 JSON 文档添加到指定索引并使其可搜索。

python example/index_documents.py

执行完上面的命令后，我们可以在 Kibana 中进行查看：

GET jobsearch/_search

打开浏览器

转到 http://localhost:5100。 下面是一些搜索的例子。

从上面的搜索结果中，我们可以看出来，尽管我们输入的词并不完全匹配文字中的描述，但是它还是给了我们最想要的最为相关的结果。这些结果是按照相关性进行排列显示的。

上面的搜索，其实在 web 里是使用了如下的搜索命令：

GET jobsearch/_search

  "_source": ["text", "title"], 
  "query": 
    "script_score": 
      "script": 
        "source": "cosineSimilarity(params.query_vector, 'text_vector')",
        "params": 
          "query_vector": [
            0.480839341878891,
            -0.3990676701068878,
            -0.1494527906179428,
            -0.6091867685317993,
            -0.014144758693873882,
            -0.053846489638090134,
            0.727445125579834,
            -0.009675377979874611,
            -0.29119399189949036,
            0.14104360342025757,
            0.2982104420661926,
            0.5848511457443237,
           ...
          ]
        
      
  

特别值得指出的是：在最新的 Elasticsearch 发布版中，我们可以使用 knn 搜索。具体例子可以参考文章 “Elasticsearch：运用 Python 实现在 Elasticsearch 上的向量搜索”。你可以尝试修改 web 里的搜索部分来完成这个练习。这里就不再展示了。