TorchV的RAG實踐分享(三):解析llama_index的資料儲存結構和召回策略過程

Knife4j發表於2024-03-27
Index

1.前言

LlamaIndex是一個基於LLM的資料處理框架，在RAG領域非常流行，簡單的幾行程式碼就能實現本地的檔案的對話功能，對開發者提供了極致的封裝，開箱即用。

本文以官方提供的最簡單的代理示例為例，分析LlamaIndex在資料解析、向量Embedding、資料儲存及召回的整個原始碼過程。

透過學習框架的原始碼也能讓開發者們在實際的企業大模型應用開發中，對RAG有一個更清晰的瞭解和認知。

本次選用的技術元件：

  • llm：OpenAI
  • Embedding：OpenAI
  • VectorDB：ElasticSearch

官方程式碼示例如下：


# 1.構建向量資料庫儲存物件例項
vector_store = ElasticsearchStore(
    index_name="my_index",
    es_url="http://localhost:9200",
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# 載入本地的資料集
documents = SimpleDirectoryReader('data').load_data()
# 構建索引
index = VectorStoreIndex.from_documents(documents,storage_context=storage_context)
# 服務物件，構建query引擎
service_context = ServiceContext.from_defaults(llm=OpenAI())
query_engine = index.as_query_engine(service_context=service_context)
# 問問題
resp=query_engine.query("住院起付線多少錢?")
# 響應答案
print(resp)

2.處理過程

2.1 資料處理過程

在資料處理的過程中，主要包含幾個核心的步驟：

  • 初始化向量儲存引擎，目前向量資料庫型別非常多，筆者本機跑了一個es的docker映象，這裡就選擇es了
  • 讀取資料，資料格式包括：PDF、WORD、TXT等等文字資料
  • 在資料讀取完成後，會對文件內容進行分割，然後Embedding(呼叫embedding模型)儲存入庫

2.1.1 處理載入不同的檔案型別(構建Document)

SimpleDirectoryReader是llamaindex提供的一個基於資料夾的讀取器類，會根據資料夾中的檔案擴充套件字尾型別自動載入資料

主要支援的檔案資料型別如下：

DEFAULT_FILE_READER_CLS: Dict[str, Type[BaseReader]] = {
    ".hwp": HWPReader,
    ".pdf": PDFReader,
    ".docx": DocxReader,
    ".pptx": PptxReader,
    ".ppt": PptxReader,
    ".pptm": PptxReader,
    ".jpg": ImageReader,
    ".png": ImageReader,
    ".jpeg": ImageReader,
    ".mp3": VideoAudioReader,
    ".mp4": VideoAudioReader,
    ".csv": PandasCSVReader,
    ".epub": EpubReader,
    ".md": MarkdownReader,
    ".mbox": MboxReader,
    ".ipynb": IPYNBReader,
}


class SimpleDirectoryReader(BaseReader):
    """Simple directory reader.

    Load files from file directory.
    Automatically select the best file reader given file extensions.

    Args:
        input_dir (str): Path to the directory.
        input_files (List): List of file paths to read
            (Optional; overrides input_dir, exclude)
        exclude (List): glob of python file paths to exclude (Optional)
        exclude_hidden (bool): Whether to exclude hidden files (dotfiles).
        encoding (str): Encoding of the files.
            Default is utf-8.
        errors (str): how encoding and decoding errors are to be handled,
              see https://docs.python.org/3/library/functions.html#open
        recursive (bool): Whether to recursively search in subdirectories.
            False by default.
        filename_as_id (bool): Whether to use the filename as the document id.
            False by default.
        required_exts (Optional[List[str]]): List of required extensions.
            Default is None.
        file_extractor (Optional[Dict[str, BaseReader]]): A mapping of file
            extension to a BaseReader class that specifies how to convert that file
            to text. If not specified, use default from DEFAULT_FILE_READER_CLS.
        num_files_limit (Optional[int]): Maximum number of files to read.
            Default is None.
        file_metadata (Optional[Callable[str, Dict]]): A function that takes
            in a filename and returns a Dict of metadata for the Document.
            Default is None.
    """

    supported_suffix = list(DEFAULT_FILE_READER_CLS.keys())
    //....

總共支援了16個檔案資料型別，整理到表格如下：

檔案型別 依賴元件 說明
hwp olefile
pdf pypdf
docx docx2txt
pptx、pptm、ppt python-pptx、transformers、torch 用到一些模型，對資料進行理解、提取
jpg、png、jpeg、 sentencepiece、transformers、torch 用到一些模型，對資料進行理解、提取
mp3、mp4 whisper 用到一些模型，對資料進行理解、提取
csv pandas
epub EbookLib、html2text
md 無 本地流直接open，讀取文字
mbox bs4、mailbox
ipynb nbconvert

整個Reader類的UML類圖設計如下：

所有檔案資料型別的Reader，透過load_data方法，最終得到該文件的Document物件集合，Document類是LlamaIndex框架的處理文件的核心類物件,從該類的結構設計來看，我們可以總結一下：

  • 核心欄位：id(文件唯一id)、text(文字內容)、embedding(向量float浮點型集合)、metadata(後設資料)
  • BaseNode提供了一個樹結構的設計，對於一篇文件而言，從多級標題劃分來看，樹結構能更好的描述一篇文件的基礎結構
  • Document提供了一些外部應用框架適配的方法，比如：LangChain、EmbedChain等等

最終構建完成所有的Document資訊後，我們可以看到下面一個結構資訊

本示例程式，使用的是一個PDF檔案，由於我們並未指定分割等策略，LlamaIndex對於PDF檔案是以Page為單位，進行切割，最終將所有的Document物件儲存進入向量資料庫

image-20240114153311229

2.1.2 構建向量資料庫索引(Index)

當本地資料集處理完成，得到一個Document集合的時候，此時，這需要構建向量資料庫的索引，主要是包含幾個過程：

  • 呼叫不同的向量資料庫中介軟體，構建集合索引，對於ES來說，則是建立Index
  • 呼叫Embedding模型(基於OpenAI提供的text-embedding-ada-002模型)，將Document物件集合中的text文字，進行向量化處理並賦值
  • 將Document集合的物件值(text、embedding、metadata)儲存進入向量資料庫

在LlamaIndex建立ES的向量索引結構中，初始情況下，核心欄位也是前面我們提到的Document類中的幾個核心欄位(id、embedding、content、metadata)，如下圖：

image-20240114164439411

但是在Document物件遍歷結束後，在資料儲存階段，考慮到後設資料的資訊，LlamaIndex會擴充metadata後設資料的欄位，如下圖：

image-20240114164647308

後設資料資訊會將文件的資訊提取出來，包括頁碼、檔案大小、檔名稱、建立日期等等資訊

最終在本地資料集的情況下，LlamaIndex建立的ES資料索引結構最終就會變成下面這種結構形式：

{
    "mappings": {
        "properties": {
            "content": {
                "type": "text"
            },
            "embedding": {
                "type": "dense_vector",
                "dims": 1536,
                "index": true,
                "similarity": "cosine"
            },
            "metadata": {
                "properties": {
                    "_node_content": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "_node_type": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "creation_date": {
                        "type": "date"
                    },
                    "doc_id": {
                        "type": "keyword"
                    },
                    "document_id": {
                        "type": "keyword"
                    },
                    "file_name": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "file_path": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "file_size": {
                        "type": "long"
                    },
                    "file_type": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "last_accessed_date": {
                        "type": "date"
                    },
                    "last_modified_date": {
                        "type": "date"
                    },
                    "page_label": {
                        "type": "text",
                        "fields": {
                            "keyword": {
                                "type": "keyword",
                                "ignore_above": 256
                            }
                        }
                    },
                    "ref_doc_id": {
                        "type": "keyword"
                    }
                }
            }
        }
    }
}

資料Index定義完成，Document物件儲存進入向量資料庫，此時，我們的資料集結構如下：

image-20240114173127581

2.2 問答獲取答案

在獲取答案的過程中，主要包含幾個核心的步驟：

  • 構建使用者查詢Query，對query進行Embedding處理，召回Topk的相似片段內容。
  • 組裝Prompt工程內容，傳送大模型獲取答案

2.2.1 召回查詢獲取TopK

首先，在RAG的查詢階段，不管是使用那個向量資料庫，根據資料庫的型別，將使用者的query語句進行Embedding後，再構建資料庫的查詢條件，如下圖：

image-20240114155247333

這裡面會包含幾個核心的引數：

  • embedding：knn查詢的浮點型向量陣列值
  • top_k:根據knn相似度查詢獲取得到的topk值數量，在這個例子中，LlamaIndex預設值是2
  • filters：過濾條件
  • alpha：語義&關鍵詞混合檢索的權重，0代表bm25演算法檢索，1則代表knn

VectorStoreQuery類結構定義如下：

@dataclass
class VectorStoreQuery:
    """Vector store query."""
    # knn搜尋的查詢Embedding浮點型陣列
    query_embedding: Optional[List[float]] = None
    # knn搜尋的top k取值
    similarity_top_k: int = 1
    doc_ids: Optional[List[str]] = None
    node_ids: Optional[List[str]] = None
    query_str: Optional[str] = None
    output_fields: Optional[List[str]] = None
    embedding_field: Optional[str] = None

    mode: VectorStoreQueryMode = VectorStoreQueryMode.DEFAULT

    # NOTE: only for hybrid search (0 for bm25, 1 for vector search)
    alpha: Optional[float] = None

    # metadata filters
    filters: Optional[MetadataFilters] = None

    # only for mmr
    mmr_threshold: Optional[float] = None

    # NOTE: currently only used by postgres hybrid search
    sparse_top_k: Optional[int] = None
    # NOTE: return top k results from hybrid search. similarity_top_k is used for dense search top k
    hybrid_top_k: Optional[int] = None

根據query的條件，會從向量資料庫中召回獲取得到topk的TextNode陣列，如下：

image-20240114175533596

2.2.2 構建Prompt傳送大模型獲取答案

最終召回到引用文件內容後，剩下的就是構建整個大模型對話的Prompt工程，來看看LlamaIndex的基礎Prompt是如何構建的

image-20240114160925472

image-20240114180819057

partial_format方法獲取得到一個基礎的Prompt模版資訊，內容如下：

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

這裡有兩個核心的引數：

  • context_str: 從向量資料庫召回查詢的知識庫引用文字資料上下文資訊，從模版的設定也是告訴大模型基於知識資訊進行回答
  • query_str：使用者提問的問題

而最終的context_str資訊，我們可以看到，如下圖：

image-20240114181236770

我們的問題是：住院起付線多少錢?

從最終knn檢索召回的文件片段來看，精準的找到了知識庫的引用內容，此時，交給大模型進行回答，獲取我們想要的答案結果。

image-20240114181429648

3.總結

好了，本文從LlamaIndex給我們提供的基礎的示例程式，基於Basic RAG的基礎架構來分析資料的處理、召回響應等過程，我們可以看到LlamaIndex框架給了我們一個很好的處理流程，從這裡我們可以總結如下：

  • 對於基礎的RAG架構，有一個很好的認知，讓開發者知道RAG是一個怎樣的處理過程
  • 底層的向量資料庫儲存結構設計和中間程式的結構設計，能夠給做RAG應用的開發人員一些啟發，流行的RAG框架在資料結構設計上是如何做的，這對於企業開發人員來說，架構&資料結構設計是有很重要的參考意義。

相關文章