michimani.net

近傍探索を使ってブログ内検索してみる

2023-02-24

最近何かと話題の ChatGPT を運営している OpenAI が公開している API と Python の近傍探索ライブラリを用いてブログ内検索を実現してみます。

具体的には、 OpenAI が提供している Embeddings API と、 Facebook Research が開発している Python の近似最近傍探索ライブラリ faiss を使って、任意のキーワードで当ブログ内の記事を検索できるようにしてみます。

概要

ブログ内検索の機能を実現するにあたってはいくつか方法があると思います。

今回は、ブログ記事と検索クエリをベクトル化して、それぞれの類似度をもとに検索を行う 近似最近傍探索 によって検索してみます。

ブログ記事および検索クエリのベクトル化には、 OpenAI が提供している Embeddings API を利用し、近似最近傍探索には Facebook Research が開発している faiss を利用します。

環境構築

faiss には GPU 対応しているものもあり、使い方について検索していると色々な環境で使っている例が出てきますが、今回は Intel Mac (MacBook Pro 2018) のローカル環境で CPU 版を使用します。

❯ system_profiler SPHardwareDataType
Hardware:

    Hardware Overview:

      Model Name: MacBook Pro
      Model Identifier: MacBookPro15,2
      Processor Name: Quad-Core Intel Core i7
      Processor Speed: 2.7 GHz
      Number of Processors: 1
      Total Number of Cores: 4
      L2 Cache (per Core): 256 KB
      L3 Cache: 8 MB
      Hyper-Threading Technology: Enabled
      Memory: 16 GB
      
❯ system_profiler SPSoftwareDataType
Software:

    System Software Overview:

      System Version: macOS 13.2 (22D49)
      Kernel Version: Darwin 22.3.0

Python の仮想環境については Anaconda で作成し、 conda コマンドで faiss をインストールします。

conda install -c pytorch faiss-cpu

今回紹介するコードについては下記に置いてあるので、 clone して

conda env create -f=nns-bs.yaml

とすれば必要な環境は作成されるはずです。

michimani/nns-blog-search: Example implementation of search in a blog using nearest neighbor search.

やってみる

手順としては、まずはブログ記事をベクトル化して近傍探索用のインデックスを作成します。そして、検索クエリをベクトル化してインデックスに対して検索を実行します。

ブログ記事をベクトル化してインデックスを作成

前提として、ブログ記事の一覧については Hugo の機能を使って下記のような index.json として生成しておきます。(contents には実際には記事の本文すべてが含まれています)

[
  {
    "contents": "Lambda 関数に対して URL を発行して外部から HTTP で実行できるようにする Function URLs ですが...",
    "permalink": "https://michimani.net/post/aws-lambda-function-urls-with-custom-domain/",
    "title": "Lambda Function URLs でカスタムドメインを使う構成を AWS CDK v2 (Go) で構築する"
  },
  {
    "contents": "何も書かないのはあれなので、雑に 2022 年の振り返りをしておきます。\n去年の振り返りの最後に書いていたこと...",
    "permalink": "https://michimani.net/post/other-retrospect-in-2022/",
    "title": "2022 年を雑に振り返る"
  },
  {
    "contents": "AWS Lambda の Runtime API について調べていたら Extension API、 Telemetry API というものの存在を知り...",
    "permalink": "https://michimani.net/post/aws-lambda-extension-using-telemetry-api/",
    "title": "Telemetry API を使う Lambda Extension を作ってみた"
  },
]

インデックスの作成処理は下記のような手順になります。

  1. faiss のインデックスを初期化
  2. index.json を読み込んで contents の配列を作成
  3. 2 で生成した string の配列をもとに Embeddings API にリクエストしてベクトル化
  4. 3 で得られたレスポンスをもとに 1 で作ったインデックスに追加
  5. インデックスをファイルに出力して永続化

シーケンスにすると下記のような感じです。

sequenceDiagram actor u as User participant m as main participant bm as Blog Contents (on memory) participant b as Blog Contents (index.json) participant im as Faiss Index (on memory) participant if as Faiss Index (as file) participant a as OpenAI Embeddings API alt: indexing m ->> b: json.load b ->> bm: [
{"contents":"...","title":"...","permalink":"..."},
{"contents":"...","title":"...","permalink":"..."}
] m ->> bm: create sentences bm ->> m: ["...", "...", "..."] m ->> im: init m ->> a: sentences a ->> m: embeddings m ->> im: add embeddings im ->> if: save as a file end alt: search u ->> m: query: 'ポエム' m ->> if: read index if ->> im: _ m ->> b: json.load b ->> bm: blog_index = [
{"contents":"...","title":"...","permalink":"..."},
{"contents":"...","title":"...","permalink":"..."}
] m ->> a: 'ポエム' a ->> m: embeddings m ->> im: search im ->> m: [[3,27,81]] m ->> u: show blog_index[3], blog_index[27], blog_index[81] end

1. faiss のインデックスを初期化

OpenAI の Embeddings API では使用する言語モデルを指定できるのですが、推奨されているのは `` というモデルです。このモデルが生成するベクトルの次元は 1536 なので、 faiss のインデックスも次元数 1536 で初期化しておきます。

import faiss

DIMENSION = 1536
NNS_INDEX_FILE = 'data/nns_index.faiss'

def init_nns_index():
    return faiss.IndexFlatL2(DIMENSION)

2. index.json を読み込んで contents の配列を作成

Embeddings API へのインプットとしては、 string または string の配列を渡すことができます。今回は string の配列を渡すので、 index.json を読み込んで contents を string の配列に詰めていきます。その際、配列の各要素の トークンの長さ (≠文字数) の上限は 8191 なので、それを超えないように string の長さを調整します。トークンの長さについては tiktoken ライブラリを使って調べることができるので、今回は各 contents の値についてトークンの長さをチェックして、超えるようであれば単純に contents の先頭 2,500 文字を切り出して使うようにします。 (マルチバイト文字が一文字のトークン長がだいたい 3 くらいになりそうだったので)

import json
import tiktoken

enc = tiktoken.get_encoding('cl100k_base')

def str_to_tokens(s: str):
    t = enc.encode(s)
    return t, len(t)


BLOG_INDEX_FILE = 'data/index.json'

def load_blog_indexes():
    with open(BLOG_INDEX_FILE) as f:
        indexes = json.load(f)
        return indexes


TOKEN_LIMIT = 8191

if __name__ == '__main__':
    sentences = []
    for bi in blog_indexes:
        sentence = bi['contents']
        _, count = str_to_tokens(sentence)
        if count > TOKEN_LIMIT:
            sentence = sentence[:2500]
        sentences.append(sentence)

3. 2 で生成した string の配列をもとに Embeddings API にリクエストしてベクトル化

OpenAI API へのリクエストについては openai ライブラリを使います。API の利用にあたっては API Key と Organization ID が必要になります。今回はそれぞれ環境変数にセットする形で使います。

import openai
import os

EMBEDDING_MODEL = 'text-embedding-ada-002'

org_id = os.getenv("OPENAI_ORGANIZATION_ID")
api_key = os.getenv("OPENAI_API_KEY")

if org_id is None or len(org_id) == 0:
    print("OPENAI_ORGANIZATION_ID is empty")
    exit

if api_key is None or len(api_key) == 0:
    print("OPENAI_API_KEY is empty")
    exit

openai.organization = org_id
openai.api_key = api_key

embeddings = openai.Embedding.create(input=sentences, model=EMBEDDING_MODEL)['data']

4. 3 で得られたレスポンスをもとに 1 で作ったインデックスに追加

Embeddings API のレスポンスを faiss のインデックスに追加します。このとき、レスポンスに含まれる各ベクトルを faiss に追加するために numpy の array に変換します。

nns_index = init_nns_index()

embeddings = np.array([x["embedding"] for x in embeddings], dtype=np.float32)

nns_index.add(embeddings)

5. インデックスをファイルに出力して永続化

インデックス作成にはそこそこコストがかかるので、作成したインデックスを再利用できるようにファイルに書き出しておきます。

NNS_INDEX_FILE = 'data/nns_index.faiss'

faiss.write_index(nns_index, NNS_INDEX_FILE)

ファイルに書き出して置けば、次回以降は faiss.read_index(NNS_INDEX_FILE) でメモリ上に再構築できます。

検索クエリをベクトル化して検索

faiss のインデックスに対して検索を行うために、検索クエリをベクトル化します。ベクトル化には Embeddings API を使います。

_, count = str_to_tokens(query)
if count > TOKEN_LIMIT:
    print('over token limit. token_count:{} query_len:{}'.format(count, len(query)))
    exit

embeddings = create_embeddings(openai_client, query)

query_embedding = np.array([embeddings[0]["embedding"]], dtype=np.float32)

faiss への検索には、検索クエリのベクトルと類似度が高い順に何件取得するか (= k) をパラメータとして渡します。結果としてはインデックスの配列が返ってきますが、このインデックスは index.json で保持している各記事のオブジェクトの配列のインデックスと一致します。なので、 faiss への検索結果をもとに index.json 内の情報に対してインデックス指定して実際の値 (タイトル、URL) を取得します。

blog_indexes = load_blog_indexes()

_, idxs = nns_index.search(query_embedding, k)
for idx in idxs[0]:
    blog_content = blog_indexes[idx]
    print('-----------\nTitle: {}\nURL: {}\n'.format(
        blog_content['title'], blog_content['permalink']))

検索例

以下、適当な文字列で検索してみた結果です。

S3 と Lambda をいい感じに使いたい

-----------
Title: Amazon S3 の仕様とユースケースについてあらためて調べてみた
URL: https://michimani.net/post/aws-about-amazon-s3/

-----------
Title: Gatsby で静的サイトを作ってみた - 自動デプロイ編 -
URL: https://michimani.net/post/deploying-gatsby/

-----------
Title: Lambda の同時起動を S3 に置いたロックファイルで制御する
URL: https://michimani.net/post/aws-lambda-single-process/

エンジニアリングに関係ない話

-----------
Title: この三年半を雑に振り返ってみる
URL: https://michimani.net/post/other-retrospect-in-2017-to-2020/

-----------
Title: [レポート] Developers.IO 2019 Tokyo に行ってきました
URL: https://michimani.net/post/event-developersio-1029-tokyo/

-----------
Title: 一般的なエンジニアとしてそれっぽいテレワーク環境を整えてみた
URL: https://michimani.net/post/gadget-for-remote-work-tele-work/

明日使える豆知識

-----------
Title: PHPerKaigi 2019 に行ってきました − 2日目 −
URL: https://michimani.net/post/event-phperkaigi-2019-day2/

-----------
Title: AWS のサービス条件がつい最近更新されていたので Beta 版とプレビュー版の扱いについて確認してみた
URL: https://michimani.net/post/aws-service-terms-update-at-20211203/

-----------
Title: スマホのメイン回線を楽天モバイルから LINEモバイルに変更して 10 ヶ月くらいたったので使用感をざっくり比較してみます
URL: https://michimani.net/post/gadget-line-mobile-review/

なんとなく近い内容の記事が検索できているような気がします。

あとがき

近傍探索を使ってブログ内検索してみた話でした。
ChatGPT が注目されていて、なにか投げれば良い感じに返してくれる というイメージが先行していますが (私もそんなイメージでしたが)、その動きに使われている Completions API だけでなく、今回使った Embeddings API やその他関連する API 、関連技術を良い感じに組み合わせればチャット以外にも使い所は色々ありそうだなと思いました。 (小並)

参考

Tags for this post: Python OpenAI
comments powered by Disqus

Other posts tagged by "Machine Learning"

other categories

Other posts tagged by "Python"

more ...

Other posts tagged by "OpenAI"

other tags
Copyright © 2016, michimani; all rights reserved.
Power by Hugo / Theme simplog by michimani