Elasticsearch 遇上 BERT:使用 Elasticsearch 和 BERT 构建搜索引擎

作者:Hironsan

编译:ronghuaiyang

导读: 强强联合,看看是否能有 1+1>2 的效果。

在这篇文章中,我们使用一个预先训练好的 BERT 模型和 Elasticsearch 来构建一个搜索引擎。Elasticsearch 最近发布了带有矢量字段的文本相似性搜索。另一方面,你可以使用 BERT 将文本转换为固定长度的向量。一旦我们通过 BERT 将文档转换成向量并存储到 Elasticsearch 中,我们就可以使用 Elasticsearch 和 BERT 搜索类似的文档。

本文使用 Elasticsearch 和 BERT 按照以下架构实现了一个搜索引擎。这里,我们使用 Docker 将整个系统划分为三个部分:application, BERT 和 Elasticsearch。这样做的目的是为了更容易地扩展每个服务。

我在这篇文章中只展示了重要的部分,整个系统都在 <span>docker-compose.yaml</span 中。在以下的 GitHub 存储库中:https://github.com/Hironsan/bertsearch

1. 下载预训练 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

2. 设置环境变量

你需要设置一个预先训练好的 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

3. 启动 Docker 容器

现在,我们使用 Docker compose 来启动 Docker 容器。这里要启动三个容器:application 容器、BERT 容器和 Elasticsearch 容器。



$ docker-compose up

注意,我建议你分配更多的内存(超过 8GB)给 Docker。因为 BERT 容器需要大内存。

4. 创建 Elasticsearch 索引

您可以使用 create index API 向 Elasticsearch 集群添加新的索引。创建索引时,你可以指定以下内容:

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

例如,如果你想创建带有“title”、“text”和“text_vector”字段的“jobsearch”索引,可以通过以下命令创建索引:

$ python example/create_index.py --index_file=example/index.json --index_name=jobsearch
# index.json
{
  "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
      }
    }
  }
}

注意:<span>text_vector</span<span>dims</span 值必须与预先训练的 BERT 模型的 dims 匹配。

5. 创建文档

一旦创建了索引,就可以为某个文档建立索引了。这里的要点是使用 BERT 将文档转换为向量。得到的向量存储在 <span>text_vector</span 字段中。让我们把你的数据转换成一个 JSON 文档:

$ python example/create_documents.py --data=example/example.csv --index_name=jobsearch
# example/example.csv
"Title","Description"
"Saleswoman","lorem ipsum"
"Software Developer","lorem ipsum"
"Chief Financial Officer","lorem ipsum"
"General Manager","lorem ipsum"
"Network Administrator","lorem ipsum"

脚本执行完成后,可以得到如下的 JSON 文档:

# documents.jsonl
{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Saleswoman", "text_vector": [...]}
{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Software Developer", "text_vector": [...]}
{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Chief Financial Officer", "text_vector": [...]}
...

6. 索引文档

将数据转换成 JSON 后,可以向指定的索引添加一个 JSON 文档,并使其可搜索。

$ python example/index_documents.py

7. 打开浏览器

打开http://127.0.0.1:5000,下面是一个搜索 job 的例子。你可以看到,“I’m looking for a lawyer” 的 query 返回的是“Legal assistant”以及 “Lawyer”。

总结

在这篇文章中,我们使用 Elasticsearch 和 BERT 实现了搜索引擎。虽然 BERT 的执行速度存在问题,但是像这样将 BERT 作为一个独立的容器来处理很容易扩展,所以我认为这个问题是可以解决的。我希望这篇文章对你有用。

英文原文:https://towardsdatascience.com/elasticsearch-meets-bert-building-search-engine-with-elasticsearch-and-bert-9e74bf5b4cf2


本文地址:https://www.6aiq.com/article/1594084088157
本文版权归作者和AIQ共有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出