【LlamaIndex】Indexにクエリした際に回答で参考にした箇所（リファレンス）を取得する方法

山本紘暉

2023.04.09

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

新規事業統括部Passregiチームの山本です。

最近、OpenAIのAPIが公開されたこともあり、効率的に情報を見つけ出す手助けをするための方法として、文章や資料に関する質問に回答できるチャットボットが採用されることが増えてきました。チャットボットが回答する際に、単に直接的な答えだけを返答するのでも良いですが、どこの文章を参考にしたか（=リファレンス）も合わせて回答すれば、ユーザは一次情報を読むことができたり、周辺の情報も合わせて見ることができるので、より安心して利用できそうです。

今回は、LlamaIndexでリファレンスした箇所を把握する方法を調べましたので、そのコードを共有します。

(補足)

文章ファイルの内容に関する質問に回答するチャットボット作成するために、ChatGPTなどのAPIのクエリにその文章の内容を入力として加える方法があります。

その際の課題として以下のようなことが挙げられます。

クエリとして入力できる文字数の制限があり、大量の文章を読ませることができない
質問に関係ない多くの文章を入れてしまうと、回答の精度が低くなる可能性がある

これを回避する方法として、以下のような手法がよく採用されています。

文章ファイルを一定の長さ内に分割し、それぞれをベクトルに変換する（Embedding・Indexing）
ユーザが質問した文章と、文章ファイルとの類似度を計算し、関連性が高い文を採用する（Retriving）
関連性高い文章と質問をあわせて、チャットボットのクエリにする

今回は、Retrivingの段階において、どの（関連性が高い）文章を採用したかの情報を取得します。

リファレンスを取得する方法

準備

Python環境を用意し、必要なライブラリをインストールします
```
pip install llama-index
```
今回はPDFファイルを読み込むため、以下のライブラリもインストールします
```
pip install PyPDF2
```
※ 補足：PDF以外のファイルも読み込むことができます。llama_index/readers/file/base.pyを見る限り、ワードファイルやパワーポイントなどの文章ファイル、画像・音声などのメディアファイルにも対応していそうです（2023/4/10時点）。注意点として、各ファイル形式に対応したライブラリをpipでインストールしておく必要があります。詳細はリファレンスなどで確認してください。（実行時に出てくるエラーメッセージに、必要なライブラリが表示されるので、それを見て対応する形でも良いかと思います）
```
DEFAULT_FILE_EXTRACTOR: Dict[str, BaseParser] = {
    ".pdf": PDFParser(),
    ".docx": DocxParser(),
    ".pptx": PptxParser(),
    ".jpg": ImageParser(),
    ".png": ImageParser(),
    ".jpeg": ImageParser(),
    ".mp3": VideoAudioParser(),
    ".mp4": VideoAudioParser(),
    ".csv": PandasCSVParser(),
    ".epub": EpubParser(),
    ".md": MarkdownParser(),
    ".mbox": MboxParser(),
}
```

OpenAIのAPIキーを作成し、環境変数として設定します（別の方法でも問題ありません）

export OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
(Windowsの場合)
set OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

フォルダ構成

今回はdata/documentsフォルダを作成してPDFをファイルを置き、それらを読み込む処理にしました。

利用したPDFファイルはこちらのものです。

https://github.com/Azure-Samples/azure-search-openai-demo/tree/main/data

以下のような構成にしてください。

test.py (下のコード)
data
  documents
    ---.pdf
    ---.pdf

コード

今回作成したコードは以下のとおりです。得られた回答の情報を、show_response関数で表示しています。

SimpleDirectoryReaderを作成する際に、file_metadataを使うことで、node.node.extra_infoにファイル名をもたせることができます
- このコードでは、filename_fnを指定することで、”file_name”にファイルパスを保持させています
リファレンスの情報は、index.queryの結果に含まれています（response.source_nodes）
- node.node.extra_infoにfile_metadataで指定した情報が含まれており、今回の場合はファイルパスを取り出すことができます
- （下のコードの処理以外にも、nodeから辿ることで必要な情報を取得できると思います）
（途中作成した、data/indexにキャッシュとして保存しています。data/documents以下の同じファイル数とファイル名が同じであれば、保存したindexを再度利用します）

import json
import os
from typing import List

from llama_index import Document, GPTSimpleVectorIndex, ServiceContext, SimpleDirectoryReader
from llama_index.data_structs.node_v2 import DocumentRelationship
from llama_index.response.schema import RESPONSE_TYPE

def print_source_of_response(response: RESPONSE_TYPE, documents: List[Document]):
    for node in response.source_nodes:
        if DocumentRelationship.SOURCE not in node.node.relationships:
            print("no source doc found")
            continue
        source_doc_id = node.node.relationships[DocumentRelationship.SOURCE]
        source_docs = [doc for doc in documents if doc.doc_id == source_doc_id]
        if len(source_docs) == 0:
            print("no source doc found")
            continue
        elif len(source_docs) == 1:
            source_doc = source_docs[0]
            print(source_doc.get_text())
        else:
            print("More than one source doc found")
            continue

def show_response(response: RESPONSE_TYPE):
    print("============================================================================================================")
    print(response)
    print("============================================================================================================")
    for node in response.source_nodes:
        if node.node.extra_info is not None:
            if "file_name" in node.node.extra_info:
                print(node.node.extra_info["file_name"])
        print(node.node.node_info)
        if node.score is not None:
            print(node.score)
        print("============================================================================================================")
        print(node.node.text)
        print("============================================================================================================")

def query_with_index(
    folderpath_documents: str,
    folderpath_index: str,
    filepath_cache_metadata: str,
    filepath_cache_index: str,
    chunk_size_limit: int,
    similarity_top_k: int,
):
    # select docments
    filename_fn = lambda filename: {"file_name": filename}
    directory_reader = SimpleDirectoryReader(
        folderpath_documents,
        file_metadata=filename_fn,
    )
    filenames = [str(input_file) for input_file in directory_reader.input_files]

    # decide whether to use cache
    use_cache = False
    if os.path.exists(filepath_cache_metadata):
        metadata = json.loads(open(filepath_cache_metadata, "r").read())
        filenames_ = metadata["filenames"]
        if sorted(filenames_) == sorted(filenames):
            use_cache = True

    # load index
    if use_cache:
        index = GPTSimpleVectorIndex.load_from_disk(filepath_cache_index)
    else:
        # read documents
        documents = directory_reader.load_data()

        # make index
        service_context = ServiceContext.from_defaults(chunk_size_limit=chunk_size_limit)
        index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)

        # save cache
        os.makedirs(folderpath_index, exist_ok=True)
        index.save_to_disk(filepath_cache_index)
        json.dump({"filenames": filenames}, open(filepath_cache_metadata, "w"), indent=2)

    # query with index
    response = index.query("What is Workplace Violence?", similarity_top_k=similarity_top_k)

    return response

def main():
    # settings for paths
    folderpath_index = os.path.join("data", "indexes")
    folderpath_documents = os.path.join("data", "documents")
    filepath_cache_metadata = os.path.join(folderpath_index, "metadata.json")
    filepath_cache_index = os.path.join(folderpath_index, "index.json")
    chunk_size_limit = 512
    similarity_top_k = 1

    response = query_with_index(
        folderpath_documents,
        folderpath_index,
        filepath_cache_metadata,
        filepath_cache_index,
        chunk_size_limit,
        similarity_top_k,
    )

    show_response(response)

if __name__ == "__main__":
    main()

結果

上のファイルを実行した結果、以下のような出力が得られました。

質問に対する回答
参考にした文章が含まれているファイルパス
参考にした文章がファイル中のどこに有るかのインデックス（文字数）
- （と思われる値。自分が調べた限りでは、これに関する情報がちょっと見つかりませんでした、、）
質問文章に対する、リファレンスの類似度
該当箇所のテキスト

Workplace violence is any act of physical aggression, intimidation, or threat of physical harm toward another individual in the workplace. This includes but is not limited to physical assault, threats of violence, verbal abuse, intimidation, harassment, bullying, stalking, and any other behavior that creates a hostile work environment.
============================================================================================================
data\documents\employee_handbook.pdf
{'start': 11375, 'end': 13475}
0.8670518862058239
============================================================================================================
Workplace Violence Prevention Program.

Purpose

The purpose of this program is to promote a safe and  healthy work environment by
preventing violence, threats, and abuse in the workplace. It is also intended to provide a
safe, secure and protected environment for our employees, customers, and visitors.

Definition of Workplace Violence

Workplace violence  is any act of physical aggression, intimidation, or threat of physical
harm toward another individual in the workplace. This includes but is not limited to
physical assault, threats of violence, verbal abuse, intimidation, harassment, bullying,
stalking, and any other behavior that creates a hostile work environment.

Prevention and Response

Contoso Electronics is committed to preventing workplace violence and will not tolerate
any acts of violence, threats, or abuse in the workplace. All employees are ex pected to
follow the company’s zero tolerance policy for workplace violence.

If an employee believes that they are in danger or are the victim or witness of workplace
violence, they should immediately notify their supervisor or Human Resources
Representat ive. Employees are also encouraged to report any suspicious activity or
behavior to their supervisor or Human Resources Representative.

In the event of an incident of workplace violence, Contoso Electronics will respond
promptly and appropriately. All inc idents will be thoroughly investigated and the
appropriate disciplinary action will be taken.

Training and Education

Contoso Electronics will provide regular training and education to all employees on
workplace violence prevention and response. This trai ning will include information on
recognizing potential signs of workplace violence, strategies for responding to incidents,
and the company’s zero tolerance policy.

We are committed to creating a safe and secure work environment for all of our employees.
By following the guidelines outlined in this program, we can ensure that our workplace is
free from violence and abuse.
Privacy

Privacy Policy

上のテキストボックス中の17行目に回答と同様の文章があり、目的どおり回答の際に使用された箇所を取得できました。こういった情報をユーザに示せば、ユーザ側でも回答が合っているかを確認できそうです。

他にあった方が追加したほうが良さそうな機能

今回は調べられなかったのですが、以下のような機能を追加すると良さそうです

インデックス（文字数）が、ファイルの何ページ目に該当するか・ページ内のどのあたりに表示されているかを調べる

まとめ

LlambaIndexを使用して、文章ファイルに関する質問をした際に、回答に使用した箇所（リファレンス）を取得するコードを作成しました。

今回実装したコードでは、リファレンスの「ファイル名・インデックス（文字数）・類似度・テキスト」を得られました。

他・補足

以下のリポジトリで、Azure Coginitive Searchを使用した、社内文章に関する質問を行えるチャットボットを構築するサンプルが公開されています。こちらでも、リファレンスを含めて答える機能があり、また、Agentを使用することでインタラクティブに質問したり、クエリをより良いものにしてから検索できるようになっているので、参考になりそうです。

https://github.com/Azure-Samples/azure-search-openai-demo