langchainとDatabricksで(私が)学ぶRAG :  RAG Fusion

# 導入

私が学ぶRAGの実質5回目です。シリーズ一覧は[こちら](https://qiita.com/isanakamishiro2/items/66411e1443cf78a2c1e8)。
今回はRAG Fusionです。

# これは何？


こちらが詳細内容です。

https://towardsdatascience.com/forget-rag-the-future-is-rag-fusion-1147298d8ad1

LangchainのBlogにも取り上げられており、
イメージとして、以下のような画像が掲載されています。

https://blog.langchain.dev/query-transformations/

![RAG Fusion](https://blog.langchain.dev/content/images/size/w1000/2023/10/1_5nG7iLyBoO-B5Tna6oDc-Q@2x.webp)

ざっくり言うと、元の問い合わせから派生となるクエリを複数生成し、各クエリの検索結果をReciprocal Rank Fusionというアルゴリズムを用いてRe-ranking(順序付け)し、関連度の高いものを抽出する考え方です。

LangchainによるTemplateは以下にあります。

https://github.com/langchain-ai/langchain/tree/master/templates/rag-fusion

というわけで、やってみましょう。

DatabricksのDBRは14.1 ML、GPUクラスタで動作を確認しています。

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

今後、使うモジュールをインストールします。
今回は```langchainhub```を追加します。

```python
%pip install -U -qq transformers accelerate langchain langchainhub faiss-cpu
# CUDA 11.8用
%pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.7/autoawq-0.1.7+cu118-cp310-cp310-linux_x86_64.whl
%pip install "databricks-feature-engineering"

dbutils.library.restartPython()
```

AutoAWQはCUDA 11.8版をWheelを指定してインストールしています。
(DatabricksのDBR ML 14台はCUDAのバージョンが11.8のため)

# Step1. Document Loading

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

```python
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

ごくごく単純なText Splitterを使って長文データをチャンキングします。
[ベーシックなRAG](https://qiita.com/isanakamishiro2/items/3fc4f9111c6be240a69e)のときと同じです。

```python
from typing import Any
import pandas as pd
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", split_text("page_content"))

# Pandas Dataframeに変換し、チャンクのリストデータを取得
pdf = df.select("chunk").toPandas()
texts = list(pdf["chunk"][0])

print(len(texts))
print(texts)
```

# Step3. Storage

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

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

https://huggingface.co/intfloat/multilingual-e5-large

```python
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},
)
```

FAISSでベクトルストアを作成します。
```python
vectorstore = FAISS.from_texts(texts, embedding)
```

# Step4. LLM Preparation

複数クエリを生成するためのLLMをロードします。
今回も、前回同様以下のモデルを利用しました。

https://huggingface.co/TheBloke/openchat_3.5-AWQ

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


```python
from transformers import AutoModelForCausalLM, AutoTokenizer
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="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)
```

```python
retriever_llm = ChatHuggingFaceModel(
    generator=generator,
    tokenizer=tokenizer,
    human_message_template="GPT4 Correct User: {}
\nGPT4 Correct Assistant: ",
    repetition_penalty=1.2,
    temperature=0.1,
    max_new_tokens=1024,
)
```

# Step5-a. RAG Fusion: Reciprocal Rank Fusion

**今回のポイントです。**

まず、Reciprocal Rank Fusionを実行するための関数を定義します。

Reciprocal Rank Fusionの解説は、以下のLINE社Blogに解説がありましたので、参考にしてください。

https://engineering.linecorp.com/ja/blog/a-new-challenge-for-line-timeline-3

>Reciprocal Rank Fusion(RRF)は、とても簡単で性能は優れている教師なし学習（unsupervised）のランク学習（learning-to-rank）で、情報検索（information retrieval）分野で検索語を入力すると、属性別に違う基準で検索される複数の文書を最終的にどの順番で組み合わせてユーザーに見せるかを決める方法です。

```python
from langchain.load import dumps, loads

def reciprocal_rank_fusion(results: list[list], k=10):
    fused_scores = {}
    for docs in results:
        # Assumes the docs are returned in sorted order of relevance
        for rank, doc in enumerate(docs):
            doc_str = dumps(doc)
            if doc_str not in fused_scores:
                fused_scores[doc_str] = 0
            fused_scores[doc_str] += 1 / (rank + k)

    reranked_results = [
        (loads(doc), score)
        for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    ]
    return reranked_results
```

# Step5-b. RAG Fusion: Query Generation

複数クエリを生成するためのChainを作成します。

まずは、Langchain Templateに基づいて、langchain hubからクエリ生成用のプロンプトテンプレートを読み込みます。

```python
from langchain import hub

prompt = hub.pull("langchain-ai/rag-fusion-query-generation")
```

promptの中身はこんな感じ。

```
ChatPromptTemplate(
    input_variables=['original_query'], 
    messages=[
        SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='You are a helpful assistant that generates multiple search queries based on a single input query.')), 
        HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['original_query'], template='Generate multiple search queries related to: {original_query}')), 
        HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='OUTPUT (4 queries):'))
    ]
)
```

クエリを4個生成するテンプレートになっています。

最後に、LLMと上記のテンプレートを用いたChainを作成します。

```python
from langchain.schema.output_parser import StrOutputParser

generate_queries = (
    prompt | retriever_llm | StrOutputParser() | (lambda x: x.split("\n"))
)
```

動作確認してみましょう。

```python
generate_queries.invoke({"original_query": "この契約において知的財産権はどのような扱いなのか？"})
```

```:出力
['1. この契約での知的財産権の譲渡条件は何ですか？',
 '2. この契約において知的財産権の利用権はどのように管理されますか？',
 '3. この契約に基づく知的財産権の拡張権はどのように決定されますか？',
 '4. この契約による知的財産権の終了条件は何ですか？']
```

類似クエリが4件作成されることを確認できました。


では、```reciprocal_rank_fusion```関数と```generate_queries```を使って、実際にRAGを実行するChainを組んでみます。

# Step6. Chain creation

Retrieverやプロンプトテンプレート等を作成して、Chainを作ります。
```retrieve_chain```がRAG Fusionを実行して必要な文書を取得するところになります。

 ```python
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough, RunnableLambda
from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import (
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# 10個の関連文書を取ってくるようにする
retriever = vectorstore.as_retriever(search_kwargs={'k': 10})


template = """You are an expert of world knowledge. I am going to ask you a question. Your response should be comprehensive and not contradicted with the following context if they are relevant. Otherwise, ignore them if they are not relevant.
{context}

Question: {question}
"""
prompt2 = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(template),
        AIMessagePromptTemplate.from_template(""),
    ]
)

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

retrieve_chain = (
    {"original_query": lambda x: x}
    | generate_queries
    | retriever.map()
    | reciprocal_rank_fusion
)

chain = (
    {
        "context": retrieve_chain,
        "question": lambda x: x,
    }
    | prompt2
    | chat_model
    | StrOutputParser()
)
```

# Step7. Run

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

```python
for s in chain.stream("この契約において知的財産権はどのような扱いなのか？"):
    print(s, end="", flush=True)
```

```:出力
この契約において、知的財産権は以下のように扱われます。

1. 知的財産権は、乙が本契約により行うこととされた全ての給付を完了する前に、乙が前項ただし書に遵守し、及び第７条第１項 の再委託先に遵守させ ることを約束する。
2. 本契約の締結時に乙が既に所有又は管理していた知的財産権（乙知的財産権）を乙が納入物に使用した場合には、甲は当該乙知的財産権を、仕様書の「目的」のため、仕様書の「納入物」の項に記載した利用方法に従い、本契約終了後の納入物の利用についても同様とする。
3. 乙は、本契約により生じる権利の全部又は一部を甲の承諾を得ずに、第三者に譲渡し、又は承継させてはならない。ただし、信用保証協会、資産の流動化に関する法律（平成１０年法律第１０５号）第２条第３項に規定する金融機関に対して債権を譲渡する場合にあっては、この限りでない。
4. 乙は、本契約に関し、第４条又は前条第２項の規定に該当したときは、甲が本契約を解除するか否かにかかわらず、かつ、甲が損害の発生及び損害額を立証することを要する。

これらの規定により、知的財産権は乙が本契約により行うこととされた全ての給付を完了する前に、乙が前項ただし書に遵守し、及び第７条第１項 の再委託先に遵守させ ることを約束する。また、乙は本契約により生じる権利の全部又は一部を甲の承諾を得ずに、第三者に譲渡し、又は承継させてはならない。
```

結果が得られました。(正しいかどうかはさておき)

RAG Fusion部単体での動作も見てみます。

```python
retrieve_chain.invoke("この契約において知的財産権はどのような扱いなのか？")
```
```:出力
[(Document(page_content='新規知的財産権は 約定の委託金額以外の追加支払なしに、納入物の引渡しと同時に乙か\nら甲に譲渡され、甲単独に帰属する。  \n５ 前項の規定にかかわらず、著作権等について は第２８条の定めに従う。  \n６ 乙は、本契約終了後であっても、知的財産権の取 扱いに関する本契約 の約定を 自ら遵\n守し、及び第７条第１項 の再委託先に遵守させ ることを 約束する。'),
  0.3909090909090909),
 (Document(page_content='納入物の改良・改変をはじめとして、あらゆる使用（利用）態様を含む 。また、本契約\nにおいて「知的財産権」とは、知的財産基本法第２条第２項所定の知的財産権をいい、\n知的財産権を受ける権利及びノウハウその他の秘密情報を含む。  \n２ 乙は、 納入物に第三者の知的財産権を利用する場合には、 第１条第２項 の規定に従い、\n乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による'),
  0.33674242424242423),
 (Document(page_content='（契約保証金）  \n第３条 甲は、本契約に係る乙が納付すべき契約保証金の納付を全額免除する。  \n \n （知的財産 権の帰属及び 使用） \n第４条 本契約の締結時に乙が既に所有又は管理していた 知的財産権（以下「 乙知的財産\n権」という。）を 乙が納入物に使用した場合には、甲は、当該乙知的財産権を、仕様書\n記載の「目的」のため、仕様書の「納入物」の項 に記載した利用方法に従い、本契約終'),
  0.32808857808857805),
 (Document(page_content='が所有し、又は管理する知的財産権の実施許諾や動産・不動産の使用許可の取得等を含\nむ。）が必要な場合には乙の費用及び責任で行うものとする。 甲の指示により、委託者\n名を明示して業務を行う場合も同様とする。  \n３ 甲は、委託業務及び納入物に関して、約定の委託金額以外の支払義務を負わない。 本\n契約終了後の納入物の利用についても同様とする。 委託金額には委託業務の遂行に必要'),
  0.3170995670995671),
 (Document(page_content='自体を含 む。）に関 して著作者人格権を行使しないことに同意する。また、乙は、当該\n著作物の著作者が乙以外の 者であるときは、当該著作者が著作者人格権を行使しないよ\nうに必要な措置をとるものとす る。 \n４ 乙は、 本条及び知的財産権の帰属等に関する本契約及び仕様書の約定を 遵守するため、\n必要な範囲で職 務発明や著作権に 関する管理規程その他の社内規程を整備 するととも'),
  0.28717948717948716),
 (Document(page_content='乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による\n納入物の利用 に必要な書面の許諾を得なければならない。なお、第三者より当該許諾に\n条件を付された場合には（以下「第三者の許諾条件」という。）、乙は、納入物に第三\n者の知的財産権を利用する前に、甲に対して第三者の 許諾条件を書面で速やかに通知し'),
  0.24429824561403507),
 (Document(page_content='しに、納入物の利用 に必要な範囲で、前項の第三者の知的財産権を 自由かつ対価の追加\n支払なしに 使用し、又は第三者に使用させることができる。  \n４ 委託業務の遂行中に納入物に関して乙（甲の同意を得て一部を再委託する場合は再委\n託先を含む。） が新たに知的財産権（以下 「新規知的財産権」という。）を取得した場\n合には、乙は、その詳細を書面にしたものを納入物に添付して甲に提出するものとする。'),
  0.19691876750700282),
 (Document(page_content='甲に譲渡され、甲単独に帰属 する。乙は、 甲が求める場合には、本項に定める著作権の\n譲渡証の作 成等、譲渡を証する書面の作成に協力しなければなら ない。  \n２ 本契約締結日現在乙、乙以外の委託事業参加者又は第三者の権利対象となる著作 物が\n納入物に含まれている場合であっても、甲は、納入物の利用のため、本契約期間中 及び'),
  0.12698412698412698),
 (Document(page_content='変更その他契約内容の変更を行うことがあり、この場合、丙は異議を申し立てない\nものとし、当該契約の変更により、譲渡対象債権の内容に影響が及ぶ場合 の対応に\nついては、専ら乙と丙の間 の協議により決定 されなければならないこと。  \n３ 第１項ただし書の規定に基づい て乙が第 三者に債権の譲渡を行った場合においては、'),
  0.11805555555555555),
 (Document(page_content='【特記事項２】  \n （暴力団関与の属性要件に基づく契約解除）  \n第４条 甲は、乙が次の各号の一に該当すると 認められ るときは、何ら の催告を要せず、\n本契約を解除することができる。  \n（１）法人等（個人、法人又は団体をいう。）が、暴力団（暴力団員による不当な行為\nの防止等に関する法律（平成３年法律第 ７７号）第２条第２号に規定する暴力団 を'),
  0.08333333333333333),
 (Document(page_content='に譲渡し、又は承継させてはならない。ただし、信用保証協会、資産の流動化に関する\n法律（平成１０年法律第１０５号）第２条第３項に規定する特定目的会社 又は中小企業\n信用保険法施行令（昭和２５年政令第３５０号）第１条の ３に規定する金融機関 に対し\nて債権を譲渡する場合にあっては、この限りでない。  \n２ 乙が本契約により行うこととされた 全ての給付を完了 する前に、乙が 前項ただし書に'),
  0.058823529411764705),
 (Document(page_content='守し、及び第７条第１項 の再委託先に遵守させ ることを 約束する。 \n７ 委託業務又は納入物に関して、第三者の知的財産権の侵害 に関する紛争 その他第三者\nとの間で何らかの紛争が発生した場合には、当該紛争の解決については乙が全責任を負\nう。 \n \n（計画変更等）  \n第５条  乙は、実施計画を変更しようとするとき（事業内容の軽微な変更の場合及び支出'),
  0.058823529411764705),
 (Document(page_content='除く。）に提供し、又はその内容を知らせること。  \n（２）甲から預託された個人情報 等について、甲が示した利用目的（ 特に明示がない場 \n合は本契約の目 的）の範囲を 超えて使用し、複製し、又は改変すること。  \n（３）委託業務に関して自ら収集し、又は作成した個 人情報について、甲が示した利用\n目的（特に明示がな い場合は本契約の目的） の範囲を超 えて使用す ること。'),
  0.058823529411764705),
 (Document(page_content='なければならない。甲は、当該第三者の許諾条件に同意できない場合には、本契約の 解\n約又は変更を含め、乙に対して協議を求めることができる。甲が当該条件に同意した場\n合、乙は、委託業務の遂行及び納入物の作成に 当たって第三者の許諾条件を遵守するこ\nとにつき全責任を負う。   \n３ 甲は、第三者の許諾条件を遵守することを条件として、 本契約終了後も 期間の制限な'),
  0.05555555555555555),
 (Document(page_content='乙に生じた損害について、何ら賠償ないし補償する ことは要しない。  \n２ 乙は、甲が第４条又は前条第２項の規定により本契約を解除した場合において、甲に\n損害が生じたときは、その損害を 賠償する ものとする。  \n３ 乙が、本契約に関 し、第４条又は前条第２項 の規定に該当したときは、甲が本契約を\n解除するか否かにかかわらず、かつ、甲が損害の発生及び損害額を立証することを要す'),
  0.05555555555555555),
 (Document(page_content='再委託に関する承認を得た又は履行体制図変更届出を行ったものとみなす。  \n２ 第７条第２項の再委託の承認を得た場合は 、その承認された範 囲内において、履行体\n制図変更届出を行ったものとみなす。  \n \n（債権譲渡の禁止）  \n第１０条 乙は、本 契約によって生じる権利の全部又は一部を甲の承諾を得ずに、第三者\nに譲渡し、又は承継させてはならない。ただし、信用保証協会、資産の流動化に関する'),
  0.05263157894736842),
 (Document(page_content='善良なる管理者の注意をもって取り扱わなければならない。  \n２ 乙は、個人情報 等を取り扱わせる業務を第三者に再委託する場合は、事前に甲の 承認\nを得るとともに、本条に定める、甲が乙に求めた個人情報 等の適切な管理のために必要\nな措置と同 様の措置を当該第三者も講ずるように求め、かつ 、当該第三者が 約定を遵守'),
  0.05263157894736842),
 (Document(page_content='り本契約が終了した後で あっても、なおその効力を有する。  \n \n（著作権等の帰属）  \n第２８条  納入物に係る 著作権（著 作権法第２ ７条及び第２８条の権利を含む。ただし、\n本契約締結日現 在、乙、乙以外 の委託事業参加 者又は第三者の権利 対象となっているも\nのを除く。以下同じ。）は、委 託金額以外の追加支払なしに、その発生と同時に乙から'),
  0.05263157894736842)]
```

各ドキュメントのスコアが追加されています。
今回は行っていませんが、スコアの上位N件だけにしぼるなどすると、より精度のよいRAGが実現できるかと思います。


# まとめ

RAG Fusionを実践してみました。
思ったよりRAG部の動作速度が速く、積極的に取り入れてもよい手法かなと思いました。

次回はStep Back Promptingを実践する予定です。