LlamaIndex(GPTIndex)を触る

2023-03-04

LlamaIndexを動かしていきます。最近名前がGPTIndexから変わったみたい。

LlamaIndexとは、LLMのPromptの作成などを外部のデータソースと連携して行えるようにしたツールキット。
https://gpt-index.readthedocs.io/en/latest/guides/primer.html
データソースにはGoogle DocsやSlackなどが使えるらしい。

メインの機能は、DocsやSlack、Webの情報を体系化したIndexとして保持して、抽出できる機能でしょう。

OpenAIのAPIを使うことが多いと思いますが、今回はHugging faceのモデルを使ってみます。
LangChainをwrapしている部分がいくつかあるので、LangChainについても知っておくと良いと思う。

Install

doc通り

pip install llama-index

このあたりも使うので入れておく。

pip install transformers sentencepiece torch langchain

とりあえず動かす

大まかな流れは、ドキュメントからIndexを作成、Indexに対してqueryを投げ、抽出する、といった感じ。

ドキュメントからIndexを作成する際に、以下の４つを使う。

llm_predictor
embed_model
prompt_helper
template

llm_predictorの作成

llm_predictorは最終的な文章生成の際に使われる。インデックスの構築、挿入、クエリの走査中にも使用される場合がある。

Huggingfaceにあるmodelを使う例。モデルは好きなものを使ってください。

transformersのpipelineをlangchainのHuggingFacePipelineでwrap、
langchainのHuggingFacePipelineをllama_indexのLLMPredictorでwrapする。

model_name = "facebook/xglm-564M"
tokenizer = XGLMTokenizer.from_pretrained(model_name)
model = XGLMForCausalLM.from_pretrained(model_name)

pipe = pipeline("text2text-generation",
                model=model,
                tokenizer=tokenizer,
                torch_dtype=torch.float16)                
llm = HuggingFacePipeline(pipeline=pipe)

llm_predictor = LLMPredictor(llm=llm)

embed_modelの作成

embed_modelはdocumentをsentense vectorにしindexを作成する際に使われる。　　 (queryによる抽出の方法がいくつかあるが、vectorの類似度から抽出する方法があるので、そこのembeddingでも使ってそう)　　

model_name = "sonoisa/sentence-bert-base-ja-en-mean-tokens"
embeddings = HuggingFaceEmbeddings(model_name = model_name)
embed_model = LangchainEmbedding(embeddings)

prompt helperの作成

tokenの制限に使われる。OpenAIのモデルを使う場合に設定するものだと思う。

prompt_helper=PromptHelper.from_llm_predictor(
    llm_predictor=llm_predictor,
    max_chunk_overlap=0,
    chunk_size_limit=512,
    embedding_limit=1023,
)

promptの作成

日本語のpromptが欲しかったので作成。

QA_PROMPT_TMPL = (
    "以下は文脈情報です。\n"
    "---------------------\n"
    "{context_str}"
    "\n---------------------\n"
    "予備知識ではなく、文脈情報が与えられます。"
    "質問: {query_str}\n"
    "回答:\n"
)
QA_PROMPT = QuestionAnswerPrompt(QA_PROMPT_TMPL)

indexの作成

これまでに作成した４つをあわせてindexを作成

llm_predictor
embed_model
prompt_helper
template

index = GPTSimpleVectorIndex(documents, 
                            llm_predictor=llm_predictor,
                            embed_model=embed_model,
                            prompt_helper=prompt_helper,
                            text_qa_template=QA_PROMPT,
                            )

query

GPTSimpleVectorIndexに渡すtext_qa_templateと同じものを、queryにも渡します。

query_str = "podについて教えてください。"
response = index.query(query_str, text_qa_template=QA_PROMPT)

print("response:", response.response, "\n")
print("source_nodes:", response.source_nodes, "\n")