将带有 OpenAI API 请求结果的 jsonl 文件加载到 pandas data.frame-解网

问：

我有一个包含大约 500k 观测值的大型数据集。它有一个字符串变量，我想为它创建一个嵌入。我使用 OpenAI API 来创建嵌入，由于有大量的观察结果，我将他们的脚本用于并行请求：

https://github.com/openai/openai-cookbook/blob/main/examples/api_request_parallel_processor.py

一切都很好。但是我正在努力将结果加载到 pandas data.frame 中。包含结果的 jsonl 文件具有以下结构，每行对应于 500k 观测值之一：

[{“model”： “text-embedding-ada-002”， “input”： “INPUT STRING NR 1”}， {“object”： “list”， “data”： [{“object”： “embedding”， “index”： 0， “embedding”： [1,2,3,4...，1536]}]， “model”： “text-embedding-ada-002-v2”， “usage”： {“prompt_tokens”： 2， “total_tokens”： 2}}]

[{“model”： “text-embedding-ada-002”， “input”： “输入字符串 NR 2}， {”object“： ”list“， ”data“： [{”object“： ”embedding“， ”index“： 0， ”embedding“： [1,2,3,4...，1536]}]， ”model“： ”text-embedding-ada-002-v2“， ”usage“： {”prompt_tokens“： 2， ”total_tokens“： 2}}]

现在，我想将这些结果读入具有以下结构的熊猫数据帧中。它应该有一个包含“INPUT STRING”的变量和 1536 个包含嵌入的其他变量。

我是python和json文件的新手。我通常使用 csv 文件和 R。

我尝试使用熊猫的read_json功能，但没有用

import pandas as pd
openai_results = pd.read_json("results.jsonl", lines=True)

但这给了我一个只有 2 个变量的数据集：例如，对于第一个观察，第一个变量包含： {“model”： “text-embedding-ada-002”， “input”： “INPUT STRING NR 1”} 和第二个变量 {“object”： “list”， “data”： [{“object”： “embedding”， “index”： 0， “embedding”： [1,2,3,4...，1536]}]， “model”： “text-embedding-ada-002-v2”， “usage”： {“prompt_tokens”： 2， “total_tokens”： 2}}

python json pandas openai-api

df = pd.read_json('your_file.json', lines=True)
df
'''
   0                                                  1
0  {'model': 'text-embedding-ada-002', 'input': '...  {'object': 'list', 'data': [{'object': 'embedd...
1  {'model': 'text-embedding-ada-002', 'input': '...  {'object': 'list', 'data': [{'object': 'embedd...
'''

访问值：

df["input"] = df[0].str["input"]
df["embedding"] = df[1].str["data"].str[0].str["embedding"] # or df["embedding"]=df[1].apply(lambda x: x["data"][0]["embedding"])
df = df[["input","embedding"]]

输出：

               input           embedding
0  INPUT STRING NR 1  [1, 2, 3, 4, 1536]
1  INPUT STRING NR 2  [1, 2, 3, 4, 1536]

如果你想分解嵌入列，那么使用 explode（）：

df = df.explode("embedding")
df
'''
               input embedding
0  INPUT STRING NR 1         1
0  INPUT STRING NR 1         2
0  INPUT STRING NR 1         3
0  INPUT STRING NR 1         4
0  INPUT STRING NR 1      1536
1  INPUT STRING NR 2         1
1  INPUT STRING NR 2         2
1  INPUT STRING NR 2         3
1  INPUT STRING NR 2         4
1  INPUT STRING NR 2      1536
'''

将带有 OpenAI API 请求结果的 jsonl 文件加载到 pandas data.frame

load jsonl File with OpenAI API request results to pandas data.frame

评论

评论