| import torch |
|
|
| def collate_fn_hard(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): |
| |
| 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']), |
| } |
|
|
