AI/LLM

SmolDocling으로 OCR 하는 방법 및 후기(PyMuPDFLoader markdown 비교)

bigc 2025. 4. 1. 15:25
반응형

회사에서 RAG를 하기 위해 표가 있는 문서를 markdown으로 추출하는 여러 기법들을 테스트해보고 있었습니다.

지금까지 해본 오픈 모델 중엔 PyMuPDFLoader이 그나마 가장 깔끔하게 텍스트, table을  markdown 형식으로 변환해주었습니다.

 

- PyMuPDFLoader 예시

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings

# 1. PDF 로드 + 마크다운으로 표 추출
loader = PyMuPDFLoader(
    "./llmissue_spri.pdf",
    mode="page",
    extract_tables="markdown"
)
pages = loader.load()

# 2. 텍스트 분할기 설정
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
)

# 3. 문서 분할
documents = text_splitter.split_documents(pages)

 

- table 부분 markdown 결과

Llama3.3 70B 모델로 위 내용 vector DB에 저장 후 질문해보면 답 잘함.

 

But 첫 번째로, 벡터 임베딩 모델을 multilingual-e5-large-instruct를 사용하여 chunk_size를 1000까지 할수 있기에, 한 table이 1000을 넘어가는 사이즈의 내용을 포함시 markdown이 되지 않는 문제가 발생.

두 번째로, 복잡한 table에 대해선 언어 모델이 답을 제대로 못함.

- ex) Q: "Google의 모델 종류에 대해 알려줘"   A: 위 표의 나눠진 3행 데이터 전부 가져오지 못함 or 웹 서칭으로 가버림.

 

 

문서 페이지를 이미지로 읽어 DocTags라는 semantic markup을 출력하는 256M Visual Language Model SmolDocling이 얼마전에 나와서 조금의 기대를 가지고 바로 테스트를 해봅니다.

 

- SmolDocling 사용 예시

https://colab.research.google.com/drive/1sKYsLWEv7YicIFDyXgSzwIrosaMxC4WV?usp=sharing#scrollTo=C66AqIxj0w4G

!pip install docling_core
!pip install flash-attn
import torch
from docling_core.types.doc import DoclingDocument
from docling_core.types.doc.document import DocTagsDocument
from transformers import AutoProcessor, AutoModelForVision2Seq
from transformers.image_utils import load_image
from pdf2image import convert_from_path
import os

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load model and processor
processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
model = AutoModelForVision2Seq.from_pretrained(
    "ds4sd/SmolDocling-256M-preview",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
).to(DEVICE)

# Convert only page 7 of the PDF
pdf_path = "KDS677015.pdf"
page7 = convert_from_path(pdf_path, first_page=7, last_page=7)[0]

# Save and reload with transformers' image loader
temp_filename = "temp_page_7.png"
page7.save(temp_filename)
image = load_image(temp_filename)
os.remove(temp_filename)  # cleanup

# Prepare prompt
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": "Convert this page to docling."}
        ]
    },
]
prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=prompt, images=[image], return_tensors="pt").to(DEVICE)

# Generate doctags
generated_ids = model.generate(**inputs, max_new_tokens=8192)
prompt_len = inputs.input_ids.shape[1]
trimmed_ids = generated_ids[:, prompt_len:]
doctags = processor.batch_decode(trimmed_ids, skip_special_tokens=False)[0].lstrip()

# Create docling document
doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
doc = DoclingDocument(name="Page7Document")
doc.load_from_doctags(doctags_doc)

# Export to Markdown
print(doc.export_to_markdown())

 

PDF를 넣어보니 예전 Docling이랑 똑같이 한국어에 대해선 아얘 노답 성능..

table 말고 일반 text 조차 대부분 오타로 가져옴.

 

아직까지는 한국어 기반 문서로 vector DB 저장시 PyMuPDFLoader 사용할 것 같습니다.

 

반응형