[开发篇]检索增强生成(RAG)知识库部分的实现Part1-文档处理

大型语言模型 (LLM) 的知识来源于预训练数据,这使得它们容易给出过时甚至不准确的答案。而检索增强生成(RAG)技术在内容生成前为 LLM 提供信息参考,有效克服了这一局限性。

根据自己的学习路线和职业规划,我参与智能体平台开发已经有一段时间了。虽然的热度相比之前有所回落,市面上也陆续出现了不少成熟的智能体工作流平台和智能体开发框架,但作为智能体平台早期阶段的核心组成部分之一,RAG 依然值得系统梳理和实践总结。因此,也算是搭上一个“末班车”,记录一下自己在实现过程中的一些思路和经验。

这篇文章主要讲讲RAG技术知识库部分文档处理的实现,主要涉及文档解析和文档切片。

以下主要介绍下思路,并提供Python代码例子(我主要使用golang开发,对于python的使用理解得还不够透彻,可能有写法不够严谨的地方,还请谅解)

文档解析

Markdown 和 Word 格式的文档通常具备较为明确的大纲结构和正文层级,整体结构性较强,能够较为自然地解析为对应的语法树。因此本文优先从这两种文档格式入手进行处理与实现,后续也会逐步补充对 PDF 等其他文档格式的解析方案。

Markdown文档

mistune 在 Markdown 和 Html 格式上有解析稳定、速度快的优点,只需要几行代码就可以实现从文档到语法树的解析

import mistune
from mistune.renderers.markdown import MarkdownRenderer

with open(filepath, "r", encoding="utf-8") as file:
    data = file.read()

render = MarkdownRenderer() # 后续如果需要支持 Html 格式,也可以将这里直接替换为内置的 HtmlRenderer
parser = mistune.create_markdown(renderer=render)

_, block_state = parser.parse(data) # block_state.tokens 即为后续所需要的语法树结构

Docx文档

使用 mammoth 将 Docx 文档转化为 Markdown 格式,即可复用 Markdown 文档的解析,便于后续的统一处理

但是需要注意额外处理两个内容:

  • Docx 文档可能包含内嵌的图片,需要提取其中的替代文本(alt)并将其上传到第三方存储来保留文档的图片内容
  • mammoth 解析出的 Markdown 内容包含锚点标签,需要移除(具体也可以调用下mammoth.convert_to_markdown通过断点验证下是否如此)
# 实现提取alt文本并上传到第三方存储的方法
# 该方法返回处理后的图片链接和替代文本字典即可,格式例如{"src": "", "alt": ""}
def image_convertor(image_name, image):
    extension = image.content_type.split("/")[-1]
    image_name = f"{image_name}.{extension}"

    try:
        with image.open() as image_bytes:
            data = image_bytes.read()  # 读取为 bytes
            data_stream = BytesIO(data)  # 转换为 BytesIO
            url = MinioClient.upload_stream(
                object_name=image_name,
                data=data_stream,
                content_type=image.content_type
            )
        alt_text = BasicCleaner.clean_string_blank(getattr(image, "alt_text") or "")
        return {"src": MinioClient.get_public_download_url(url), "alt": alt_text}
    except Exception as e:
        logger.error("docx图片处理失败", exc_info=True)
        return {"src": "", "alt": ""}

# 用于移除 mammoth 解析 Markdown内容包含的锚点标签
import mammoth

def remove_anchor_tags(text: str) -> str:
    return re.sub(r'<a\s+id="[^"]*"\s*></a>', '', text)

def image_convertor_with_doc_id(image):
    image_name = f"{doc_id}/{uuid.uuid4()}"
    return image_convertor(image_name, image)

with open(filepath, "rb") as data:
    md_result = mammoth.convert_to_markdown(
        content,
        convert_image=mammoth.images.img_element(image_convertor_with_doc_id) # 传递 convert_image 参数即可实现自定义的图片元素处理方法
    )
    input_content = remove_anchor_tags(md_result.value) # 这里即为 Markdown 文档的内容,可以走 Markdown 文档解析,复用先前的解析方法

文档切片

目前主要实现了两种方案的文档切片方式,层级切片和递归切片。层级切片适用于结构性较强的 Markdown 和 Word 格式。递归切片则比较通用,除了 Markdown 和 Word 格式,也适用于PDF格式。

层级切片

层级切片(结构化切片):指定最大层级数,按照标题、段落直接进行切片。能够保留文档的自然语义结构,但是也依赖于传入的文档需要有良好的层级结构

通过下面的方法得到每个切片Segment的标题、原始内容、清洗内容以及层级

BLOCK_KEY_TYPE = "type"  # 块类型Key
BLOCK_KEY_CHILDREN = "children"  # 块子项Key
BLOCK_KEY_RAW = "raw"  # 块内容Key
BLOCK_KEY_ATTRS = "attrs"  # 块属性Key
BLOCK_ATTRS_KEY_LEVEL = "level"  # 标题层级Key

BLOCK_TYPE_HEADING = "heading"  # 标题
BLOCK_TYPE_BLANK_LINE = "blank_line"  # 空白行, 直接跳过

SEGMENT_BLANK_LINE = "\n"

def parse_by_level(self, doc_id: int,  content: str, max_level: int = 3) -> List[Segment]:
    segments: List[Segment] = []
    current_segment_title: str = ""  # 当前切片标题
    current_segment_tokens: List[str] = []  # 当前切片token, 用于生成切片清洗内容
    current_segment_blocks: List[Dict[str, Any]] = []  # 当前切片块, 用于生成切片原始内容
    current_segment_level = 0  # 当前切片层级, 通常根据标题层级有1-6

    _, block_state = self.parser.parse(content)
    for block in block_state.tokens:
        # 无法识别类型, 跳过处理
        if BLOCK_KEY_TYPE not in block:
            continue
        # 空行需要判断上一个块结尾是否为换行符, 否则需要添加一个换行符
        if block[BLOCK_KEY_TYPE] == BLOCK_TYPE_BLANK_LINE:
            if len(current_segment_tokens) > 0 and not current_segment_tokens[-1].rstrip(' \t\r').endswith(
                    SEGMENT_BLANK_LINE):
                current_segment_tokens.append(SEGMENT_BLANK_LINE)
            else:
                continue
        elif block[BLOCK_KEY_TYPE] == BLOCK_TYPE_HEADING:
            # 标题处理
            # 超过分片的最大层级
            level = self.handle_heading_block_level(block)
            if level > max_level:
                self.handle_block_raw(current_segment_tokens, block)
                continue
            # 遇到标题则暂存之前的切片
            if len(current_segment_title) > 0 or len(current_segment_tokens) > 0:
                self.handle_new_segment(segments, current_segment_title, current_segment_tokens,
                                        current_segment_blocks, block_state, current_segment_level)
            # 并开始新的切片
            block_children = block.get(BLOCK_KEY_CHILDREN, [])
            if len(block_children) > 0:
                current_segment_title = BasicCleaner.clean_string_blank(block_children[0].get(BLOCK_KEY_RAW, ''))
            current_segment_tokens = []
            current_segment_blocks = []
            current_segment_level = level
        else:
            # 其他文本处理
            is_handle = self.handle_block_raw(current_segment_tokens, block)
            # 记录无法处理的内容
            if not is_handle:
                logger.warning(f"该类型无法处理, type: {block[BLOCK_KEY_TYPE]}, block: {block}")
        # 加入块
        current_segment_blocks.append(block)

    # 跳出循环后处理剩余内容
    if len(current_segment_title) > 0 or len(current_segment_tokens) > 0:
        self.handle_new_segment(segments, current_segment_title, current_segment_tokens,
                                current_segment_blocks, block_state, current_segment_level)

    return segments

# 保存新切片对象
def handle_new_segment(
        self,
        segments: List[Segment],
        title: str,
        tokens: List[str],
        blocks: List[Dict[str, Any]],
        block_state: BlockState,
        level: int
):
    content = self.render(blocks, block_state)
    segments.append(Segment(
        title=title, clean_content="".join(tokens), origin_content=content, level=level)
    )

def handle_heading_block_level(block: Dict[str, Any]) -> int:
    level = 0
    attrs = block.get(BLOCK_KEY_ATTRS, None)
    if isinstance(attrs, Dict):
        level = attrs.get(BLOCK_ATTRS_KEY_LEVEL, 0)
    return level

# 处理提取markdown块内容
def handle_block_raw(raw_segments: List[str], block: Dict[str, Any]) -> bool:
    is_handle = False
    # 无法处理的内容
    if BLOCK_KEY_RAW not in block and BLOCK_KEY_CHILDREN not in block:
        return is_handle
    # 处理内容
    if BLOCK_KEY_RAW in block:
        is_handle = True
        segment = BasicCleaner.clean_string_blank(block.get(BLOCK_KEY_RAW, ''))
        if len(segment) > 0:
            raw_segments.append(segment)
    # 存在子块则向下遍历
    block_children = block.get(BLOCK_KEY_CHILDREN, [])
    if len(block_children) > 0:
        for child in block_children:
            if MarkdownParser.handle_block_raw(raw_segments, child):
                is_handle = True
    return is_handle

递归切片

递归切片:先将文章内容拆分为段落,再通过分割符号将段落拆分成句子,将句子递归组合成接近指定长度的切片,每个切片之间包含重叠的内容以便于保持内容语义的关联性和完整性,但是配置复杂、速度较慢

这个实现还是比较权威的,可以参考下eino项目:
eino的递归切片实现

关联阅读

文档切片在 RAG 系统中的方法详解:原理、优劣与实战代码

发布者

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注