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']),
    }