LlamaIndexでAutoGPTQモデルを使う（vicuna-13B-v1.5-GPTQ）
npakaさんの記事を見て、LlamaIndexでもAutoGPTQモデルを使いたいと思い、試してみました。ポイントは下記2点です。
- GPTQモデルをLlamaIndexに渡す
-
text_splitterをトークン数で分割するよう設定する
環境
Google Colab 無料版のT4インスタンス（VRAM 15GB）で動作確認しています。
前準備
モデルの指定
vicuna-13B, multilingual-e5-baseの組み合わせで、VRAM使用量は11GB～15GB程度でした。
vicuna-7BならローカルPCのRTX3060 12GBでも動作しました。
vicuna-7B + multilingual-e5-smallなら8GB以下でも推論できましたが、複数回質問していると8GB以上になることもありました。トークン数を絞れば安定するのかも知れません。
llm_model_name = "TheBloke/vicuna-13B-v1.5-GPTQ"
embed_model_name = "intfloat/multilingual-e5-base"
multilingual-e5はtext-embedding-ada-002に近い性能が出るので、baseにしてもそれほど問題は無さそうです。largeに至ってはadaを超えています。
環境の確認
!nvidia-smi
!python -V
Sun Aug 6 02:36:47 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17 Driver Version: 525.105.17 CUDA Version: 12.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 69C P0 31W / 70W | 10473MiB / 15360MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
+-----------------------------------------------------------------------------+
Python 3.10.12
Google Driveをマウント
from google.colab import drive
drive.mount('/content/drive')
google_drive = "drive/MyDrive/data"
!mkdir {google_drive}
必要なモジュールをインストール
auto-gptqはビルドに時間がかかるのでwhlを直接指定しています。
環境によって書き換えてください。
!pip install llama-index sentence_transformers https://github.com/PanQiWei/AutoGPTQ/releases/download/v0.3.2/auto_gptq-0.3.2+cu118-cp310-cp310-linux_x86_64.whl
主要なモジュールのバージョン
!pip list | egrep "auto-gptq|langchain|llama-index|sentence-transformers|torch|transformers"
auto-gptq 0.3.2+cu118
langchain 0.0.253
llama-index 0.7.19
sentence-transformers 2.2.2
torch 2.0.1+cu118
torchaudio 2.0.2+cu118
torchdata 0.6.1
torchsummary 1.5.1
torchtext 0.15.2
torchvision 0.15.2+cu118
transformers 4.31.0
各インスタンスの初期化
LLMの初期化
モデルの読み込みについてはモデルカードに記載があります。
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
from transformers import pipeline
from langchain.llms import HuggingFacePipeline
# トークナイザーの初期化
tokenizer = AutoTokenizer.from_pretrained(
llm_model_name,
use_fast=True,
)
# LLMの読み込み
model = AutoGPTQForCausalLM.from_quantized(
llm_model_name,
device_map="auto",
use_safetensors=True,
)
# パイプラインの作成
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=512,
# 指定すると英語で回答する率が上がってしまう
# repetition_penalty=1.15,
)
# LLMの初期化
llm = HuggingFacePipeline(pipeline=pipe)
AutoGPTQForCausalLMは、transformers.pipeline > HuggingFacePipeline > service_context の順でLlamaIndexに渡すことができます。
注意点として、モデルカードのサンプルにあるrepetition_penalty=1.15を指定すると英語で回答する確率が上がってしまいます。プロンプトを工夫すれば指定しても大丈夫かも知れません。
埋め込みモデルの初期化
multilingual-e5は、クエリの先頭に"query: "を付け足すと精度が上がるそうです。
【参考】npakaさんの記事（note）
from langchain.embeddings import HuggingFaceEmbeddings
from llama_index import LangchainEmbedding
from typing import Any, List
# query付きのHuggingFaceEmbeddings
class HuggingFaceQueryEmbeddings(HuggingFaceEmbeddings):
def __init__(self, **kwargs: Any):
super().__init__(**kwargs)
def embed_documents(self, texts: List[str]) -> List[List[float]]:
return super().embed_documents(["query: " + text for text in texts])
def embed_query(self, text: str) -> List[float]:
return super().embed_query("query: " + text)
# 埋め込みモデルの初期化
embed_model = LangchainEmbedding(
HuggingFaceQueryEmbeddings(model_name=embed_model_name)
)
モデルカードのサンプルコードにも記載があります。
# Each input text should start with "query: " or "passage: ", even for non-English texts.
# For tasks other than retrieval, you can simply use the "query: " prefix.
ノードパーサーの初期化
multilingal-e5の最大入力トークン数は512なので、チャンクはそれ以下に分割する必要があります。さらに、クエリの先頭に"query: "という文字列を追加しているので、その分のトークンを引いておく必要があります。
TextSplitter.from_huggingface_tokenizerを使うと、Tokenizerのトークン数でカウントしてくれます（デフォルトは文字数）[1]。
# 埋め込み用トークナイザーを初期化
# from transformers import AutoTokenizer
# embed_tokenizer = AutoTokenizer.from_pretrained(
# embed_model_name,
# use_fast=True
# )
# special tokenの確認
print(tokenizer.all_special_tokens)
print(tokenizer.all_special_ids)
# ['<s>', '</s>', '<unk>', '<unk>']
# [1, 2, 0, 0]
# "query: "のトークン数を確かめる
tokenizer.encode("query: ")
# [1, 2346, 29901, 29871] ※最初は <s> というspecial tokenなので実質3トークン
from langchain.text_splitter import RecursiveCharacterTextSplitter
from llama_index.node_parser import SimpleNodeParser
text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
tokenizer,
chunk_size=512-3, # チャンクの最大トークン数（実際はTokenizerが違うのでここまでする意味はない）
chunk_overlap=20, # オーバーラップの最大トークン数
separators=["\n= ", "\n== ", "\n=== ", "\n\n", "\n", "。", "「", "」", "！", "？", "、", "『", "』", "(", ")"," ", ""],,
)
node_parser = SimpleNodeParser(text_splitter=text_splitter)
WikipediaReaderで読み込んだドキュメントは、"== 見出し ==" のような形式で見出しがつくので、それに合わせてseparators=を設定しています。
なお、e5-base-multilingual-4096という、4Kコンテキスト長のモデルも存在します。
これとvicuna-13B-v1.5-16Kを組み合わせれば、巨大なコンテキスト長も扱うことができます[2]。
サービスコンテキストの初期化
from llama_index import ServiceContext
service_context = ServiceContext.from_defaults(
llm=llm,
embed_model=embed_model,
node_parser=node_parser,
)
インデックスの作成
ドキュメントの読み込み
pages=に記事タイトルのリストを渡せば複数のページを読み込むこともできます。
from llama_index import download_loader
WikipediaReader = download_loader("WikipediaReader")
loader = WikipediaReader()
documents = loader.load_data(pages=["ステンレス鋼"], lang="ja")
インデックスの作成
from llama_index import VectorStoreIndex
index = VectorStoreIndex.from_documents(
documents,
service_context=service_context,
)
質問応答
クエリエンジンの作成
query_engine = index.as_query_engine(
similarity_top_k=3
)
クエリ用関数
torch.cuda.empty_cache()で、クエリのたびになるべくVRAMを開放しています。
logging.basicConfig(stream=sys.stdout, level=logging.WARNING, force=True)
import torch
def query(question):
print(f"Q: {question}")
response = query_engine.query(question).response.strip()
print(f"A: {response}\n")
torch.cuda.empty_cache()
質問
questions = [
"ステンレス鋼の定義を教えてください。",
"ステンレス鋼が発明されたのはいつですか？",
"ステンレス鋼のスクラップの回収率はどのくらいですか？",
"ステンレス鋼を切断する方法にはどんなものがありますか？",
"ステンレス鋼をめっきするとどんなメリットがありますか？",
"オーステナイトは強磁性体ですか？",
"ステンレス製の流し台のメリットを教えてください。"
]
for question in questions:
query(question)
回答
かなり正しく回答できています。驚くべき性能です。
今回は「日本語で回答してください。」なしでも、日本語で回答してくれています。
Q: ステンレス鋼の定義を教えてください。
A: ステンレス鋼とは、鉄にクロムが一定量以上添加された錆びにくい合金の一種で、高合金鋼または特殊鋼に位置づけられます。現在の国際的な定義では、ステンレス鋼は「クロム含有量が 10.5 % 以上、炭素含有量が 1.2 % 以下の合金鋼」とされています。
Q: ステンレス鋼が発明されたのはいつですか？
A: The answer is: ステンレス鋼が発明されたのは、20世紀初頭の1910年代のことである。
Q: ステンレス鋼のスクラップの回収率はどのくらいですか？
A: 答え: ステンレス鋼のスクラップの回収率は、その内、その原料の約 60 % がステンレス鋼スクラップを利用できているという調査結果があることから、高いと考えられます。
Q: ステンレス鋼を切断する方法にはどんなものがありますか？
A: 答え: ステンレス鋼を切断する方法には、溶断、せん断、ウォータージェット切断、アーク切断、プラズマ切断、レーザー切断があります。
Q: ステンレス鋼をめっきするとどんなメリットがありますか？
A: 答え: ステンレス鋼をめっきすることで、耐食性、装飾性、導電性の向上が期待できます。めっきの方法としては、電気めっき、溶融めっき、ガス還元法による溶融めっきなどがあります。ステンレス鋼に施工する際には、めっきの前処理が必要となることがあります。また、ステンレス鋼と競合する他材料としては、塗装・めっき・ホーローなどの表面処理を施した鋼、ポリプロピレンのような樹脂材料、アルミニウムやチタンなどの他金属材料があります。
Q: オーステナイトは強磁性体ですか？
A: オーステナイトは、一般的な鉄鋼材料とは異なり、常磁性材料であることが特徴です。強磁場中でもごくわずかにしか磁化しないため、オーステナイト系は非磁性材料です。ただし、オーステナイト系は、加工誘起マルテンサイト変態が起こると磁性を帯びるようになることがあります。
Q: ステンレス製の流し台のメリットを教えてください。
A: ステンレス製の流し台のメリットは、耐久性があり、メンテナンスしやすいことです。ステンレス製流し台本体は、板材からプレス成形で造られる。台所の天板でも、ステンレス鋼が選択肢の一つで、エンボス仕上げや着色処理による外観を良くしたものも採用されている。また、鍋やフライパンなどでもステンレス製が使われている。ただし、ステンレス鋼は熱伝導があまりよくないので、ステンレス鋼でアルミを挟み込んだ三層構造クラッド鋼などにして対策される。IH調理器用には、磁性のあるフェライト系や普通鋼と複合させた、ステンレスクラッド鋼が使われる。業務用の厨房は、流し台、テーブル、ケース類に至るまで、清潔さを保つために清浄しやすいステンレス鋼が全面的に使われている。魔法瓶の水筒もステンレス鋼を使った製品で、ステンレス鋼管のプレス成形で造られる。魔法瓶水筒の場合は、ステンレス鋼の熱伝導の悪さを逆に有効活用している事例といえる。
ストレージコンテキストの保存／読み込み
うまくいった場合は保存しておきます。
保存
import os
storage_dir = os.path.join(google_drive, "wiki_stainless")
storage_context = index.storage_context.persist(storage_dir)
読み込み
保存したインデックスを読み込む場合。
service_contextを忘れずに指定しましょう。
from llama_index import StorageContext, load_index_from_storage
storage_dir = os.path.join(google_drive, "wiki_stainless")
storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
index = load_index_from_storage(
storage_context,
service_context=service_context
)
参考リンク
-
EmbeddingモデルのTokenizerを使用したほうが良さそうなのですが、何故かLLMのTokenizerを使ったほうが結果が良かったので一旦そのままにしています。
Embedding用のTokenizerで分割したチャンクをLLMのTokenizerでトークン化すると1000近くになったので、与えるコンテキストが長すぎることが原因のような気がします。LLMのトークン数で言うと512、Embeddingなら300（× similarity_top_k=3）程度が今のところ妥当なようです。
Vicunaは日本語の長いコンテキストが苦手なのかも知れません。 ↩︎ -
残念ながら今のところうまく行っていません。これもVicunaが長い日本語のコンテキストをうまく解釈できないことが原因のような気がします。
長いコンテキストを扱うときは一度英語に変換したほうがよいかも知れません。 ↩︎
Discussion