LlamaIndexでClaude3を使ったRAGを試してみる

!pip install llama-index llama-index-llms-anthropic llama-index-callbacks-arize-phoenix

from google.colab import userdata
import os

os.environ["ANTHROPIC_API_KEY"] = userdata.get('ANTHROPIC_API_KEY')
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

import phoenix as px
from llama_index.core import set_global_handler

px.launch_app()
set_global_handler("arize_phoenix")

🌍 To view the Phoenix app in your browser, visit https://XXXXXXXXXX-colab.googleusercontent.com/

!wget https://d.line-scdn.net/stf/linecorp/ja/csr/dataset_.zip
!unzip dataset_.zip

import pandas as pd

df = pd.read_excel("dataset_.xlsx")
df.drop(columns=["ID","カテゴリ2","出典","<参考>UMカテゴリタグ","<参考>ＵＭサービスメニュー\n（標準的な行政サービス名称）"], inplace=True)
df.rename(columns={
    'サンプルID': 'ID',
    'サンプル 問い合わせ文': 'Question',
    'サンプル 応答文': 'Answer',
    'カテゴリ1': 'Category',
}, inplace=True)
df

from llama_index.core.schema import TextNode

nodes = []
for idx, row in df.iterrows():
    id = row["ID"]
    text = row["Answer"].replace("\n\n","\n")
    node = TextNode(text=text, id_=id)
    nodes.append(node)

from llama_index.core import VectorStoreIndex
from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

index = VectorStoreIndex(nodes, embed_model=embed_model, show_progress=True)

from llama_index.core.prompts import PromptTemplate
from llama_index.llms.anthropic import Anthropic
from xml.dom.minidom import Document

# コンテキストをXML形式に変換
def format_context_fn(**kwargs):
    context_str = kwargs.get("context_str", "")
    doc = Document()
    
    # ルート要素 <documents> を作成
    root_element = doc.createElement("documents")
    doc.appendChild(root_element)
    
    for section in context_str.split("\n\n"):
        # <document> 要素を作成
        document_element = doc.createElement("document")
        root_element.appendChild(document_element)
        
        # セクション内の各行を処理
        lines = section.split("\n")
        for line in lines:
            text_node = doc.createTextNode(line)
            document_element.appendChild(text_node)

    # 整形済みのXMLを文字列として返す。XML宣言は省略。
    pretty_xml_as_str = doc.toprettyxml(indent="  ")
    # XML宣言を削除
    return '\n'.join(pretty_xml_as_str.split('\n')[1:])  # 最初の行（XML宣言）を除外

# QAテンプレート
text_qa_prompt = """\
コンテキスト情報は以下です。

{context_str}

与えられたコンテキスト情報を元に、事前知識を使用することなく、質問に答えてください。

質問: {query_str}
回答: \
"""
text_qa_template = PromptTemplate(text_qa_prompt, function_mappings={"context_str": format_context_fn})

llm = Anthropic(model="claude-3-sonnet-20240229", temperature=0.1)

query_engine = index.as_query_engine(llm=llm, similarity_top_k=5, response_mode="compact")
query_engine.update_prompts({"response_synthesizer:text_qa_template": text_qa_template})

response = query_engine.query("母子手帳を受け取りたいのですが、手続きを教えてください。")
print(response)

母子手帳を受け取るための手続きは以下の通りです。

1. 妊娠届を提出する際に、診断を受けた病院名と医師名を記入する必要があります。

2. 妊娠届の内容を確認した上で、その場で母子手帳がお渡しされます。

3. 母子手帳は、○○市役所本庁舎△△階××課窓口、◎◎出張所、その他の場所で受け取ることができます。

母子手帳の受け取り場所の詳細については、自治体のウェブサイトで確認できます。診断書は不要ですが、本人確認できるものを持参する必要があります。