Build RAG system

# Wikipediaの天文関連情報を使ったRAGシステム
#　LLMは、OpenAI（いわゆるChatGPT）を使い、RAGはFAISSとas_retrieverを使う。

WIKI_DB = "../20240813_RAG_MakeDB/wiki_vs.db"
import os
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

model_name = "intfloat/multilingual-e5-large"
model_path = f"/workdir/models/{model_name}"

embeddings = HuggingFaceEmbeddings(
    model_name = model_path,
    model_kwargs = {'device':'cuda:0'},
)

# 事前に構築したベクトルデータベースを読み込む
if os.path.exists(WIKI_DB):
    db = FAISS.load_local(WIKI_DB, embeddings, allow_dangerous_deserialization=True)
else:
    print("="*80)
    print("You need to make vector database")
    print("="*80)

# 検索器を構築
# ベクトル検索で2文書を得る（k=2)
retriever = db.as_retriever(search_kwargs={'k':4})

# OpenAIのAPIキーを設定する
import os

os.environ['OPENAI_API_KEY'] = 'xxxxxxxxxxxxxxxxx'

# LLMを設定
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o")

# テンプレートを準備

template = """
以下の「コンテキスト情報」を元に「質問」に回答してください。
なお、コンテキスト情報に無い情報は回答に含めないでください。
また、コンテキスト情報から回答が導けない場合は「分かりません」と回答してください。

前提条件：{context}

質問：{question} 
"""

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template(template)

# LangChain（LCEL）でChainを構築

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableParallel, RunnablePassthrough

prompt = ChatPromptTemplate.from_template(template)
output_parser = StrOutputParser()
setup_and_retrieval = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
)
chain = setup_and_retrieval | prompt | llm | output_parser

question = "r過程について教えてください"
print(chain.invoke(question))

question = "ナンシー・グレース・ローマン宇宙望遠鏡について教えてください"
print(chain.invoke(question))

question = "B2FH論文について教えてください"
print(chain.invoke(question))

# Wikipediaの天文関連情報を使ったRAGシステム
#　LLMは、Llama-3-ELYZA-JP-8Bを使い、RAGはFAISSとas_retrieverを使う。

# Llama-3-ELYZA-JP-8Bをモデルとして準備

import sys
from langchain_community.llms import LlamaCpp

llm = LlamaCpp(
    model_path="/workdir/models/Llama-3-ELYZA-JP-8B-q4_k_m.gguf", 
    n_gpu_layers=-1,
    n_ctx=2048,
    f16_kv=True,
    verbose=False,
)

#　Wikipediaの天文関連情報を使ったRAGシステム
#　LLMは、OpenAI（いわゆるChatGPT）を使い、RAGは、BM25Retrieverを使う。

from janome.tokenizer import Tokenizer

t = Tokenizer()

def my_preprocess_func(text):
    keywords = []
    for token in t.tokenize(text):
        pos = token.part_of_speech.split(',')[0]
        if (pos in ["名詞", "動詞", "形容詞"]):
            keywords.append(token.surface)
    return keywords

WIKI_DB = "../20240813_RAG_MakeDB/wiki_kw.pkl"
import os
import pickle
from langchain_community.retrievers import BM25Retriever

# 事前に構築したキーワードのデータベースを読み込む
if os.path.exists(WIKI_DB):
    with open(WIKI_DB, 'rb') as f:
        retriever = pickle.load(f)
else:
    print("="*80)
    print("You need to make keyword database")
    print("="*80)

# Wikipediaの天文関連情報を使ったRAGシステム
#　LLMは、Llama-3-ELYZA-JP-8Bを使い、RAGはBM25Retrieverを使う。

# テンプレートを設定
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

template = """
{question}
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", "あなたはユーザの質問に回答する優秀なアシスタントです。以下の質問に可能な限り丁寧に回答してください。"),
    ("user", template)
])

output_parser = StrOutputParser()

# LangChain(LCEL）によりChainを作成
chain = prompt | llm | output_parser