import torch


def batchify(batch):
    """Gather a batch of individual examples into one batch."""

    # batch is a list of vectorized examples
    batch_size = len(batch)
    max_source_len = max([len(ex['source_ids']) for ex in batch])

    source_ids = torch.LongTensor(batch_size, max_source_len).zero_()
    source_mask = torch.LongTensor(batch_size, max_source_len).zero_()
    source_type_ids = torch.LongTensor(batch_size, max_source_len).zero_()
    source_pos_ids = None
    label = torch.LongTensor(batch_size)

    for bidx in range(batch_size):
        example = batch[bidx]
        label[bidx] = example['label']
        ex_len = len(example['source_ids'])
        source_ids[bidx, :ex_len].copy_(torch.LongTensor(example['source_ids']))
        source_type_ids[bidx, :ex_len].copy_(torch.LongTensor(example['type_ids']))
        source_mask[bidx, :ex_len].copy_(torch.ones(ex_len))

    return {
        'ids': [ex['id'] for ex in batch],
        'batch_size': batch_size,
        'sent_id': [ex['sent_id'] for ex in batch],
        'sentence': [ex['sentence'] for ex in batch],
        'ques_id': [ex['ques_id'] for ex in batch],
        'question': [ex['question'] for ex in batch],
        'source_ids': source_ids,
        'source_pos_ids': source_pos_ids,
        'source_type_ids': source_type_ids,
        'source_mask': source_mask,
        'label': label,
        'raw_label': [ex['raw_label'] for ex in batch]
    }