Chunk text in chunks of N words and calculate the average embedding from all chunks iterating over a pandas df

chunk_and_average_embeddings.py

import numpy as np
from tqdm import tqdm
tqdm.pandas()

def chunk_text(text, max_words=100):
    words = text.split()
    chunks = [' '.join(words[i:i + max_words]) for i in range(0, len(words), max_words)]
    
    # Convert the list of embeddings to a NumPy array and change its dtype to float32
    embeddings_array = np.array(model.encode(chunks)).astype('float32')
    
    # Calculate the mean vector
    mean_vector = np.mean(embeddings_array, axis=0)
    
    return mean_vector

df["mean_embeddings"] = df["text"].progress_apply(lambda x: chunk_text(x, 100))