LlamaIndexのデータ形式に関するまとめ

本記事では，LlamaIndexでRAGを使用する際の，データ構造に関する情報をまとめます．対象読者は以下の通りです．
* Retriever，Generator，ベクトル検索などのRAGの基礎がなんとなくわかっている方
* Llama indexのチュートリアルをとりあえずやったがデータ形式がいまいちわからない方
* 独自のRAGシステムを作成する際，情報源ファイルに適したRetrieverを作成したい方

# RAGの情報源になるまでの流れ
わかりやすい図を見つけたので，引用させていただきます．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2718862/45f8cc23-e2ae-409c-8331-6e7032255821.png)

引用元: https://medium.com/@mojo885/llamaindex-saddle-up-and-control-your-llama-5131dde3a1fb


1. PDF，Word，Markdownなどさまざまな形式のテキストファイル
1. Document：さまざまなファイルからテキストとメタデータを抽出した形式
1. Node：Documentを細かくした形式．チャンクと同義
1. Embedding：Nodeを埋め込みモデルでベクトル化した形式
1. Index：Nodeの集合体

本記事では，これらのデータがどのように変換されていくのかをまとめていきます．
なお，PDFファイルを情報源とする場合を例に挙げます．

# データ形式
## Document
Documentオブジェクトは，PDFやWordファイルなどのから作成されるオブジェクトであり，以下の要素で構成される．
* テキスト（`text`）：PDFファイルなどから抽出されたテキスト
* ドキュメントID（`doc_id`）：一つのPDFファイルから多くのドキュメントが作成されるため，その識別子
* メタデータ（`extra_info`）：PDFファイルのタイトルやページ数，パスなどのメタデータが保存される

つまり，データ検索を行う際にはメタデータをテキスト共に使用することが多くある．特に参照元を表示することの多いRAGでは，メタデータとテキストをまとめて管理することに大きな意味があるのだろう．

該当するクラスをソースコードより抽出する．
```python
class Document(Node):
    """Generic interface for a data document.

    This document connects to data sources.
    """

    def __init__(self, **data: Any) -> None:
        """Keeps backward compatibility with old 'Document' versions.

        If 'text' was passed, store it in 'text_resource'.
        If 'doc_id' was passed, store it in 'id_'.
        If 'extra_info' was passed, store it in 'metadata'.
        """
        if "doc_id" in data:
            value = data.pop("doc_id")
            if "id_" in data:
                msg = "'doc_id' is deprecated and 'id_' will be used instead"
                logging.warning(msg)
            else:
                data["id_"] = value

        if "extra_info" in data:
            value = data.pop("extra_info")
            if "metadata" in data:
                msg = "'extra_info' is deprecated and 'metadata' will be used instead"
                logging.warning(msg)
            else:
                data["metadata"] = value

        if "text" in data:
            text = data.pop("text")
            if "text_resource" in data:
                text_resource = (
                    data["text_resource"]
                    if isinstance(data["text_resource"], MediaResource)
                    else MediaResource.model_validate(data["text_resource"])
                )
                if (text_resource.text or "").strip() != text.strip():
                    msg = (
                        "'text' is deprecated and 'text_resource' will be used instead"
                    )
                    logging.warning(msg)
            else:
                data["text_resource"] = MediaResource(text=text)

        super().__init__(**data)
        ...
```

## Node
Nodeオブジェクトは，Documentオブジェクトから複数作成されるオブジェクトである．Nodeオブジェクトの特徴は以下の通りである．

* 1Node = 1chunk
* テキストやメタデータで構成される
* Documentオブジェクトのメタデータが継承される
* 他のNodeやIndexとの関係性を保存

イメージとしては，Documnetクラスは1つまるまるのPDFファイルやPDF1ページまるまるなどのある程度まとまったデータを統一するために存在し，NodeクラスはそれぞれがRAGのRetrieverで取得されるアイテムとなる（1Node=1チャンク）．つまり，RAGの文脈でよく出てくる「チャンキングアルゴリズムの工夫」はLlama indexの文脈では「DocumentオブジェクトからNodeオブジェクトを作成する際の工夫」と考えることができる．

Nodeオブジェクトを作成する際は以下のように記述する．

```python
from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo

node1 = TextNode(text="<text_chunk>", id_="<node_id>")
node2 = TextNode(text="<text_chunk>", id_="<node_id>")
# set relationships
node1.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(
    node_id=node2.node_id
)
node2.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(
    node_id=node1.node_id
)
nodes = [node1, node2]
```

このようにNodeオブジェクトは，どのNode同士に関連性があるのかを保持している．

## Embedding
Embeddingは，Nodeを埋め込みモデルでベクトル化したものである．
1つのNodeから一つのEmbeddingが作成される．
後述のIndexは，複数のNodeで構成されるが，それぞれのNodeはそれぞれの埋め込みベクトルを持つ．

## Index
Indexは日本語で索引という意味である．Llama-indexにおけるIndexとは，複数のドキュメントオブジェクト（テキストとメタデータ）から構成されるデータ構造である．ここでは，最も使用頻度の高い`Vector Index`についてまとめるが，実は`Tree Index`や`Keyword Table Index`なども存在するらしい．以下の図がVector Store Indexのイメージ図である．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2718862/3d09993a-7e05-4b52-ae18-deff362eaf53.png)

Vector Store Index は複数のNodeをまとめたデータ構造である．それぞれのNodeは先ほどまとめた様に，分割されたテキストチャンク，メタデータ，そして埋め込みベクトルを保有している．RAGの情報源としてこのVector Store Index を使用する際は，埋め込まれたクエリと各Nodeの埋め込みベクトルが比較され，関連度の高いNodeが検索に引っかかる．

# データの変形に使用する機構
これ以降，PDFファイルからIndexを作成するまでの手順を順々に追っていく．なお，いくつかコードを記載するが，実際にLlama indexを使用する際はもっと抽象化されたメソッドが用意されており，あくまで仕組みの理解のためであることを注意してもらいたい．

## PDF → Document
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2718862/74b347f5-ce25-4770-8baa-833f67b24eb3.png)


PDFからDocumentオブジェクトを作成する選択肢として，Llama indexでは，PDFReaderクラスというものが存在する．他にもLlama HubというサイトでさまざまなドキュメントReaderが公開されているが，今回はPDFReaderクラスを見てみる．

```python
class PDFReader(BaseReader):
    """PDF parser."""

    def __init__(self, return_full_document: Optional[bool] = False) -> None:
        """
        Initialize PDFReader.
        """
        self.return_full_document = return_full_document

    @retry(
        stop=stop_after_attempt(RETRY_TIMES),
    )
    def load_data(
        self,
        file: Path,
        extra_info: Optional[Dict] = None,
        fs: Optional[AbstractFileSystem] = None,
    ) -> List[Document]:
        """Parse file."""
        if not isinstance(file, Path):
            file = Path(file)

        try:
            import pypdf
        except ImportError:
            raise ImportError(
                "pypdf is required to read PDF files: `pip install pypdf`"
            )
        fs = fs or get_default_fs()
        with fs.open(str(file), "rb") as fp:
            # Load the file in memory if the filesystem is not the default one to avoid
            # issues with pypdf
            stream = fp if is_default_fs(fs) else io.BytesIO(fp.read())

            # Create a PDF object
            pdf = pypdf.PdfReader(stream)

            # Get the number of pages in the PDF document
            num_pages = len(pdf.pages)

            docs = []

            # This block returns a whole PDF as a single Document
            if self.return_full_document:
                metadata = {"file_name": file.name}
                if extra_info is not None:
                    metadata.update(extra_info)

                # Join text extracted from each page
                text = "\n".join(
                    pdf.pages[page].extract_text() for page in range(num_pages)
                )

                docs.append(Document(text=text, metadata=metadata))

            # This block returns each page of a PDF as its own Document
            else:
                # Iterate over every page

                for page in range(num_pages):
                    # Extract the text from the page
                    page_text = pdf.pages[page].extract_text()
                    page_text = page_text + "\n\n\n"
                    page_label = pdf.page_labels[page]

                    metadata = {"page_label": page_label, "file_name": file.name}
                    if extra_info is not None:
                        metadata.update(extra_info)

                    docs.append(Document(text=page_text, metadata=metadata))

            return docs
```

これを見てみると内部的には，pypdfというライブラリを呼び出し，1ページずつDocumentオブジェクトを作成するか，PDFの全てのページを一つのドキュメントにするかを引数の`return_full_document`によって選ぶことができる．いずれにせよ，Documentクラスのリストを返す．

また，以下のようにカスタマイズすることで，自分の所望のセパレーター（`"\n\n\n"`とか．．．）をtextに追加することができ，のちの`TextSplitter`と合わせて使うことで，好きな箇所でチャンクを分割することができると考えられる．

```python
test = text + "好きなセパレータ"
Document(text=text, metadata=metadata)
```

## Document → Node
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2718862/505bc59e-9c03-4044-b03b-3f8843cf1f95.png)


DocumentからNodeを作成する際，以下の手順が必要になる．
1. DocumentのTextSplitterによるチャンク化
1. チャンクからNodeを作成

TextSplitterとは，指定したルールに従ってテキストを分割するためのオブジェクトである．例えば，「##」や「．」でテキストを分割するなどのルールをユーザー側が指定できる．チャンキングのルールを独自にカスタマイズしたい場合は，TextSplitterをカスタマイズすることになる．

Documentオブジェクトからテキストチャンクを作るには，以下の様な処理が行われる．

```Python
from llama_index.core.node_parser import SentenceSplitter

text_parser = SentenceSplitter(
    chunk_size=1024,
    # separator=" ",
)

text_chunks = []
# maintain relationship with source doc index, to help inject doc metadata in (3)
doc_idxs = []
for doc_idx, page in enumerate(doc):
    page_text = page.get_text("text")
    cur_text_chunks = text_parser.split_text(page_text)
    text_chunks.extend(cur_text_chunks)
    doc_idxs.extend([doc_idx] * len(cur_text_chunks))
```


## Node → Embedding
ここでは，OpenAIの公開している埋め込みモデルによる埋め込みを示す．各Nodeオブジェクトに，テキストを埋め込んだベクトルを格納していることがわかる．

```python
from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding()
for node in nodes:
    node_embedding = embed_model.get_text_embedding(
        node.get_content(metadata_mode="all")
    )
    node.embedding = node_embedding
```


## Embedding → Index
これまでの処理で，Nodeは，テキスト，メタデータ，埋め込みベクトルを保有しているので，あとはVector_storeに渡すだけで，Vector Store Index の完成である．
```python
vector_store.add(nodes)
```

なお，これらの処理を抽象化して，ドキュメントのリストからインデックスを一撃で作成するには，以下のように記述する．ついでに，RAGによる質問に対する回答が得られる例も掲載する．

```python
from llama_index import VectorStoreIndex, SimpleDirectoryReader
# ドキュメントを読み込む
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
# インデックスを作成
index = VectorStoreIndex.from_documents(documents)
# クエリエンジンを作成
query_engine = index.as_query_engine()
# クエリを実行
response = query_engine.query("What did the author do growing up?")
print(response)
```

# まとめ
ここまで読んでいただきありがとうございました．簡単なRAGシステムを作ってみるだけなら，最後の抽象化された5行だけで良いのですが，正確な出力を求めるほど，データ構造に関する知識やチャンキングに関する知識が必要とされる傾向があったので，一度勉強も兼ねてまとめてみました．


# 参考資料
https://medium.com/@mojo885/llamaindex-saddle-up-and-control-your-llama-5131dde3a1fb

https://github.com/run-llama/llama_index

https://docs.llamaindex.ai/en/stable/module_guides/loading/documents_and_nodes/

https://docs.llamaindex.ai/en/stable/examples/low_level/ingestion/#2-use-a-text-splitter-to-split-documents