小样本·信息抽取 关系分类 复现-3 常用代码集-attention mask& ids

有些场景下，我们得到的mask的长度或者input_ids的长度并不是一样的。
有的时候，ids的长度和句子的长度一致，但是不同句子的长度是不一致的。
在使用PLM encoder时，需要统一到相同的长度在做考虑。

import torch
from transformers import AutoTokenizer, AutoModel

ids=[[1, 2, 3, 4], [3, 4, 2],[0,1]]

def get_atten_mask( ids):
    len_ids = [len(i) for i in ids]
    max_len = max(len_ids)
    attention_mask = torch.zeros(len(ids), max_len)
    for j in range(len(len_ids)):
        attention_mask[j, :len_ids[j]] = 1
    return attention_mask

'''mask=get_atten_mask(ids)
print(mask)'''

max_ids=[len(i) for i in ids]
max_len=max(max_ids)
idss=[]
for id in ids:
    id_=id+[0]*(max_len-len(id))
    idss.append(torch.tensor(id_))
id_inp=torch.tensor([i.cpu().detach().numpy() for i in idss]).cuda()
print(id_inp)