Wieghted average embedding in Python with numpy

weighted_mean_embedding.py

import numpy as np

# takes an arbitrary number of vectors/weights as input 
def weighted_mean_embedding(vectors, weights):
    if len(vectors) != len(weights):
        raise ValueError("Number of vectors must be equal to the number of weights")

    # Convert lists to NumPy arrays and cast to float
    vectors = [np.array(vector, dtype=float) for vector in vectors]

    weighted_sum = np.zeros_like(vectors[0], dtype=float)
    total_weight = sum(weights)

    for vector, weight in zip(vectors, weights):
        weighted_sum += vector * weight

    weighted_mean = weighted_sum / total_weight
    
    return weighted_mean.tolist()

# Example usage:
list1 = [1, 1, 1, 1]
list2 = [4, 5, 6, 7]
list3 = [0,0,0,0]

weights = [1, 0.5, 1]

weighted_mean_embedding([list1, list2, list3], weights)

# [1.2, 1.4, 1.6, 1.8]