ファイルサイズの大きいJSONファイルの読み込み

はじめに

非常に多くの行数のあるJSONファイルをPythonで読み込む方法について備忘録として残します。

以下の、izumi-lab/llm-japanese-dataset の一部を取り出して保存してみたいと思います。
拡張子を手動で.jsonlから.jsonに変更しました。

huggingface.co

あらかじめ、こちらのJSONファイルををダウンロードしているとします。

また、本記事で利用したコードは以下のページにアップロードしています。

github.com

Pythonコード

モジュールのインポート

import json
import codecs

ジェネレータ関数を作成する

def read_json_file_in_batches(file_path, batch_size):
    with open(file_path, 'r', encoding='utf-8') as file:
        batch = []
        for line in file:
            batch.append(json.loads(line))
            if len(batch) == batch_size:
                yield batch
                batch = []
        # すべてを読み取ったあと、batchにまだ中身があればyieldする
        if batch:
            yield batch

保存するファイル名を指定する

output_file_path = 'output.json'

練習のため、次の日本語を英語に翻訳してください という文章を含む行が現れるまで読み込みます。また、それまでの行の情報をlistに格納していきます。

# 指定した文字列が見つかったかどうかを示すflagを用意する    
flag = False
in_file_path = 'data-cc-by-sa.json'
batch_size = 10
search_string = "次の日本語を英語に翻訳してください"
output = []
for json_batch in read_json_file_in_batches(in_file_path, batch_size):
    # json_batchは最大でbatch_size行のJSONオブジェクトを含むリストである
    # 1行ずつ読み込んでいく
    for json_object in json_batch:
        # 指定した文字列が含まれるか判定する
        if any(search_string in value for value in json_object.values())==True:
            print('指定した文字列が含まれます')
            print(json_object)
            flag = True
        
        # 指定した文字列が見つかるまでリストに文章を追加していく
        output.append(json_object)
        # 指定した文列が含まれていた場合ループから抜け出す
        if flag:
            print('BREAK')
            break
    # 指定した文列が含まれていた場合ループから抜け出す
    if flag:
        break

ファイルの書き込み

fw = codecs.open(output_file_path , 'w', 'utf-8')
json.dump(output, fw, ensure_ascii=False, indent=0)
# 書き込みオブジェクトを閉じる
fw.close()

以下のように 次の日本語を英語に翻訳してください が含まれる文章までのデータを書き込むことができました。