作者: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