langchainとDatabricksで(私が)学ぶRAG : Step Back Prompting

# 導入

私が学ぶRAGの実質6回目です。シリーズ一覧は[こちら](https://qiita.com/isanakamishiro2/items/66411e1443cf78a2c1e8)。
今回はStep Back PromptingによるRAGの改善です。

# これは何？

langchianのBlog内に説明のセクションがあります。

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

また、npaka大先生が説明＆実践をされています。

https://note.com/npaka/n/n55f276ad2988

先生の記事からの抜粋で言うと、Step back promptingは以下のアルゴリズムとなります。

> (1) ユーザーの元の質問に基づいて、ステップバック質問を生成
> (2) 元の質問とステップバック質問の両方を情報収集
> (3) 取得した両方の情報に基づいて回答を生成

複数質問を生成して回答を生成する、というのは[前回のRAG Fusion](https://qiita.com/isanakamishiro2/items/f4387443b86723eecf36)と同様ですが、
生成する質問は「ステップバック質問」という、元質問に比べてより抽象度・概念性の高い質問を生成します。

Colabでの実践も大先生の記事中にありますので、是非確認ください。

<br>


・・・あれ、私のこの記事いらないんじゃない？

<br>

という、若干存在意義を疑いつつ、やってみます。

なお、Step Back Promptingに関するLangchainのTemplateは以下にあります。

https://github.com/langchain-ai/langchain/tree/master/templates/stepback-qa-prompting

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

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

使うモジュールをインストールします。

```python
%pip install -U -qq transformers accelerate langchain 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

[準備編](https://qiita.com/isanakamishiro2/items/66411e1443cf78a2c1e8)で保管した特徴量を取得します。
読み込んだデータは、```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])
```

# 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)
```

# Step5. Question Gen Chain for Step-back

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

ステップバック質問を生成するためのChainを作成します。

まずは、Langchainのchat modelを作成。

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

次にプロンプトテンプレートを準備。
こちらはLangchainのtemplateを基に作成しました。

```python
from langchain.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate

# Few Shot Examples
examples = [
    {
        "input": "Could the members of The Police perform lawful arrests?",
        "output": "what can the members of The Police do?",
    },
    {
        "input": "Jan Sindel’s was born in what country?",
        "output": "what is Jan Sindel’s personal history?",
    },
]
# We now transform these to example messages
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}"),
    ]
)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are an expert at world knowledge. Your task is to step back "
            "and paraphrase a question to a more generic step-back question, which "
            "is easier to answer. Here are a few examples:",
        ),
        # Few shot examples
        few_shot_prompt,
        # New question
        ("user", "{question}"),
        ("ai", ""),
    ]
)
```

FewShotで事例を与えつつ、ステップバック質問を生成するシステムプロンプトを含めた内容となっています。なるほど、プロンプト構築の参考になる。。。

<br>

最後に、Chainを[LCEL](https://python.langchain.com/docs/expression_language/)で作成。

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

question_gen = prompt | gen_model | StrOutputParser()
```

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

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

```:出力
'この契約に関連する知的財産権について、どのような制度が存在していますか？'
```

より抽象度・一般性の高い質問が生成されます。

ちなみに、以下の質問だとこうなります。可愛い存在全体を聞く質問を生成していますね。
```python
question_gen.invoke({"question": "まどか☆マギカで一番可愛いのは誰？"})

# <出力結果>
# まどか☆マギカのキャラクターの中で誰が可愛いとされていますか？
```

では、このChainを使ってRAGを実行してみましょう。

# Step6. Chain creation

元質問とステップバック質問からコンテキストを取得し、元質問について回答させるChainを構築します。


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

retriever = vectorstore.as_retriever()

response_prompt_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.

{normal_context}
{step_back_context}

Original Question: {question}
Answer:"""

response_prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(response_prompt_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,
)

chain = (
    {
        # 元質問から関連するコンテキストを抽出
        "normal_context": RunnableLambda(lambda x: x["question"]) | retriever,
        # ステップバック質問から関連するコンテキストを抽出
        "step_back_context": question_gen | retriever,
        # 元質問を設定
        "question": lambda x: x["question"],
    }
    | response_prompt
    | chat_model
    | StrOutputParser()
)
```

個人の感想ですが、LCELに慣れると、このパイプ表現が結構読みやすく感じています。
(R言語におけるTidyverseの感じに似ている）

# Step7. Run

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

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

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

1. 乙が既に所有又は管理していた知的財産権（乙知的財産権）を乙が納入物に使用した場合、甲は仕様書に記載の「目的」のため、仕様書の「納入物」の項に記載した利用方法に従って本契約終了後の知的財産権を譲渡します。
2. 乙は、納入物に第三者の知的財産権を利用する場合には、第１条第２項の規定に従って乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による知的財産権を取得します。
3. 新規知的財産権は約定の委託金額以外の追加支払なしに、納入物の引渡しと同時に乙から甲に譲渡され、甲単独に帰属します。
4. 著作権等については、第２８条の定めに従い扱われます。
5. 乙は、本契約終了後であっても知的財産権の取扱いに関する本契約の約定を自ら遵守し、及び第７条第１項の再委託先に遵守させることを約束します。
```

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

ちなみに、```astream_log```メソッドを使って実行のログを見たところ、上の処理で生成されたステップバック質問は、「この契約に関連する知的財産権について、どのような制度が存在していますか？」であり、取得されたコンテキストは以下のようになります。

```python
[Document(page_content='（契約保証金）  \n第３条 甲は、本契約に係る乙が納付すべき契約保証金の納付を全額免除する。  \n \n （知的財産 権の帰属及び 使用） \n第４条 本契約の締結時に乙が既に所有又は管理していた 知的財産権（以下「 乙知的財産\n権」という。）を 乙が納入物に使用した場合には、甲は、当該乙知的財産権を、仕様書\n記載の「目的」のため、仕様書の「納入物」の項 に記載した利用方法に従い、本契約終'),
Document(page_content='が所有し、又は管理する知的財産権の実施許諾や動産・不動産の使用許可の取得等を含\nむ。）が必要な場合には乙の費用及び責任で行うものとする。 甲の指示により、委託者\n名を明示して業務を行う場合も同様とする。  \n３ 甲は、委託業務及び納入物に関して、約定の委託金額以外の支払義務を負わない。 本\n契約終了後の納入物の利用についても同様とする。 委託金額には委託業務の遂行に必要'),
Document(page_content='納入物の改良・改変をはじめとして、あらゆる使用（利用）態様を含む 。また、本契約\nにおいて「知的財産権」とは、知的財産基本法第２条第２項所定の知的財産権をいい、\n知的財産権を受ける権利及びノウハウその他の秘密情報を含む。  \n２ 乙は、 納入物に第三者の知的財産権を利用する場合には、 第１条第２項 の規定に従い、\n乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による'),
Document(page_content='新規知的財産権は 約定の委託金額以外の追加支払なしに、納入物の引渡しと同時に乙か\nら甲に譲渡され、甲単独に帰属する。  \n５ 前項の規定にかかわらず、著作権等について は第２８条の定めに従う。  \n６ 乙は、本契約終了後であっても、知的財産権の取 扱いに関する本契約 の約定を 自ら遵\n守し、及び第７条第１項 の再委託先に遵守させ ることを 約束する。')]
```

ちなみに、元質問である「この契約において知的財産権はどのような扱いなのか？」で得られたコンテキストはこちら。

```python
[Document(page_content='納入物の改良・改変をはじめとして、あらゆる使用（利用）態様を含む 。また、本契約\nにおいて「知的財産権」とは、知的財産基本法第２条第２項所定の知的財産権をいい、\n知的財産権を受ける権利及びノウハウその他の秘密情報を含む。  \n２ 乙は、 納入物に第三者の知的財産権を利用する場合には、 第１条第２項 の規定に従い、\n乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による'),
Document(page_content='（契約保証金）  \n第３条 甲は、本契約に係る乙が納付すべき契約保証金の納付を全額免除する。  \n \n （知的財産 権の帰属及び 使用） \n第４条 本契約の締結時に乙が既に所有又は管理していた 知的財産権（以下「 乙知的財産\n権」という。）を 乙が納入物に使用した場合には、甲は、当該乙知的財産権を、仕様書\n記載の「目的」のため、仕様書の「納入物」の項 に記載した利用方法に従い、本契約終'),
Document(page_content='が所有し、又は管理する知的財産権の実施許諾や動産・不動産の使用許可の取得等を含\nむ。）が必要な場合には乙の費用及び責任で行うものとする。 甲の指示により、委託者\n名を明示して業務を行う場合も同様とする。  \n３ 甲は、委託業務及び納入物に関して、約定の委託金額以外の支払義務を負わない。 本\n契約終了後の納入物の利用についても同様とする。 委託金額には委託業務の遂行に必要'),
Document(page_content='新規知的財産権は 約定の委託金額以外の追加支払なしに、納入物の引渡しと同時に乙か\nら甲に譲渡され、甲単独に帰属する。  \n５ 前項の規定にかかわらず、著作権等について は第２８条の定めに従う。  \n６ 乙は、本契約終了後であっても、知的財産権の取 扱いに関する本契約 の約定を 自ら遵\n守し、及び第７条第１項 の再委託先に遵守させ ることを 約束する。')]
```

順番はことなりますが、今回のケースで言えばどちらの質問でも同じコンテキストが取得されていました。

取得件数やしきい値を変えたりすると、違う結果になるかもしれません。

# まとめ

Step Back Promptingを実践しました。
ステップバック質問という考え方自体に面白味があると感じます。

実務においては、適切な問い合わせ自体を作ることが難しい、という場面にも直面するので、こういった内容を取り入れながら応答性能を改善していきたいですね。実装自体はシンプルですし。

次回はSemi Structuredなデータに対するRAGを予定しています。（こちらは別物にする可能性があります）