本記事ではElasticsearchとBERTを組み合わせて類似文書検索を行う方法について紹介します。Elasticsearchでは最近、ベクトルに対する類似文書検索の機能が実装されました。一方、BERTを使うことでテキストを固定長のベクトルに変換することができます。つまり、BERTを使ってテキストをベクトルに変換すれば、Elasticsearchを使って類似文書検索ができるということになります。
本記事では以下のアーキテクチャでElasticsearchとBERTを組み合わせた検索システムを実現します。Dockerを使ってアプリケーション、BERT、Elasticsearchのコンテナを分けることでそれぞれをスケールしやすくする狙いがあります。記事中では重要な部分のみ提示しますが、システム全体はdocker-composeのファイルとして記述しこちらのリポジトリに置いてるので、参照してください。
実装は以下の手順で行います。
- 学習済みモデルのダウンロード
- 環境変数の設定
- Dockerコンテナの起動
- インデックスの作成
- ドキュメントの作成
- ドキュメントのインデックス
- ブラウザを開く
Step 1: 学習済みモデルのダウンロード
まずは、以下のリストから学習済みのBERTのモデルをダウンロードします。
| BERT-Base, Uncased | 12-layer, 768-hidden, 12-heads, 110M parameters |
| BERT-Large, Uncased | 24-layer, 1024-hidden, 16-heads, 340M parameters |
| BERT-Base, Cased | 12-layer, 768-hidden, 12-heads , 110M parameters |
| BERT-Large, Cased | 24-layer, 1024-hidden, 16-heads, 340M parameters |
| BERT-Base, Multilingual Cased (New) | 104 languages, 12-layer, 768-hidden, 12-heads, 110M parameters |
| BERT-Base, Multilingual Cased (Old) | 102 languages, 12-layer, 768-hidden, 12-heads, 110M parameters |
| BERT-Base, Chinese | Chinese Simplified and Traditional, 12-layer, 768-hidden, 12-heads, 110M parameters |
$ 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
Step 2: 環境変数の設定
次に、ダウンロードしたBERTのモデルとElasticsearchのインデックス名を環境変数に設定します。
$ export PATH_MODEL=./cased_L-12_H-768_A-12 $ export INDEX_NAME=jobsearch
Step 3: Dockerコンテナの起動
Docker composeを使ってDockerコンテナを起動します。ここで起動するコンテナは、アプリケーションコンテナ、BERTコンテナ、Elasticsearchコンテナの3つです。詳細についてはdocker-compose.yamlを確認してください。
$ docker-compose up
注意としては、Dockerに割り当てるメモリの設定をある程度大きくしておかないと、BERTがいつまで経っても使用可能にならない点です。8GB程度は確保しておきましょう。
Step 4: インデックスの作成
コンテナを起動したら、Elasticsearchのインデックスを作成します。たとえば、title、text、text_vectorの3つのフィールドを持つ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
}
}
}
}
ここで、text_vectorフィールドのdimsの値は学習済みBERTモデルの次元数と等しくなることに注意してください。
Step 5: ドキュメントの作成
インデックスを作成したら、ドキュメントを格納します。ここでのポイントは、BERTを使ってテキストをベクトルに変換することです。BERTで得られたベクトルはtext_vectorフィールドに格納されます。以下ではCSVファイルを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": [...]} ...
Step 6: ドキュメントのインデックス
データをJSONに変換したら、指定したインデックスに格納します。
$ python example/index_documents.py
Step 7: ブラウザを開く
データをインデックスしたら、ブラウザでhttp://127.0.0.1:5000を開きます。以下の画像は求人検索を行っている例です。「I'm looking for a lawyer」というクエリに対して、「Legal assistant」や「Lawyer」の求人を返していることが確認できます。
ソースコード
おわりに
今回はElasticsearchとBERTを使って類似文書検索を行ってみました。BERTの実行速度には課題がありますが、今回のようにBERTを独立したコンテナとして扱うことでスケールさせやすくなるので、課題は解決できると思います。今後、BERTのようなモデルを検索システムに使うときに本記事の内容が役に立てば幸いです。
