Merge PyTorch predictions to original dataframe

I have obtained predictions from my PyTorch model as a tensor with the following shape (torch.Size([2958, 96])). My original dataset has 2958 qids and some of them have a max number of documents of 96 (min is 47). The predictions have padded the missing one with -1s. The shape of my original dataframe is (221567, 7).

I would like to merge the predictions from the PyTorch model or tensor back to this dataframe, using the qid. Every row of the tensor represents a qid, while every column represents the rank of that particular document (based on the order of documents in each qid).

Below is a min example (after converting the tensor to a df):

tensor = {'0': ['3', '1','2'],'1': ['2', '1','2'],'2': ['2', '1','-1']}
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)

Notice that for qid==2, there are only 2 rows, and the tensor has a '-1' in row 2 and column 2 because of this. Also, the order of the tensor is correct in the sense that it matches the order of the items in the dataframe. This is the target output:

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)

EDIT: I fixed an incorrect column (2 instead of 3, and made the last y_pred -1.

1. First convert he tensor to a dataframe.

2. Then reshape the dataframe to match the original shape by stacking it, and drop the the second level of the index and creat a new index.

3. Finally merge the original dataframe and the reshaped y_pred dataframe to the qid column.

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')

Use merged_df will create an extra column index, you can use the drop() method if you don't want it:

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

Thank for your posted answers. I might have not specified the question correctly but I was able to use your suggestions. The final step was (after following your suggestions for reshaping y_pred), that instead of using merge I had to simply assign the column like so:

merged_df['y_pred'] = y_pred .values

Thanks again!!!!