Google Colab で Llama 2 + LlamaIndex の QA を試す

# パッケージのインストール
!pip install llama-index
!pip install transformers accelerate bitsandbytes sentence_transformers

# 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

import logging
import sys

# ログレベルの設定
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)

from llama_index import SimpleDirectoryReader

# ドキュメントの読み込み
documents = SimpleDirectoryReader("data").load_data()

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
from langchain.llms import HuggingFacePipeline
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",
)

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

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

from langchain.embeddings import HuggingFaceEmbeddings
from llama_index import LangchainEmbedding

# 埋め込みモデルの準備
embed_model = LangchainEmbedding(
    HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
)

from langchain.text_splitter import RecursiveCharacterTextSplitter
from llama_index.node_parser import SimpleNodeParser

# チャンクの分割
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=514,  # チャンクの最大文字数
    chunk_overlap=20,  # オーバーラップの最大文字数
)

# ノードパーサーの準備
node_parser = SimpleNodeParser(text_splitter=text_splitter)

from llama_index import ServiceContext

# サービスコンテキストの準備
service_context = ServiceContext.from_defaults(
    llm=llm,
    embed_model=embed_model,
    node_parser=node_parser,
)

from llama_index import VectorStoreIndex

# インデックスの作成
index = VectorStoreIndex.from_documents(
    documents,
    service_context=service_context,
)

Adding chunk: 結束バンド...
Adding chunk: 後藤ひとりは友達を作れない陰キャでいつも一人で過ごしていたが、...
Adding chunk: ことを知る。紆余（うよ）曲折の末、...
Adding chunk: 文化祭ライブ...
Adding chunk: 夏休みに入り、...
Adding chunk: る。そしてひとりたちは、...
Adding chunk: デモ審査...
Adding chunk: 未確認ライオットに参加するためには、...
Adding chunk: れる結束バンドだったが、...
Adding chunk: 後藤ひとり(ごとうひとり)...
Adding chunk: 秀華高校に通う女子。...
Adding chunk: ものだが、他人と合わせるセッションの経験が皆無なため、...
Adding chunk: 伊地知虹夏(いじちにじか)...
Adding chunk: 下北沢高校に通う女子。...
Adding chunk: いてきており、...
Adding chunk: 山田リョウ(やまだりょう)...
Adding chunk: 下北沢高校に通う女子。...
Adding chunk: も及ばず、成績はつねに赤点。...
Adding chunk: 喜多郁代(きたいくよ)...
Adding chunk: 秀華高校に通う女子。...
Adding chunk: 以上いる。...
Adding chunk: 廣井きくり(ひろいきくり)...
Adding chunk: 実力派サイケデリックロックバンド「SICKHACK(シックハック)」に所属する女性。...
Adding chunk: のシャワーを使ったり、タダ飯にありつこうとしているため、...
Adding chunk: 後藤ふたり(ごとうふたり)...
Adding chunk: 清水イライザ(しみずいらいざ)...
Adding chunk: 岩下志麻(いわしたしま)...
Adding chunk: 長谷川あくび(はせがわあくび)...
Adding chunk: 佐々木次子(ささきつぐこ)...
Adding chunk: 大槻ヨヨコ(おおつきよよこ)...
Adding chunk: 人気メタルバンド「SIDEROS(シデロス)」のリーダーを務める少女。...
Adding chunk: 佐藤愛子(さとうあいこ)...
Adding chunk: フリーライターとして活動する女性。...
Adding chunk: に参加させるきっかけをつくった。...
Adding chunk: ひとりの父(ひとりのちち)...
Adding chunk: 2号(にごう)...
Adding chunk: 吉田銀次郎(よしだぎんじろう)...
Adding chunk: PAさん(ぴーえーさん)...
Adding chunk: 伊地知星歌(いじちせいか)...
Adding chunk: 完熟マンゴー仮面(かんじゅくまんごーかめん)...
Adding chunk: ギターヒーロー...

from llama_index.prompts.prompts import QuestionAnswerPrompt

# QAテンプレートの準備
qa_template = QuestionAnswerPrompt("""Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

{context_str}

Question: {query_str}
Helpful Answer:""")

# クエリエンジンの作成
query_engine = index.as_query_engine(
    similarity_top_k=3,
    text_qa_template=qa_template,
)

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

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

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

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

Context information is below.
---------------------
{context_str}
---------------------
Given the context information and not prior knowledge, answer the question: {query_str}

Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

{context_str}

Question: {query_str}
Helpful Answer: