大型 DataFrame 的高效嵌入计算

Question

鉴于DataFrame：

    id  articleno           target
0   1   [607303]            607295
1   1   [607295]            607303
2   2   [243404, 617953]    590448
3   2   [590448, 617953]    243404

对于每一行，通过查找字典列表中的每个项目来计算平均文章嵌入：

embeddings = {"607303": np.array([0.19, 0.25, 0.45])
             ,"607295": np.array([0.77, 0.76, 0.55])
             ,"243404": np.array([0.35, 0.44, 0.32])
             ,"617953": np.array([0.23, 0.78, 0.24])
             ,"590448": np.array([0.67, 0.12, 0.10])}

例如，为了澄清，对于第三行（索引 2），243404 和 617953 的文章嵌入分别是 [0.35, 0.44, 0.32] 和 [0.23, 0.78, 0.24]。平均文章嵌入计算为所有元素的元素相加除以文章数量，因此：([0.35, 0.44, 0.32]+[0.23, 0.78, 0.24])/2=[0.29, 0.61, 0.28]。

预期输出：

    id  dim1     dim2     dim3      target
0   1   0.19     0.25     0.45      607295
1   1   0.77     0.76     0.55      607303
2   2   0.29     0.61     0.28      590448
3   2   0.45     0.45     0.17      243404

实际上，我的DataFrame 有数百万行，articleno 中的列表可以包含更多项目。因此，迭代行可能太慢，可能需要更有效的解决方案（可能是矢量化的）。

而且，维数（嵌入大小）是事先知道的，但是是几百，所以列数； dim1, dim2, dim3, ... dimN 应该是动态的，基于嵌入的尺寸 (N)。

Answer 1

在上一个问题中，您在articleno 列表中分离元素，然后从articleno 列表中删除target。现在，如果您想访问articleno 列表中的元素，您需要再次加倍努力将它们分开。

为了说明我的意思，这里有一种方法可以生成两个问题的输出，同时添加最少的额外代码：

# construct the embeddings dataframe:
embedding_df = pd.DataFrame(embeddings).T.add_prefix('dim')

# aggregation dictionary
agg_dict = {'countrycode':'first','articleno':list}

# taking mean over embedddings
for i in embedding_df.columns: agg_dict[i] = 'mean'

new_df = df.explode('articleno')

(new_df.join(new_df['articleno'].rename('target'))
    .query('articleno != target')
    .merge(embedding_df, left_on='articleno', right_index=True)  # this line is extra from the previous question
    .groupby(['id','target'], as_index=False)
    .agg(agg_dict)
)

输出：

   id  target countrycode         articleno  dim0  dim1  dim2
0   2  590448          US  [617953, 617953]  0.23  0.78  0.24
1   2  617953          US  [590448, 590448]  0.67  0.12  0.10

---

现在，如果您不关心最终输出中的 articleno 列，您甚至可以像这样在降低内存/运行时间的同时简化代码：

total_embeddings = g[embedding_df.columns].sum()
article_counts = g['id'].transform('size')

new_df[embedding_df.columns] = (total_embeddings.sub(new_df[embedding_df.columns])
                                  .div(article_counts-1, axis=0)
                               )

你会得到相同的输出。