File size: 1,766 Bytes
aa45001 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | import torch
def collate_fn_hard(batch):
# Each stored row is a pre-batch.
item = batch[0]
return {
'anchor_input_ids': torch.tensor(item['anchor_input_ids']),
'anchor_attention_mask': torch.tensor(item['anchor_attention_mask']),
'positive_input_ids': torch.tensor(item['positive_input_ids']),
'positive_attention_mask': torch.tensor(item['positive_attention_mask']),
'negative_input_ids': torch.tensor(item['negative_input_ids']),
'negative_attention_mask': torch.tensor(item['negative_attention_mask']),
'anchor_labels': torch.tensor(item['anchor_labels']),
'positive_labels': torch.tensor(item['positive_labels']),
'negative_labels': torch.tensor(item['negative_labels']),
'hard_negative_input_ids': torch.tensor(item['hard_negative_input_ids']),
'hard_negative_attention_mask': torch.tensor(item['hard_negative_attention_mask']),
'hard_negative_labels': torch.tensor(item['hard_negative_labels'])
}
def collate_fn(batch):
# Each stored row is a pre-batch.
item = batch[0]
return {
'anchor_input_ids': torch.tensor(item['anchor_input_ids']),
'anchor_attention_mask': torch.tensor(item['anchor_attention_mask']),
'positive_input_ids': torch.tensor(item['positive_input_ids']),
'positive_attention_mask': torch.tensor(item['positive_attention_mask']),
'negative_input_ids': torch.tensor(item['negative_input_ids']),
'negative_attention_mask': torch.tensor(item['negative_attention_mask']),
'anchor_labels': torch.tensor(item['anchor_labels']),
'positive_labels': torch.tensor(item['positive_labels']),
'negative_labels': torch.tensor(item['negative_labels']),
}
|