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

  • 研究表明,使用 Camelot 库可以有效处理 PDF 文件中的表格数据,适合 RAG 开发。
  • 提取表格后,将每行数据转换为带列名的字符串,便于 RAG 系统检索。
  • 可能需要额外安装 Ghostscript,具体取决于 PDF 的复杂性。
AI教我做事之RAG开发-4 如何处理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′) 或手动检查结果。
AI教我做事之RAG开发-4 如何处理PDF文件中的表格数据

RA/SD 衍生者AI训练营。发布者:稻草人,转载请注明出处:https://www.shxcj.com/ai%e6%95%99%e6%88%91%e5%81%9a%e4%ba%8b%e4%b9%8brag%e5%bc%80%e5%8f%91-4-%e5%a6%82%e4%bd%95%e5%a4%84%e7%90%86pdf%e6%96%87%e4%bb%b6%e4%b8%ad%e7%9a%84%e8%a1%a8%e6%a0%bc%e6%95%b0%e6%8d%ae/

(0)
上一篇 1天前
下一篇 1天前

相关推荐

发表回复

登录后才能评论
本文授权以下站点有原版访问授权 https://www.shxcj.com https://www.2img.ai https://www.2video.cn