LangChainでRAGを試す｜CALM2 + FAISS + RetrievalQAの使用例

!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7
!pip install langchain sentence_transformers
!pip install faiss-gpu 

import requests

def get_wikipedia_page(keyword):
  url = "https://ja.wikipedia.org/w/api.php"
  params = {
      "action": "query",
      "list":"search",
      "srsearch": keyword,
      "srlimit":1,
      # "srnamespace": 1,
      "format": "json",
  }

  response = requests.get(url, params=params)
  data = response.json()
  id = data['query']['search'][0]['pageid']

  params = {
    "action": "query",
    "format": "json",
    "pageids": id,
    "prop": "extracts",
    "exlimit": 1,
    "explaintext": True,
    "redirects": 1,
  }

  response = requests.get(url, params=params)
  data = response.json()
  page = next(iter(data['query']['pages'].values()))

  return page.get("extract", "")

wiki_txt = get_wikipedia_page("ワンパンマン")
len(wiki_txt), wiki_txt[:1000]

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer


model = AutoModelForCausalLM.from_pretrained("cyberagent/calm2-7b-chat", 
                                             device_map="auto", 
                                             torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained("cyberagent/calm2-7b-chat")
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

model.eval()

from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.faiss import FAISS

# CALM2は32000トークンまで対応しているので、5000文字と少し大きめのサイズで区切っている
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=5000,
    chunk_overlap=500,
)
texts = text_splitter.split_text(wiki_txt)

embedding=HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")

index = FAISS.from_texts(
    texts=texts,
    embedding=embedding
)
index.save_local("storage")

txt = "ワンパンマンとはなんですか？"
docs = index.similarity_search(txt)

txt = "アクエリオンのあらすじを教えてください"

docs = index.similarity_search(txt)[0]

prompt = f"""USER: あなたは以下の情報に詳しいアシスタントです。わかりやすい日本語で答えてください。
以下の情報が質問に関連していない場合は、情報を利用せずに質問に答えてください。

情報：
{docs.page_content}

質問: {txt}
ASSISTANT:"""

input_ids = tokenizer.encode(prompt, return_tensors="pt")

print(f"\n文字数={len(prompt)}, トークン数={len(input_ids[0])}")

print("アシスタント：", end=" ")
with torch.no_grad():
  tokens = model.generate(
    input_ids = input_ids.to(model.device),
    max_new_tokens=1000,
    do_sample=True,
    temperature=0.7,
    top_k = 50, 
    streamer=streamer,
  )

from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema import StrOutputParser
from langchain import PromptTemplate
from transformers import pipeline


pipe = pipeline(
    'text-generation',
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=1000,
    do_sample=True,
    temperature=0.7,
    top_k = 50, 
)
local_llm = HuggingFacePipeline(pipeline=pipe)

base_template = """USER: あなたは以下の情報に詳しいアシスタントです。わかりやすい日本語で答えてください。
以下の情報が質問に関連していない場合は、情報を利用せずに質問に答えてください。

情報：
{context}

質問: {question}
ASSISTANT:"""

prompt = PromptTemplate.from_template(base_template)

qa_chain = RetrievalQA.from_chain_type(
    llm=local_llm,
    chain_type="stuff",
    chain_type_kwargs={"prompt": prompt},
    retriever=index.as_retriever(search_kwargs={"k": 2}),
    verbose=True,
)
res = qa_chain.invoke('アクエリオンのアニメシリーズはいくつありますか？')
print(res['result'])

res = qa_chain.invoke('ワンパンマンの敵のランクを教えてください')
print(res['result'])

res = qa_chain.invoke('無限パンチとはなんですか？')
print(res['result'])