保存带 numpy.ndarray 的 dataframe

# 创建 DataFrame
texts = ["hi", "hello", "you"]
embeddings = [np.random.randn(10,) for i in range(3)]
df = pd.DataFrame({"text":texts, "embedding":embeddings})

type(df.embedding.values[0]) # numpy.ndarray

# 保存数据
import pickle
save_path = "./result/tmp/text_embedding_df.pkl"
# 假设你的DataFrame叫做 df
with open(save_path, 'wb') as f:
    pickle.dump(df, f)

# 读取数据
with open(save_path, 'rb') as f:
    df = pickle.load(f)
type(df.embedding.values[0]) # numpy.ndarray

需要注意的是，pickle文件是以二进制形式存储的，因此在写入和读取文件时，需要以二进制模式（‘wb’或’rb’）打开文件。同时，pickle有其特定的序列化方式，适合保存大量数据，并且能保存很多Python的内置类型

参考：
chatGPT 提示词：“我有一个 pandas 的 dataframe, 字段包括 text 和 embedding。 text 数据类型为 str, embedding 的数据类型为 numpy.ndarray。如何保存数据，随后如何读取，需要将embedding 的数据类型保持为 numpy.ndarray”