本記事では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 ファイルが得られます。
{"_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のようなモデルを検索システムに使うときに本記事の内容が役に立てば幸いです。
あわせて読む
hironsan.hatenablog.com
参考文献
