English SDK for Apache Sparkで日本語を使って頑張ってみる

弊社内でも盛り上がりを見せているEnglish SDK for Apache Spark。

https://qiita.com/maroon-db/items/89f7a1aae11a112f9700

誰でもOpenAI APIでgpt-4を使えるようになったので、日本語で色々頑張ってみます。

https://openai.com/blog/gpt-4-api-general-availability

# 準備

ライブラリのインストール。

```py
%pip install pyspark-ai
dbutils.library.restartPython()
```

環境設定。

```py
import os

os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
os.environ["GOOGLE_API_KEY"] = "<Google Cloud PlatformのAPIキー>"
os.environ["GOOGLE_CSE_ID"] = "<Google CSEのID>"
```

モデルを確認します。

```py
import os
import openai
openai.api_key = os.getenv("OPENAI_API_KEY")
openai.Model.list()
```

ありました。

```
      "id": "gpt-4",
      "object": "model",
      "created": 1687882411,
      "owned_by": "openai",
```

`gpt-4`を指定します。

```py
from langchain.chat_models import ChatOpenAI
from pyspark_ai import SparkAI

# If 'gpt-4' is unavailable, use 'gpt-3.5-turbo' (might lower output quality)
llm = ChatOpenAI(model_name='gpt-4', temperature=0)

# Initialize SparkAI with the ChatOpenAI model
spark_ai = SparkAI(llm=llm, verbose=True)

# Activate partial functions for Spark DataFrame
spark_ai.activate()
```

ちなみに`SparkAI`を呼び出す際の引数に`verbose=True`を指定すると、処理過程が確認できて便利です。で、上のコメントの`Activate partial functions for Spark DataFrame`をスルーしていたのですが、以下で理解できました。

# データの読み込み

普通にSparkデータフレームを読み込みます。

```py
df = spark.read.table("takaakiyayoi_catalog.covid.covid_cases")
display(df)
```
![Screenshot 2023-07-07 at 17.33.14.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c411a6d1-c2ac-5fd2-3909-92e887595dbd.png)

本当は`spark_ai.create_df`を以下のように使ってやりたいのですが、以下の理由から現時点では断念しています。

```py
df = spark_ai.create_df("2023年の日本の都道府県別、月別covid-19感染者数")
```

- 指定した文字列でGoogle検索を行います。ヒットページにテーブルの記述があればそれをパースしてデータフレームを構成する仕組みのようです。
- 当然、日本語のページがヒットするのですが、意外と`TABLE`タグでデータを記述しているサイトがありません。大抵GUIに埋もれていたり、csvファイルへのリンクなどになっています。

このため、現時点では別のデータソースからデータフレームを組み立てる方が確実です。

# プロット

で、先ほどの`Activate partial functions for Spark DataFrame`ですが、書いてある通りに**Sparkデータフレームで一部の関数を有効化したよ**ということでした。つまり、**SparkデータフレームにEnglish SDKのメソッドが追加**されているのです。

こちらにも説明があります。

https://github.com/databrickslabs/pyspark-ai/issues/49

> the English SDK can seamlessly integrate with the existing Spark APIs. This means that you use the the SDK on any Spark DataFrame. You can load your saved source in to `df`
> English SDKは既存のSpark APIとシームレスに連携できます。これは、すべてのSparkデータフレームでEnglsh SDKを使うことができることを意味します。保存したソースを`df`にロードすることができます。

つまり、上で`spark_ai.activate()`時点で、Sparkデータフレームで`.ai...`というメソッドを呼び出せるようになっています。

```py
df.ai.plot("都道府県別の折れ線グラフ")
```
![Screenshot 2023-07-07 at 17.41.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c680f351-5f1f-990e-ff57-e25965151bc8.png)

やったー。

次は円グラフ。しかし、ここでつまづきました。

```py
df.ai.plot("都道府県別の合計感染者数の円グラフ、上位5つを表示、それ以外は「その他」にまとめる。")
```

これだと、背後で行うデータフレームの結合処理で非推奨の`append`を使ってしまい、処理に失敗します。`concat`でなくてはいけません。なので、プロンプトエンジニアリングのノリで以下のように。

```py
df.ai.plot("都道府県別の合計感染者数の円グラフ、上位5つを表示、それ以外は「その他」にまとめる。まとめる際にはappendではなくconcatを使ってください。")
```

通りました。
![Screenshot 2023-07-07 at 17.43.48.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/47f20535-3e11-600d-fd1b-78e67289122f.png)

# フィルタリング

`東京`ではマッチしなくてダメでした。

```py
tokyo_df = df.ai.transform("Tokyoの2023年のデータ")
display(tokyo_df)
```

ちゃんと、年でもフィルタリングしてくれます。
![Screenshot 2023-07-07 at 17.52.04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/64c07f19-0ad7-f452-b511-9d1e644ad10c.png)

```py
tokyo_df.ai.plot("月別の合計の棒グラフ")
```
![Screenshot 2023-07-07 at 17.53.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/926bf29c-47d2-2d06-f046-f644a435c3ac.png)

なんというか、半分プログラミング、半分プロンプトエンジニアリングという取り組みになりますね。面白い。

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)