Dataikuで実施するRAG構築 4 - granite-vision-3.1-2b-preview とvllmでPDFをOCR処理 -
前回の記事
・今回は、Windows(WSL2)のローカル環境です。
・また、使用ブラウザーはChromeです。現在はEdgeをメインに使用していたため、最初は登録が間違っていたのか？などで数日無駄に過ごしてしまいました。推奨環境はきちんと確認しましょう。
・あくまで最低限の動作確認です。
使用するモデル
ibm-granite/granite-vision-3.1-2b-previewとは、
ibmさんが提供しているOCR性能が高そうな画像認識モデルです。
対応言語は英語のみとなっていますが、モデルサイズが小さい事と、Dataikuでvllmを使ってみたい事もあったため、試してみる事にしました。
性能表
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の処理を高速化してくれるライブラリーです。
最近ではflash-attentionのv3にも対応したため、さらに早くなっているようです。
（そこまで詳しくはありません。）
対応したモデルであれば
export VLLM_USE_V1=1
という設定でさらに高速化する事ができるモデルもあるようです。
ここでダウンロードしたモデルは
echo $HF_HOME
で表示されるパスに格納されます。
また、読み込んだモデルがどのくらいVRAMを必要とするのかは、nvidia-smiコマンドで確認することができます。2bサイズモデルはさすがに軽いですね！ミドルレンジのGPUでも使えそうです。
LLM推論モデルを自PCでやる場合、途中で処理を停止させた場合のVRAM開放などがよく分かっておらず、使用してないのに占有し続ける状況があります。
その場合はカーネルをシャットダウンしたり、最悪はwslをシャットダウンなどで対応などしています。。。
処理実施後には、VRAM使用量が爆増しています。現時点、よく分かっておらず、、、。
ノートブックで、Kernel→Shutdownにすれば減らすことができます。
補足
あまり話題にならないのですが?SG Langはvllmをさらに高速化させたり、去年の前半段階時点でJSON出力形式にも対応していたため、お勧めです。
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