import torch


def collate_dataset(batch, test=False):
    """
    Preprocessing for given batch.
    It will use for the Torch DataLoader(collate_fn=collate_dataset).
    
    Args: 
        batch(Torch.Tensor)
    
    return:
        List(Torch.Tensor)
    """
    drug_batch = []
    protein_batch = []
    y_batch = []

    for data in batch:
        drug_batch.append(data[0])
        protein_batch.append(data[1])
        if not test:
            y_batch.append(data[2])
    if not test:
        return [torch.LongTensor(drug_batch), torch.LongTensor(protein_batch), torch.FloatTensor(y_batch)]
    else:
        return [torch.LongTensor(drug_batch), torch.LongTensor(protein_batch)]