LangChain でより高い vector 検索精度が期待できる HyDE 仮説をやってみる

# OpenAI APIs
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
# LLMChain, HyDE
from langchain.chains import LLMChain, HypotheticalDocumentEmbedder
# カスタムテンプレート用(Chat特化)
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import (
    # メッセージテンプレート
    ChatPromptTemplate,
    # System メッセージテンプレート
    SystemMessagePromptTemplate,
    # assistant メッセージテンプレート
    AIMessagePromptTemplate,
    # user メッセージテンプレート
    HumanMessagePromptTemplate,
)
from langchain.schema import (
    # それぞれ GPT-3.5-turbo API の assistant, user, system role に対応
    AIMessage,
    HumanMessage,
    SystemMessage
)

# env に読み込ませるAPIキーの類
import key
# 環境変数にAPIキーを設定
import os
os.environ["OPENAI_API_KEY"] = key.OPEN_API_KEY

# 基本的なオブジェクト生成
base_embeddings = OpenAIEmbeddings()
llm = ChatOpenAI()

# HyDE モジュールに同梱されている web_search prompt を使って HyDE を作る
embeddings = HypotheticalDocumentEmbedder.from_llm(llm, base_embeddings, "web_search")

# 埋め込み query
result = embeddings.embed_query("Where is the Taj Mahal?")

print(result)
#[-0.008099346421658993, 0.003965758252888918, ...... ,-0.028891606256365776]

# 複数の埋め込みを結合（精度が上がる？）
multi_llm = ChatOpenAI(n=4, best_of=4)
embeddings = HypotheticalDocumentEmbedder.from_llm(llm, base_embeddings, "web_search")
result = embeddings.embed_query("Where is the Taj Mahal?")
# print(result)

# 日本語で回答するように独自prompt を使用する LLMChainを作成
prompt_template = "Please answer questions from users in Japanese like a management consultant."
user_template = """Question: {question}
Answer:"""
system_message_prompt = SystemMessagePromptTemplate.from_template(prompt_template)
human_message_prompt = HumanMessagePromptTemplate.from_template(user_template)

chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
chat_prompt.input_variables = ["question"]

# カスタムプロンプトを使ったLLMChain から HyDE を生成
llm_chain = LLMChain(llm=llm, prompt=chat_prompt)
embeddings = HypotheticalDocumentEmbedder(llm_chain=llm_chain, base_embeddings=base_embeddings)

# 日本語でembedding してみる
result = embeddings.embed_query("経営戦略の重要な判断基準を3つ教えてください。")

# Vector 格納 / FAISS
from langchain.vectorstores import FAISS
# テキストファイルを読み込む
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter

loader = TextLoader('reiwa4-1.txt')
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=0, separator="\n")
docs = text_splitter.split_documents(documents)

# HyDE の embeddings を使って vector indexを作成
db = FAISS.from_documents(docs, embeddings)

query = "A社の主力商品は？"

docs = db.similarity_search(query)

print(docs[0].page_content)

from langchain.chains.question_answering import load_qa_chain
chain = load_qa_chain(ChatOpenAI(), chain_type="stuff")
print("----------------")
print(chain({"input_documents": docs, "question": query}, return_only_outputs=True))