はじめに

nlp4j-local-search-embedding という Python パッケージを PyPI に公開しました。

このパッケージは、テキストを Embedding に変換し、ローカル環境でベクトル検索を行うためのシンプルなライブラリです。

pip install nlp4j-local-search-embedding==0.1.0

GitHub:

https://github.com/oyahiroki/nlp4j-local-search-embedding

PyPI:

https://pypi.org/project/nlp4j-local-search-embedding/0.1.0/

何ができるか

通常のキーワード検索では、検索語と文書内の語が一致しているかが重要になります。

一方、Embedding を使った検索では、検索語と文書の意味的な近さを使って検索できます。

たとえば、次のような検索ができます。

検索クエリ: 自転車

文書:
- 自動車が走っています。
- 自転車を買います
- サイクリングが趣味です
- Go shopping by bicycle.

この場合、単なる文字列一致だけではなく、

自転車
サイクリング
bicycle

のような意味的に近い文書を検索できます。

インストール

Google Colab やローカル環境で、以下のようにインストールできます。

pip install nlp4j-local-search-embedding==0.1.0

Google Colab の場合は、セルで次のように実行します。

!pip install -q nlp4j-local-search-embedding==0.1.0

内部では Java/Lucene を利用しているため、Java が利用できることも確認できます。

!java -version

Google Colab では、以下のように OpenJDK が利用できる環境で動作確認できました。

openjdk version "17.0.19"
OpenJDK Runtime Environment
OpenJDK 64-Bit Server VM

最小サンプル

以下は、Embedding 検索を行う最小サンプルです。

from nlp4j_local_search_embedding import SemanticSearch

app = SemanticSearch("ja")

app.add({
    "doc1": "自動車が走っています。",
    "doc2": "自転車を買います",
    "doc3": "サイクリングが趣味です",
    "doc4": "Go shopping by bicycle.",
})

app.commit()

results = app.search("自転車", limit=10)

print("=== Search results ===")
print(f"number of results: {len(results)}")

for i, result in enumerate(results):
    print(f"result[{i}].id: {result.id}")
    print(f"result[{i}].text: {result.text}")
    print(f"result[{i}].score: {result.score}")
    print(f"result[{i}].metadata: {result.metadata}")
    print("---")

実行結果

実行すると、以下のような結果になります。

=== Search results ===
number of results: 4
result[0].id: doc2
result[0].text: 自転車を買います
result[0].score: 0.9391793012619019
result[0].metadata: {}
---
result[1].id: doc3
result[1].text: サイクリングが趣味です
result[1].score: 0.9307862520217896
result[1].metadata: {}
---
result[2].id: doc4
result[2].text: Go shopping by bicycle.
result[2].score: 0.9053230285644531
result[2].metadata: {}
---
result[3].id: doc1
result[3].text: 自動車が走っています。
result[3].score: 0.8970762491226196
result[3].metadata: {}
---

検索クエリは日本語の 自転車 ですが、英語の bicycle を含む文書も検索されています。

これは、多言語 Embedding モデルを利用しているためです。

英語文書の例

英語の文書でも同じように検索できます。

from nlp4j_local_search_embedding import SemanticSearch

app = SemanticSearch("en")

app.add({
    "doc1": "Kyoto is a historic city in Japan.",
    "doc2": "Tokyo is the capital city of Japan.",
    "doc3": "Python is a popular programming language.",
    "doc4": "Nintendo is a video game company headquartered in Kyoto.",
})

app.commit()

results = app.search("an old Japanese capital", limit=10)

print("=== Search results ===")
print(f"number of results: {len(results)}")

for i, result in enumerate(results):
    print(f"result[{i}].id: {result.id}")
    print(f"result[{i}].text: {result.text}")
    print(f"result[{i}].score: {result.score}")
    print(f"result[{i}].metadata: {result.metadata}")
    print("---")

この例では、an old Japanese capital という検索クエリに対して、Kyoto is a historic city in Japan. が上位に来ることが期待できます。

メタデータ付き文書の追加

文書にはメタデータも付与できます。

from nlp4j_local_search_embedding import SemanticSearch

documents = [
    {
        "id": "doc1",
        "text": "Kyoto is a historic city in Japan with many temples and shrines.",
        "metadata": {
            "category": "city",
            "country": "Japan",
        },
    },
    {
        "id": "doc2",
        "text": "Nintendo is a video game company headquartered in Kyoto.",
        "metadata": {
            "category": "company",
            "country": "Japan",
        },
    },
    {
        "id": "doc3",
        "text": "Python is widely used for data science and machine learning.",
        "metadata": {
            "category": "technology",
        },
    },
]

app = SemanticSearch("en")
app.add(documents)
app.commit()

results = app.search("a Japanese game company", limit=3)

for result in results:
    print(result.id)
    print(result.text)
    print(result.score)
    print(result.metadata)
    print("---")

検索結果には、文書ID、本文、スコア、メタデータが含まれます。

設計の考え方

nlp4j-local-search-embedding は、単体ですべてを行う巨大なライブラリではなく、nlp4j-local-search の上に載る Embedding 検索用の薄いレイヤーとして作っています。

構成としては、以下のようなイメージです。

text documents
    |
    v
Embedding model
    |
    v
vectors
    |
    v
nlp4j-local-search vector index
    |
    v
semantic search results

ベースとなる nlp4j-local-search は、ローカル環境で Lucene を使った検索を行うパッケージです。

一方、nlp4j-local-search-embedding は、テキストを Embedding に変換し、ベクトル検索を簡単に使えるようにするためのパッケージです。

キーワード検索との違い

キーワード検索では、基本的には検索語が文書に含まれているかどうかが重要です。

たとえば、自転車 で検索した場合、自転車を買います は見つかりますが、サイクリングが趣味です や Go shopping by bicycle. は単純な文字列一致では見つけにくいです。

Embedding 検索では、文書とクエリをベクトル化して意味的な近さで検索するため、表記が違っても意味が近い文書を探しやすくなります。

自転車
サイクリング
bicycle

このような関係を扱えるのが Embedding 検索の利点です。

内部で利用しているモデル

デフォルトでは、以下のモデルを利用します。

intfloat/multilingual-e5-large

E5 系のモデルでは、文書と検索クエリに異なる prefix を付ける使い方が一般的です。

passage: <document text>
query: <search query>

nlp4j-local-search-embedding では、この prefix 付与を内部で行うため、ユーザーは通常のテキストをそのまま追加・検索できます。

app.add({
    "doc1": "Kyoto is a historic city in Japan.",
})

results = app.search("old Japanese capital")

Google Colab で試す

Google Colab では、次の2セルで試すことができます。

インストール

!pip install -q nlp4j-local-search-embedding==0.1.0

実行

from nlp4j_local_search_embedding import SemanticSearch

app = SemanticSearch("ja")

app.add({
    "doc1": "自動車が走っています。",
    "doc2": "自転車を買います",
    "doc3": "サイクリングが趣味です",
    "doc4": "Go shopping by bicycle.",
})

app.commit()

results = app.search("自転車", limit=10)

print("=== Search results ===")
print(f"number of results: {len(results)}")

for i, result in enumerate(results):
    print(f"result[{i}].id: {result.id}")
    print(f"result[{i}].text: {result.text}")
    print(f"result[{i}].score: {result.score}")
    print(f"result[{i}].metadata: {result.metadata}")
    print("---")

初回実行時は、Embedding モデルのダウンロードとロードが行われるため、少し時間がかかります。

今後の予定

今後は、以下のような機能を追加していく予定です。

インデックスの保存と読み込み
より大きな文書集合での検索
RAG 用パッケージとの連携
nlp4j-local-search-rag との接続
検索結果の整形
Google Colab 用サンプルノートブック

まとめ

nlp4j-local-search-embedding を使うと、Python から簡単にローカル Embedding 検索を試すことができます。

pip install nlp4j-local-search-embedding==0.1.0

最小コードは次のようになります。

from nlp4j_local_search_embedding import SemanticSearch

app = SemanticSearch("ja")

app.add({
    "doc1": "京都は日本の歴史的な都市です。",
    "doc2": "東京は日本の首都です。",
})

app.commit()

results = app.search("古い都", limit=10)

for result in results:
    print(result.id, result.score, result.text)

キーワード検索だけでは拾いにくい、意味的に近い文書をローカル環境で検索したい場合に使いやすいライブラリを目指しています。

参考

以下、Google Colab で実際に実行した例です。

以上です