使用 zip 函数处理字典在 Python 中保持相同的结构

Question

我有一个字典，其中每个键都映射到一个数组列表，除了键“引用”，它只是一个整数数组。

cls["input_ids"], cls["attention_masks"], cls["labels"], cls["reference"]

一个键的每一行都链接到其他键的行（这是一个 Bert 分词器的修改输出）

我想通过参考值过滤掉一些行并保持输出相同的字典结构，现在我设法做到这一点的唯一方法是这样的： 我放了一些随机数据来给出一个想法

cls= {"input_ids":[[22,22,22],[33,33,33]], "attention_masks":[[22,22,22],[33,33,33]], "reference":[1,0], "labels":[[[22,22,22],[33,33,33]]]}
mcp = {"input_ids":[], "attention_masks":[], "reference":[], "labels":[]}

        for el in zip(cls["input_ids"], cls["attention_masks"], cls["reference"], cls["labels"]):
            if el[2] == 1:
                mcp["input_ids"].append(el[0])
                mcp["attention_masks"].append(el[1])
                mcp["reference"].append(el[2])
                mcp["labels"].append(el[3])

但我真的不喜欢这段代码，我想知道是否有更漂亮的方法来做到这一点。

Answer 1

假设您的值都具有相同的长度，我建议使用pandas，它是为这种修改而制作的。

import pandas as pd

cls= {"input_ids":[[22,22,22],[33,33,33]], 
      "attention_masks":[[22,22,22],[33,33,33]], 
      "reference":[1,0], 
      "labels":[[22,22,22],[33,33,33]] # I assume this is what you meant
      }

# Turn the data into a dataframe which is sth like a table
df = pd.DataFrame(cls)

这是df 的样子：

>>> df
      input_ids attention_masks  reference        labels
0  [22, 22, 22]    [22, 22, 22]          1  [22, 22, 22]
1  [33, 33, 33]    [33, 33, 33]          0  [33, 33, 33]

您可以访问df['reference'] 之类的值并过滤数据，例如：

>>> df[df['reference']==1]
      input_ids attention_masks  reference        labels
0  [22, 22, 22]    [22, 22, 22]          1  [22, 22, 22]

如果你需要它作为字典：

>>> df[df['reference']==1].to_dict()
{'input_ids': {0: [22, 22, 22]},
 'attention_masks': {0: [22, 22, 22]},
 'reference': {0: 1},
 'labels': {0: [22, 22, 22]}}