LangChainとVertex AI PaLMを用いた大規模ドキュメントでのスケーラブルな質問応答

!pip install google-cloud-aiplatform langchain==0.0.323 chromadb==0.3.26 pypdf

from langchain.document_loaders import PyPDFLoader  
from langchain.llms import VertexAI
from langchain.chains.question_answering import load_qa_chain
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import VertexAIEmbeddings

vertex_llm_text = VertexAI(model_name="text-bison@001")
vertex_embeddings = VertexAIEmbeddings(model_name="textembedding-gecko@001")  

pdf_url = "https://services.google.com/fh/files/misc/practitioners_guide_to_mlops_whitepaper.pdf"
pdf_loader = PyPDFLoader(pdf_file)
pages = pdf_loader.load_and_split()

prompt_template = """提供されたコンテキストを使用して、できるだけ正確に質問に回答してください。
コンテキストに回答が含まれていない場合は、「コンテキストに回答はありません」と言ってください。\n\n
コンテキスト:\n{context}?\n
質問:\n{question}\n
回答:
"""
 
prompt = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

stuff_chain = load_qa_chain(vertex_llm_text, chain_type="stuff", prompt=prompt)

question = "実験とは何ですか？"
context = "\n".join(str(p.page_content) for p in pages[:7])
stuff_answer = stuff_chain(
    {"input_documents": pages[7:10], "question": question}, return_only_outputs=True
)

try:
    print(stuff_chain(
        {"input_documents": pages[7:], "question": question}, 
        return_only_outputs=True))
except Exception as e:  
    print("The code failed since it won't be able to run inference on such a huge context")

question_prompt_template = """
提供されたコンテキストを使用して、できるだけ正確に質問に回答してください。\n\n
コンテキスト:\n{context}\n
質問:\n{question}\n  
回答:
"""
question_prompt = PromptTemplate(
    template=question_prompt_template, input_variables=["context", "question"]
)
 
combine_prompt_template = """抽出されたコンテンツと質問を考慮して、最終的な回答を作成してください。
コンテキストに回答が含まれていない場合は、「コンテキストに回答はありません」と言ってください。\n\n
サマリー:\n{summaries}?\n
質問:\n{question}\n
回答:  
"""
combine_prompt = PromptTemplate(
    template=combine_prompt_template, input_variables=["summaries", "question"]
)

map_reduce_chain = load_qa_chain(
    vertex_llm_text, 
    chain_type="map_reduce",
    return_intermediate_steps=True,
    question_prompt=question_prompt,
    combine_prompt=combine_prompt,
)

map_reduce_outputs = map_reduce_chain({"input_documents": pages, "question": question})

for doc, out in zip(
    map_reduce_outputs["input_documents"], map_reduce_outputs["intermediate_steps"]
):
    print(f"Page: {doc.metadata['page']}")
    print(f"Answer: {out}")

refine_prompt_template = """
オリジナルの質問は：\n{question}\n\n
提供された回答は：\n{existing_answer}\n\n
次のコンテキストで必要に応じて既存の回答を洗練してください：\n{context_str}\n\n
取得したコンテンツと質問を考慮して、最終的な回答を作成してください。
コンテキストに回答が含まれていない場合は、"コンテキストに回答が利用できません"と答えてください。\n\n
"""
refine_prompt = PromptTemplate(
    input_variables=["question", "existing_answer", "context_str"],
    template=refine_prompt_template,
)

refine_chain = load_qa_chain(
    vertex_llm_text,
    chain_type="refine",
    return_intermediate_steps=True,
    question_prompt=initial_question_prompt,
    refine_prompt=refine_prompt,
)

refine_outputs = refine_chain({"input_documents": pages, "question": question})

for doc, out in zip(
    refine_outputs["input_documents"], refine_outputs["intermediate_steps"]
):
    print(f"Page: {doc.metadata['page']}")
    print(f"Answer: {out}")

vector_index = Chroma.from_documents(pages, vertex_embeddings).as_retriever()

docs = vector_index.get_relevant_documents(question)

map_reduce_embeddings_outputs = map_reduce_chain(
    {"input_documents": docs, "question": question}
)
print(map_reduce_embeddings_outputs["output_text"])