RetrievalQA.py

# MIT License
# This code will run on VRAM 12GB+ GPU such as T4, RTX 3060
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from langchain import PromptTemplate


model_id = "elyza/ELYZA-japanese-Llama-2-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたはプロの研究者です。あなたが得意な専門分野に関する文章を正確に理解し、答えることに努めてください。"
text = "これから与えられる英文で記述された論文の一部をよく読み、ユーザーからの質問に日本語で答えてください。まず、ユーザからの質問はこちらです。\n{question}\n英文で記述された論文の一部はこちらです。\n\n{context}\n\n次のユーザからの質問に日本語で答えてください。{question}"
template = "{bos_token}{b_inst} {system}{prompt} {e_inst} ".format(
    bos_token=tokenizer.bos_token,
    b_inst=B_INST,
    system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
    prompt=text,
    e_inst=E_INST,
)

loader = PyPDFLoader("/path/to/pdf")
documents = loader.load()

print(len(documents))
document=""
for doc in documents:
    document+=doc.page_content
text=document.replace("\n","")

custom_text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 300,
    chunk_overlap  = 30,
    length_function = len,
    separators = ['.']
)

texts=custom_text_splitter.split_text(text)
texts= custom_text_splitter.create_documents(texts)
print(len(texts))
embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
db = FAISS.from_documents(texts, embeddings)
retriever = db.as_retriever(search_kwargs={"k": 3})
model = AutoModelForCausalLM.from_pretrained(model_id,load_in_8bit=True)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
PROMPT = PromptTemplate(template=template,input_variables=["question","context"],template_format="f-string")
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(
    llm=HuggingFacePipeline(pipeline=pipe),
    retriever=retriever,
    chain_type="stuff",
    return_source_documents=False,
    chain_type_kwargs=chain_type_kwargs)

result = qa.run("この論文の学術的貢献を要約しなさい。")
print(result)

例
"以下の英文をよく読み、ユーザーからの質問に日本語で答えてください。

...

ユーザー: この論文の学術的貢献を要約しなさい。
システム: この論文は、Relation Inversionという新しい概念を提案し、その概念が既存の文法構造に適用できることを示しています。"
https://arxiv.org/abs/2303.13495

例
"以下の英文で記述された論文の一部をよく読み、ユーザーからの質問に日本語で答えてください。
...
ユーザからの質問は次のとおりです。この論文の学術的貢献を要約しなさい。
学術的貢献は以下の通りです。

• 新しい問題、関係逆転問題を提案し、関係の強化学習の新しい方向を示しました。
• ReVerseフレームワークを提案し、関係の強化学習の効率性を向上させました。
• 関係の強化学習の問題解決能力を評価する新しい評価指標を提案しました。"
  https://arxiv.org/abs/2303.13495

コンテキスト長に合わせてパラメータ調整しました。
"以下の英文で記述された論文の一部をよく読み、ユーザーからの質問に日本語で答えてください。
...
ユーザからの質問は次のとおりです。この論文の学術的貢献を要約しなさい。
この論文では、関係構造を学習するための新しいディープラーニングモデルを提案しています。このモデルは、関係構造を考慮したイメージとテキストのデータを用いて、関係構造を学習することができます。このモデルは、関係構造を学習することで、イメージとテキストのデータの関係性を説明することができるようになります。"
https://arxiv.org/abs/2303.13495

検索が正常にできていなかったため、修正したところ正常に動きました。
"あなたはプロの研究者です。あなたが得意な専門分野に関する文章を正確に理解し、答えることに努めてください。
これから与えられる英文で記述された論文の一部をよく読み、ユーザーからの質問に日本語で答えてください。まず、ユーザからの質問はこちらです。
この論文の学術的貢献を要約しなさい。
英文で記述された論文の一部はこちらです。
....
次のユーザからの質問に日本語で答えてください。この論文の学術的貢献を要約しなさい。

与えられた英文で記述された論文の一部をよく読み、ユーザーからの質問に回答いたします。

論文の学術的貢献を要約すると以下の通りです。

この論文では、深層学習の新しいアーキテクチャーであるTransformerを提案しました。Transformerは、従来のRNNやCNNと比較してパラメーター数が少なく、トレーニング時間が短いという特徴があります。このアーキテクチャーは、自然言語処理の他、画像や音声の処理にも応用することができます。

また、この論文では、Transformerのパラメーターを最適化するための新しいアルゴリズムを提案しています。このアルゴリズムは、トレーニング時間と正確性のバランスを取ることができるようになっています。"
https://arxiv.org/abs/2212.04089