Google Colab で Llama 2 + LangChain の RetrievalQA を試す

# パッケージのインストール
!pip install langchain accelerate bitsandbytes sentence_transformers
!pip install faiss-gpu

# HuggingFaceのログイン
!huggingface-cli login

    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|
    
    To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens .
Token: 
Add token as git credential? (Y/n) n
Token is valid (permission: read).
Your token has been saved to /root/.cache/huggingface/token
Login successful

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

# ドキュメントの読み込み
with open("bocchi.txt", encoding="utf-8") as f:
    test_all = f.read()

# チャンクの分割
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=514,
    chunk_overlap=20,
)
texts = text_splitter.split_text(test_all)

# チャンクの確認
print(len(texts))
for text in texts:
    print(text[:10].replace("\n", "\\n"), ":", len(text))

# インデックスの作成
index = FAISS.from_texts(
    texts=texts,
    embedding=HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large"),
)

41
結束バンド: : 6
後藤ひとりは友達を作 : 513
ことを知る。紆余（う : 232
文化祭ライブ: : 7
夏休みに入り、後藤ひ : 513
る。そしてひとりたち : 64
デモ審査: : 5
未確認ライオットに参 : 513
れる結束バンドだった : 92
後藤ひとり(ごとうひ : 14
秀華高校に通う女子。 : 513
ものだが、他人と合わ : 362
伊地知虹夏(いじちに : 14
下北沢高校に通う女子 : 513
いてきており、最近は : 123
山田リョウ(やまだり : 14
下北沢高校に通う女子 : 513
も及ばず、成績はつね : 151
喜多郁代(きたいくよ : 12
秀華高校に通う女子。 : 513
以上いる。またその社 : 208
廣井きくり(ひろいき : 14
実力派サイケデリック : 513
のシャワーを使ったり : 99
後藤ふたり(ごとうふ : 277
清水イライザ(しみず : 497
岩下志麻(いわしたし : 316
長谷川あくび(はせが : 323
佐々木次子(ささきつ : 200
大槻ヨヨコ(おおつき : 15
人気メタルバンド「S : 508
佐藤愛子(さとうあい : 13
フリーライターとして : 513
に参加させるきっかけ : 123
ひとりの父(ひとりの : 312
2号(にごう): 美 : 413
吉田銀次郎(よしだぎ : 471
PAさん(ぴーえーさ : 221
伊地知星歌(いじちせ : 397
完熟マンゴー仮面(か : 236
ギターヒーロー: 後 : 398

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
import torch

# トークナイザーとモデルの準備
model_id = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float16,
    quantization_config=bnb_config, 
    device_map="auto",
)

from langchain.llms import HuggingFacePipeline

# パイプラインの準備
pipe = pipeline(
    "text-generation",
    model=model, 
    tokenizer=tokenizer,
    max_new_tokens=256
)

# LLMの準備
llm = HuggingFacePipeline(pipeline=pipe)

from langchain.chains import RetrievalQA

# 質問応答チェーンの作成
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=index.as_retriever(search_kwargs={"k": 3}),
)

qa_chain.run("日本語で回答してください。後藤ひとりの得意な楽器は？")

後藤ひとりの得意な楽器はギターです。

# 入力
inputs = [
    "後藤ひとりの得意な楽器は？",
    "後藤ひとりの妹の名前は？",
    "後藤ひとりが加入したバンド名は？",
    "ギターヒーローの正体は？",
    "喜多郁代の髪の色は？",
    "伊地知虹夏が通う学校の名前は？",
    "山田リョウの趣味は？",
    "廣井きくりが所属するバンド名は？",
    "ライブハウス「STARRY」の店長の名前は？",
    "ぼっちちゃんが文化祭で披露した演奏法は？",
]

# チェーンの実行
for input in inputs:
    print(qa_chain.run("日本語で回答してください。" + input))