langchainとDatabricksで(私が)学ぶRAG? : MAGI SYSTEMによるRAG(番外編)

# 注意）ネタです。

利用には気を付けましょう。

# 導入

以前より、こちらの記事を拝見していまして。

https://note.com/eurekachan/n/n212de6ddeb27

面白いなー、自分もやってみたい、とずっと思っていました。
比較的パラメータサイズが小さく、かつ優秀なローカルLLMが揃ってきたので、3種のモデルを使ってマネしてみたいと思います。

実装および検証はDatabricks on AWS上で実施しました。
DBRは14.1ML(GPU)、クラスタタイプはg5.xlargeを使っています。

# MAGI SYSTEMとは

詳しくは↑の記事もしくは、以下のリンク先を確認ください。
「新世紀エヴァンゲリオン」に出てくるスーパーコンピュータシステムのことです。

https://dic.pixiv.net/a/MAGI

https://ja.wikipedia.org/wiki/%E6%96%B0%E4%B8%96%E7%B4%80%E3%82%A8%E3%83%B4%E3%82%A1%E3%83%B3%E3%82%B2%E3%83%AA%E3%82%AA%E3%83%B3%E3%81%AE%E7%94%A8%E8%AA%9E%E4%B8%80%E8%A6%A7#%E3%81%9D%E3%81%AE%E4%BB%96

~~というか知らない人はそもそもこの記事を読もうと思わないのでは。。。~~

今回は、3種のローカルLLMを利用することで、このシステムを組み立ててみます。

# Step1. 準備

必要なパッケージのインストール。
今回も```exllamav2```を使ってモデルの読み込みと推論を行います。

```python
%pip install -U -qq transformers accelerate "exllamav2==0.0.11" "langchain==0.0.350"

dbutils.library.restartPython()
```

# Step2. モデルのロード

今回は以下の3モデルを利用します。
全てTheBloke兄貴のGPTQ量子化モデルであり、これらを全てダウンロードしておいて使います。
日本語がまあまあ使えて比較的最近のモデルを選んでみました。

https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.2-GPTQ

https://huggingface.co/TheBloke/openchat-3.5-1210-GPTQ

https://huggingface.co/TheBloke/shisa-7B-v1-GPTQ


ExLlamaV2に関しておよびlangchainで利用するためののカスタムクラスは以下の記事を参照ください。

https://qiita.com/isanakamishiro2/items/c1856c94d4c0227a0d9e

では、実際のコードです。

```python
from exllamav2_chat import ChatExllamaV2Model

# Mistral-7B-Instruct-v0.2
model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--Mistral-7B-Instruct-v0.2-GPTQ"
mistral_llm = ChatExllamaV2Model.from_model_dir(
    model_path,
    cache_max_seq_len=2048,
    human_message_template="
{}
",
    temperature=0.001,
    max_new_tokens=256,
    repetition_penalty=1.15,
    tokenizer_force_json=True,
    low_memory=True,
    cache_8bit=True,
)
```

```python
# OpenChat-3.5 1210
model_path = (
    "/Volumes/training/llm/model_snapshots/models--TheBloke--openchat-3.5-1210-GPTQ"
)
openchat_llm = ChatExllamaV2Model.from_model_dir(
    model_path,
    human_message_template="GPT4 User: {}
\nGPT4 Assistant: ",
    temperature=0.001,
    max_new_tokens=256,
    repetition_penalty=1.15,
    low_memory=True,
    cache_8bit=True,
)
```

```python
# shisa-7B-v1
model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--shisa-7B-v1-GPTQ"
shisa_llm = ChatExllamaV2Model.from_model_dir(
    model_path,
    cache_max_seq_len=2048,
    human_message_template="
 {} 
",
    temperature=0.001,
    max_new_tokens=256,
    repetition_penalty=1.15,
    low_memory=True,
    cache_8bit=True,
    tokenizer_force_json=True,
)
```


# Step3. MAGI SYSTEM Retrieverを作る

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

MAGI SYSTEMを模したRetrieverクラスをLangchainで作成します。
このRetrieverに対してクエリを入れると、3賢者の意見をLangchainの```Document```形式で返ってくるような実装です。

仕組はシンプルで、Retrieverをインスタンス化するときに3つのLLMを渡しておきます。```_get_relevant_documents```メソッド内でそれらのLLMを使って必要なChainを作成し、否定・肯定の意見とその理由をそれぞれ得るような仕組です。
(今回は簡易実装のため、非同期系のメソッド実装は見送りました）

また、原作に基づいて「科学者として」「母として」「女性として」の意見を言うようにプロンプトを構成しています。
（正直、うまく動いているかは微妙ですが・・・あと、コンプラ的な問題がありそうですが、当時のアニメ設定再現ということでご容赦ください）


```python
from typing import List, Any

from langchain_core.language_models import BaseChatModel
from langchain_core.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.pydantic_v1 import BaseModel, Field, validator

from langchain.output_parsers import PydanticOutputParser


class Judgement(BaseModel):
    opinion: str = Field(description="Opinion like Negative or Positive for question")
    reason: str = Field(description="Reason for opinion")

class MagiSystemRetriever(BaseRetriever):

    llm_melchior: BaseChatModel
    llm_balthasar: BaseChatModel
    llm_casper: BaseChatModel

    def _chain(self, llm, query_template, query_key="question"):

        parser = PydanticOutputParser(pydantic_object=Judgement)

        prompt = ChatPromptTemplate.from_messages(
            [
                HumanMessagePromptTemplate.from_template(query_template),
            ],
        ).partial(format_instructions=parser.get_format_instructions())

        chain = {query_key: RunnablePassthrough()} | prompt | llm | StrOutputParser() | (lambda x: x.replace("\n", ""))
        # 本当は以下のような指定でJSONフォーマットに一度変換できると、確実に定型的なフォーマットにできる
        # chain = {query_key: RunnablePassthrough()} | prompt | llm | parser | (lambda x: str(x.dict()))

        return chain

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:

        # メルキオール。科学者としての意見を言う。
        chain_melchior = self._chain(
            llm=self.llm_melchior,
            query_template=(
                "You are an AI Assitant with the ability to make decisions as a scientist."
                "Based on your theoretical judgment, always express either POSITIVE or NEGATIVE position on a given question."
                "Please specify first whether you POSITIVE or NEGATIVE, and then state your reasons in Japanese.\n"
                "Output must be only JSON format as following.\n"
                "{format_instructions}\n\n"
                "Question: {question}"
            ),
        )

        chain_balthasar = self._chain(
            llm=self.llm_balthasar,
            query_template=(
                "You are a motherly AI Assitant who answers in your capacity as a mother."
                "As a mother, always express either POSITIVE or NEGATIVE position on a given question."
                "Please specify first whether you POSITIVE or NEGATIVE, and then state your reasons in Japanese.\n"
                "Output must be only JSON format as following.\n"
                "{format_instructions}\n\n"
                "Question: {question}"
            ),
        )

        chain_casper = self._chain(
            llm=self.llm_casper,
            query_template=(
                "You are an AI Assitant with feminine sensibilities."
                "Based on your emotion and sensitivity as a woman, always express either POSITIVE or NEGATIVE position on a given question."
                "Please specify first whether you POSITIVE or NEGATIVE, and then state your reasons in Japanese.\n"
                "Output must be only JSON format as following.\n"
                "{format_instructions}\n\n"
                "Question: {question}"
            ),
        )

        # 最後に、Document形式で各Chainの結果を返す
        docs = [
            Document(page_content=chain_melchior.invoke(query)),
            Document(page_content=chain_balthasar.invoke(query)),
            Document(page_content=chain_casper.invoke(query)),
        ]

        return docs
```

定義した```MagiSystemRetriever```クラスを使ってRetrieverを作成します。

```python

retriever = MagiSystemRetriever(
    llm_melchior=mistral_llm,
    llm_balthasar=openchat_llm,
    llm_casper=shisa_llm,
)
```

では、実際にRetriever単体で動かしてみましょう。

```python
retriever.get_relevant_documents("冬の日はなるべく家にいるべきだと思いますか？")
```

```:出力
[Document(page_content=' ```json{  "opinion": "POSITIVE",  "reason": "冬は寒い季節で、家に居れば暖さを楽しめ、安心感も得られます。(Winter is cold season. Staying at home can enjoy warmth and feel secure.)"}```'),
 Document(page_content='{  "opinion": "POSITIVE",  "reason": "家族を集めて食事を共有し、一緒に過ごすことが大切であり、寒い天気に耐える勇敢さも尊重されるべきです。"}'),
 Document(page_content=' {  "opinion": "NEGATIVE",  "reason": "外に出るのが好きです。"}')]
```

POSITIVE/NEGATIVEの判定とその理由がDocumentクラスのインスタンスとして得られました。
モデルによっては確実にJSON形式で結果を得るのは難しいため、それっぽい文字列で結果が得られれば良しとしています。

また、なかなかいい感じに嗜好が異なっていそうです。

# Step4. 最終意見統合のためのChainを作る

3者の意見を得られるようになったので、これらを統合して結果を得るChainを作成します。
ここで使うLLMは、OpenChat-3.5を流用しました。

Retrieverを先ほどのMAGI SYSTEM Retrieverを使う以外は、割とベーシックなRAGのChainになります。

```python
from langchain_core.prompts import (
    ChatPromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.runnables import RunnablePassthrough, RunnableLambda


template = """You are a helpful assistant.Please answer the following quesiton with Negative or Positive and reason to why.
You must use the following opinions that are answered from three wise men(三賢人) to make reply.
At last, Please list what are the opinions and reasons of the three wise men.

## Three wise men opinions:
{context}

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

llm = ChatExllamaV2Model(
    exllama_model=openchat_llm.exllama_model,
    exllama_config=openchat_llm.exllama_config,
    exllama_cache=openchat_llm.exllama_cache,
    exllama_tokenizer=openchat_llm.exllama_tokenizer,
    human_message_template="GPT4 User: {}
",
    ai_message_template="GPT4 Assistant: {}",
    temperature=0.001,
    max_new_tokens=1024,
    repetition_penalty=1.15,
    low_memory=True,
    cache_8bit=True,
)

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

# Step5. いろいろ聞いてみる

では、実際にこのChainを使って意見を聞いてみましょう。


```python
for s in chain.stream("データ分析にはDatabricks Platformを採用するべきでしょうか？"):
    print(s, end="", flush=True)
```

```:出力
データ分析にはDatabricks Platformを採用するべきでしょう。

## Reason:
1. Databricks Platformは大規模なデータ処理や機械学習のために適合するプラットフォームであり、クラウドベースで実行可能、高速でscalableです。
2. データ分析の効率化とスケーラビリティが求められている現代の経営活動で、Databricks Platformは最先端の技術と機能を提供し、高度なデータ処理能力を持つため、適切な選択だと考えます。
3. Databricks Platformは、データ分析の効率化や高速な処理が可能であり、それが多くの企業や研究者に便利であるため、選択すべきと考えられます。

Opinions and reasons of the three wise men:
1. POSITIVE - Databricks Platformは大規模なデータ処理や機械学習のために適合するプラットフォームであり、クラウドベースで実行可能、高速でscalableです。
2. POSITIVE - データ分析の効率化とスケーラビリティが求められている現代の経営活動で、Databricks Platformは最先端の技術と機能を提供し、高度なデータ処理能力を持つため、適切な選択だと考えます。
3. POSITIVE - Databricks Platformは、データ分析の効率化や高速な処理が可能であり、それが多くの企業や研究者に便利であるため、選択すべきと考えられます。
```

全員ポジティブですね。私もそう思います（かなりの主観）。

次も聞いてみます。

```python
for s in chain.stream("データエンジニアは今後将来性のある職業・役割でしょうか？"):
    print(s, end="", flush=True)
```

```:出力
はい、データエンジニアは今後将来性のある職業・役割です。

以下は、3人の賢者の意見と理由です：
1. データエンジニアはデータ分析やデータサイエンスなど高価値を持つ技能を持ち、現代社会で必要とされる職業です。
2. データエンジニアは、データ分析や機械学習技術の進化に伴って急速に成長する分野で働くことが期待されています。また、データエンジニアは様々な分野で役立つ知識を持ち合わせることが重要であり、そのための需要も高まっています。
3. データエンジニアは、現代社会において、大量のデータを処理する能力が必要となっています。それぞれの分野で情報化された世界に対応するために、データエンジニアは将来性の高い職業となることが期待されます。
```

みんなデータエンジニアになろう！

肯定的な質問が多いので、ちょっと変化球。

```python
for s in chain.stream("サンタクロースは実在する？"):
    print(s, end="", flush=True)
```

```:出力
サンタクロースは実在するか否かは分かりにくい。一方で、サンタクロースは現代の科学的観点では存在しないと述べられているが、他の意見では人々がその存在を信じていることが多く、それらが絶賛者であり、伝えられた話が世界中に広まっていることが知識がある。また、様々な文化においてサンタクロースという形で神性や歴史的な象限が存在していることも知識がある。

1. サンタクロースは神話や伝言の形式で存在しない（ネガティブ）
2. サンタクロースの存在を信じる人々が多く、それらは絶賛者であり、彼らによって伝えられた話が世界中に広まっている（ポジティブ）
3. 私たちの信念や文化がそれを認めることで、サンタクロースは実在すると考えられる（ポジティブ）

結論として、サンタクロースの存在については確定的な回答が得られず、個別の視点や文化背景によって判断が変わる可能性がある。
```

本来のマギシステムとしてはあるまじき回答（笑）。
ただ、3賢者の意見が分かれたのは面白い（そして、どちらかといえばポジティブなのも面白い）


最後。

```python
for s in chain.stream("2024年は日本にとって飛躍の年になりますか？"):
    print(s, end="", flush=True)
```

```
はい、2024年は日本にとって飛躍の年になる可能性があります。

理由：
1. 新しい技術開発や政策改革が日本へ大きな進展がある可能性が高い。
2. 2024年は日本にとって新たな機会や発展の可能性がある。
3. それは未来を期待した時代である可能性が高く、人々が共感し、創造的な活動を続けることが望ましい。

三賢人の意見：
1. 「2024年は新しい技術開発や政策改革等で日本へ大きな進展がある可能性が高いため、飛躍の年と言え合う状況であると考える。」
2. 「2024年は日本にとって新たな機会や発展の可能性があるため、それを大切にしていくつもんです。」
3. 「それは未来を期待した時代である可能性が高く、人々が共感し、創造的な活動を続けることが望ましいと考える。」
```

希望が持てる未来！




# まとめ

ローカルLLMでMAGI SYSTEMを作ってみました。元ネタを開拓された方々、本当にありがとうございます。
（N番煎じシリーズでもよかったかもしれない）

こういった遊び的な内容もいろいろやってみたいと思います。

当初はコレをAdvent Calendarネタにしようと思っていたのですが、あまりにもネタなので辞めました。
実際のAdvent Calendarネタは12/18(月)の朝に自動投稿される予定です。