将 PyTorch 预测合并到原始数据帧

Question

我已经从我的 PyTorch 模型中获得了具有以下形状的张量 (torch.Size([2958, 96])) 的预测。我的原始数据集有 2958 个 qid，其中一些的最大文档数为 96（最小为 47）。预测用 -1 填充了缺失的那个。我的原始数据框的形状是 (221567, 7)。

我想使用 qid 将来自 PyTorch 模型或张量的预测合并回这个数据框。张量的每一行代表一个 qid，而每一列代表该特定文档的排名（基于每个 qid 中文档的顺序）。

下面是一个最小的例子（在将张量转换为 df 之后）：

tensor = {'0': ['3', '1','2'],'1': ['2', '1','2'],'3': ['2', '1','2']}
y_pred = pd.DataFrame(tensor)

data = {'qid': ['0', '0','0','1', '1','1','2', '2'],'irrelevant_col': ['foo', 'foo','foo','foo', 'bar','bar','bar', 'bar']}
original_df = pd.DataFrame(data)

注意对于 qid==2，只有 2 行，因此张量在第 2 行和第 2 列中有一个“-1”。此外，张量的顺序在某种意义上是正确的，因为它与数据框中项目的顺序相匹配。这是目标输出：

target = {'qid': ['0', '0','0','1', '1','1','2', '2'],'irrelevant_col': ['foo', 'foo','foo','foo', 'bar','bar','bar', 'bar'],'y_pred': ['3', '2','2','1', '1','1','2', '2']}
target_df = pd.DataFrame(target)

Answer 1

1. 首先将张量转换为数据框。

2. 然后通过堆叠重塑数据框以匹配原始形状，并删除索引的第二级并创建新索引。

3. 最后将原始数据框和改造后的

   y_pred

   数据框合并到

   qid

   列。

y_pred = pd.DataFrame(tensor).replace(-1, np.nan)

y_pred = y_pred.stack().reset_index(level=1, drop=True).to_frame('y_pred').reset_index()

merged_df = original_df.merge(y_pred, on='qid', how='left')

使用

merged_df

会创建一个额外的列索引，如果你不想的话可以使用

drop()

方法：

merged_df = merged_df.drop('index', axis=1)

Answer 2

import pandas as pd

tensor = {'0': ['3', '1', '2'], '1': ['2', '1', '2'], '3': ['2', '1', '2']}
y_pred = pd.DataFrame(tensor)

data = {'qid': ['0', '0', '0', '1', '1', '1', '2', '2'], 'irrelevant_col': ['foo', 'foo', 'foo', 'foo', 'bar', 'bar', 'bar', 'bar']}
original_df = pd.DataFrame(data)

y_pred_long = y_pred.stack().reset_index()
y_pred_long.columns = ['rank', 'qid', 'y_pred']
y_pred_long = y_pred_long[y_pred_long['y_pred'] != '-1']

merged_df = original_df.merge(y_pred_long, on=['qid'], how='left')
merged_df.drop(columns=['rank'], inplace=True)

print(merged_df)

运行此代码将为您提供所需的输出，您可以使用它将 PyTorch 预测与原始 DataFrame 合并。