【Python】gpt-indexのsave_to_disk関数でTypeError: Object of type float32 is not JSON serializableになったときの対処

はじめに

GPT Index でのHuggingFaceの埋め込みモデルの利用 | npakaにしたがって作ったindexをsave_to_diskで保存しようとしたときに以下のエラーが出ました。対応方法を考えます。

index.save_to_disk("index.json")

TypeError: Object of type float32 is not JSON serializable

環境

% pip list | grep gpt-index
gpt-index                0.2.3
% python -V
Python 3.10.2

（gpt-indexの更新が非常に活発（記事投稿日2023/1/20時点での最新が0.2.8）なので、すぐにout-datedな情報になるかもしれません。）

解決方法

保存

save_to_disk関数の内部で辞書をjson.dumpに渡す処理があり、そのとき渡す辞書の要素にfloat32型が含まれていることが原因です。これをfloat型に置き換えてから保存します。

def save_index(index, save_path):
  # embedding_dictの変換
  index_struct_dict = index.index_struct.to_dict()
  embdict = index_struct_dict["embedding_dict"]
  embdict = {k:list(map(float, v)) for k, v in  embdict.items()} # floatに変換
  index_struct_dict["embedding_dict"] = embdict

  # 元のsave_to_diskと同じ方法で保存
  out_dict = {
    "index_struct": index_struct_dict,
    "docstore": index.docstore.to_dict(),
  }
  
  with open(save_path, "w") as f:
    #json.dump(out_dict, f)
    json.dump(out_dict, f, ensure_ascii = False, indent = 2) # ついでなので日本語をそのまま出力する。

save_index(index, "index.json")

エラーなく実行できたらファイルの中身を確認します。

index.json

{
  "index_struct": {
    "text": null,
    "doc_id": null,
    "embedding": null,
    "nodes_dict": {
      "6916722896900258078": {
        "text": "結束バンド\n後藤ひとり...

読み込み

load_from_diskでindex.jsonを読み込みます。
その際、index作成時と同じ引数（documentsはおそらく不要）を与えるようにしてください。

# ライブラリimportを済ませておく

# 埋め込みモデルの準備
embed_model = LangchainEmbedding(HuggingFaceEmbeddings(
    model_name="oshizo/sbert-jsnli-luke-japanese-base-lite"
))

# インデックスの生成
index = GPTSimpleVectorIndex.load_from_disk(
    "index.json", 
    # documents=documents,
    prompt_helper=PromptHelper(
        max_input_size=4000,  # LLM入力の最大トークン数
        num_output=256,  # LLM出力のトークン数
        chunk_size_limit=2000,  # チャンクのトークン数
        max_chunk_overlap=0,  # チャンクオーバーラップの最大トークン数
        separator="。"  # セパレータ
    ),
    embed_model=embed_model,  # 埋め込みモデル
    verbose=True
)

応答を確認します
※OpenAIの課金が発生しますのでご注意ください。

response = index.query("文化祭で何が起きた？")
print(response)

文化祭では、結束バンドが初の依頼にワクワクしながら参加し、カオス状態のイベント会場で最善を尽くしてライブを行い、大盛り上がりのうちに終了し、大成功を収めた。

いい感じです。

補足：原因の詳細

エラーコードの通り、float32の型をjsonで保存しようとしていることが原因です。
今回はOpenAIではなくoshizoさん作成の日本語埋め込みモデルを利用しているので、それによってできた埋め込みがfloat32型なのだと思います。

ソースコードを読んで、save_to_diskの処理を見ると、以下のdictをjson.dumpに渡していることがわかります。

        out_dict: Dict[str, dict] = {
            "index_struct": self.index_struct.to_dict(),
            "docstore": self.docstore.to_dict(),
        }

このfloat32を普通のfloatに変換すればエラーが出なくなると期待できます。

index.index_struct.to_dict()の中身を確認します。

index_struct_dict = index.index_struct.to_dict()
print(list(index_struct_dict.keys()))

['text', 'doc_id', 'embedding', 'nodes_dict', 'id_map', 'embedding_dict']

更に見ていくと、embedding_dictのvalueがnumpy.float32型のlistであることがわかります。

embdict = index_struct_dict["embedding_dict"]
values = list(embdict.values()) # keyとvalueのうちvalueのリストを取得
print(type(values[0]))
print(type(values[0][0]))

<class 'list'>
<class 'numpy.float32'>

これをfloat型に変換してから保存すれば解決できると期待されます。

その他の試行錯誤

最初、indexをまるごとpickleで保存しようとしたのですが、失敗しました。

import pickle
with open("test.json", "wb") as f:
  pickle.dump(index, f)

---------------------------------------------------------------------------
TypeError                                 
Traceback (most recent call last)
Cell In[146], line 3
      1 import pickle
      2 with open("test.json", "wb") as f:
----> 3   pickle.dump(index, f)

TypeError: cannot pickle 'builtins.CoreBPE' object

参考

gpt-index | GitHub

GPT Index でのHuggingFaceの埋め込みモデルの利用 | npaka

gpt-index(0.2.5)をOpenAI APIなし＆日本語で動かす | oshizo