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

wufei123 2025-01-26 阅读:74 评论:0

让我们探索一种结合ocr和llm技术分析图像的方法。虽然这不是专家级方案，但它源于实际应用中的类似方法，更像是一个便捷的周末项目，而非生产就绪代码。让我们开始吧！目标: 构建一个简单的管道，用于处理图像（或PDF），利用OCR提取文本，再...

让我们探索一种结合ocr和llm技术分析图像的方法。虽然这不是专家级方案，但它源于实际应用中的类似方法，更像是一个便捷的周末项目，而非生产就绪代码。让我们开始吧！

目标:

构建一个简单的管道，用于处理图像（或PDF），利用OCR提取文本，再用LLM分析文本以获取有价值的元数据。这对于文档自动分类、来信分析或智能文档管理系统非常有用。我们将使用一些流行的开源工具，简化流程。

前提:

本文假设您已熟悉Hugging Face Transformers库。如不熟悉，请参考Hugging Face Transformers快速入门。

所需库:

我们将使用torch、transformers、pymupdf和rich库。rich用于提升控制台输出的可读性。

PHP

import json
import time
import fitz

import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

from rich.console import Console
console = Console()

图像准备:

我们将使用Hugging Face官网首页作为测试样本。为了模拟实际应用场景，我们假设输入为PDF，需要将其转换为PNG格式以便模型处理：

PHP

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("正在将PDF转换为PNG...", spinner="monkey"):
    with open(output_png_file, "wb") as f:
        f.write(img)

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

OCR处理:

我测试过多种OCR方案，最终选择ucaslcl/got-ocr2_0 (https://www.php.cn/link/c29568179ca4f6e62552f14d69b810b2) 效果最佳。

PHP

tokenizer = AutoTokenizer.from_pretrained(
    "ucaslcl/got-ocr2_0",
    device_map="cuda",
    trust_remote_code=True,
)
model = AutoModelForCausalLM.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()


def run_ocr(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"耗时: {final_time:.2f} 秒")

    return res

result_text = None
with console.status(
    "正在进行OCR处理...",
    spinner="monkey",
):
    #  此处需要调整，根据got-ocr2_0的实际接口进行修改
    result_text = run_ocr(
        lambda: model.generate(
            input_ids=tokenizer(output_png_file, return_tensors="pt").input_ids.cuda(),
            max_length=512
        ),
        "纯文本OCR结果",
    )

got-ocr2_0支持多种输出格式，包括HTML。此处展示了纯文本输出示例：

PHP

hugging face- the al community building the future.  https: / / hugging face. co/  search models, datasets, users. . .  following 0
...

LLM分析:

我们将使用meta-llama/llama-3.2-1b-instruct进行文本分析。我们将进行文本分类、情感分析等基本操作。

PHP

prompt = f"分析以下文档文本，并生成包含以下字段的JSON元数据：标签(tags)、语言(language)、机密性(confidentiality)、优先级(priority)、类别(category)和摘要(summary)。只提供JSON数据，开头为{{，不包含任何解释。 文档文本：{result_text}"

model_id = "meta-llama/llama-3.2-1b-instruct"
pipe = pipeline(
    "text-generation",
    model=model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
messages = [
    {"role": "system", "content": "你将分析提供的文本并生成JSON输出。"},
    {"role": "user", "content": prompt},
]

with console.status(
    "正在使用LLM分析文本...",
    spinner="monkey",
):
    outputs = pipe(
        messages,
        max_new_tokens=2048,
    )

result_json_str = outputs[0]["generated_text"].strip()
#  移除可能存在的代码块标记
if result_json_str.startswith("```json") and result_json_str.endswith("```"):
    result_json_str = result_json_str[7:-3].strip()

parsed_json = json.loads(result_json_str)

with console.status(
    "正在保存结果到文件...",
    spinner="monkey",
):
    with open("result.json", "w") as f:
        json.dump(parsed_json, f, indent=4)

示例输出：

PHP

{
  "tags": ["Hugging Face", "AI", "machine learning", "models", "datasets"],
  "language": "en",
  "confidentiality": "public",
  "priority": "normal",
  "category": "technology",
  "summary": "This text describes Hugging Face, a platform for AI models and datasets."
}