Zenn
🏆

LlamaIndexのOpenAI Pydantic Program：構造化データ生成方法

2024/08/18に公開
Pydantic
構造化データ
LlamaIndex
OpenAI API
tech

はじめに

OpenAIの新しいAPIを使用して、LlamaIndexを介して構造化データを生成する方法を紹介します。このガイドでは、ユーザーがPydanticオブジェクトを指定するだけで、簡単に構造化データを抽出できることを示します。

主に以下の設定を紹介します：

  • Albumオブジェクトへの抽出（複数のSongオブジェクトを含むことができます）

環境設定

まず、必要なライブラリをインストールしましょう。Google Colabで実行している場合は、以下のコマンドを実行してください。

# LlamaIndexとOpenAI関連のライブラリをインストール
%pip install llama-index-llms-openai
%pip install llama-index-program-openai
%pip install llama-index

from google.colab import userdata
import os
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

1. Albumオブジェクトへの抽出

1.1 基本的な実装（docstringなし）

まず、出力スキーマを定義します。

# 必要なライブラリをインポート
from pydantic import BaseModel
from typing import List

# Songクラスを定義
class Song(BaseModel):
    title: str
    length_seconds: int

# Albumクラスを定義
class Album(BaseModel):
    name: str
    artist: str
    songs: List[Song]

次に、OpenAI Pydantic Programを定義します。

# OpenAIPydanticProgramをインポート
from llama_index.program.openai import OpenAIPydanticProgram

# プロンプトテンプレートを定義
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""

# OpenAIPydanticProgramを初期化
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True
)

プログラムを実行して構造化出力を取得します。

# プログラムを実行
output = program(
    movie_name="The Shining", description="Data model for an album."
)

# 結果を表示
print(output)

1.2 docstringを含む実装

モデルにdocstringを追加することで、より詳細な情報を提供できます。

# docstringを含むSongクラスを定義
class Song(BaseModel):
    """Data model for a song."""
    title: str
    length_seconds: int

# docstringを含むAlbumクラスを定義
class Album(BaseModel):
    """Data model for an album."""
    name: str
    artist: str
    songs: List[Song]

OpenAI Pydantic Programの定義と実行は前述と同様です。

# プロンプトテンプレートを定義
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""

# OpenAIPydanticProgramを初期化
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True
)

# プログラムを実行
output = program(movie_name="The Shining")

# 結果を表示
print(output)

2. 部分的な中間Pydanticオブジェクトのストリーミング

完全なJSONが生成されるまで待つ代わりに、利用可能になった時点で中間的なPydanticオブジェクトをストリームすることができます。

# 必要なライブラリをインポート
from pydantic import BaseModel, Field

# CharacterInfoクラスを定義
class CharacterInfo(BaseModel):
    """Information about a character."""
    character_name: str
    name: str = Field(..., description="Name of the actor/actress")
    hometown: str

# Charactersクラスを定義
class Characters(BaseModel):
    """List of characters."""
    characters: list[CharacterInfo] = Field(default_factory=list)

# OpenAIPydanticProgramを初期化
from llama_index.program.openai import OpenAIPydanticProgram

prompt_template_str = "Information about 3 characters from the movie: {movie}"

program = OpenAIPydanticProgram.from_defaults(
    output_cls=Characters,
    prompt_template_str=prompt_template_str
)

# 部分的なオブジェクトをストリーミング
for partial_object in program.stream_partial_objects(movie="Harry Potter"):
    print(partial_object)

3. Albumのリスト抽出（並列関数呼び出し）

OpenAIの最新の並列関数呼び出し機能を使用して、単一のプロンプトから複数の構造化データを同時に抽出できます。

# OpenAIをインポート
from llama_index.llms.openai import OpenAI

# プロンプトテンプレートを定義
prompt_template_str = """\
Generate 4 albums about spring, summer, fall, and winter.
"""

# OpenAIPydanticProgramを初期化（並列関数呼び出しを有効化）
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    llm=OpenAI(model="gpt-3.5-turbo-1106"),
    prompt_template_str=prompt_template_str,
    allow_multiple=True,
    verbose=True,
)

# プログラムを実行
output = program()

# 結果を表示
print(output)

まとめ

このガイドでは、OpenAI Pydantic Programを使用して構造化データを生成する方法を詳しく紹介しました。単純なオブジェクトから再帰的な構造まで、さまざまなユースケースに対応できることがわかりました。この手法を活用することで、自然言語からの構造化データ抽出がより簡単かつ効率的になります。

OpenAI Pydantic Programは、データ抽出、情報整理、そして複雑な構造のモデリングに非常に有用なツールです。特に、並列関数呼び出しやストリーミング機能を使用することで、より柔軟で効率的なデータ生成が可能になります。

この技術を使用することで、以下のような利点があります：

  1. 構造化データの迅速な生成
  2. 複雑なデータモデルの簡単な定義と抽出
  3. リアルタイムでの部分的なデータ取得
  4. 並列処理による効率的なデータ生成

ぜひ、あなたのプロジェクトでOpenAI Pydantic Programを活用し、データ処理のワークフローを改善してみてください。

📒ノートブック

https://colab.research.google.com/drive/1jzXgjqbqjj0tJEw0xJ8spJZ04c6MRDvH?usp=sharing

参考サイト

https://docs.llamaindex.ai/en/stable/examples/output_parsing/openai_pydantic_program/

Discussion