Analisis Dokumen Pantas dan Kotor: Menggabungkan GOT-OCR dan LLama dalam Python

Mari kita teroka cara untuk melakukan analisis OCR LLM untuk imej. Adakah ini cara terbaik yang diberikan oleh pakar dengan pengalaman berdekad-dekad? Tidak juga. Tetapi ia datang dari seseorang yang mengambil pendekatan yang sama dalam kehidupan sebenar. Fikirkan ini sebagai versi projek hujung minggu dengan coretan praktikal dan bukannya kod sedia pengeluaran. Mari kita mendalami!

Apakah matlamat kami di sini?

Kami akan membina saluran paip ringkas yang boleh mengambil imej (atau PDF), mengekstrak teks daripadanya menggunakan OCR, dan kemudian menganalisis teks tersebut menggunakan LLM untuk mendapatkan metadata yang berguna. Ini mungkin berguna untuk mengkategorikan dokumen secara automatik, menganalisis surat-menyurat masuk atau membina sistem pengurusan dokumen pintar. Kami akan melakukannya menggunakan beberapa alatan sumber terbuka yang popular dan memastikan perkara itu agak mudah.

Dan ya, semua perkara di bawah menganggap anda sudah cukup selesa dengan pengubah HF. Jika tidak, lihat https://huggingface.co/docs/transformers/en/quicktour - nampaknya tempat yang kukuh untuk bermula. Walaupun saya tidak pernah melakukannya dan hanya belajar daripada contoh. Saya akan melakukannya... akhirnya.

Pakej apa yang kita perlukan?

Kami akan menggunakan obor dan transformer untuk mengangkat berat, ditambah pymupdf dan kaya untuk menjadikan hidup kami lebih mudah dengan beberapa keluaran konsol mesra pengguna (Saya suka orang kaya, jadi pada asasnya kami menggunakannya untuk keseronokan).

import json
import time
import fitz

import torch
from transformers import AutoModel, AutoTokenizer, pipeline

from rich.console import Console
console = Console()

Sediakan imej

Mula-mula, imej apakah yang patut kita gunakan sebagai input? Memandangkan kami menggunakan Muka Memeluk di sini untuk tugas utama, mari gunakan halaman pertama halaman web terkemuka mereka sebagai subjek ujian kami. Ia adalah calon yang baik dengan kedua-dua teks dan pemformatan yang rumit - sesuai untuk meletakkan OCR kami melalui langkahnya.

Untuk penyelesaian yang lebih realistik, mari kita anggap input kita ialah PDF (kerana katakanlah, itulah perkara yang mungkin akan anda hadapi di dunia nyata). Kami perlu menukarnya kepada format PNG untuk diproses oleh model kami:

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)

Lakukan OCR sebenar di sini

Saya telah bermain-main dengan pelbagai penyelesaian OCR untuk tugasan ini. Sudah tentu, terdapat tesseract dan banyak pilihan lain di luar sana. Tetapi untuk kes ujian saya, saya mendapat keputusan terbaik dengan GOT-OCR2_0 (https://huggingface.co/stepfun-ai/GOT-OCR2_0). Oleh itu, mari kita teruskan dengan itu:

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()

Apa yang berlaku di sini? Nah, AutoModel dan AutoTokenizer lalai, satu-satunya bahagian yang cukup istimewa ialah kami menyediakan model untuk menggunakan cuda. Dan ini bukan pilihan. Model ini memerlukan sokongan CUDA untuk dijalankan.

Sekarang kita telah menentukan model kita, mari kita benar-benar meletakkannya untuk berfungsi pada fail yang disimpan. Selain itu, kami akan mengukur masa dan mencetaknya. Berguna bukan sahaja untuk membandingkan dengan model yang berbeza, tetapi juga untuk memahami sama ada kes penggunaan anda boleh menunggu begitu lama (walaupun ia sangat cepat untuk kes kami):

import json
import time
import fitz

import torch
from transformers import AutoModel, AutoTokenizer, pipeline

from rich.console import Console
console = Console()

Dan inilah yang kami dapat daripada imej asal kami:

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)

^ semua teks, tiada pemformatan, tetapi ia disengajakan.

GOT-OCR2_0 cukup fleksibel - ia boleh mengeluarkan dalam format yang berbeza, termasuk HTML. Berikut ialah beberapa cara lain yang boleh anda gunakan:

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()

Akhirnya cuba LLM

Kini tiba bahagian yang menyeronokkan - memilih LLM. Terdapat perbincangan yang tidak berkesudahan tentang mana yang terbaik, dengan artikel di mana-mana sahaja anda melihat. Tetapi biarlah mudah: apakah yang LLM semua orang dan anjing mereka pernah dengar? Llama. Jadi kami akan menggunakan Llama-3.2-1B untuk memproses teks.

Apa yang boleh kita perolehi daripada teks tersebut? Fikirkan perkara asas seperti klasifikasi teks, analisis sentimen, pengesanan bahasa, dll. Bayangkan anda sedang membina sistem untuk mengkategorikan dokumen yang dimuat naik secara automatik atau mengisih faks masuk untuk farmasi.

Saya akan melangkau mendalami kejuruteraan segera (itu adalah artikel lain dan saya tidak percaya saya akan menulis apa-apa), tetapi inilah idea asasnya:

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",
    )

Sebenarnya, adakah saya melakukan sesuatu yang menggelikan hati di sini dengan segera/kandungan? Beritahu saya. Agak baru dalam "kejuruteraan segera" dan tidak mengambil serius lagi.

Model kadangkala membalut hasil dalam blok kod turun, jadi kita perlu mengendalikannya (jika sesiapa tahu cara yang lebih bersih, saya semua mendengar):

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

Dan inilah yang biasanya kami peroleh sebagai output:

# 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',
)

Kesimpulannya

Kami telah membina saluran paip kecil yang boleh mengambil PDF, mengekstrak teksnya menggunakan beberapa OCR yang cukup bagus, dan kemudian menganalisis teks tersebut menggunakan LLM untuk mendapatkan metadata yang berguna. Adakah ia sedia pengeluaran? Mungkin tidak. Tetapi ia adalah titik permulaan yang kukuh jika anda ingin membina sesuatu yang serupa. Perkara yang menarik ialah bagaimana kami menggabungkan alatan sumber terbuka yang berbeza untuk mencipta sesuatu yang berguna - daripada pengendalian PDF kepada OCR kepada analisis LLM.

Anda boleh melanjutkan ini dengan mudah. Mungkin menambah pengendalian ralat yang lebih baik, sokongan untuk berbilang halaman atau mencuba LLM yang berbeza. Atau mungkin menghubungkannya dengan sistem pengurusan dokumen. Harap anda akan. Ia mungkin satu tugas yang menyeronokkan.

Ingat, ini hanyalah satu cara untuk melakukannya - mungkin terdapat berpuluh-puluh pendekatan lain yang mungkin berfungsi lebih baik untuk kes penggunaan khusus anda. Tetapi mudah-mudahan, ini memberi anda titik permulaan yang baik untuk percubaan anda sendiri! Atau tempat yang sesuai untuk mengajar saya dalam ulasan cara ia dilakukan.

Atas ialah kandungan terperinci Analisis Dokumen Pantas dan Kotor: Menggabungkan GOT-OCR dan LLama dalam Python. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!