TorchV的RAG實踐分享(三):解析llama_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框架在數據結構設計上是如何做的,這對於企業開發人員來說,架構&數據結構設計是有很重要的參考意義。
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章