Блог

Вот некоторые данные, которые я сгенерировал:

 import numpy as np
import pandas as pd
import scipy
import scipy.spatial

df = pd.DataFrame(
    {
        "item_1": np.random.randint(low=0, high=10, size=1000),
        "item_2": np.random.randint(low=0, high=10, size=1000),
    }
)
embeddings = {item_id: np.random.randn(100) for item_id in range(0, 10)}


def get_distance(item_1, item_2):
    arr1 = embeddings[item_1]
    arr2 = embeddings[item_2]
    return scipy.spatial.distance.cosine(arr1, arr2)
 

Я хотел бы применить get_distance к каждой строке. Я могу это сделать:

 df.apply(lambda row: get_distance(row["item_1"], row["item_2"]), axis=1)
 

Но это было бы очень медленно для больших наборов данных.

Есть ли способ вычислить косинусное сходство вложений, соответствующих каждой строке, без использования DataFrame.apply ?

Для версии scipy

 %%timeit
df.apply(lambda row: get_distance(row["item_1"], row["item_2"]), axis=1)
# 38.3 ms ± 84 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
 

Как бы то ни было, я добавил numba с дополнительными сложностями

Размышляя о выделении памяти и широковещательной передачи numpy с использованием tmp, я использовал для циклов

Также стоит рассмотреть возможность передачи аргументов, возможно, вы можете передавать векторы вместо словаря.

Кроме того, первый запуск выполняется медленно из-за компиляции

Также вы можете сделать это параллельно с numba

 @nb.njit((nb.float64[:, ::100], nb.float64[:, ::100]))
def cos(a, b):
    norm_a = np.empty((a.shape[0],), dtype=np.float64)
    norm_b = np.empty((b.shape[0],), dtype=np.float64)
    cos_ab = np.empty((a.shape[0],), dtype=np.float64)

    for i in nb.prange(a.shape[0]):
        sq_norm = 0.0
        for j in range(100):
            sq_norm  = a[i][j] ** 2
        norm_a[i] = sq_norm ** 0.5
    
    for i in nb.prange(b.shape[0]):
        sq_norm = 0.0
        for j in range(100):
            sq_norm  = b[i][j] ** 2
        norm_b[i] = sq_norm ** 0.5
    
    for i in nb.prange(a.shape[0]):
        dot = 0.0
        for j in range(100):
            dot  = a[i][j] * b[i][j]
        cos_ab[i] = 1 - dot / (norm_a[i] * norm_b[i])
    return cos_ab
 

 %%timeit
cos(item_1_embedded, item_2_embedded)
# 218 µs ± 1.23 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
 

Прямое использование векторизованных numpy операций намного быстрее:

 item_1_embedded = np.array([embeddings[x]for x in df.item_1])
item_2_embedded = np.array([embeddings[x]for x in df.item_2])
cos_dist = 1 - np.sum(item_1_embedded*item_2_embedded, axis=1)/(np.linalg.norm(item_1_embedded, axis=1)*np.linalg.norm(item_2_embedded, axis=1))
 

(Эта версия работает в 771 µs среднем на моем ПК, по сравнению 37.4 ms с версией для DataFrame.apply , что делает чистую версию numpy примерно в 50 раз быстрее).

Вы можете векторизовать вызов cosine с numpy.vectorize помощью . Наблюдается небольшое увеличение скорости (34 мс против 53 мс).

 vec_cosine = np.vectorize(scipy.spatial.distance.cosine)
vec_cosine(df['item_1'].map(embeddings),
           df['item_2'].map(embeddings))
 

выход:

 array([0.90680875, 0.90999454, 0.99212814, 1.12455852, 1.06354469,
       0.95542037, 1.07133003, 1.07133003, 0.        , 1.00837058,
       0.        , 0.93961103, 0.8943738 , 1.04872436, 1.21171375,
       1.04621226, 0.90392229, 1.0365102 , 0.        , 0.90180297,
       0.90180297, 1.04516879, 0.94877277, 0.90180297, 0.93713404,
...
       1.17548653, 1.11700641, 0.97926805, 0.8943738 , 0.93961103,
       1.21171375, 0.91817959, 0.91817959, 1.04674315, 0.88210679,
       1.11806218, 1.07816675, 1.00837058, 1.12455852, 1.04516879,
       0.93713404, 0.93713404, 0.95542037, 0.93876964, 0.91817959])