🤗

Dataikuで実施するRAG構築 4 - granite-vision-3.1-2b-preview とvllmでPDFをOCR処理 -

2025/02/10に公開

前回の記事
https://zenn.dev/sousouplanet/articles/7f484ff053388c
・今回は、Windows(WSL2)のローカル環境です。
・また、使用ブラウザーはChromeです。現在はEdgeをメインに使用していたため、最初は登録が間違っていたのか?などで数日無駄に過ごしてしまいました。推奨環境はきちんと確認しましょう。
・あくまで最低限の動作確認です。

使用するモデル

ibm-granite/granite-vision-3.1-2b-previewとは、
ibmさんが提供しているOCR性能が高そうな画像認識モデルです。
対応言語は英語のみとなっていますが、モデルサイズが小さい事と、Dataikuでvllmを使ってみたい事もあったため、試してみる事にしました。
https://huggingface.co/ibm-granite/granite-vision-3.1-2b-preview

性能表

0 環境構築

【アドミニストレーション】→【Code Envs】
いつもと同じく、Dataikuさんを使用して、Python環境を準備します。

【Code Envs】→【上記で作成した環境を選択】→【Pakages to install】
使用したいパッケージをpipの個所に記載し、UPDATEを押します。

git+https://github.com/huggingface/transformers

torch
torchvision
torchaudio
vllm==0.6.6
huggingface_hub
Pillow
PyMuPDF
glob
flash-attn

1 使用するデータを作成・選択

Dataiakuのフロー画面にて、データセットでフォルダーを作成し、

適用対象のPDFを準備します。

2 Pythonレシピの準備

オリジナルコードを適用するため、フロー画面にて作成したデータセットを選択し、Pythonを選択します。

出力先のデータセットフォルダーを作成し、レシピ(Python)を作成しましょう。


左側のInputsの中にあるものが、最初に作成したPDFを入れたデータセットです。
このFolder:LkS11ZMyというIDを使ってデータにアクセスできるようになります。
右側でいうと

main_pdf_datasets = dataiku.Folder("LKS11ZMy")

の部分になります。
また、作成した出力先が、前の画面で作成したlocal_ocr_vllmという出力先のデータセット
pandasのDataFrameを元にしたDataiku管理下のデータセットになるようです。

一つだけしか選択できないかと思いきや、右上のInputs/Outputs画面で入出力できるデータセットを追加することができます。

こうやってDataiku上で管理することにより、自動的にフロー図の方も更新されるようです。(便利!


加工する対象は、【local_ocr_vllm_df = ... 】の個所になります。

上側タブのAdvancedより、先ほど作成したPython環境を選択しましょう。
pandasのDataFrameです。

Python environment

Selectoin behavior → Select an environment
Envirionment → Local_vllm
逐次適用させるために保存を押すのを忘れずに!!

EDIT IN NOTEBOOKを押せば、Notebook環境に移ることもできます。
ここでも、使用するKernelを作成したPython環境を選択します。

データ操作方法

Pythonレシピ作成時の提示されるコードを見ると、

main_pdf_datasets.get_info()

という行があります。
これを表示させてみると、path情報が表示されています。つまり、このpathにアクセスさえすれば元データをオリジナルのpythonコードで扱えるという事です。
この入出力情報をどうやってアクセスできるのかが分からず、コンペ終了間際の今気づくのでした。。。
下記のようにget_info()のpathを使えば、globでのアクセスが確認できました。

3 Pythonコード

今回の入力対象はFolderなので、加工処理し、Outputへ出力するためにpandas.DataFrameに変換する必要があります。
ここでは、globを使用しDataFrameに変換します。

3.1 テキストを抽出する場合 ( not OCR

import os
import re
from glob import glob
import pymupdf
from glob import glob

# Read recipe inputs
main_pdf_datasets = dataiku.Folder("LkS11ZMy")
main_pdf_datasets_info = main_pdf_datasets.get_info()

input_pdf_list = glob(main_pdf_datasets_info["path"] + "/*.pdf")

必用になるライブラリーを追加し、pdfのパスをリストとして取得します。

def extract_text_from_pdfs(pdf_paths):
    """
    PDFパスのリストを受け取り、テキストを抽出します。
    リストが空の場合はエラーを発生させ、1つの場合はソートせずに処理します。
    """

    if not pdf_paths:
        raise ValueError("PDFパスのリストが空です。")

    extracted_text = {}
    extracted_files = []

    if len(pdf_paths) == 1:
        # リストに 1 つの要素しかない場合、ソートせずに直接処理
        pdf_path = pdf_paths[0]
        try:
            with pymupdf.open(pdf_path) as doc:
                text = ""
                for page in doc:
                    text += page.get_text()
                extracted_text[os.path.basename(pdf_path)] = text
                extracted_files.append(os.path.basename(pdf_path))
        except Exception as e:
            print(f"Error processing {pdf_path}: {e}")
            return [], {} # エラー発生時は空のリストと辞書を返す
        return extracted_files, extracted_text
    else:
        # リストに複数の要素がある場合、ソートして処理
        def numerical_sort_key(filename):
            """ファイル名から数字の部分を抽出し、ソートキーとして使用します。"""
            match = re.search(r"(\d+)", filename)
            if match:
                return int(match.group(1))
            else:
                return float('inf')  # 数字がない場合は無限大を返す

        sorted_pdf_paths = sorted(pdf_paths, key=lambda path: numerical_sort_key(os.path.basename(path)))

        for pdf_path in sorted_pdf_paths:
            try:
                with pymupdf.open(pdf_path) as doc:
                    text = ""
                    for page in doc:
                        text += page.get_text()
                    extracted_text[os.path.basename(pdf_path)] = text
                    extracted_files.append(os.path.basename(pdf_path))
            except Exception as e:
                print(f"Error processing {pdf_path}: {e}")

        return extracted_files, extracted_text

pdfパスリストを与え、pymupdf処理させる関数を作成しました。

extract_text_from_pdfs(input_pdf_list[0:1])

全ファイルを処理すると時間がかかるため、ここでは1つのファイルのみを選択します。

pymupdfを使用し、単純にテキストを抽出する場合はこれでよさそうです。

3.2 OCR処理

今回のメインはOCR処理であるため、PDF処理するための関数を準備します。
vllm用のライブラリーなどを追加します。

import os
import re
from glob import glob
import pymupdf

from typing import List, Dict, Tuple, Union

import pymupdf
from PIL import Image
from huggingface_hub import hf_hub_download
from vllm import LLM, SamplingParams
import gc
from tqdm import tqdm

OCR処理用のプロンプトを準備します。

system_prompt = "<|system|>\nあなたは優秀な日本語OCR処理システムです。OCR処理してください。.\n"
question = "日本語として自然になるようにレイアウトを適切に判断し、OCR処理しmarkdown方式にしてください。"

vllmに対して画像やプロンプトを渡す関数を準備します。

def perform_ocr_on_image(image: Image, system_prompt: str, question: str, temperature: float, max_tokens: int, model: LLM) -> str:
    """
    画像に対してOCRを実行し、結果を返します。

    Args:
        image: OCR処理を行うPIL Imageオブジェクト。
        system_prompt: LLMのシステムプロンプト。
        question: LLMに与える質問。
        temperature: サンプリングパラメータの温度。
        max_tokens: 生成されるトークンの最大数。
        model: vllmのLLMインスタンス

    Returns:
        生成されたテキスト (OCRの結果)。
    """

    sampling_params = SamplingParams(
        temperature=temperature,
        max_tokens=max_tokens,
    )

    image_token = "<image>"
    prompt = f"{system_prompt}<|user|>\n{image_token}\n{question}\n<|assistant|>\n"

    inputs = {
        "prompt": prompt,
        "multi_modal_data": {
            "image": image,
        }
    }

    outputs = model.generate(inputs, sampling_params=sampling_params)
    return outputs[0].outputs[0].text

pdfとページ番号、プロンプトなどを与えて処理を実行する関数を作成します。
ここでは、ひとつのpdfファイルに対して処理を実施するようにしています。

def extract_text_and_ocr_from_pdf(
    pdf_path: str,
    target_pages: List[int],  # ページ番号 のリスト
    system_prompt: str,
    question: str,
    temperature: float = 0.2,
    max_tokens: int = 64,
) -> List[Dict]:
    """
    PDFファイルを受け取り、指定されたページに対してOCRを実行します。

    Args:
        pdf_path: PDFファイルのパス。
        target_pages: OCRを行うページ番号のリスト
        system_prompt: OCRモデルのシステムプロンプト
        question: OCRモデルに与える質問。
        temperature:  サンプリングパラメータの温度 (デフォルト: 0.2)。
        max_tokens: 生成されるトークンの最大数 (デフォルト: 64)。

    Returns:
        OCRの結果を格納した辞書のリスト。各辞書は以下のキーを持つ:
            - file_name: ファイル名
            - page_number: ページ番号
            - ocr_text: OCRで抽出されたテキスト
    """

    if not pdf_path:
        raise ValueError("PDFパスが空です。")

    results = []
    model = None  # モデルのスコープを広げるため、ここで定義

    try:
        # GraniteVision モデルの初期化
        model_path = "ibm-granite/granite-vision-3.1-2b-preview"
        model = LLM(
            model=model_path,
            limit_mm_per_prompt={"image": 1},
        )


        file_name = os.path.basename(pdf_path)

        # 全体の処理に対するプログレスバー
        with tqdm(total=len(target_pages), desc=f"Processing {file_name}", unit="page") as pbar:
            try:
                with pymupdf.open(pdf_path) as doc:

                    for page_number in target_pages:
                        if not (0 < page_number <= doc.page_count): # ページ番号は1から始まる
                            print(f"Warning: Page {page_number} is out of range for {file_name}. Skipping.")
                            continue

                        page = doc.load_page(page_number - 1) # pymupdfのページ番号は0から始まる
                        # ピクセル値が大きいとOOMになるため、画質を調整
                        pix = page.get_pixmap(matrix=pymupdf.Matrix(1, 1))
                        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

                        try:
                            ocr_text = perform_ocr_on_image(img, system_prompt, question, temperature, max_tokens, model)
                            results.append(
                                {
                                    "file_name": file_name,
                                    "page_number": page_number,
                                    "ocr_text": ocr_text,
                                }
                            )
                        except Exception as ocr_error:
                            print(f"Error during OCR on {file_name} page {page_number}: {ocr_error}")
                            results.append(
                                {
                                    "file_name": file_name,
                                    "page_number": page_number,
                                    "ocr_text": f"OCR Failed: {ocr_error}",
                                }
                            ) # OCR失敗時のエラーを格納
                        finally:
                            pbar.update(1) # ページ処理完了後にプログレスバーを更新


            except Exception as e:
                print(f"Error processing {pdf_path}: {e}")

    except KeyboardInterrupt:
        print("Interrupted by user. Releasing resources...")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
    finally:
        if model is not None:
            del model
            gc.collect()  # 明示的なガベージコレクション
            print("Model released from memory.")

    return results

3.3 OCR処理の実行

作成した関数にpdfファイルを一つ与え、先頭の5ページを処理させてみます。

ocr_results = extract_text_and_ocr_from_pdf(
        input_pdf_list[1],
        [1,2,3,4,5],
        system_prompt,
        question,
    )

モデルのダウンロードや、モデルのロードを含め、そこそこ時間がかかります。

初回のロードこそ時間はかかるものの、vllmを使用しているため高速な処理を実施することができます。
※ vllmとは、KVCacheを上手い事やる(キャッシュを再利用して早いよ!)というイメージでLLMの処理を高速化してくれるライブラリーです。
https://docs.vllm.ai/en/latest/getting_started/installation/index.html
最近ではflash-attentionのv3にも対応したため、さらに早くなっているようです。
(そこまで詳しくはありません。)
対応したモデルであれば

export VLLM_USE_V1=1

という設定でさらに高速化する事ができるモデルもあるようです。
ここでダウンロードしたモデルは

echo $HF_HOME

で表示されるパスに格納されます。
https://huggingface.co/docs/huggingface_hub/main/en/package_reference/environment_variables
また、読み込んだモデルがどのくらいVRAMを必要とするのかは、nvidia-smiコマンドで確認することができます。2bサイズモデルはさすがに軽いですね!ミドルレンジのGPUでも使えそうです。

LLM推論モデルを自PCでやる場合、途中で処理を停止させた場合のVRAM開放などがよく分かっておらず、使用してないのに占有し続ける状況があります。
その場合はカーネルをシャットダウンしたり、最悪はwslをシャットダウンなどで対応などしています。。。

処理実施後には、VRAM使用量が爆増しています。現時点、よく分かっておらず、、、。
ノートブックで、Kernel→Shutdownにすれば減らすことができます。

補足

あまり話題にならないのですが?SG Langはvllmをさらに高速化させたり、去年の前半段階時点でJSON出力形式にも対応していたため、お勧めです。
https://github.com/sgl-project/sglang

3.4 結果

INFO 02-10 07:34:24 config.py:510] This model supports multiple tasks: {'generate', 'score', 'classify', 'embed', 'reward'}. Defaulting to 'generate'.
INFO 02-10 07:34:24 llm_engine.py:234] Initializing an LLM engine (v0.6.6) with config: model='ibm-granite/granite-vision-3.1-2b-preview', speculative_config=None, tokenizer='ibm-granite/granite-vision-3.1-2b-preview', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=16384, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=ibm-granite/granite-vision-3.1-2b-preview, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=False, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"splitting_ops":["vllm.unified_attention","vllm.unified_attention_with_output"],"candidate_compile_sizes":[],"compile_sizes":[],"capture_sizes":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":256}, use_cached_outputs=False, 
INFO 02-10 07:34:26 model_runner.py:1094] Starting to load model ibm-granite/granite-vision-3.1-2b-preview...
INFO 02-10 07:34:26 weight_utils.py:251] Using model weights format ['*.safetensors']
Loading safetensors checkpoint shards:   0% Completed | 0/2 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  50% Completed | 1/2 [13:50<13:50, 830.62s/it]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [16:21<00:00, 430.48s/it]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [16:21<00:00, 490.50s/it]

INFO 02-10 07:50:48 model_runner.py:1099] Loading model weights took 5.5532 GB
INFO 02-10 07:50:50 worker.py:241] Memory profiling takes 1.45 seconds
INFO 02-10 07:50:50 worker.py:241] the current vLLM instance can use total_gpu_memory (23.99GiB) x gpu_memory_utilization (0.90) = 21.59GiB
INFO 02-10 07:50:50 worker.py:241] model weights take 5.55GiB; non_torch_memory takes -0.02GiB; PyTorch activation peak memory takes 1.11GiB; the rest of the memory reserved for KV Cache is 14.94GiB.
INFO 02-10 07:50:50 gpu_executor.py:76] # GPU blocks: 12241, # CPU blocks: 3276
INFO 02-10 07:50:50 gpu_executor.py:80] Maximum concurrency for 16384 tokens per request: 11.95x
INFO 02-10 07:50:50 model_runner.py:1415] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
Capturing CUDA graph shapes: 100%|██████████| 35/35 [00:13<00:00,  2.52it/s]
INFO 02-10 07:51:04 model_runner.py:1535] Graph capturing finished in 14 secs, took 0.00 GiB
INFO 02-10 07:51:04 llm_engine.py:431] init engine (profile, create kv cache, warmup model) took 16.18 seconds

Processing 14.pdf:   0%|          | 0/5 [00:00<?, ?page/s]
Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Processed prompts: 100%|██████████| 1/1 [00:04<00:00,  4.74s/it, est. speed input: 1061.46 toks/s, output: 108.42 toks/s]
Processing 14.pdf:  20%|██        | 1/5 [00:04<00:19,  4.86s/page]
Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Processed prompts: 100%|██████████| 1/1 [00:03<00:00,  3.39s/it, est. speed input: 1615.82 toks/s, output: 110.75 toks/s]
Processing 14.pdf:  40%|████      | 2/5 [00:08<00:12,  4.23s/page]
Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Processed prompts: 100%|██████████| 1/1 [00:04<00:00,  4.52s/it, est. speed input: 1210.82 toks/s, output: 113.76 toks/s]
Processing 14.pdf:  60%|██████    | 3/5 [00:13<00:08,  4.42s/page]
Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Processed prompts: 100%|██████████| 1/1 [00:04<00:00,  4.52s/it, est. speed input: 1210.23 toks/s, output: 113.70 toks/s]
Processing 14.pdf:  80%|████████  | 4/5 [00:17<00:04,  4.47s/page]
Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Processed prompts: 100%|██████████| 1/1 [00:04<00:00,  4.51s/it, est. speed input: 1212.65 toks/s, output: 113.93 toks/s]
Processing 14.pdf: 100%|██████████| 5/5 [00:22<00:00,  4.48s/page]
Model released from memory.

HF_HOMEをCドライブではない場所に置いてるせいか?設定が悪いのか?
モデルサイズのわりにロードが20分近くかかってしまいましたが、回答は一ページに対して5秒という速度で処理ができています。

結果としては、雑のプロンプト・設定とモデルサイズに応じた内容でしょうか・・・?
様々な工夫の余地があるかと思います。

local_ocr_vllm_df = pd.DataFrame(ocr_results)

# Write recipe outputs
local_ocr_vllm = dataiku.Dataset("local_ocr_vllm")
local_ocr_vllm.write_with_schema(local_ocr_vllm_df)

あとは、処理した結果をDataFrameにして、dataikuのデータセット先に書き込むだけです。
以下、出力されたデータセットデータ

LLMは非常に難しいですが、面白いですね。

Discussion