こんにちは、CCCMKホールディングスAIエンジニアの三浦です。
ファイルを読み込み、その内容について回答が出来るLLM Agentを作ることが時々あるのですが、毎度ファイル読み込み処理を作るのが面倒だな、と感じることがありました。PDFやPPTXファイルからテキストデータを抽出したいというニーズはLLM Agentに限らず色々なシーンで発生するので、WebAPIにして様々なアプリケーションから利用できるようにしてみようと考えました。
稼働する環境はdatabricksのModel Servingで構築しました。Model Servingを利用することで、databricksを利用できる環境さえあれば追加でサーバーなどを構築する必要なくAPIを提供することが可能です。
今回の記事ではModel ServingでPDF/PPTXファイルからテキストデータを抽出するAPIの構築から提供するまでの手順をソースコードを交えてご紹介したいと思います。
Model Servingに任意のPythonコードをデプロイするために
databricksのModel ServingはMLflowで管理されたモデルをデプロイするものですが、MLflowのPyFuncというモデルを使うことで任意のPythonコードや関数をデプロイすることが出来ます。
PyFuncを使う場合、上のドキュメントにあるようにモデルロード時に1度だけ実行されるload_context関数とリクエストが来た際に実行されるpredict関数を実装します。あとはPyFuncモデルを通常のモデルと同様にlog_modelでExperimentに記録し、register_modelでUnity Catalogに登録したのちModel Serving Endpointにデプロイする手順を取ります。
PythonでPDFとPPTXの内容をテキスト化する
肝心のPDFやPPTXの内容をテキストデータで取得する処理は、PDFはPyMuPDFを、PPTXはpython-pptxというライブラリを使って実装しました。
・PyMuPDFのドキュメント
・python-pptxのドキュメント
どちらも軽量で、CPU環境で動かすことが出来ました。ただOCR機能はついていないので、その機能を付けようとするともしかしたらもう少しリソースが必要になるかもしれません。
具体的な実装
ユーティリティ関数
PyFuncモデルで使用する関数をPythonファイルに以下の内容で実装し、保存しておきます(とりあえず"utils.py"という名前で作成)。
import base64 import io from typing import Any, Dict, Optional, Tuple, List import fitz from pptx import Presentation def decode_base64_to_bytes(file_base64: str) -> bytes: """base64文字列をbytesに変換する""" return base64.decodebytes(file_base64.encode("utf-8")) def detect_ext(filename: str) -> str: """ファイル名から拡張子を判定する""" if not filename or "." not in filename: return "" return filename.lower().rsplit(".",1)[-1] def extract_pdf_text(file_bytes: bytes)->Tuple[str, dict]: """PDF形式のファイルbytesからテキストを抽出する""" doc = fitz.open(stream=file_bytes, filetype="pdf") pages = [] for i, page in enumerate(doc): text = page.get_text("text") pages.append(text) text = "==============================\n".join(pages) # ページ区切り meta = {"pages": len(doc), "filetype": "pdf"} doc.close() return text.strip(), meta def extract_pptx_text(file_bytes: bytes)->Tuple[str, dict]: """PPTX形式のファイルbytesからテキストを抽出する""" prs = Presentation(io.BytesIO(file_bytes)) pages = [] for si, slide in enumerate(prs.slides): pages.append(f"[Slide {si+1}]") for shape in slide.shapes: if hasattr(shape, "text"): pages.append(shape.text) text = "==============================\n".join(pages) meta = {"pages": len(prs.slides), "filetype": "pptx"} return text.strip(), meta
- decode_base64_to_bytes: リクエストで受け取るbase64文字列に変換されたファイルデータをbytesに変換
- detect_ext: リクエストで受け取るファイル名からファイル拡張子を抽出
- extract_pdf_text: PDF形式のbytesからテキストデータとメタデータを抽出
- extract_pptx_text: PPTX形式のbytesからテキストデータとメタデータを抽出
モデル記録用のNotebook
databricksのWorkspaceで先ほど作成した"utlis.py"と同階層に新しくNotebookを作成し、PyFuncモデルの定義と記録用の処理を書いていきます。
まず必要なライブラリをインストールします。
%pip install mlflow>=3.2 pymupdf python-pptx dbutils.library.restartPython()
Model Servingにデプロイするモデルを以下のように実装します。Model Servingでは、リクエストはList形式で受け付けるため、predictのmodel_inputはListが格納されることを想定して実装する必要があります。
このモデルはfile_base64とfilenameをリクエストで受け取ります。それぞれ対象のファイルデータのbase64文字列、対象のファイル名が格納されることを想定しています。
import base64 from utils import ( decode_base64_to_bytes, detect_ext, extract_pdf_text, extract_pptx_text ) class TextExtractorPyFunc(mlflow.pyfunc.PythonModel): def load_context(self, context): """モデルロード時に1度だけ実行される。今回は使用しない。""" pass def predict(self, context, model_input:list[Dict[str, str]]) -> list[Dict[str, Dict]]: """リクエスト処理""" b64 = model_input[0].get("file_base64") if not b64: raise ValueError("file_base64 is required") filename = model_input[0].get("filename","") ext = detect_ext(filename) file_bytes = decode_base64_to_bytes(b64) if ext == "pdf": text, meta = extract_pdf_text(file_bytes) elif ext == "pptx": text, meta = extract_pptx_text(file_bytes) else: raise ValueError(f"Unsupported file type: {ext}") return [{"text": text, "meta": meta}]
MLflowのExperimentへの記録からUnity Catalogへの登録まで、次のコードで進めることが出来ます。log_modelを実行する際にモデルへの入力例を与えるため、サンプルのPDFファイルを読み込んでinput_exampleに渡すようにしました。
import base64 import mlflow # input_example用にサンプルファイルを読み込む with open("sample.pdf", "rb") as f: b64 = base64.b64encode(f.read()).decode("utf-8") input_example={ "file_base64": b64, "filename":"sample.pdf" } model_name = "text_extractor" with mlflow.start_run() as run: #モデルの記録 model_info = mlflow.pyfunc.log_model( name=model_name, python_model=TextExtractorPyFunc(), code_paths=["utils.py"], # ユーティリティ関数のPythonファイル input_example=[input_example] ) #モデルのUnity Catalogへの記録 model_uri = model_info.model_uri reg = mlflow.register_model(model_uri, f"db.schema.{model_name}")
ちなみにモデルのUnity Catalogへの登録はGUIでも出来るのですが、何か問題があった時の原因が掴みにくいのと、登録までに時間がかかるように感じられるため、コードで実現したほうが良いと思います。
Unity Catalogへの登録が完了したら、あとはModel Servingへのデプロイに進むことが出来ます。こちらの手順は前回の記事と同様ですので、そちらをご覧ください。
Model Servingにデプロイしたテキスト抽出APIを使ってみる
Model ServingのAPIはアクセストークンを利用して接続が出来ます。
以下はdatabricksの同じWorkspaceから接続を試みた際に作成したコードです。
import requests # xxxはWorkspace固有の文字列を設定 endpoint = "https://adb-xxxxxxxxxxx.azuredatabricks.net/serving-endpoints/text-extractor/invocations" # PDFファイルの読み込み with open("test.pdf", "rb") as f: b64_pdf = base64.b64encode(f.read()).decode("utf-8") input_data = { "inputs": [{ "file_base64": b64_pdf, "filename":"test.pdf" }] } headers = { # トークンを取得 "Authorization": f"Bearer {dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()}" } response = requests.request(method="POST", headers=headers, url=endpoint, json=input_data) print(response.json())
この処理をツールにしてLLM Agentに渡すことでPDFやPPTXのファイルの中身を参照して回答出来るアプリなどを開発することが出来ます。
まとめ
ということで今回はdatabricksのModel ServingでPDFやPPTXファイルの中身をテキスト化するAPIを構築した話をご紹介しました。Model Servingに今回のように色々な機能を実装していくことでAgentのツールをどんどん拡張していくとこも可能になりそうです。今度は画像解析などのもう少し処理にリソースが必要になりそうな機能もModel Servingで動かしてみたいと思います。
