3行で始める文章検索 ― txtai入門

こんにちは
AIチームの戸田です

今回は文章検索をひとつのパッケージで実現できるtxtaiを紹介したいと思います。

LLMは強力ですが、手元の社内資料や議事録などの固有データに関する質問には回答することができません。そこで重要になるのがRAG（Retrieval Augmented Generation）です。 固有データを検索し、関連コンテキストを取り出してPromptとしてLLMに渡すことで、固有データに関する質問に対しても回答することができるようになります。実務での活用において、現在も多くの企業がこの仕組みを使っています。

RAGは非常に強力なアプローチなのですが、固有データを保存するベクトルDBや文章を機械的に扱えるようにするためのEmbeddingなど構成が複雑になりがちで、管理コストも高くなってしまう、という課題もあります。

txtaiはall-in-one AI frameworkを自称しており上記のようなRAGによる検索システムを構築する上でややこしくなる各モジュールをひとつのパッケージとして提供しています。

本記事ではtxtaiの簡単な使い方と類似ライブラリと比較したメリットを紹介します。

txtai

txtaiはPythonベースのフレームワークで、ベクトル検索やグラフ分析、リレーショナルデータベースを組み合わせたEmbeddings Databaseを中心に構成されています。

今回は扱いませんがyamlでワークフローとして定義をすることも可能なようで、応用の幅も広そうです。

環境構築

Hugging Face Transformers/Sentence Transformers、FastAPIで構成されており、Python 3.10以上の環境が必要です。今回はGoogle Colaboratoryで試します。

pipでインストールできます。

pip install txtai

基本的な使い方

以下に簡単なベクトルの類似度を利用した質問応答の例を示します。

import txtai

embeddings = txtai.Embeddings()
embeddings.index(["カーナビ", "ラーメン", "ダンゴムシ"])
print(embeddings.search("食べ物は？", 1))  # 第2引数の1は近傍サンプルを最大何件返すか
# [(1, 0.33019936084747314)]  # 出力: (index, cos類似度)

indexが1、つまり"ラーメン"が最も近いという予測でもっともらしいです。import文を除くと3行で文章検索が書けてしまいました。ちなみにEmbeddingsではモデルを指定することができるのですが、デフォルトだと sentence-transformers/all-MiniLM-L6-v2 が使われるようです。

作成したindexは以下の方法でローカルファイルとして読み書きできます。

embeddings.save("{folder_name}")

embeddings = Embeddings()
embeddings.load("{folder_name}")

これがtxtaiの最も基礎的な使い方ですが、ここから検索手法をカスタマイズすることもできます。

1. SQL

txtaiは自然文クエリだけでなく、SQLでのクエリも組み合わせることができます。

embeddings = txtai.Embeddings(content=True, objects=True)
embeddings.index([{
    "text": "カーナビ", "n_char": 4,
    "text": "ラーメン", "n_char": 4,
    "text": "ダンゴムシ",  "n_char": 5,
    "text": "カレーライス", "n_char": 6,
}])

query = """
SELECT text, n_char
FROM txtai
WHERE similar('食べ物は？') AND n_char >= 6
"""
print(embeddings.search(query, 2))
# [{'text': 'カレーライス', 'n_char': 6}]

indexにはテキストに加え、文字数(n_char)も一緒に登録します。基礎的な使い方の部分と同じ質問を投げかけていますが6文字以上でフィルタリングしているので、"ラーメン"はヒットせず、indexに追加した"カレーライス"のみがヒットします。

2. 表層検索

ベクトル検索は優秀ですが、品番など単語の表層が重要になるケースもあります。txtaiはBM25による表層検索もサポートしており、こういったケースにも対応できます。

例えば以下の例を見てみましょう。

embeddings = txtai.Embeddings()
embeddings.index(["iphone 13", "iphone 14", "iphone 15"])
print(embeddings.search("15日に買ったiphone 13", 1))
# [(2, 0.4083036184310913)]

この場合"iphone 13"をヒットさせたいのですが、デフォルトのベクトル検索だと15日の方に引っ張られてしまい、indexが2、つまり"iphone 15"がヒットしてしまっています。

ここで、keyword=Trueで表層検索をすると、求める挙動通りに"iphone 13"をヒットさせることが確認できました。・

embeddings = txtai.Embeddings(keyword=True)
embeddings.index(["iphone 13", "iphone 14", "iphone 15"])
print(embeddings.search("15日に買ったiphone 13", 1))
# [(0, 0.7487186789512634)]

ちなみに今回は試しませんでしたが、txtai.Embeddingsでhybrid=Trueを指定すると、ベクトル検索と表層検索のハイブリッド検索も行うことができるようです。

RAG

txtaiは実際にRAGを構築することもできます。以下に例を示します。Generate部分はTinySwallow-1.5B-Instructを使用しました。GPTのようなAPI経由で一般によく使われるモデルも使用できるようなのですが、HuggingFaceに公開されているOpen Weightなモデルと簡単に連携できることがアピールされていたことが理由です。

# 架空の企業の就業規則
data = [
  "勤務時間は原則9:30-18:30（休憩60分）で、フレックスはコアタイム11:00-15:00。",
  "月の所定労働時間は営業日数×8時間。半休は午前（9:30-13:30）と午後（14:30-18:30）。",
  "リモート勤務は週3日まで、申請は前日18:00までにWorkPortalで行う。",
  "年次有給休暇は入社日に10日付与し、以後毎年4月1日に11日、12日、…と1日ずつ加算され最大20日。",
  "病気休暇は年度あたり5日まで（有給とは別枠）。遅刻は30分単位で控除される。",
  "残業申請は当日中にSlackの#overtimeへ理由と見込み時間を投稿し、翌営業日までに精算する。"
]

# IDをつける。なんでも良いが今回はuuid5
embeddings = txtai.Embeddings(content=True, autoid="uuid5")
embeddings.index(data)

# RAG用のPrompt Template
template = """
  Contextに応じて質問に答えてください。

  Question:
  {question}

  Context:
  {context}
"""

# Create and run RAG instance
rag = txtai.RAG(
    embeddings,
    "SakanaAI/TinySwallow-1.5B-Instruct",
    template=template,
output="reference")

# RAGの実行
resp = rag("フレックスのコアタイムは何時から何時まで？")
print(resp)
# {'answer': 'フレックスのコアタイムは **11:00 - 15:00** です。', 'reference': '20316a1f-6aa0-5680-97ba-b5e2e7b8461b'}

# 参照先の確認
i = resp["reference"]
ref = embeddings.search(f"select id, text from txtai where id = '{i}'")
print(ref)
# [{'id': '20316a1f-6aa0-5680-97ba-b5e2e7b8461b', 'text': '勤務時間は原則9:30-18:30（休憩60分）で、フレックスはコアタイム11:00-15:00。'}]

質問に対して正しいリファレンスをもとに答えられたことがわかるかと思います。

他のライブラリとの比較

文章検索周りでよく聞くライブラリと比較したメリットを紹介したいと思います。

LangChain

LLM周りでは一番よく聞くライブラリなのではないでしょうか？正直、今回試したようなことはLangChainでも実現可能なのですが、コード量が少なくシンプルに書けるという点に関してはtxtaiが優れていると思います。一方でHyDEのような複雑なRAGを実装しようとするならばLangChainの方が柔軟性が高くて良いかもしれません。類似するライブラリとしてLlamaIndexなどが挙げられます

Pinecone

いわゆるベクトルDBです。txtaiはpip install txtai だけで動き、データはローカルファイル（SQLite等）として管理されるため、インフラ知識がなくても動かせることがメリットだと思います。類似ライブラリだとQdrantなどがあります。

FAISS

Metaが開発したベクトル検索ライブラリです。最近はあまり使っていませんが、以前Chat Bot開発をしている頃はお世話になりました。非常に高速なのですが、前処理などを自分で書く必要があり、そこをよしなに済ませてくれるtxtaiはサクッと試したい時に便利だと思いました。類似するものだとElasticsearchが代表的なものだと思います。

おわりに

本記事では文章検索を簡単に実装できるtxtaiを紹介しました。

非常に少ないコード量で書くことができ、軽いPoCなどでは活躍するのではないかと考えています。

今回は試しませんでしたが、API化を簡単にできたり、マルチモーダル対応があったりと、まだ興味深い機能があるので、また機会をみて試した上で紹介したいと思います。

最後までお読みいただき、ありがとうございました！