Average embeddings in pandas with numpy.py

average_embeddings.py

import pandas as pd 
import numpy as np

# Useful when you indxed document chunks and now need to create average embeddings per document
df = pd.read_parquet("chunk_embeddings.parquet") # ~3GB file

# Convert the embeddings column to a NumPy array
df['embeddings'] = df['embeddings'].apply(np.array)

# Group by filename and calculate the mean of the embeddings
result_df = df.groupby('filename')['embeddings'].mean().reset_index()

# If you want to keep the embeddings as lists instead of NumPy arrays, you can convert them back
result_df['embeddings'] = result_df['embeddings'].apply(list)

result_df

# takes 1 min for a df with 1.323.054 chunks to result_df with 68.492 rows on my old i7 with 32gb Ram


# in case you run into any issues like this: os error list index overflow 
# use this polars workaround: https://github.com/apache/arrow/issues/14229#issuecomment-1826259855

# import polars as pl
# df = pl.read_parquet("chunk_embeddings.parquet")
# df = df.to_pandas()
# del df["__index_level_0__"] 
# and continue with the normal pandas logic