Function collate_fn

predict_downstream_condition.py:118–136

(batch_input)

    return wf - wf.mean()

def collate_fn(batch_input):
    input_ids = pad_sequence([torch.tensor(
        [tokenizer.cls_token_id] + d['source'] + [tokenizer.mask_token_id] * (args.seq_len - len(d['source']) - 1)
    ) for d in batch_input], batch_first=True)

    attention_mask = torch.ones_like(input_ids)

    target_mask = torch.stack([torch.cat([
        torch.zeros(len(d['source']) + 1), torch.ones(input_ids.size(1) - len(d['source']) - 1)
    ]) for d in batch_input])
    target_start = torch.tensor([len(d['source']) + 1 for d in batch_input]).long()

    assert input_ids.size() == attention_mask.size() == target_mask.size()
    return {
        'input_ids': input_ids.repeat(1, MBR_size).view(-1, input_ids.size(-1)),
        'attention_mask': attention_mask.repeat(1, MBR_size).view(-1, input_ids.size(-1)),
        'target_mask': target_mask.repeat(1, MBR_size).view(-1, input_ids.size(-1)),
        'target_start': target_start
    }

test_data = Dataloaders[task_name](tokenizer=tokenizer).my_load(splits=['test'])[0]
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, collate_fn=collate_fn, num_workers=4, pin_memory=True)