首頁 > 後端開發 > Python教學 > 快速而骯髒的文檔分析:在 Python 中結合 GOT-OCR 和 LLama

快速而骯髒的文檔分析:在 Python 中結合 GOT-OCR 和 LLama

DDD
發布: 2025-01-09 12:14:41
原創
744 人瀏覽過

讓我們來探索一種對影像進行 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 測試。

That landing page of HF that was mentioned above

為了更現實的解決方案,我們假設我們的輸入是 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中文網其他相關文章!

本網站聲明
本文內容由網友自願投稿,版權歸原作者所有。本站不承擔相應的法律責任。如發現涉嫌抄襲或侵權的內容,請聯絡admin@php.cn
熱門教學
更多>
最新下載
更多>
網站特效
網站源碼
網站素材
前端模板