どんな内容？

前回、この記事にて、記事推薦モデルをBERT+AllenNLPを用いて作成しました。

今回、この記事では、faissとFastAPIを更に用いることで、任意の日本語検索クエリに対して、記事を推薦するモデルを作っていきます。

今回のコードは以下にすべて用意しています。

構成について

非常にシンプルな構成になっています。

今回は、検索クエリ（例えば”iPhone 最新”等）をリクエストボディとして受け取り、それをFastAPI経由で予測モデルに渡すようにします。予測結果がjsonとして返ってきます。

実装について

前回の記事で、学習モデルを用いて、任意の文字列（タイトル、検索クエリ）に対して埋め込みを吐き出す部分の実装を行いました。

今回は、この実装と、近似近傍探索ライブラリfaissを用いることで、
①[準備]探索対象となる全記事の埋め込み変換
②記事埋め込みのfaissへの格納
③実際の検索

この順に実装していきたいと思います。

faissとは？

近似近傍探索ライブラリの一種です。C++ベースの実装ではあるものの、pythonのインターフェースも用意しており、pythonからも大変使いやすくなっています。

近似近傍探索ライブラリには、annoyなどもあるようです。

[準備]探索対象となる全記事の埋め込み変換

前回用意したデータセットのイテレータを用いて、train及びdevデータの各記事タイトルを埋め込みへと変換します。

emb_dumper.py

import pickle
from allennlp.data import (
    DataLoader,
    DatasetReader,
    Instance,
    Vocabulary,
    TextFieldTensors,
)
from allennlp.predictors import Predictor
from allennlp.common.util import JsonDict
from tqdm import tqdm
import os

class EmbeddingEncoder(Predictor):
    def predict(self, sentence: str) -> JsonDict:
        # This method is implemented in the base class.
        return self.predict_json({"context": sentence})

    def _json_to_instance(self, json_dict: JsonDict) -> Instance:
        context = json_dict["context"]
        return self._dataset_reader.text_to_instance(mention_uniq_id=None,
                                                     data={'title': context})

class ArticleKB:
    def __init__(self, model, dsr, config):
        self.predictor = EmbeddingEncoder(model, dsr)
        self.dsr = dsr
        self.train_mention_ids, self.dev_mention_ids, self.mention_id2data = \
            dsr.train_mention_ids, dsr.dev_mention_ids, dsr.mention_id2data
        self.config = config

        self._dump_dir_maker()
        if not os.path.exists(self.config.dump_emb_dir+'kbemb.pkl'):
            mention_idx2emb = self._article_emb_iterator_from_train_and_dev_dataset()
            with open(self.config.dump_emb_dir+'kbemb.pkl', 'wb') as f:
                pickle.dump(mention_idx2emb, f)
            self.mention_idx2emb = mention_idx2emb
        else:
            with open(self.config.dump_emb_dir+'kbemb.pkl', 'rb') as g:
                self.mention_idx2emb = pickle.load(g)


    def _article_emb_iterator_from_train_and_dev_dataset(self):
        print('=== emb making from train and dev')

        mention_id2emb = {}

        for mention_id in tqdm(self.train_mention_ids + self.dev_mention_ids):
            its_article_title_emb = self.predictor.predict(
                self.mention_id2data[mention_id]['title']
            )['encoded_embeddings']
            mention_id2emb.update({mention_id: its_article_title_emb})

        return mention_id2emb

    def _dump_dir_maker(self):
        if not os.path.exists(self.config.dump_emb_dir):
            os.mkdir(self.config.dump_emb_dir)

Predictor クラスを継承することで、今回実装したモデルのforwardメソッドの返り値を受け取ることが可能になります。

今回実装したモデルのforward部分では


    def forward(self, context,
                mention_uniq_id: torch.Tensor = None,
                label: torch.Tensor = None):
        emb = self.mention_encoder(context)
        scores = self.linear_for_classify(emb)
        probs = softmax(scores, dim=1)
        output = {}
        if label is not None:
            loss = self.loss(scores, label)
            self.accuracy(probs, label)
            output['loss'] = loss
            output['logits'] = scores
            output['probs'] = probs
            output['mention_uniq_id'] = mention_uniq_id

        output['encoded_embeddings'] = emb
        return output

としており、ラベルが存在しないテキストを引数として受け取った場合でもencoded_embeddingsに埋め込みが格納されます。

記事埋め込みのfaissへの格納

上記のようにして各記事の埋め込みを用意した後、faissへの格納を行います。と言っても、行列に変換してfaissに渡してあげるだけで可能です。

ただし、faissへ格納する行列のインデックスと、実際の記事のインデックスとを混合しないように注意して実装します。

import faiss
import numpy as np

class ArticleTitleIndexerWithFaiss:
    def __init__(self, config, mention_idx2emb, dsr, kbemb_dim=768):
        self.config = config
        self.kbemb_dim = kbemb_dim
        self.article_num = len(mention_idx2emb)
        self.mention_idx2emb = mention_idx2emb
        self.dsr = dsr
        self.search_method_for_faiss = self.config.search_method_for_faiss
        self._indexed_faiss_loader()
        self.KBmatrix, self.kb_idx2mention_idx = self._KBmatrixloader()
        self._indexed_faiss_KBemb_adder(KBmatrix=self.KBmatrix)

    def _KBmatrixloader(self):
        KBemb = np.random.randn(self.article_num, self.kbemb_dim).astype('float32')
        kb_idx2mention_idx = {}
        for idx, (mention_idx, emb) in enumerate(self.mention_idx2emb.items()):
            KBemb[idx] = emb
            kb_idx2mention_idx.update({idx: mention_idx})

        return KBemb, kb_idx2mention_idx

    def _indexed_faiss_loader(self):
        if self.search_method_for_faiss == 'indexflatl2':  # L2
            self.indexed_faiss = faiss.IndexFlatL2(self.kbemb_dim)
        elif self.search_method_for_faiss == 'indexflatip':  #
            self.indexed_faiss = faiss.IndexFlatIP(self.kbemb_dim)
        elif self.search_method_for_faiss == 'cossim':  # innerdot * Beforehand-Normalization must be done.
            self.indexed_faiss = faiss.IndexFlatIP(self.kbemb_dim)

    def _indexed_faiss_KBemb_adder(self, KBmatrix):
        if self.search_method_for_faiss == 'cossim':
            KBemb_normalized_for_cossimonly = np.random.randn(self.article_num, self.kbemb_dim).astype('float32')
            for idx, emb in enumerate(KBmatrix):
                if np.linalg.norm(emb, ord=2, axis=0) != 0:
                    KBemb_normalized_for_cossimonly[idx] = emb / np.linalg.norm(emb, ord=2, axis=0)
            self.indexed_faiss.add(KBemb_normalized_for_cossimonly)
        else:
            self.indexed_faiss.add(KBmatrix)

    def _indexed_faiss_returner(self):
        return self.indexed_faiss

    def search_with_emb(self, emb):
        _, faiss_search_candidate_result_kb_idxs = self.indexed_faiss.search(
            np.array([emb]).astype('float32'),
            self.config.how_many_top_hits_preserved)
        top_titles = []
        for kb_idx in faiss_search_candidate_result_kb_idxs[0]:
            mention_idx = self.kb_idx2mention_idx[kb_idx]
            candidate_title = self.dsr.mention_id2data[mention_idx]['title']
            top_titles.append(candidate_title)

        return top_titles

一つ前のセクションで作成した、mention_idx2emb を、_indexed_faiss_KBemb_adderメソッドを用いて格納しています。

今回は探索手法としてL2距離をデフォルトでは用いています。コサイン距離での探索を行う場合は、格納時にノルムの正規化が必要となるので注意しましょう。

実際の検索

main.py に予測のサンプルを追記してあります。

    # load kb
    article_kb_class = ArticleTitleIndexerWithFaiss(
        config=config, mention_idx2emb=mention_idx2emb, dsr=dsr, kbemb_dim=768
    )
    top_titles = article_kb_class.search_with_emb(
        emb=emb_dumper.predictor.predict('iPhoneとパソコン')['encoded_embeddings'])

予測部分では、再びPredictor.predictメソッドを再利用しています。

"iPhoneとパソコン"の推薦記事検索結果は、トップ順に以下のようになりました。似たワードが出ていることから、うまく推薦できているように見えますね。

[
'Android版の「LINE」に、新プラットフォーム「LINE Channel」提供サービスとして「LINE占い」を先行公開！仮想通貨「LINEコイン」も導入【Androidアプリ】', 
'ウィルコム、世界的な色彩規格のリーディングカンパニーであるPANTONE社とのコラボPHS「PANTONE WX01SH」を発表！', 
'Microsoft、Windows Phone 7.x搭載端末向けの最新「Windows Phone 7.8」を発表！Windows Phone 8の機能を先取り', 
'ウィルコム初の「PANTONE」ブランドのPHS音声端末「PANTONE PHS WX01SH」を写真で紹介【レポート】', 
'Google、iOS向けWebブラウザー「Chrome」とクラウドストレージ「Google Drive」の専用アプリを提供開始！Android向け「Chrome」も正式版に'
]

FastAPIを用いたAPI化

更に、このモデルを用いて、レコメンドエンジンをAPI化することを考えます。

冒頭の構成図をもう一度再掲します。

uvicorn, FastAPI経由で検索クエリをモデルに渡し、推薦記事のタイトルがjsonで返ってくるAPIを作製します。

FastAPI上での予測モデル呼び出し

import uvicorn
from fastapi import FastAPI, Query, Path, Body
from pydantic import BaseModel, Field
from main import main

app = FastAPI()
article_kb_class, emb_dumper = main()

class Article(BaseModel):
    title: str

@app.get("/hello")
def hello():
   return {"Hello": "World!"}

@app.post("/request/")
async def predict_nearest_title(article: Article):
    top_titles = article_kb_class.search_with_emb(
        emb=emb_dumper.predictor.predict(article.title)['encoded_embeddings'])
    return top_titles

リクエストボディはBaseModelを継承することで定義が出来ます。

今回は/requestへ検索クエリをリクエストボディとして渡し、予測結果が返ってくるのを実際に見てみます。

uvicornによるFastAPIの起動

# 実験が一回でも終了していれば、1行目は省略可能
$ python3 main.py -num_epochs 10
$ uvicorn api:app --reload --host 0.0.0.0 --port 8000

FastAPIは各メソッドを自動でドキュメント化してくれるので、今回は実際に/docs/を見てみましょう。

docsの確認

起動しているサーバ下でlocalhost:8000/docs を見てみましょう。

(ローカルでない他のサーバで起動している場合、適宜事前にポートフォワーディングしておきましょう。)

実際にrequestメソッドがPOSTメソッドとして定義されていることが確認できます。

docsインターフェース経由での予測確認

実際にここでも予測とその結果を見てみましょう。先程Articleクラスで渡すtitleを定義しました。ここでもそのように渡してみましょう。

Executeを押して、結果を見てみましょう。

"食べ"に関連するワードを含む記事が推薦されていることが確認できます。キン肉マンの"肉"も反応しているのは、予測モデルの改良が必要そうですね。

"Mac iPhone" でもクエリを与えてみます。

[
  "ウィルコム初の「PANTONE」ブランドのPHS音声端末「PANTONE PHS WX01SH」を写真で紹介【レポート】",
  "Google、iOS向けWebブラウザー「Chrome」とクラウドストレージ「Google Drive」の専用アプリを提供開始！Android向け「Chrome」も正式版に",
  "渋谷駅の駅ビルで通り魔",
  "今日の運勢は？タロットで占おう！「タロットby Hangame」【Androidアプリ】",
  "5000以上のタイトルが見放題！使ってわかった楽しすぎる「dマーケット VIDEOストア Powered by BeeTV」の秘密に迫る"
]

テクノロジー関連の記事が出ていますね。

まとめ

BERTを用いて記事推薦モデルを2記事に渡って作りました。実際には予測のパフォーマンスを高速化するために、検索クエリの事前バッチ処理のスケジュール化なども考えられます。

コードは

にあります。よければクローンして遊んでみて下さい。