lansinuote / huggingface_toturials Goto Github PK

labels = torch.LongTensor(labels)

def collate_fn(data):
    # print('data: %s' % data)
    
    # 取每个评论的第一个字符作为数据, 第二作为标签
    # 猫和老鼠的DVD,我在当当网已买过10余次了。除了做为礼物送给亲朋好有的孩子外，...
    
    sents =  [i[0] for i in data] # 猫 
    labels = [i[1] for i in data] # 和
    
    # print('sents: %s sents长度: %s' % (sents, len(sents)))
    # print('labels: %s' % labels)    
    
    # 编码
    data = token.batch_encode_plus(batch_text_or_text_pairs=sents,
                                   truncation=True,
                                   padding='max_length',
                                   max_length=500,
                                   return_tensors='pt',
                                   return_length=True)
    
    labels_data = token.batch_encode_plus(batch_text_or_text_pairs=labels,
                                           truncation=True,
                                           padding='max_length',
                                           max_length=500,
                                           return_tensors='pt',
                                           return_length=True)
    
    # input_ids:编码之后的数字
    # attention_mask:是补零的位置是0,其他位置是1
    input_ids = data['input_ids']
    attention_mask = data['attention_mask']
    token_type_ids = data['token_type_ids']
    
    n_labels = torch.LongTensor(labels_data['input_ids'])

    #print(data['length'], data['length'].max())

    return input_ids, attention_mask, token_type_ids, n_labels

# 数据加载器
loader = torch.utils.data.DataLoader(dataset=dataset['train'],
                                     batch_size=16,
                                     collate_fn=collate_fn,
                                     shuffle=True,
                                     drop_last=True)


for i, values in enumerate(loader):
    input_ids, attention_mask, token_type_ids, labels = values
    
    # print(i, values)
    print(i, input_ids.shape, 
          attention_mask.shape, 
          token_type_ids.shape, 
          labels.shape,
         )
    
    if i >= 5:
        break
    

print(len(loader))

from transformers import AdamW


# 训练
optimizer = AdamW(model.parameters(), lr=5e-4)
loss = torch.nn.CrossEntropyLoss()


model.train()


for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(loader):
    
    input_ids = input_ids.to(device)
    attention_mask = attention_mask.to(device)
    token_type_ids = token_type_ids.to(device)
    labels = labels.to(device)
        
    out = model(input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids)
    
    # 梯度下降
    l = loss(out, labels)
    optimizer.zero_grad()
    l.backward()
    optimizer.step()
    

    if i % 5 == 0:
        out = out.cpu()
        labels = labels.cpu()
            
        out = out.argmax(dim=1)
        accuracy = (out == labels).sum().item() / len(labels)
        print(i, l.item(), accuracy)

    if i == 300:
        break