AI Agent RAG Grok AI教我做事

AI教我做事之RAG开发-4 如何处理PDF文件中的表格数据

Posted by Ramendeus on March 11, 2025

AI教我做事之RAG开发-4 如何处理PDF文件中的表格数据

  • 研究表明，使用 Camelot 库可以有效处理 PDF 文件中的表格数据，适合 RAG 开发。
  • 提取表格后，将每行数据转换为带列名的字符串，便于 RAG 系统检索。
  • 可能需要额外安装 Ghostscript，具体取决于 PDF 的复杂性。

AI教我做事之RAG开发-4 如何处理PDF文件中的表格数据

1 提取表格的步骤

首先，确保已安装 Camelot 和 pandas 库，可以通过以下命令安装：

pip install camelot-py pandas

2 完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233

#from langchain.document_loaders import PyPDFLoader
from langchain_community.document_loaders import PyPDFLoader

from langchain.text_splitter import RecursiveCharacterTextSplitter

#from langchain_community.llms import Ollama
from langchain_ollama import OllamaLLM
#from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
#from langchain.vectorstores import FAISS
from langchain_community.vectorstores import FAISS

#from langchain.llms import HuggingFacePipeline
from langchain_community.llms import HuggingFacePipeline

from langchain.chains import RetrievalQA
from transformers import pipeline
from langchain.prompts import PromptTemplate
from pathlib import Path
import fitz  # PyMuPDF，用于PDF处理
import pytesseract  # 用于OCR
from PIL import Image  # 图像处理
import camelot
import pandas as pd
import io
import os
import torch

from langchain.docstore.document import Document

# 从PDF文件中将表格数据，以一行一行的数据导出。
# 这会将每行数据转换为如“姓名：Alice, 年龄：30”的格式，适合 RAG 系统使用。
def extract_data_from_pdf(pdf_path):
    """
     从PDF文件中将表格数据，以一行一行的数据导出。
    参数：
        pdf_path: PDF文件路径
    返回：
        提取的文本内容字典列表
    """
    tables = camelot.read_pdf(pdf_path)
    if len(tables) > 0:
        first_table = tables[0]
        df = first_table.dataframe
        documents = []
        for index, row in df.iterrows():
            row_str = ', '.join([f"{col}: {value}" for col, value in zip(df.columns, row)])
            documents = documents + [row_str]
    return documents

def extract_images_from_pdf(pdf_path, output_dir="extracted_images"):
    """
    从PDF中提取图像并保存到指定目录
    参数：
        pdf_path: PDF文件路径
        output_dir: 提取图像的保存目录
    返回：
        提取的图像文件路径列表
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 打开PDF文件
    pdf_document = fitz.open(pdf_path)
    image_paths = []

    # 遍历每一页
    for page_num in range(len(pdf_document)):
        page = pdf_document[page_num]
        image_list = page.get_images(full=True)  # 获取页面中的所有图像

        # 遍历页面中的每张图像
        for img_index, img in enumerate(image_list):
            xref = img[0]  # 图像的xref引用
            base_image = pdf_document.extract_image(xref)  # 提取图像数据
            image_bytes = base_image["image"]  # 图像的字节数据
            image_ext = base_image["ext"]  # 图像扩展名（如jpeg, png）

            # 保存图像
            image_filename = f"{output_dir}/image_page{page_num+1}_{img_index}.{image_ext}"
            with open(image_filename, "wb") as image_file:
                image_file.write(image_bytes)
            image_paths.append(image_filename)

    pdf_document.close()
    return image_paths

def extract_text_from_image(image_path):
    """
    使用OCR从图像中提取文本
    参数：
        image_path: 图像文件路径
    返回：
        提取的文本内容
    """
    # 打开图像
    image = Image.open(image_path)
    # 使用pytesseract进行OCR
    text = pytesseract.image_to_string(image)
    return text

def process_pdf_images_for_rag(pdf_path):
    """
    处理PDF中的图像并提取内容，用于RAG
    参数：
        pdf_path: PDF文件路径
    返回：
        包含图像内容的字典列表
    """
    # 提取图像
    image_paths = extract_images_from_pdf(pdf_path)
    image_contents = []

    # 对每张图像进行OCR并提取文本
    for img_path in image_paths:
        text = extract_text_from_image(img_path)
        image_contents.append({
            "image_path": img_path,
            "text_content": text.strip()
        })

    return image_contents

# 1. 加载和处理文档
def load_documents(file_path):
    # 使用PyPDFLoader加载PDF文档
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    return documents

def split_document(documents):
    # 分割文档为小块，便于检索,每个块约500字符，适合短文档；对于长文档可调整到1000
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = text_splitter.split_documents(documents)
    return docs

# 2. 创建向量存储
def create_vector_store(docs):
    # 使用HuggingFace的嵌入模型生成向量,all-MiniLM-L6-v2是一个轻量高效的嵌入模型，适合快速开发；生产环境可考虑更强的模型如all-mpnet-base-v2。
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    
    # 使用FAISS创建向量存储
    
    vector_store = FAISS.from_documents(docs, embeddings)
    return vector_store

# 3. 初始化语言模型
def initialize_llm():
    # 使用Hugging Face的生成模型（这里用distilgpt2作为示例）,Demo中使用distilgpt2（轻量但生成质量一般）；建议替换为gpt-neo-1.3B或LLaMA（需本地部署）。
    # text_generation_pipeline = pipeline("text-generation", model="distilgpt2", max_length=200)
    # llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
    myLLM = OllamaLLM(model="llama2")
    return myLLM

# 4. 创建RAG链
def create_rag_chain(vector_store, llm):
# 自定义提示模板，只要求简洁回答
    prompt_template = """根据以下上下文，直接用中文回答问题，不要添加多余内容。
上下文: {context}
问题: {question}
回答: """
    
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    
    # 创建RetrievalQA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=False,  # 不返回来源文档
        chain_type_kwargs={"prompt": PROMPT}  # 使用自定义提示
    )
    return qa_chain

# 5. 主函数
def main():
    pytesseract.pytesseract.tesseract_cmd = 'tesseract.exe'
    pdf_file1 = Path('testdata') / 'RAG-TestImage-3.pdf'
    
    # 处理PDF中的图像
    results = process_pdf_images_for_rag(pdf_file1)

    # 打印结果
    myStrings = []
    for result in results:
        print(f"Image: {result['image_path']}")
        print(f"Extracted Text: {result['text_content']}")
        myStrings.append(result['text_content'])
        print("-" * 50)

    # 加载并分割文档
    print("加载和分割文档...")
    # 将字符串内容认为是documents
    documents = [Document(page_content=s) for s in myStrings]
    # 分割文档（这是为了保持原来的逻辑和架构）
    docs=split_document(documents)

    print("处理表格数据...")
    #  提取后，documents 列表可能为：
    # "姓名: Alice, 年龄: 30"
    # "姓名: Bob, 年龄: 25"
    pdf_file2 = Path('testdata') / 'RAG-TestTable-4.pdf'
    sheet_documents=extract_data_from_pdf(pdf_file2)
    sheet_doc=split_document(sheet_documents)

    print("合并Doc数据...")
    # 将两部分文档叠加
    allDocs=docs+sheet_doc
    # 创建向量存储
    print("创建向量存储...")
    vector_store = create_vector_store(allDocs)
    
    # 初始化语言模型
    print("初始化语言模型...")
    llm = initialize_llm()
    
    # 创建RAG链
    print("创建RAG链...")
    rag_chain = create_rag_chain(vector_store, llm)
    
    # 示例查询
    #query = "What is the main topic of the document?"
    query = "字形绘梦是哪个公司的产品?"
    print(f"查询: {query}")
    
    # 执行RAG查询,获取精准答案
    #result = rag_chain({"query": query})
    result = rag_chain.invoke({"query": query})
    print(f"回答: {result['result'].strip()}")

if __name__ == "__main__":
    main()

3 代码解释

  1. 读取 PDF：
    1. camelot.read_pdf(‘example.pdf’) 使用默认的 ‘lattice’ 方法（适合有清晰网格线的表格）读取 PDF，返回一个 TableList 对象，包含所有提取的表格。
    2. 如果 PDF 表格无网格线，可尝试 camelot.read_pdf(‘example.pdf’, flavor=’stream’)。
    3. “姓名: Alice, 年龄: 30”
    4. “姓名: Bob, 年龄: 25”
  2. 检查和提取：
    1. len(tables) 检查找到的表格数量，若为 0，则提示无表格。
    2. tables[0].dataframe 将第一个表格转换为 pandas DataFrame，便于后续处理。
  3. 准备 RAG 数据：
    1. 遍历 DataFrame 的每一行，使用 iterrows() 获取行数据。
    2. 通过 zip(df.columns, row) 结合列名和值，创建如“姓名：Alice, 年龄：30”的字符串。
    3. 所有字符串存储在 documents 列表中，供 RAG 系统进一步处理，如创建嵌入或存储在向量数据库。

考虑与优化

  • 多表格处理：若 PDF 包含多个表格，可循环遍历 tables，为每个表格重复上述步骤，生成多个文档列表。
  • 复杂表格：对于合并单元格或无网格线的表格，Camelot 可能需要调整参数（如 flavor=’stream’），或结合其他工具如 Tabula-py。
  • 输出格式：Camelot 支持导出为 CSV、JSON 等格式（如 first_table.to_csv(‘table.csv’)），这为 RAG 系统的灵活集成提供了便利。

潜在挑战与解决方案

  • 依赖问题：Camelot 可能需要 Ghostscript，若安装失败，可参考 [Camelot Py Anaconda.org](https://anaconda.org/conda-forge/camelot-py) 使用 conda 安装。
  • 提取准确性：对于复杂 PDF，提取结果可能不完美，可通过调整 pages 参数（如 pages=’1-3′) 或手动检查结果。

AI教我做事之RAG开发-4 如何处理PDF文件中的表格数据

CATALOG
    FEATURED TAGS
    流式 Streaming WebRTC RAG AI教我做事 AI Agent Grok AIGC LLM Ollama 字形绘梦 微信小程序 GPT WebUI Lite 个人大模型 Deepseek HayStack 绘声 C# LlamaIndex SD
    FRIENDS