跳到正文
八一菜刀
返回

TorchV的RAG实践分享(三):解析llama_index的数据存储结构和召回策略过程

1.前言

LlamaIndex是一个基于LLM的数据处理框架,在RAG领域非常流行,简单的几行代码就能实现本地的文件的对话功能,对开发者提供了极致的封装,开箱即用。

本文以官方提供的最简单的代理示例为例,分析LlamaIndex在数据解析、向量Embedding、数据存储及召回的整个源码过程。

通过学习框架的源码也能让开发者们在实际的企业大模型应用开发中,对RAG有一个更清晰的了解和认知。

本次选用的技术组件:

官方代码示例如下:


# 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 数据处理过程

在数据处理的过程中,主要包含几个核心的步骤:

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个文件数据类型,整理到表格如下:

文件类型依赖组件说明
hwpolefile
pdfpypdf
docxdocx2txt
pptx、pptm、pptpython-pptx、transformers、torch用到一些模型,对数据进行理解、提取
jpg、png、jpeg、sentencepiece、transformers、torch用到一些模型,对数据进行理解、提取
mp3、mp4whisper用到一些模型,对数据进行理解、提取
csvpandas
epubEbookLib、html2text
md本地流直接open,读取文本
mboxbs4、mailbox
ipynbnbconvert

整个Reader类的UML类图设计如下:

所有文件数据类型的Reader,通过load_data方法,最终得到该文档的Document对象集合,Document类是LlamaIndex框架的处理文档的核心类对象,从该类的结构设计来看,我们可以总结一下:

最终构建完成所有的Document信息后,我们可以看到下面一个结构信息

本示例程序,使用的是一个PDF文件,由于我们并未指定分割等策略,LlamaIndex对于PDF文件是以Page为单位,进行切割,最终将所有的Document对象存储进入向量数据库

image-20240114153311229

2.1.2 构建向量数据库索引(Index)

当本地数据集处理完成,得到一个Document集合的时候,此时,这需要构建向量数据库的索引,主要是包含几个过程:

在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 问答获取答案

在获取答案的过程中,主要包含几个核心的步骤:

2.2.1 召回查询获取TopK

首先,在RAG的查询阶段,不管是使用那个向量数据库,根据数据库的类型,将用户的query语句进行Embedding后,再构建数据库的查询条件,如下图:

image-20240114155247333

这里面会包含几个核心的参数:

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信息,我们可以看到,如下图:

image-20240114181236770

我们的问题是:住院起付线多少钱?

从最终knn检索召回的文档片段来看,精准的找到了知识库的引用内容,此时,交给大模型进行回答,获取我们想要的答案结果。

image-20240114181429648

3.总结

好了,本文从LlamaIndex给我们提供的基础的示例程序,基于Basic RAG的基础架构来分析数据的处理、召回响应等过程,我们可以看到LlamaIndex框架给了我们一个很好的处理流程,从这里我们可以总结如下:


分享文章:

上一篇
创业:大模型RAG系统三个月的开发心得和思考
下一篇
TorchV的RAG实践分享(二):基于ElasticSearch的混合检索实战&原理分析