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

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

1 提取表格的步骤

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

pip install camelot-py pandas

2 完整代码

#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 使用 conda 安装。
• 提取准确性：对于复杂 PDF，提取结果可能不完美，可通过调整 pages 参数（如 pages=’1-3′) 或手动检查结果。

RA/SD 衍生者AI训练营