RAG工程落地：处理文档中表格数据

import pymupdf4llmfrom pathlib import Path
# 设置参数pdf_path = r"D:Testmuxuedata2caiwubaogao.pdf"  # 替换为您的 PDF 文件路径output_md = r"D:Testmuxuedata2caiwubaogao.md"           # 输出的 Markdown 文件名image_dir = r"D:Testmuxuedata2images"              # 图片保存目录dpi = 300                         # 图片分辨率image_format = "png"              # 图片格式，可选 "png"、"jpg" 等
# 创建图片保存目录Path(image_dir).mkdir(parents=True, exist_ok=True)
# 转换 PDF 为 Markdown，并提取图片md_text = pymupdf4llm.to_markdown(    doc=pdf_path,    write_images=True,    image_path=image_dir,    image_format=image_format,    dpi=dpi)
# 保存 Markdown 内容到文件with open(output_md, "w", encoding="utf-8") as f:    f.write(md_text)
print(f"Markdown 内容已保存到 {output_md}")print(f"图片已保存到目录 {image_dir}")

import mammothimport os
def docx_to_markdown_with_images(docx_path, output_md_path=None, image_dir="images"):    os.makedirs(image_dir, exist_ok=True)
    def save_image(image):        image_name = image.alt_text.replace(" ", "_") if image.alt_text else "image"        ext = {            "image/png": ".png",            "image/jpeg": ".jpg",            "image/gif": ".gif"        }.get(image.content_type, ".bin")        filename = f"{image_name}{ext}"        image_path = os.path.join(image_dir, filename)
        # 避免重名        counter = 1        base_name = filename.rsplit(".", 1)[0]        while os.path.exists(image_path):            filename = f"{base_name}_{counter}{ext}"            image_path = os.path.join(image_dir, filename)            counter += 1
        # 读取图片数据，保存 —— **改这里！**        with image.open() as img_file:            with open(image_path, "wb") as out_file:                out_file.write(img_file.read())
        # 返回 Markdown 中图片的路径，注意替换成相对路径或 URL 时修改这里        return {"src": image_path.replace("\", "/")}
    with open(docx_path, "rb") as docx_file:        result = mammoth.convert_to_markdown(            docx_file,            convert_image=mammoth.images.img_element(save_image)        )        markdown_text = result.value
    if output_md_path:        with open(output_md_path, "w", encoding="utf-8") as f:            f.write(markdown_text)
    return markdown_text
# 示例markdown = docx_to_markdown_with_images(    r"D:Testmuxuedata2caiwubaogao.docx",    output_md_path=r"D:muxuedata2caiwubaogao.md",    image_dir=r"D:Testmuxuedata2images")print(markdown)

from llama_index.core import SimpleDirectoryReader, VectorStoreIndexfrom llama_index.core.node_parser import MarkdownNodeParser
# 加载 Markdown 文档documents = SimpleDirectoryReader(input_dir=r"D:TestRAGTestdatamarkdown", required_exts=[".md"]).load_data()
# 创建 Markdown 节点解析器node_parser = MarkdownNodeParser.from_defaults(    include_metadata=True,            # 包含元数据    include_prev_next_rel=True,       # 包含前后节点关系    header_path_separator="/")
# 将文档解析为节点列表nodes = node_parser.get_nodes_from_documents(documents)

from llama_index.core.llms.mock import MockLLMfrom llama_index.core.node_parser.relational.markdown_element import MarkdownElementNodeParserfrom llama_index.core.schema import Document, TextNode, IndexNode
# 示例markdown文本，包含文本、标题和表格md_text = """# 第一章这是第一章的内容。| 年份 | 收益 || ---- | ---- || 2020 | 12000 || 2021 | 15000 |## 第二节这是第二节的内容。| 产品 | 数量 | 价格 || ---- | ---- | ---- || A    | 10   | 5    || B    | 20   | 8    |"""
# 构建Document对象doc = Document(text=md_text)
# 初始化MarkdownElementNodeParserparser = MarkdownElementNodeParser(llm=MockLLM())
# 解析为节点nodes = parser.get_nodes_from_documents([doc])
# 输出每个节点的类型和内容for i, node in enumerate(nodes):    print(f"Node (i): 类型: {type(node).__name__}")    print(f"内容: {getattr(node, 'text', getattr(node, 'table', ''))}n")

'''markdown中表格数据的切割和查询'''from llama_index.core import VectorStoreIndex, Settings, SimpleDirectoryReaderfrom llama_index.llms.openai_like import OpenAILikefrom llama_index.embeddings.openai_like import OpenAILikeEmbeddingfrom llama_index.core.node_parser.relational.markdown_element import (    MarkdownElementNodeParser,)from llama_index.core.llms.mock import MockLLM
# ================== 初始化模型 ==================def init_models():    """初始化模型并验证"""    # Embedding模型
    embed_model = OpenAILikeEmbedding(        model_name="BAAI/bge-m3",        api_base="https://api.siliconflow.cn/v1",        api_key="sk-xxx",        embed_batch_size=10,    )
    llm = OpenAILike(        model="DeepSeek-ai/DeepSeek-V3",        api_base="https://api.siliconflow.cn/v1",        api_key="sk-xxx",        context_window=128000,        is_chat_model=True,        is_function_calling_model=False,    )

    Settings.embed_model = embed_model    Settings.llm = llm
    # 验证模型    test_embedding = embed_model.get_text_embedding("测试文本")    print(f"Embedding维度验证：{len(test_embedding)}")
    return embed_model, llm
init_models()
# load documents, split into chunksdocuments = SimpleDirectoryReader(r"D:Testmuxuedata2", required_exts=[".md"]).load_data()
# 2. 强大的分割器
node_parser = MarkdownElementNodeParser(llm=MockLLM())nodes = node_parser.get_nodes_from_documents(documents)
index = VectorStoreIndex(nodes)
from llama_index.core.query_engine import CitationQueryEngine
query_engine = CitationQueryEngine.from_args(    index,    similarity_top_k=3,    # here we can control how granular citation sources are, the default is 512    citation_chunk_size=512,)
res = query_engine.query("股本增减变动幅度多大?请使用中文回答")print(res.response)              # LLM 输出回答print("------来源---------------")for node in res.source_nodes:    print("相关片段：", node.text)    print("片段分数：", node.score)    print("片段元数据：", node.metadata)    print("="*40)