Pyserini（Faiss）を使ってお手軽Entity検索をやってみた！

2023年8月30日 15:08

こんにちは。
リサーチャーの勝又です。
私はレトリバで自然言語処理、とくに要約や文法誤り訂正に関する研究の最新動向の調査・キャッチアップなどを行っております。

今回の記事では、Pyseriniという情報検索の研究で使われるPythonライブラリの簡単な使い方、拡張方法について紹介します。

Pyseriniとは

近年、Large Language Model（LLM）の流行に伴い、Retrieval-augmented Language Modelのように、情報検索技術の需要は高まっていると思います。
たとえば、LangChainではRetrieverとしてさまざまなライブラリをサポートしていて、
これらを利用して検索結果を取得、その検索結果をLLMに加えることでより効果的な出力を行うことが期待されます。

このように情報検索技術に触れる機会が増えると、より最新の検索手法を触ってみたくなりませんか？

PyseriniではBM25のような古典的な手法や、Faissを利用したDense Vectorによる検索、
uniCOIL[1]やSPLADE[2]といったSparse Vectorによる検索を動かすことができます。

Pyseriniは以下の画像のようにAnserini（Luceneベースの情報検索ツール）やFaissのインターフェイスとなっていて、
非常に簡単に動かすことも可能ですし、拡張に関しても簡単に行うことができます。

今回の記事ではLuke[3]を利用したEntity検索をPyseriniで実現しようと思います。

Pyseriniを使ってEntityの検索を実施する

今回はPyseriniを利用してEntity検索（Entity Linking）を行ってみます。

Entity検索について

Pyseriniの準備の前に、まずは今回試すEntity検索について説明します。
ここでのEntity検索とは、ずばりEntity Linkingのことです。

Entity Linkingとはテキスト中の固有表現（mention）を、なんらかの知識データ（Knowledgeデータ; WikiDataやDBPediaなど）に対して対応付けるタスクです。
以下の図であれば、春の嵐という固有表現に対して、WikiDataのQ1515513（Gertrud）を対応付けることになります。

Entity検索では、分散表現を利用してEntity Linkingを行いたいと思います。

分散表現はstudio-ousia/luke-japanese-base-lite [link]を利用して作成します。
mentionデータはMewsli-9[4]と呼ばれる、WikiNewsに基づくデータを使用します。
KnowledgeデータはWikiDataを使用します。
ただし、今回はPyseriniの拡張がちゃんと動くことを確認できれば良いので、WikiDataすべてではなく、Mewsli-9中に出現するEntityのみ使用しています。

Pyseriniの準備について

Pyseriniは前述した通り、AnseriniやFaissなどに依存しています。
pip install pyseriniでAnseriniの方は問題ないのですが、FaissについてはCPU版とGPU版のどちらを入れるかがユーザーによって変わってくるのでこちらは別途導入する必要があります。

この辺りがちょっと大変なので、今回はDockerfileを作成しました。

FROM nvidia/cuda:11.4.3-cudnn8-devel-ubuntu20.04

RUN ln -sf /usr/share/zoneinfo/Asia/Tokyo /etc/localtime

# install conda
RUN apt-get update && apt-get install -y wget bzip2 ca-certificates \
    libglib2.0-0 libxext6 libsm6 libxrender1 git mercurial subversion

RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-py310_23.5.2-0-Linux-x86_64.sh \
    && /bin/bash Miniconda3-py310_23.5.2-0-Linux-x86_64.sh -b -p /opt/conda \
    && rm Miniconda3-py310_23.5.2-0-Linux-x86_64.sh

ENV PATH /opt/conda/bin:$PATH

# install conda update and build tools
RUN conda update -y -q conda \
    && conda install -y -q conda-build

# install python library
# install faiss-gpu
RUN conda install -y -c pytorch -c nvidia faiss-gpu=1.7.4 mkl=2021 blas=1.0=mkl \
    && conda install -c conda-forge openjdk=11

# install pytorch for cuda 11.4
# RUN conda install -y -c pytorch -c nvidia pytorch pytorch-cuda=11.4
RUN pip install torch

# install requirements
# このrequirements.txtの中身はpyseriniのみでok
COPY requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt

# set workdir
WORKDIR /workspace

以降はこのDockerコンテナ内で作業を実行します。

Pyseriniを動かす流れ

PyseriniのUsageにも書いてありますが、
Dense Vectorの場合、一連の流れは以下の通りです。

検索対象（今回はKnowledge）をEncodeする
Indexingの実施
mentionをEncodeして検索する

2番についてはPyseriniそのままで問題ないのですが、1番と3番についてはLuke用にちょっと工夫してあげる必要があります。
今回はLuke用にPyseriniの各種クラスを継承してみようと思います。

それでは、前述の流れにしたがって、Lukeを利用するための拡張について説明します。

Encodeを行う

最初に検索対象のベクトル表現を作成するためのEncoderを作成します。
作成します、といっても基本的にはhuggingface/transformersを使えばいいのでそこまで大変ではありません。

今回は以下のようなEncoderを用意しました。

from typing import Optional

import numpy as np
from sklearn.preprocessing import normalize
from transformers import AutoModel, AutoTokenizer
from pyserini.encode import DocumentEncoder


class LukeDocumentEncoder(DocumentEncoder):
    def __init__(self, model_name: str, tokenizer_name: Optional[str] = None, device: str = "cuda:0", l2_norm=False):
        self.device = device
        self.model = AutoModel.from_pretrained(model_name)
        self.model.to(self.device)
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name if tokenizer_name else model_name, use_fast=True)
        except:
            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name if tokenizer_name else model_name, use_fast=False)

        self.has_model = True
        self.l2_norm = l2_norm

    def encode(self, texts, span, max_length=256, **kwargs) -> np.ndarray:
        """Lukeを用いて入力Entityのベクトル化を行う"""
        tokenizer_kwargs = {
            "max_length": max_length,
            "truncation": True,
            "padding": "longest",
            "return_tensors": "pt",
        }

        inputs = self.tokenizer(text=texts, entity_spans=span, **tokenizer_kwargs)
        inputs.to(self.device)

        outputs = self.model(**inputs)
        entity_vector = outputs.entity_last_hidden_state.detach().cpu().numpy()
        batch_size = entity_vector.shape[0]
        entity_vector = entity_vector.reshape([batch_size, -1])
        if self.l2_norm:
            entity_vector = normalize(entity_vector, norm='l2', axis=1)
        return entity_vector

上記のようなEncoderを用意したら、実際にEncodeを行いEntityのベクトル表現を作成してみます。

from pyserini.encode import FaissRepresentationWriter

encoder = LukeDocumentEncoder("studio-ousia/luke-japanese-base-lite")
embedding_writer = FaissRepresentationWriter("/path/to/output/encoded", dimension=768)

# batch size 1 のサンプルデータ
data = {
    "start": [0],
    "end": [3],
    "contexts": ["渋谷区（しぶやく）は、東京都の区部南西部に位置する特別区。"],
    "entity-name": ["渋谷区"],
    "id": ["Q193638"],
}

with embedding_writer:
    spans = [[(start, end)] for start, end in zip(data["start"], data["end"])]
    kwargs = {
        "texts": data["contexts"],
        "span": spans,
    }
    embeddings = encoder.encode(**kwargs)
    data["vector"] = embeddings
    embedding_writer.write(data)

Pyseriniの便利な実装として、FaissRepresentationWriterというものがあり、今回はそれを使ってベクトル表現を保存しています。
上記の場合ですと、/path/to/output/encoded/docidに文書ID（上記の例だとQ193638）が保存され、/path/to/output/encoded/indexにベクトル表現が保存されます。

Indexingを実施する

Faissの特徴として、多様なIndex作成手法があると思います。
Pyseriniを通して、先ほどEncodeしたベクトル表現からIndexを作成してみます。

たとえばHNSWの場合は以下のようにしてIndexingを実施できます。

$ python -m pyserini.index.faiss \
  --input /path/to/output/encoded \
  --output /path/to/output/index \
  --hnsw

このようにして、Indexが/path/to/output/index以下に作成されます。

Pyseriniを利用したIndex作成手法についてはこちらをご確認ください。

Searcherを作成する

作成したIndexに対して検索を行うSearcherを準備します。
Encoderの時と同様にPyseriniのSearcherを継承してLukeに必要な部分だけオーバーライドします。

from typing import Tuple, List, Union

import faiss
import numpy as np
from pyserini.search import FaissSearcher, DenseSearchResult, PRFDenseSearchResult

class LukeFaissSearcher(FaissSearcher):
    """Faiss searcher for Luke.

    This code is based on the following code:
    https://github.com/castorini/pyserini/blob/b56d04a823d8fd063614524dec799ef84db0cac1/pyserini/search/faiss/_searcher.py#L379
    """
    def search(
        self, query: str, span: List[Tuple[int ,int]], k: int = 10, threads: int = 1, return_vector: bool = False
    ) -> Union[List[DenseSearchResult], Tuple[np.ndarray, List[PRFDenseSearchResult]]]:
        emb_q = self.query_encoder.encode(query, span)
        assert len(emb_q) == self.dimension
        emb_q = emb_q.reshape((1, len(emb_q)))
        faiss.omp_set_num_threads(threads)

        if return_vector:
            distances, indexes, vectors = self.index.search_and_reconstruct(emb_q, k)
            vectors = vectors[0]
            distances = distances.flat
            indexes = indexes.flat
            return emb_q, [PRFDenseSearchResult(self.docids[idx], score, vector)
                           for score, idx, vector in zip(distances, indexes, vectors) if idx != -1]
        else:
            distances, indexes = self.index.search(emb_q, k)
            distances = distances.flat
            indexes = indexes.flat
            return [DenseSearchResult(self.docids[idx], score)
                    for score, idx in zip(distances, indexes) if idx != -1]

そして、Searcherで利用するQueryEncoder（今回はmentionをベクトル化する部分）を用意すれば準備は完了です。
基本的な部分は先ほど作成したLukeDocumentEncoderで問題ないのですが、pyseriniの都合でちょっと修正する必要があります。
具体的には下記の通りです。

from typing import Optional

import numpy as np
from pyserini.search import QueryEncoder

class LukeQueryEncoder(LukeDocumentEncoder, QueryEncoder):
    def __init__(self, model_name: str, tokenizer_name: Optional[str] = None, device: str = "cuda:0", l2_norm=False):
        super().__init__(model_name, tokenizer_name, device, l2_norm)

    def encode(self, texts, span, max_length=256, **kwargs) -> np.ndarray:
        entity_vector = super().encode(texts, span, max_length, **kwargs)
        return entity_vector.flatten()

このように、Searcherで使用するEncoderはpyserini.search.QueryEncoderを継承する必要があり、改めてLukeQueryEncoderを作成しました。

Entity検索してみる

ここまで準備したものを使い、Entityの検索を行いたいと思います。

encoder = LukeQueryEncoder("studio-ousia/luke-japanese-base-lite")
searcher = LukeFaissSearcher("/path/to/output/index", encoder)

query_mention = {
    "mention": "屋久島",
    "context": "同9時30分までの1時間に90mm、降り始めからの合計が319mmとなる豪雨を記録するなど、種子島や屋久島は局地的な豪雨となった。\n\n一方気象庁は22日、東日本と西日本を中心に5月の連休明けからの日照時間が",  # 今回はmentionの周囲50文字を取り出しています
    "start": 50,
    "end": 53,
}
kwargs = {
    "query": query_mention["context"],
    "span": [[(query_mention["start"], query_mention["end"])]],
}
results = seacher.search(**kwargs)

for i in range(0, 10):  # top-10
    entity_id = results[i].docid
    print(f"{i+1:2} {entity_id} {results[i].score:.5f}")

このようにすることで、作成したIndexに対して検索を行うことができます。

実際にMewsli-9とWikiDataを利用して検索した結果（top-3）が次の通りです。

3例ともある程度ジャンルとしては似ているものが出力されていることがわかります（棋士、力士はちょっとわかりませんが。。。[5]）。
よりTopic Specificなベクトル表現を作成したい場合は、LukeをもとにContrastive Learningを施したUCTopic[6]がございますので、そちらも検証してみるとおもしろいと思います。

まとめ

今回はPyseriniを使用したEntity検索の実装について簡単に紹介しました。

ベクトル表現としてLukeを利用する都合上、EncoderやSearcherを工夫する必要がありましたが、実はPyseriniにはhuggingface/transformersのAutoModelに対応したEncoderの実装もあるので、大抵の場合は実装済みのものを使えば検索まで行うことが可能です。

Pyseriniを使うことで、Faissを比較的簡単に（文書IDの管理など含む）使うことができます。
また今回触ってはいませんが、冒頭に述べた通り、Sparse Vectorを利用した検索なども対応しており、こちらも簡単に使用することができます。

気になった方はぜひPyseriniで遊んでみてください。

注釈

A Few Brief Notes on DeepImpact, COIL, and a Conceptual Framework for Information Retrieval Techniques. Jimmy Lin, Xueguang Ma. [paper]
SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking. Thibault Formal, Benjamin Piwowarski, Stéphane Clinchant. [paper]
LUKE: Deep Contextualized Entity Representations with Entity-aware Self-attention. Ikuya Yamada, Akari Asai, Hiroyuki Shindo, Hideaki Takeda, Yuji Matsumoto. [paper]
Entity Linking in 100 Languages. Jan A. Botha, Zifei Shan, Daniel Gillick. [paper]
自分の話になりますが、以前同義語検索をベクトル検索で行おうとして際にLukeを試したのですが、そちらでもあまりうまくいきませんでした [ブログ記事] [NLP2023]。特徴量として有効なベクトル表現は、必ずしも類似度計算ではそうならない、ということかとは思うのですが。。。（NLP2023ではLukeによる類似度では細かい粒度が区別できない例が散見されました。）
UCTopic: Unsupervised Contrastive Learning for Phrase Representations and Topic Mining. Jiacheng Li, Jingbo Shang, Julian McAuley. [paper]