快速而肮脏的文档分析：在 Python 中结合 GOT-OCR 和 LLama

让我们探索一种对图像进行 OCR LLM 分析的方法。这会是拥有数十年经验的专家给出的最佳方法吗？并不真地。但它来自于现实生活中采取类似方法的人。将此视为带有实用片段的周末项目版本，而不是可用于生产的代码。让我们深入挖掘吧！

我们的目标是什么？

我们将构建一个简单的管道，可以获取图像（或 PDF），使用 OCR 从中提取文本，然后使用 LLM 分析该文本以获得有用的元数据。这对于自动分类文档、分析传入的信件或构建智能文档管理系统非常方便。我们将使用一些流行的开源工具来完成此操作，并使事情相对简单。

是的，下面的所有内容都假设您已经对高频变压器非常熟悉。如果没有，请查看 https://huggingface.co/docs/transformers/en/quicktour - 似乎是一个不错的起点。虽然我从来没有这样做过，只是从例子中学习。我最终会做到的。

我们需要什么包？

我们将使用 torch 和 Transformer 来完成繁重的工作，再加上 pymupdf 和 rich 来通过一些用户友好的控制台输出让我们的生活更轻松（我喜欢 rich，所以基本上我们使用它是为了好玩）。

import json
import time
import fitz

import torch
from transformers import AutoModel, AutoTokenizer, pipeline

from rich.console import Console
console = Console()

准备图像

首先，我们应该使用什么图像作为输入？由于我们在这里使用 Hugging Face 作为主要工作，因此我们使用其主要网页的第一页作为我们的测试主题。它是文本和复杂格式的良好候选者 - 非常适合我们的 OCR 测试。

为了更现实的解决方案，我们假设我们的输入是 PDF（因为让我们面对现实，这就是您在现实世界中可能会处理的内容）。我们需要将其转换为 PNG 格式以便我们的模型进行处理：

INPUT_PDF_FILE = "./data/ocr_hf_main_page.pdf"
OUTPUT_PNG_FILE = "./data/ocr_hf_main_page.png"

doc = fitz.open(INPUT_PDF_FILE)
page = doc.load_page(0)
pixmap = page.get_pixmap(dpi=300)
img = pixmap.tobytes()

with console.status("Converting PDF to PNG...", spinner="monkey"):
    with open(OUTPUT_PNG_FILE, "wb") as f:
        f.write(img)

在这里进行真正的 OCR

我已经尝试过各种 OCR 解决方案来完成这项任务。当然，还有超正方体和许多其他选择。但对于我的测试用例，我使用 GOT-OCR2_0 (https://huggingface.co/stepfun-ai/GOT-OCR2_0) 获得了最佳结果。那么让我们直接开始吧：

tokenizer = AutoTokenizer.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    device_map="cuda",
    trust_remote_code=True,
)
model = AutoModel.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    pad_token_id=tokenizer.eos_token_id,
)
model = model.eval().cuda()

这是怎么回事？好吧，默认的 AutoModel 和 AutoTokenizer，唯一足够特别的部分是我们正在设置模型以使用 cuda。这不是可选的。该模型需要 CUDA 支持才能运行。

现在我们已经定义了模型，让我们实际将其应用于保存的文件。此外，我们还将测量时间并打印出来。不仅可以与不同的模型进行比较，还可以了解您的用例等待这么长时间是否可行（尽管对于我们的情况来说非常快）：

import json
import time
import fitz

import torch
from transformers import AutoModel, AutoTokenizer, pipeline

from rich.console import Console
console = Console()

这是我们从原始图像中得到的结果：

INPUT_PDF_FILE = "./data/ocr_hf_main_page.pdf"
OUTPUT_PNG_FILE = "./data/ocr_hf_main_page.png"

doc = fitz.open(INPUT_PDF_FILE)
page = doc.load_page(0)
pixmap = page.get_pixmap(dpi=300)
img = pixmap.tobytes()

with console.status("Converting PDF to PNG...", spinner="monkey"):
    with open(OUTPUT_PNG_FILE, "wb") as f:
        f.write(img)

^ 所有文本，没有格式，但这是故意的。

GOT-OCR2_0 非常灵活 - 它可以以不同的格式输出，包括 HTML。以下是您可以使用它的一些其他方法：

tokenizer = AutoTokenizer.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    device_map="cuda",
    trust_remote_code=True,
)
model = AutoModel.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    pad_token_id=tokenizer.eos_token_id,
)
model = model.eval().cuda()

最后尝试一下LLM

现在是有趣的部分 - 选择法学硕士。关于哪一个最好的讨论一直没完没了，随处可见相关文章。但让我们保持简单：每个人和他们的狗都听说过的法学硕士是什么？骆驼。所以我们将使用 Llama-3.2-1B 来处理文本。

从文字中我们能得到什么？考虑文本分类、情感分析、语言检测等基本内容。想象一下，您正在构建一个系统来自动对上传的文档进行分类或对药房收到的传真进行排序。

我将跳过对即时工程的深入研究（这是另一篇文章，我不相信我会写任何文章），但基本思想如下：

def run_ocr_for_file(func: callable, text: str):
    start_time = time.time()
    res = func()
    final_time = time.time() - start_time

    console.rule(f"[bold red] {text} [/bold red]")
    console.print(res)
    console.rule(f"Time: {final_time} seconds")

    return res

result_text = None
with console.status(
    "Running OCR for the result file...",
    spinner="monkey",
):
    result_text = run_ocr_for_file(
        lambda: model.chat(
            tokenizer,
            OUTPUT_PNG_FILE,
            ocr_type="ocr",
        ),
        "plain texts OCR",
    )

顺便问一下，我是不是在用提示/内容做一些非常愚蠢的事情？让我知道。对“即时工程”还很陌生，还没有足够认真地对待它。

模型有时会将结果包装在 markdown 代码块中，因此我们需要处理它（如果有人知道更简洁的方法，我洗耳恭听）：

Hugging Face- The Al community building the future.  https: / / hugging face. co/  Search models, datasets, users. . .  Following 0
All Models Datasets Spaces Papers Collections Community Posts Up votes Likes New Follow your favorite Al creators Refresh List black-
forest- labs· Advancing state- of- the- art image generation Follow stability a i· Sharing open- source image generation models
Follow bria a i· Specializing in advanced image editing models Follow Trending last 7 days All Models Datasets Spaces deep see k- a
i/ Deep Seek- V 3 Updated 3 days ago· 40 k· 877 deep see k- a i/ Deep Seek- V 3- Base Updated 3 days ago· 6.34 k· 1.06 k 2.39 k
TRELLIS Q wen/ QV Q- 72 B- Preview 88888888888888888888 888888888888888888 301 Gemini Co der 1 of 3 2025-01-01,9:38 p. m

这是我们通常得到的输出：

# format texts OCR:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='format',
)

# fine-grained OCR:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='ocr',
  ocr_box='',
)
# ... ocr_type='format', ocr_box='')
# ... ocr_type='ocr', ocr_color='')
# ... ocr_type='format', ocr_color='')

# multi-crop OCR:
# ... ocr_type='ocr')
# ... ocr_type='format')

# render the formatted OCR results:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='format',
  render=True,
  save_render_file = './demo.html',
)

总结

我们构建了一个小管道，可以获取 PDF，使用一些非常好的 OCR 提取其文本，然后使用 LLM 分析该文本以获得有用的元数据。生产就绪了吗？可能不会。但如果您想要构建类似的东西，那么这是一个坚实的起点。最酷的是我们如何结合不同的开源工具来创建有用的东西 - 从 PDF 处理到 OCR 到 LLM 分析。

您可以轻松扩展它。也许添加更好的错误处理、对多个页面的支持，或者尝试不同的法学硕士。或者也许将其连接到文档管理系统。希望你会的。这可能是一项有趣的任务。

请记住，这只是一种方法 - 可能还有许多其他方法可能更适合您的特定用例。但希望这能为您自己的实验提供一个良好的起点！或者是在评论中教我如何完成的完美地方。

以上是快速而肮脏的文档分析：在 Python 中结合 GOT-OCR 和 LLama的详细内容。更多信息请关注PHP中文网其他相关文章！