import torch


def pad_collate_textonly(samples, pad_idx=1, pad_first=False):
    """For each batch, fill in all texts with the pad_idx, in order to make all
    lines have the same length as the longest text of this batch.
    A sample is a batch of dataset (a batch of tuples of x and y). s[0] is x,
    s[1] is y.
    When pad_first=True, the text is at the end, when pad_first=False, the text
    is in the beginning.
    """
    max_len = max([len(s[0]) for s in samples])
    # res is a matrix
    res = torch.zeros(len(samples), max_len).long() + pad_idx
    for i, s in enumerate(samples):
        if pad_first:
            res[i, -len(s[0]) :] = torch.LongTensor(s[0])
        else:
            res[i, : len(s[0])] = torch.LongTensor(s[0])
    return res, torch.tensor([s[1] for s in samples])