import torch


def collate_fn_rank(samples):
    """
    Creates a batch out of samples.
    Assumes samples is a list of sample of the form:
    Sample: ([(premise + ending, # of tokens for permise), ... 4x]
    The first item must be the correct ending.

    Returns:
        (premise tensor, endings tensor)
    """
    # Maximum length of tokens
    max_len = max(len(x) for sample in samples for x, _ in sample)
    
    # Mask the premise as well as padding
    x_mask = torch.ByteTensor([
        [[0] * prem_len + [1] * (len(x) - prem_len) + [0] * (max_len - len(x)) for x, prem_len in sample]
        for sample in samples 
    ])

    # Pad all token lists to same length
    x = torch.LongTensor([[x + [0] * (max_len - len(x)) for x, _ in sample] for sample in samples])
    return x, x_mask