利用大模型提取PDF文件内容并结构化
用大模型提取pdf文件内容并格式化的全过程
利用大模型提取PDF文件内容并结构化
今天刷Github的时候看到一个有趣的仓库：https://github.com/getomni-ai/zerox。这个仓库提供了一个大模型提取PDF文件的方法，而我则是加了一个函数能将提取出来的内容结构化，下面跟我一起了解这个仓库它是如何工作的，而内容结构化又是如何实现的。
仓库工作流程：
该仓库的运行流程如图所示，这里以python为例，仓库使用的是pdf2image来将pdf文件转换成图片，而在使用这个库之前，亲测自己的电脑上需要安装有poppler,不然就会出现TypeError: 'NoneType' object is not iterable这样的错误，没法将文件转换成图片，出师未捷身先死hhh，后来我在他们的issue中找到了解决办法：
1、安装 Poppler：
-
在 Windows 上： 从 Poppler for Windows 下载 Poppler 文件，然后将 zip 文件解压到一个文件夹中（例如，C:\poppler）。
-
在 macOS 上： 您可以通过 Homebrew 安装 Poppler：
brew install poppler
- 在 Linux（Debian/Ubuntu）上： 使用包管理器安装 Poppler：
sudo apt-get install poppler-utils
2、将 Poppler 添加到系统 PATH：
如果你使用的是 Windows，你需要将 Poppler 安装目录中的 bin 文件夹添加到系统的 PATH 中。步骤如下：
- 右键点击“此电脑”或“我的电脑”，选择“属性”。
- 点击“高级系统设置”。
- 在系统属性窗口中，点击“环境变量”按钮。
- 在系统变量中找到 Path 变量，点击“编辑”。
- 点击“新建”，添加 Poppler bin 目录的路径（例如，C:\poppler\bin）。
- 点击“确定”以关闭所有窗口。
3、验证 Poppler 安装：
安装 Poppler 并将其添加到 PATH 后，您可以通过在终端（命令提示符或 shell）中运行以下命令来验证其是否正确设置：
pdftoppm -h
这应该会显示 pdftoppm 的帮助信息，如果您看到这些信息，说明 Poppler 已正确安装并添加到 PATH 中。
然后根据他们仓库提供的python脚本来配置模型，配置完模型后，要记得将其他的无关的参数删除，比如我配置了gpt-4o-mini，后面的gemini、claude等等那些都需要删除，配置完便可以运行脚本。
不过目前测试下来，这个仓库对于本地模型的支持不是很好，我本地用ollama跑模型，启动脚本的时候一直提醒我“该模型不是视觉模型”，后面我在后端起了一个one-api接口，接入了本地的ollama，测试才算正常。
下面是官方仓库提供的python脚本：
from pyzerox import zerox
import os
import json
import asyncio
### 模型设置（仅使用视觉模型）参考：https://docs.litellm.ai/docs/providers ###
## 可选的额外模型参数，某些模型可能需要
kwargs = {}
## 用于视觉模型的系统提示
custom_system_prompt = None
# 可覆盖
# custom_system_prompt = "对于以下 PDF 页面，执行某些操作..." ## 示例
###################### OpenAI 示例 ######################
model = "gpt-4o-mini" ## OpenAI 模型
os.environ["OPENAI_API_KEY"] = "" ## 你的 API 密钥
###################### Azure OpenAI 示例 ######################
model = "azure/gpt-4o-mini" ## "azure/<你的部署名称>" -> 格式 <提供者>/<模型>
os.environ["AZURE_API_KEY"] = "" # "你的 Azure API 密钥"
os.environ["AZURE_API_BASE"] = "" # "https://example-endpoint.openai.azure.com"
os.environ["AZURE_API_VERSION"] = "" # "2023-05-15"
###################### Gemini 示例 ######################
model = "gemini/gpt-4o-mini" ## "gemini/<gemini_model>" -> 格式 <提供者>/<模型>
os.environ['GEMINI_API_KEY'] = "" # 你的 Gemini API 密钥
###################### Anthropic 示例 ######################
model = "claude-3-opus-20240229"
os.environ["ANTHROPIC_API_KEY"] = "" # 你的 Anthropic API 密钥
###################### Vertex AI 示例 ######################
model = "vertex_ai/gemini-1.5-flash-001" ## "vertex_ai/<模型名称>" -> 格式 <提供者>/<模型>
## 获取凭据
## 运行 ##
# !gcloud auth application-default login - 运行此命令以将 Vertex 凭据添加到你的环境
## 或者 ##
file_path = 'path/to/vertex_ai_service_account.json'
# 加载 JSON 文件
with open(file_path, 'r') as file:
vertex_credentials = json.load(file)
# 转换为 JSON 字符串
vertex_credentials_json = json.dumps(vertex_credentials)
vertex_credentials = vertex_credentials_json
## 额外参数
kwargs = {"vertex_credentials": vertex_credentials}
###################### 其他提供者的参考： https://docs.litellm.ai/docs/providers ######################
# 定义主要的异步入口点
async def main():
file_path = "https://omni-demo-data.s3.amazonaws.com/test/cs101.pdf" ## 支持本地文件路径和文件 URL
## 处理部分页面或全部页面
select_pages = None ## None 表示全部，也可以是整数或整数列表（1 索引）
output_dir = "./output_test" ## 保存合并后的 Markdown 文件的目录
result = await zerox(file_path=file_path, model=model, output_dir=output_dir,
custom_system_prompt=custom_system_prompt, select_pages=select_pages, **kwargs)
return result
# 运行主函数：
result = asyncio.run(main())
# 打印 Markdown 结果
print(result)
Markdown输出结果
我使用了经典的《attention is all you need》来做测试，输出结果如下：
文本
公式
表格
结构化
而在输出markdown之后，因为有一级标题二级标题正文等内容，所以很容易对其进行结构化，下面是对markdown内容结构化存储的一种方案：
代码如下：
import json
def markdown_to_json(markdown_file, json_file, indent=2):
with open(markdown_file, 'r', encoding='utf-8') as md_file:
markdown_content = md_file.read()
json_data = []
current_item = {"一级标题": None, "二级标题": None, "三级标题": None, "content": ""}
current_level = 0
for line in markdown_content.split('\n'):
if line.startswith('#'):
# 处理新标题
level = len(line.split()[0])
title = line.strip('# ')
if current_item["content"].strip():
json_data.append(current_item)
current_item = {"一级标题": None, "二级标题": None, "三级标题": None, "content": ""}
if level == 1:
current_item["一级标题"] = title
current_item["二级标题"] = None
current_item["三级标题"] = None
elif level == 2:
current_item["二级标题"] = title
current_item["三级标题"] = None
elif level == 3:
current_item["三级标题"] = title
current_level = level
elif line.strip():
# 非空内容行
if current_item["content"]:
current_item["content"] += "\n"
current_item["content"] += line.strip()
elif current_item["content"].strip():
# 空行，但前面有内容，开始新的项
json_data.append(current_item)
current_item = {
"一级标题": current_item["一级标题"],
"二级标题": current_item["二级标题"] if current_level >= 2 else None,
"三级标题": current_item["三级标题"] if current_level >= 3 else None,
"content": ""
}
# 添加最后一项
if current_item["content"].strip():
json_data.append(current_item)
# 写入JSON文件
with open(json_file, 'w', encoding='utf-8') as json_file:
json.dump(json_data, json_file, ensure_ascii=False, indent=indent)
print(f"已将Markdown文件 '{markdown_file}' 转换为JSON文件 '{json_file}'")
使用方法：
markdown_to_json("markdown's path", "json's output path")
输出的json文件长这样：
{
"一级标题": "Figure 1: The Transformer - model architecture.",
"二级标题": "3.1 Encoder and Decoder Stacks",
"三级标题": null,
"content": "**Encoder:** The encoder is composed of a stack of \\( N = 6 \\) identical layers. Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed-forward network. We employ a residual connection [1] around each of the two sub-layers, followed by layer normalization \\( \\text{LayerNorm}(z + \\text{Sublayer}(z)) \\). That is, the output of each sub-layer is itself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension \\( d_{model} = 512 \\)."
},
{
"一级标题": "Figure 1: The Transformer - model architecture.",
"二级标题": "3.1 Encoder and Decoder Stacks",
"三级标题": null,
"content": "**Decoder:** The decoder is also composed of a stack of \\( N = 6 \\) identical layers. In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization. We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with the fact that the output embeddings are offset by one position, ensures that the predictions for position \\( i \\) can depend only on the known outputs at positions less than \\( i \\)."
},
至此，PDF文件的内容结构化存储就完成啦。
这些结构化的文件可以用作入库，结构化数据的好处在于可以减少数据噪音，用在各种RAG系统上可以让检索到的内容质量更高。后面再出点RAG相关的文章，点击关注，敬请期待。