今回の内容は、モデル等によってはきちんと動かないかもしれません。
ローカルLLMでは結構不安定な印象。

導入

私が学ぶRAGの実質3回目です。準備編はこちら。

今回はSelf Queryingを使ったRAGを実践します。

langchainの公式Docは以下です。

これは何？

上の公式Docより、ざっくり和訳。

自己クエリ型レトリーバとは、その名の通り、自分自身でクエリを実行する機能を持つレトリーバのことである。具体的には、任意の自然言語クエリが与えられると、レトリーバーはクエリ構築LLMチェーンを使って構造化クエリを記述し、その構造化クエリを基礎となるVectorStoreに適用する。
これにより、レトリーバーはユーザが入力したクエリを保存文書の内容との意味的類似性比較に使用するだけでなく、保存文書のメタデータに関するユーザクエリからフィルタを抽出し、そのフィルタを実行することができる。

端的に言うと、問い合わせ文言の中身から、Vectorstoreへのフィルタ処理を作り出すような仕組です。例えば、「ユーザ評価スコアが3以上の英語を抽出して」みたいな自然言語クエリをLLMを使って解釈し、user_score==3みたいなフィルタ処理を作成した上で関連文書を取り出します。

今回は上記公式ドキュメントの内容を簡易化して実施してみます。

Step0. モジュールインストール

使うモジュールをインストールします。
今回もベクトルストアにChromaを利用。また、larkパッケージが必要です。

%pip install -U transformers accelerate autoawq=="0.1.5" langchain chromadb lark
%pip install databricks-feature-engineering 

dbutils.library.restartPython()

Step1. Document Loading

準備編で保管した特徴量を取得します。
読み込んだデータは、docsという名前のビューから参照できるようにします。

from databricks.feature_engineering import FeatureEngineeringClient

fe = FeatureEngineeringClient()

feature_name = "training.llm.sample_doc_features"
df = fe.read_table(name=feature_name)

df.createOrReplaceTempView("docs")

Step2. Splitting and add attribute

ごくごく単純なText Splitterを使って長文データをチャンキングします。
また、フィルタ条件ように、仮想の重要度であるimportanceという列を追加します。
中身は1-5までのランダムな数字を割り当てます。

from typing import Any
import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql.functions import pandas_udf
from langchain.text_splitter import RecursiveCharacterTextSplitter


class JapaneseCharacterTextSplitter(RecursiveCharacterTextSplitter):
    """句読点も句切り文字に含めるようにするためのスプリッタ"""

    def __init__(self, **kwargs: Any):
        separators = ["\n\n", "\n", "。", "、", " ", ""]
        super().__init__(separators=separators, **kwargs)


@pandas_udf("array<string>")
def split_text(texts: pd.Series) -> pd.Series:

    # 適当なサイズとオーバーラップでチャンク分割する
    text_splitter = JapaneseCharacterTextSplitter(chunk_size=200, chunk_overlap=40)
    return texts.map(lambda x: text_splitter.split_text(x))


# チャンキング
df = spark.table("docs")
df = df.withColumn("chunk", F.explode(split_text("page_content")))

# importance列を追加
df = df.withColumn("importance", F.floor(F.rand() * 5) + 1)

# 必要な列のみに限定
df = df.select("chunk", "importance")

display(df)

df.createOrReplaceTempView("chunked_docs")

こんな感じのデータになります。

Step3. Storage

チャンクデータに埋め込み(Embedding)を行い、ベクトルストアへデータを保管します。

まずはEmbedding用のモデルをロード。
以下のモデルをダウンロードしたものを利用します。

import torch
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS

device = "cuda" if torch.cuda.is_available() else "cpu"
embedding_path = "/Volumes/training/llm/model_snapshots/models--intfloat--multilingual-e5-large"

embedding = HuggingFaceEmbeddings(
    model_name=embedding_path,
    model_kwargs={"device": device},
)

Chromaを使ってベクトルストアを作成します。
この際、metadataにimportance列を設定します。

from langchain.schema.document import Document
from langchain.vectorstores import Chroma
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever

from pprint import pprint

pdf = spark.table("chunked_docs").select("chunk", "importance").toPandas()
docs = [
    Document(page_content=row["chunk"], metadata={"importance": row["importance"]})
    for index, row in pdf.iterrows()
]

vectorstore = Chroma.from_documents(docs, embedding)

Step4. Model for retriever loading

Self-queryingを実行するためのLLMをロードします。
このLLMを使って、クエリからフィルタ条件を抽出/フィルタ実行可能な形式への変換を行います。
どういったLLMがこの処理に向いているのかわからないのですが、今回は以下のモデルを利用しました。

事前にダウンロードしてあるものを利用します。

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, AwqConfig
from transformers_chat import ChatHuggingFaceModel

model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--openchat_3.5-AWQ"

generator = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_path)

retriever_llm = ChatHuggingFaceModel(
    generator=generator,
    tokenizer=tokenizer,
    human_message_template="GPT4 User: {}<|end_of_turn|>GPT4 Assistant:",
    repetition_penalty=1.2,
    temperature=0.1,
    max_new_tokens=1024,
)

Step5. Creating our self-querying retriever

フィルタをかけるmetadata情報の作成、ならびにSelfQueryRetrieverを作成します。
ここがコアの部分であり、公式Docのこの部分に対応する内容です。

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever

metadata_field_info = [
    AttributeInfo(
        name="importance",
        description="文章の重要度",
        type="integer",
    ),
]
document_content_description = "委託契約書フォーマット（概算契約）"
retriever = SelfQueryRetriever.from_llm(
    retriever_llm,
    vectorstore,
    document_content_description,
    metadata_field_info,
    enable_limit=True,
)

テスト。

# サンプル1
retriever.invoke("I want to check a document importance less than 3")

出力

[Document(page_content='含む場合には、その製造工程を含む。）を行う場合には、以下を実施しなければならな\nい。 \n（１）各工程において、 当省の意図しない変更や機密 情報の窃取等が行われないことを \n保証する管理が、一貫し た品質保証体制の下でなされていること。また、具体的な\n管理手順や品質保証体制を証明する書類等を提出すること 。', metadata={'importance': 1})]

日本語もOK。(ただ、"以上"という言葉の判定に疑問。。。)

# サンプル2
retriever.invoke("重要度が3以上の契約文書を10件取得したい")

出力

[Document(page_content='ェブサイト等のサーバへ自動的にアクセスが 発生する機能が仕様に反して組み\n込まれていないことを、ＨＴＭＬソ ースを表示させるなどして確認すること。   \n（２）提供するウェブサイト又はアプリケーションが脆弱性を含まないこと 。 \n（３）実行プログラムの形式以 外にコンテ ンツを提供 する手段がない場合を除き、実行\nプログラム形式で コンテン ツを提供しな いこと。', metadata={'importance': 5}),
 Document(page_content='サービスを調達する際は、「政府情報システムのためのセキュリティ評価制度（ＩＳＭ\nＡＰ）」において登録 されたサービスから調達することを 原則とする こと。 \n３ 乙は、ウェブサイトの構築又はアプリケーション ・コンテ ンツ（アプリ ケーションプ\nログラム、ウェブコン テンツ等の総称を いう。以下同じ。）の開発・作成を行う場合に', metadata={'importance': 5}),
 Document(page_content='（外部公開ウ ェブサイトにおける情報セキュリテ ィ対策） \n第２６条の２ 乙は、外部公開ウェ ブサイト（以 下「ウェブサイト」という。）を構築又\nは運用するプラットフォームとして、乙が管理責任を有するサーバ等がある場合には、\n当該ウェブサイト又は当該サーバ等で利用するＯＳ、ミドルウェア等のソフトウェアの\n脆弱性情報を収集し、セキュリティ 修正プログラムが提供されている場合には業務影響', metadata={'importance': 5}),
 Document(page_content='ルメディアサービスを含む） を利用す る場合には、 これらのサービスで要機密情 報を扱\nってはならず、第２６条第８項に掲げる規程等に定める 不正アクセス対策を実施するな\nど規程等を 遵守しなければならない。 また、外部サービスを利用する場合は、その利用\n状況を管理 しなければならない 。なお、乙は、委託業務を実施するに当たり、クラウド', metadata={'importance': 4}),
 Document(page_content='ない。 \n２ 前項の規定にかかわらず、概算払財務大臣協議が整ったときは、乙は委託業務の完了 \n前に委託業務に必要な経費として様式第８により作成した概算払請求書を提出するこ\nとができる。この場合において、甲は、当該請求に対し支払うことが適当であると判断\nしたときは、 支払を行うことができる。  \n \n（遅延利息）', metadata={'importance': 4}),
 Document(page_content='年法律第５４号。以下「独占禁 止法」という。）第３条又は 第８条第１号の規定に\n違反する行為を行ったことにより、次 のイからハ までのいずれかに該当 することと\nなったとき  \nイ 独占禁止法第 ６１条第１項に規定する排除措置命令が確定したとき 。 \nロ 独占禁止法第６２条第１項に規定する課徴金納付命令が確定したとき 。', metadata={'importance': 5}),
 Document(page_content='対策基準」（平成１８･０３･２４シ第１号） （以下「規程等」と総称する。） に基づく\n情報セキュリティ対策を講じなければならない。  \n９ 乙は、当省 又は内閣官房内閣サイバーセキュリティセン ターが必要に応じて 実施する\n情報セキュリティ監査 、マネジメント監査又はペネトレー ションテスト を受け入れると\nともに、指摘事項への対応を行 わなければならない 。', metadata={'importance': 5}),
 Document(page_content='ばならない。  \n \n（支払うべき 金額の確定） \n第１５条  甲は、第１ ３条第１項の確認及 び納入物の引渡しを受けた後、前条の規定によ\nり提出された実績報告書の内容の審査及び必要に応じて 現地調査 を行い、委託業務の実\n施に要した経費の証ひょう、帳簿等の調査により支払うべき金額を確定し、これを乙に\n通知しなければならない。支払うべき金額を修正すべき事由が判明した場合も、同様と\nする。', metadata={'importance': 5}),
 Document(page_content='ム対策ソフトウェアを用いてスキャン を行い、不 正プログラ ムが含まれていない\nことを確認すること。  \n②アプリケーシ ョンプログラ ムを提供する場 合には、当該アプリ ケーションの仕様\nに反するプログラムコードが含まれていないことを確認すること。  \n③提供するウェブサイト又は アプリケー ション・コンテン ツにおいて、当省外のウ', metadata={'importance': 4}),
 Document(page_content='の防止等に関する法律（平成３年法律第 ７７号）第２条第２号に規定する暴力団 を\nいう。以下同じ。）であるとき 、又は法人 等の役員 等（個人である場合はその者、\n法人である場合 は役員又は 支店若しくは 営業所（常時 契約を締 結する事務所をい\nう。）の代表者、 団体である場合は代表者、理事 等、その他経営に実質的に関与し', metadata={'importance': 4})]

このように、クエリの中に条件を入れると、それを基にmetadataに対してフィルタをかけて文書を取得できます。

Step6. Chain creation

RAGを実行するChainを作成します。
LLM(ChatModel)はRetrieverで利用しているものと同じにしましたが、よりQAに特化したものや長コンテキストに対応したものなど、別のモデルを使ってもいいと思います。

from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import (
   AIMessagePromptTemplate,
   HumanMessagePromptTemplate,
)
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough, RunnableLambda, RunnableParallel
from autoawq_chat import load_pretrained_model, ChatAutoAWQ
from operator import itemgetter

template = """次のテキストの内容を使って質問に回答してください。
### テキスト
{context}

### 質問
{question}

"""
prompt = ChatPromptTemplate.from_messages(
   [
       HumanMessagePromptTemplate.from_template(template),
       AIMessagePromptTemplate.from_template(""),
   ]
)

chat_model = ChatHuggingFaceModel(
   generator=generator,
   tokenizer=tokenizer,
   human_message_template="GPT4 User: {}<|end_of_turn|>",
   ai_message_template="GPT4 Assistant: {}",
   repetition_penalty=1.2,
   temperature=0.1,
   max_new_tokens=1024,
)

chain = (
   RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
   | prompt 
   | chat_model 
   | StrOutputParser()
)

Step7. Run

準備は整いましたので、ストリーミング出力で実行してみます。

for s in chain.stream("重要度が3未満の契約文書を5件ピックアップし、そこに含まれる内容を要約してください。"):
    print(s, end="", flush=True)

出力

以下は、重要度が3未満の契約文書から要約した内容です。

1. 独占禁止法の排除措置命令書、課徴金納付命令書、通知文書
* 独占禁止法第61条第1項の排除措置命令、独占禁止法第62条第1項の課徴金納付命令、独占禁止法第7条の4第7項又は7条の7第3項の課徴金納付命令を命じない旨の通知文書
* 乙は前項第2号又は3号のいずれかに該当するときは、速やかに当該処分等に係る関係書類を甲に提出しなければならない。
1. 再委託
* 乙は再委託（委託業務の一部を第三者に委託することをいい、 請負その他委託 の形式を問わない。以下同じ。）してはならない。ただし、事業全体の企画及び立案並びに根幹に関わる執行管理以外の業務を再委託する場合であって、当該再委託が次の各号のいずれかに該当するときは、この限りでない。
1. 特記事項2（暴力団関与の属性要件に基づく契約解除）
* 甲は、乙が次の各号の一に該当すると認められるときは、何らの催告を要せず、本契約を解除することができる。
* （1）法人等（個人、法人又は団体をいう。）が、暴力団（暴力団員による不当な行為 の防止等に関する法律（平成3年法律第77号）第2条第2号に規定する暴力団を）に関与すると認められる場合
* （2）法人等が、暴力団による不当な行為の防止等に関する法律（平成3年法律第77号）第2条第2号に規定する暴力団による不当な行為を行うと認められる場合

ちなみに、この回答の基となったRetrieverで取得した文書は以下。

retriever.get_relevant_documents("重要度が3未満の契約文書を5件ピックアップし、そこに含まれる内容を要約してください。")

出力

[Document(page_content='含む場合には、その製造工程を含む。）を行う場合には、以下を実施しなければならな\nい。 \n（１）各工程において、 当省の意図しない変更や機密 情報の窃取等が行われないことを \n保証する管理が、一貫し た品質保証体制の下でなされていること。また、具体的な\n管理手順や品質保証体制を証明する書類等を提出すること 。', metadata={'importance': 1}),
 Document(page_content='存在するバージョンのＯ Ｓ、ソフトウェア等の利用を強制す るなどの情報セキュリ\nティ水準を低 下させる設定変更を ＯＳ、ソフトウェア等の利用者に要求することが\nないよう、ウェブサイト又はア プリケーション・コンテンツの提供方式を定めて開\n発すること。  \n（６）当省外へのアクセスを自動 的に発生させる機能やサービス利 用者その他の者に関', metadata={'importance': 1}),
 Document(page_content='（１）独占禁止法第 ６１条第１項 の排除措置命令書  \n（２）独占禁止法第 ６２条第１項 の課徴金納付命令書  \n（３）独占禁止法第７条の ４第７項又は第７条の７ 第３項の課徴金納付命 令を命じない\n旨の通知文 書 \n２ 乙は、前項第２号又は３号のいずれかに 該当することとなったときは 、速やかに 、当\n該処分等に係る関係書類 を甲に提出しなければ ならない 。', metadata={'importance': 1}),
 Document(page_content='（再委託）  \n第７条 乙は、再委託（委託業務の一部を第三者に委託することをいい、 請負その他委託\nの形式を問わない 。以下同じ。）してはならない。ただし、 事業全体の企画及び立案並\nびに根幹に関わる執行管理 以外の業務を再委託する場合であって、 当該再委託が次の各\n号のいずれかに該当する ときは、この限りでない。', metadata={'importance': 2}),
 Document(page_content='【特記事項２】  \n （暴力団関与の属性要件に基づく契約解除）  \n第４条 甲は、乙が次の各号の一に該当すると 認められ るときは、何ら の催告を要せず、\n本契約を解除することができる。  \n（１）法人等（個人、法人又は団体をいう。）が、暴力団（暴力団員による不当な行為\nの防止等に関する法律（平成３年法律第 ７７号）第２条第２号に規定する暴力団 を', metadata={'importance': 2})]

まとめ

Self-queryingを使ったRAGを実践してみました。

ローカルLLMだとモデル性能によっては正常に処理してくれなさそうなので、素直にOpenAIなどのモデルを使う方が確実です。ただ、最近のLLMは比較的少ないパラメータでも性能が高いものが多く、こういったシーンでも活用できるものが増えてきました。
ケースバイケースで使い分けるのがよいと思います。

個人的に、Self-queryingについては少し使うシチュエーションを選びそうな気がしていて、適切なユースケースをより理解していきたいと考えています。日々勉強。

なお、langchain templatesでは以下などでサンプルを確認することができます。

次回はMultiQueryRetrieverの予定です。