Unstructuredライブラリを使ったpdf読み込み処理の改善

LIFULLにはkeelaiいうAIプロジェクトがあり、私はその開発にコミットしています。

keelaiは「[OpenAI Assistants APIを使わずに無限にスケールする汎用AI(仮)を開発した](https://www.lifull.blog/entry/2023/11/16/170000)」にあるように、[Function calling](https://platform.openai.com/docs/guides/function-calling)をうまく使って複数の外部機能を組み合わせて自律的にタスクを解決するような設計です。今のところ特に社内用Slack botとして最もよく使われています。

> keelaiの基本的なコンセプトを私達はマルチエージェントと呼んでいて、サブタスクを解決するために自律的に動くエージェントを複数組み合わせて協調させることで無限にスケールすることを目指します。
>
> 現在では一般的なLLMのユースケースに加えて、例えば以下のようなユースケースにも対応しています。
>
> * Webから最新のコンテンツを取得して、社内情報と突合しながら新しいコンテンツを生成する
> * 社内のテーブルスキーマに応じたSQLの生成とバリデーション
> * 社内のデザインガイドラインに準拠した画像の生成


上の文章を読むと華々しいプロジェクトに見えます。しかし社内基盤として提供していると、想定できていなかったエラーケースが多々あり、そのエラーを根気強く潰していくような地味な仕事も多いです。

この記事は、keelaiの「URLを開いてpdfを読み込む」処理で問題が発生して、それを地道に解決していった内容をまとめます。


## unstructuredライブラリについて

URLの中身が全部テキストファイルとは限らず、様々なファイル形式があります。それに対応するため、keelaiではunstructuredを利用しています。

https://pypi.org/project/unstructured/

> The unstructured library provides open-source components for ingesting and pre-processing images and text documents, such as PDFs, HTML, Word docs, and many more. The use cases of unstructured revolve around streamlining and optimizing the data processing workflow for LLMs. 
> （和訳）
> `unstructured`ライブラリは、PDF、HTML、Word文書などの画像やテキスト文書を取り込み、前処理するためのオープンソースコンポーネントを提供しています。このライブラリのユースケースは、大規模言語モデル（LLM）のデータ処理ワークフローを効率化し、最適化することを中心に展開されています。 

当初の読み込み処理の実装はこんな感じでした。この `unstructured.partition.auto` というモジュールが、pdf以外でも、よくあるファイル形式（MSオフィスのファイル等）ならほぼ対応しています。かなり便利です。

```python
def read_unstructured_content(content: bytes) -> str:
    try:
        elements = unstructured.partition.auto.partition(file=io.BytesIO(content))
        return "\n\n".join([e.text for e in elements])
    # 文字コードがUTF-8じゃない場合の処理
    except UnicodeDecodeError:
        encoding = chardet.detect(content).get("encoding")

        if encoding is None:
            return "Failed to read content with unknown encoding"

        elements = unstructured.partition.auto.partition(file=io.BytesIO(content.decode(encoding).encode("utf-8")))
        return "\n\n".join([e.text for e in elements])
```


## 改善点①: unstructuredでモデルを事前ダウンロードする

あるとき、keelaiのサポート用チャンネルに次のような声が寄せられていました。外部企業や省庁の公開するpdfを要約する用途で使っていたようです。

> PDFの読み込みができるはずなのに、執拗に「PDFを読む能力がない」と主張されてしまうのですが、どうしたらいいでしょ。。スレッド切り替えるのが手っ取り早いですかね。

内容をよく確認すると、unstructured.partition.autoの内部で `requests` がタイムアウトするエラーが出ていました。

```
ReadTimeoutError("HTTPSConnectionPool(host='cdn-lfs.hf.co', port=443): Read timed out. (read timeout=10)")
```

なぜオンメモリのファイルを読み込んでいる処理で外部通信のTimeOutしてるのか不審だったんですが、実は初回実行時にOCRのモデルをダウンロードする処理が走るようです。開発中は問題に気づかなかったのですが、本番運用ではPodの再起動後に必ず発生するようでした。

というわけでコードを追って、Dockerのビルド時に該当のモデル（yolox）をダウンロードする処理を足しました。[`python -c` のオプション](https://docs.python.org/ja/3/using/cmdline.html)ってこういうときに使えるんですね。

```Dockerfile
COPY --from=builder /opt/builder/requirements.txt /opt/bot/requirements.txt
RUN bash -c "pip install --no-cache-dir --upgrade --no-deps -r <(sed -e '\| @ file:///|d' -e '\|-e file:///|d' /opt/bot/requirements.txt)"
RUN python -m nltk.downloader punkt averaged_perceptron_tagger
# これが足した処理
RUN python -c "import unstructured_inference.models.base;unstructured_inference.models.base.get_model('yolox')"
```


## 改善点②: 各要素のカテゴリを区別できるようにする

改善点①でTimeOutせずに該当のpdfを処理できるようになったのですが、それでもAIがpdfの内容を把握できず、要領を外した返答をしてしまっていました。

特に企業が発表するpdfは画像が多く、 `"\n\n".join([e.text for e in elements])` と安易に文字列結合すると悪影響が大きいようです。動作確認では文章の多い技術文書を使うことが多かったため、こういう挙動も想定できていませんでした。

例えば、pdfの1ページ目にスマホの画像があって適当な文字列が表示されていた場合、それとタイトルが混在してしまっているようなケースがありました。


```python
def read_unstructured_content(content: bytes) -> str:
    try:
        elements = unstructured.partition.auto.partition(file=io.BytesIO(content))
        return json.dumps(
            [
                {
                    "type": e.category,
                    "text": e.text,
                }
                for e in elements
            ],
            ensure_ascii=False,
        )
    except UnicodeDecodeError:
        encoding = chardet.detect(content).get("encoding")

        if encoding is None:
            return "Failed to read content with unknown encoding"

        elements = unstructured.partition.auto.partition(file=io.BytesIO(content.decode(encoding).encode("utf-8")))
        return json.dumps(
            [
                {
                    "type": e.category,
                    "text": e.text,
                }
                for e in elements
            ],
            ensure_ascii=False,
        )
```

というわけで、json形式で返答するようにしました。 "type" というカラムに[ElementType](https://docs.unstructured.io/open-source/concepts/document-elements#element-type)の区別を残して "Image" とか "Title" とか区別できるようにしています。

ちなみに `ensure_ascii=False` は非ASCII文字がエスケープされ、AIが日本語をきちんと読み込めなくなるので要注意です。詳しくは[こちら](https://ja.stackoverflow.com/questions/95082/json-dump%E9%96%A2%E6%95%B0%E3%81%AEensure-ascii%E5%BC%95%E6%95%B0%E3%81%AF-%E3%81%AA%E3%81%9C%E3%83%87%E3%83%95%E3%82%A9%E3%83%AB%E3%83%88%E3%81%A7true%E3%81%AA%E3%81%AE%E3%81%A7%E3%81%97%E3%82%87%E3%81%86%E3%81%8B)。


## 改善点③: ファイルサイズの上限を設ける

改善点②まででリリースしたのですが、しばらく経った後、サポート用チャンネルに別の方から次のような相談が寄せられました。

> 200ページ以上のpdfの要約を依頼したところ、実行中のまま早数時間経過しています。
> 何らかのコマンドで停止したほうが良いでしょうか？

実際は「実行中のまま」ではなく、Pythonの処理がCPUを使い切って落ちてしまっていて、エラーメッセージの返答ができていなかっただけでした。調査の結果、CPUを増強して読み込めるようにしてもモデル（gpt-4o）のtoken上限を超えてしまうようでした。

そこで次のように「15MB以上のファイルは上限を超えている」と返答するようにしました。該当のpdfを試した限りでは50ページ程度まで処理できる計算です。

```python
def read_unstructured_content(content: bytes) -> str:
    try:
        if len(content) > 15728640:
            return "Content too large to read. Please reduce size to 15MB or less."

        elements = unstructured.partition.auto.partition(file=io.BytesIO(content))
        return json.dumps(
            [
                {
                    "type": e.category,
                    "text": e.text,
                }
                for e in elements
            ],
            ensure_ascii=False,
        )
    except UnicodeDecodeError:
        encoding = chardet.detect(content).get("encoding")

        if encoding is None:
            return "Failed to read content with unknown encoding"

        elements = unstructured.partition.auto.partition(file=io.BytesIO(content.decode(encoding).encode("utf-8")))
        return json.dumps(
            [
                {
                    "type": e.category,
                    "text": e.text,
                }
                for e in elements
            ],
            ensure_ascii=False,
        )
```

これは[ダクトテーププログラマー](https://gihyo.jp/dev/clip/01/orangenews/vol55/0005)の精神の逃げの一手です。CPU増強 & pdf自体を要約してtoken制限内に収めるような対応がきれいではあるんですが、ユーザーの目的に沿った要約をするにはそれなりの工夫が必要そうなのと、そういう処理を足してもOpenAIの新モデルが出て不要になることも経験上多いためです。とはいえ要望が出てきたらしっかり対応しようと思ってます。


## まとめ

この記事では、内製AIであるkeelaiプロジェクトの「URLを開いてpdfを読み込む」処理に関する問題点と解決の試行錯誤についてまとめました。このように、keelaiがより信頼性高くつかってもらえるように地道な改善を重ねています。

また、サポートチャンネルを通して不具合や要望が集まってくるようになりました。開発チームの中では「pdfやパワポはとりあえず読み込めればいっか」くらいの温度感で実装していたのですが、実際には「企業分析を目的にpdfを要約する」ような意外なユースケースでも使われるようになっています。こうした意外なフィードバックを得られるのが楽しいです。

最後に社内向けの宣伝です。もしLIFULL社員がこの記事を読んでいたら、生成AI関係で困りごとや相談がある際に、ぜひお気軽に相談してもらえると嬉しいです！