Pandas AIをAzure OpenAI Serviceで動かす

はじめに

Pandas AIはPandasにAI機能を追加するPythonライブラリです。Pandas AIを使うと、自然言語によってpandas dataframeから特定のデータを抽出したり、dataframeからグラフを描画したりと、pandasの操作が圧倒的に楽になります。

公式リポのサンプルがOpenAI社APIを使っていたので、Azure OpenAI Serviceの場合はどうやって接続するか簡単に検証してまとめました。

使い方

準備

$ pip install pandasai

実行

コードは下記になります。公式リポのサンプルをベースにしています。
国別のGDPと幸福度のテーブルがあり、そこからChatGPT (gpt-35-turbo)に最も幸福な国を5つ抽出させています。

import os
import pandas as pd
from pandasai import PandasAI
from pandasai.llm.azure_openai import AzureOpenAI

api_base = "<AZURE_OPENAI_API_BASE>"
api_key = "<AZURE_OPENAI_API_KEY>"

# Sample DataFrame
df = pd.DataFrame({
    "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
    "gdp": [19294482071552, 2891615567872, 2411255037952, 3435817336832, 1745433788416, 1181205135360, 1607402389504, 1490967855104, 4380756541440, 14631844184064],
    "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12]
})

# Instantiate a LLM
llm = AzureOpenAI(api_base = api_base,
                  api_version = "2023-03-15-preview",
                  api_token = api_key,
                  deployment_name = "gpt-35-turbo")

pandas_ai = PandasAI(llm, conversational=False)
print(pandas_ai(df, prompt='Which are the 5 happiest countries?'))

このように結果が返ってきます。

api_baseとapi_keyに皆さんのAzure OpenAIリソースのURLエンドポイントとAPIキーが代入される形にします。
コード実装を見ると、指定なしだとそれぞれ環境変数のAZURE_OPENAI_ENDPOINTとAZURE_OPENAI_KEYを読み込む形になっているようです。

計算させることもできます。

グラフを描画させることもできます。

参考

gventuri/pandas-ai: Pandas AI is a Python library that integrates generative artificial intelligence capabilities into Pandas, making dataframes conversational
pandas-ai/azure_openai.py at main · gventuri/pandas-ai · GitHub