문서 추출을 Python 변환으로 배포하기

확인한 전략을 Python 변환으로 배포하면 미디어 집합에 있는 모든 문서의 모든 페이지를 배치로 추출합니다. 추출 코드 자체를 처음부터 작성하기보다 Document Intelligence의 배포 도구로 저장소를 생성합니다. 생성 템플릿은 선택한 전략과 같은 구성을 사용합니다.

저장소 실행

src/myproject/document_extraction/my_extraction.py의 @transform.using에서 출력 데이터셋을 지정하고 빌드를 시작합니다. 현재 템플릿은 경량 변환을 사용합니다. 이전 Spark 기반 템플릿은 Spark 준비 비용이 크므로 경량 방식으로 마이그레이션하는 것이 권장됩니다.

템플릿은 미리보기 모드를 지원하지 않습니다. 미리보기에서 오류가 발생할 수 있으므로 이 템플릿의 동작을 확인할 때 실제 빌드 경로와 혼동하지 않습니다.

기본은 비증분 처리여서 실행마다 모든 문서를 처리합니다. @incremental(...) 행의 주석을 해제하면 입력 미디어 집합에 새로 추가된 문서만 처리하고 결과를 출력 데이터셋에 추가합니다.

프롬프트와 이미지 전처리

생성형 AI의 프롬프트는 src/myproject/document_extraction/prompts.py에 있습니다. 템플릿을 직접 수정하면 앱에서 시험한 결과와 배치 결과가 달라질 수 있으므로 Document Intelligence에서 수정하고 확인한 뒤 다시 배포합니다.

VisionLLMDocumentsExtractorInput은 사용자 프롬프트만 조정하며 시스템 프롬프트는 고정입니다. 레이아웃 방식의 VisionLLMLayoutDocumentsExtractorInput은 시스템 프롬프트만 조정합니다. 레이아웃 사용자 프롬프트에는 구조를 보존하는 JSON 스키마가 있으므로 바꾸면 추출 성공률이 크게 낮아질 수 있습니다.

회전된 글자처럼 이미지 변환이 필요하면 별도 변환 파이프라인에서 이미지를 처리하고 새 미디어 집합에 저장한 뒤 그 집합을 Document Intelligence 입력으로 사용합니다.

오류 재처리와 동시성

레이아웃 생성형 추출은 모델이 정해진 JSON 스키마를 충족해야 합니다. 잘못된 JSON이나 스키마 불일치는 ERROR_RESPONSE_JSON_PARSING을 발생시킵니다. 고성능 모델도 약 5%의 추출에서 실패할 수 있다는 운영 예가 있으며 보장된 실패율은 아닙니다. 실패한 행에도 레이아웃 모델만의 추출 결과가 담긴 유효한 layoutInfo가 남습니다.

실패한 미디어 항목만 다시 처리하려면 filter_on_media_items에 항목 ID 목록을 전달하고 @incremental을 제거하여 이미 처리한 행으로 판단되지 않게 합니다.

THREAD_NUMBER는 페이지별 추출 스레드 수입니다. 기본은 20이며 VLM 용량이 충분한 개발 환경에서 시험한 최대값은 300입니다. 실제 허용 한도라는 뜻은 아닙니다. 용량이 부족한 환경에서 과도하게 높이면 처리량 한도 오류와 재시도가 같은 모델을 쓰는 다른 작업의 용량까지 소비할 수 있습니다. 값을 조정할 때 사용량을 관측합니다. 빌드 상세의 Telemetry에서 메시지가 aip_workflows로 시작하는 항목을 필터링하면 추출 로그를 찾을 수 있습니다.

페이지별 청크와 문서별 청크

추출 출력은 페이지당 한 행입니다. DocumentChunker.create_chunks_per_document는 같은 문서의 페이지를 하나의 Markdown 문자열로 합친 뒤 분할합니다. create_chunks_per_row는 각 행을 독립적으로 분할합니다.

chunking_result = chunker.create_chunks_per_row(
    extraction_df,
    chunk_mode="markdown",       # "recursive" for plain text, "markdown" for markdown text
    content_column="extractionResult",
    id_column="media_item_rid",  # used as prefix for chunk_id
    chunk_size=8192,
    chunk_overlap=0,
    thread_number=20,
    strip_markdown=False,        # set True to remove ```markdown and ``` wrappers before chunking
)

chunk_size와 chunk_overlap은 청크 크기와 겹침을 정합니다. 임베딩 모델의 문맥 한도가 있으므로 임베딩 전에 분할하는 것이 권장됩니다. 파이프라인 구조를 유지하면서 분할을 사실상 생략하려면 create_chunks_per_row의 크기를 sys.maxsize처럼 매우 큰 값으로 둘 수 있지만 임베딩 입력 한도는 별도로 충족해야 합니다.

전통적 추출 템플릿

아래 코드는 추출 구성별 생성 형태를 보여 줍니다. 데이터셋, 미디어 집합 및 모델 식별자는 설명용 값이며 실제 배포 도구가 해당 환경의 값으로 생성합니다. 미디어 항목이 문서가 아니거나 페이지 수가 없으면 ValueError가 발생합니다. 각 페이지를 병렬 처리하고 RID, 미디어 참조, 페이지 번호와 결과를 출력합니다.

Raw text

전자 PDF의 내장 텍스트를 UTF-8로 읽습니다.

import polars as pl
from concurrent.futures import ThreadPoolExecutor
from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from transforms.mediasets.utils._constants import MEDIA_ITEM_RID, MEDIA_REFERENCE, PATH


THREAD_NUMBER = 20


# @incremental(v2_semantics=True)  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
)
def extract(media_input, output):
    """
    Extracts content from pdf documents with raw text extraction
    """
    media_refs = pl.from_pandas(
        media_input.list_media_items_by_path_with_media_reference().pandas(),
        schema_overrides={MEDIA_ITEM_RID: pl.String, MEDIA_REFERENCE: pl.String, PATH: pl.String},
    )

    def process_batch(batch_df: pl.DataFrame) -> pl.DataFrame:
        def create_page_tasks(row):
            media_item_rid = row[MEDIA_ITEM_RID]
            metadata = media_input.get_media_item_metadata(media_item_rid).document
            if metadata is None:
                raise ValueError(f"Media item {media_item_rid} is not a document")
            if metadata.pages is None:
                raise ValueError(f"Media item {media_item_rid} has no page count")
            return [(row, page_num) for page_num in range(metadata.pages)]

        def process_single_page(task):
            row, page_num = task
            media_item_rid = row[MEDIA_ITEM_RID]
            media_reference = row[MEDIA_REFERENCE]
            extraction_result = media_input.transform_document_to_text_raw(
                media_item_rid, page_num
            ).read().decode("utf-8")
            return {
                "media_item_rid": media_item_rid,
                "media_reference": media_reference,
                "page_num": page_num,
                "extraction_result": extraction_result
            }

        all_tasks = []
        for row in batch_df.iter_rows(named=True):
            all_tasks.extend(create_page_tasks(row))

        with ThreadPoolExecutor(max_workers=THREAD_NUMBER) as executor:
            results = list(executor.map(process_single_page, all_tasks))

        return pl.DataFrame(results)

    extracted_data = media_refs.lazy().map_batches(
        process_batch,
        schema={
            "media_item_rid": pl.String,
            "media_reference": pl.String,
            "page_num": pl.Int64,
            "extraction_result": pl.String,
        },
        streamable=True,
    )

    output.write_dataframe(extracted_data)

OCR

레이아웃을 보존하지 않는 OCR 텍스트를 추출합니다.

import polars as pl
from concurrent.futures import ThreadPoolExecutor
from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from transforms.mediasets.utils._constants import MEDIA_ITEM_RID, MEDIA_REFERENCE, PATH


THREAD_NUMBER = 20


# @incremental(v2_semantics=True)  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
)
def extract(media_input, output):
    """
    Extracts content from pdf documents with OCR text extraction
    """
    media_refs = pl.from_pandas(
        media_input.list_media_items_by_path_with_media_reference().pandas(),
        schema_overrides={MEDIA_ITEM_RID: pl.String, MEDIA_REFERENCE: pl.String, PATH: pl.String},
    )

    def process_batch(batch_df: pl.DataFrame) -> pl.DataFrame:
        def create_page_tasks(row):
            media_item_rid = row[MEDIA_ITEM_RID]
            metadata = media_input.get_media_item_metadata(media_item_rid).document
            if metadata is None:
                raise ValueError(f"Media item {media_item_rid} is not a document")
            if metadata.pages is None:
                raise ValueError(f"Media item {media_item_rid} has no page count")
            return [(row, page_num) for page_num in range(metadata.pages)]

        def process_single_page(task):
            row, page_num = task
            media_item_rid = row[MEDIA_ITEM_RID]
            media_reference = row[MEDIA_REFERENCE]
            extraction_result = media_input.transform_document_to_text_ocr_output_text(
                media_item_rid, page_num
            ).read().decode("utf-8")
            return {
                "media_item_rid": media_item_rid,
                "media_reference": media_reference,
                "page_num": page_num,
                "extraction_result": extraction_result
            }

        all_tasks = []
        for row in batch_df.iter_rows(named=True):
            all_tasks.extend(create_page_tasks(row))

        with ThreadPoolExecutor(max_workers=THREAD_NUMBER) as executor:
            results = list(executor.map(process_single_page, all_tasks))

        return pl.DataFrame(results)

    extracted_data = media_refs.lazy().map_batches(
        process_batch,
        schema={
            "media_item_rid": pl.String,
            "media_reference": pl.String,
            "page_num": pl.Int64,
            "extraction_result": pl.String,
        },
        streamable=True,
    )

    output.write_dataframe(extracted_data)

Layout-aware OCR

V2 응답을 JSON 문자열로 저장합니다. pages 배열에는 해당 페이지 블록이 있으며 format=TEXT일 때 각 블록의 text는 type과 text 필드를 가진 태그된 값입니다.

import json
import polars as pl
from concurrent.futures import ThreadPoolExecutor
from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from transforms.mediasets.utils._constants import MEDIA_ITEM_RID, MEDIA_REFERENCE, PATH


THREAD_NUMBER = 20


# @incremental(v2_semantics=True)  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
)
def extract(media_input, output):
    """
    Extracts content from pdf documents with layout-aware OCR extraction
    """
    media_refs = pl.from_pandas(
        media_input.list_media_items_by_path_with_media_reference().pandas(),
        schema_overrides={MEDIA_ITEM_RID: pl.String, MEDIA_REFERENCE: pl.String, PATH: pl.String},
    )

    def process_batch(batch_df: pl.DataFrame) -> pl.DataFrame:
        def create_page_tasks(row):
            media_item_rid = row[MEDIA_ITEM_RID]
            metadata = media_input.get_media_item_metadata(media_item_rid).document
            if metadata is None:
                raise ValueError(f"Media item {media_item_rid} is not a document")
            if metadata.pages is None:
                raise ValueError(f"Media item {media_item_rid} has no page count")
            return [(row, page_num) for page_num in range(metadata.pages)]

        def process_single_page(task):
            row, page_num = task
            media_item_rid = row[MEDIA_ITEM_RID]
            media_reference = row[MEDIA_REFERENCE]
            extraction_result = media_input.transform_media_item(media_item_rid, str(page_num), {
                "type": "documentToText",
                "documentToText": {
                    "operation": {
                        "type": "extractLayoutAwareTextV2",
                        "extractLayoutAwareTextV2": {
                            "pageRange": {
                                "startPageInclusive": page_num,
                                "endPageExclusive": page_num + 1
                            },
                            "config": {
                                "mode": "SCAN",
                                "format": "TEXT",
                                "languages": [{"type": "language", "language": "ENG"}]
                            }
                        }
                    }
                }
            })
            extraction_result = json.dumps(extraction_result.json())
            return {
                "media_item_rid": media_item_rid,
                "media_reference": media_reference,
                "page_num": page_num,
                "extraction_result": extraction_result
            }

        all_tasks = []
        for row in batch_df.iter_rows(named=True):
            all_tasks.extend(create_page_tasks(row))

        with ThreadPoolExecutor(max_workers=THREAD_NUMBER) as executor:
            results = list(executor.map(process_single_page, all_tasks))

        return pl.DataFrame(results)

    extracted_data = media_refs.lazy().map_batches(
        process_batch,
        schema={
            "media_item_rid": pl.String,
            "media_reference": pl.String,
            "page_num": pl.Int64,
            "extraction_result": pl.String,
        },
        streamable=True,
    )

    output.write_dataframe(extracted_data)

생성형 AI 템플릿

기본 VLM

전처리 없이 페이지에서 Markdown을 추출합니다. USER_PROMPT는 배포 도구가 생성한 프롬프트 모듈에서 가져옵니다.

from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from aip_workflows.document_intelligence.transforms import VisionLLMDocumentsExtractorInput
from .prompts import USER_PROMPT


THREAD_NUMBER = 20


# @incremental(v2_semantics=True, snapshot_inputs=["extractor"])  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
    extractor=VisionLLMDocumentsExtractorInput(
        "ri.language-model-service..language-model.anthropic-claude-xxx-sonnet"
    ),
)
def extract(media_input, output, extractor):
    """
    Extracts content from pdf documents as markdown.
    """
    extracted_data = extractor.create_extraction(
        media_input, with_ocr=False, prompt=USER_PROMPT, thread_number=THREAD_NUMBER
    )
    output.write_dataframe(extracted_data)

OCR 전처리 결합

with_ocr=True로 OCR 결과도 모델에 제공합니다.

from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from aip_workflows.document_intelligence.transforms import VisionLLMDocumentsExtractorInput
from .prompts import USER_PROMPT


THREAD_NUMBER = 20


# @incremental(v2_semantics=True, snapshot_inputs=["extractor"])  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
    extractor=VisionLLMDocumentsExtractorInput(
        "ri.language-model-service..language-model.anthropic-claude-xxx-sonnet"
    ),
)
def extract(media_input, output, extractor):
    """
    Extracts content from pdf documents as markdown.
    """
    extracted_data = extractor.create_extraction(
        media_input, with_ocr=True, prompt=USER_PROMPT, thread_number=THREAD_NUMBER
    )
    output.write_dataframe(extracted_data)

레이아웃 추출

include_layout_info="no_overlay"로 레이아웃 정보를 함께 반환합니다.

from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from aip_workflows.document_intelligence.transforms import VisionLLMLayoutDocumentsExtractorInput
from .prompts import SYSTEM_PROMPT


THREAD_NUMBER = 20


# @incremental(v2_semantics=True, snapshot_inputs=["extractor"])  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
    extractor=VisionLLMLayoutDocumentsExtractorInput(
        "ri.language-model-service..language-model.anthropic-claude-xxx-sonnet"
    ),
)
def extract(media_input, output, extractor):
    """
    Extracts content from pdf documents as markdown.
    """
    extracted_data = extractor.create_extraction(
        media_input,
        include_layout_info="no_overlay",
        system_prompt=SYSTEM_PROMPT,
        thread_number=THREAD_NUMBER
    )
    output.write_dataframe(extracted_data)

표 영역 자르기

include_layout_info="crop_tables"로 표 영역을 잘라 표 추출에 사용합니다.

from transforms.api import Output, incremental, transform
from transforms.mediasets import MediaSetInput
from aip_workflows.document_intelligence.transforms import VisionLLMLayoutDocumentsExtractorInput
from .prompts import SYSTEM_PROMPT


THREAD_NUMBER = 20


# @incremental(v2_semantics=True, snapshot_inputs=["extractor"])  # uncomment this line if incremental is needed
@transform.using(
    output=Output("ri.foundry.main.dataset.abc"),
    media_input=MediaSetInput("ri.mio.main.media-set.abc"),
    extractor=VisionLLMLayoutDocumentsExtractorInput(
        "ri.language-model-service..language-model.anthropic-claude-xxx-sonnet"
    ),
)
def extract(media_input, output, extractor):
    """
    Extracts content from pdf documents as markdown.
    """
    extracted_data = extractor.create_extraction(
        media_input,
        include_layout_info="crop_tables",
        system_prompt=SYSTEM_PROMPT,
        thread_number=THREAD_NUMBER
    )
    output.write_dataframe(extracted_data)

청크 분할과 임베딩

다음 예시는 추출 데이터셋을 읽어 문서별로 합치고 Markdown 청크를 만든 뒤 임베딩을 생성합니다. 입력 결과 열이 extraction_result인 전통적 템플릿을 연결한다면 예시의 content_column="extractionResult"를 실제 열 이름에 맞춰야 합니다. 임베딩이 필요 없으면 embedder 선언과 임베딩 호출을 제거하고 chunking_result를 출력합니다.

from transforms.api import Input, Output, incremental, transform
from aip_workflows.document_intelligence.transforms import DocumentChunker, DocumentEmbedderInput

THREAD_NUMBER = 20


# @incremental(v2_semantics=True, snapshot_inputs=["embedder"])  # uncomment this line if incremental is needed
@transform.using(
    extraction_input=Input("ri.foundry.main.dataset.abc"),  # typically the output dataset from the extraction transform
    output=Output("ri.foundry.main.dataset.xyz"),
    embedder=DocumentEmbedderInput("ri.language-model-service..language-model.text-embedding-3-large"),
)
def chunk_and_embed(extraction_input, output, embedder):
    extraction_df = extraction_input.polars(lazy=True)
    chunker = DocumentChunker()
    chunking_result = chunker.create_chunks_per_document(
        extraction_df,
        chunk_mode="markdown",  # "recursive" for raw text, "markdown" for markdown text
        content_column="extractionResult",  # content column name
        id_column="media_item_rid",  # id of the document, used to combine content (e.g. from different pages) of the single document
        page_column="page_num",  # page number column name
        chunk_size=8192,
        chunk_overlap=0,
        thread_number=THREAD_NUMBER,
        strip_markdown=True,  # when True, removes ```markdown prefix and ``` suffix from the content before chunking
    )
    embedding_result = embedder.create_embeddings(
        chunking_result,
        content_column="chunk_content",
        thread_number=THREAD_NUMBER,
    )
    output.write_dataframe(embedding_result)