import torch


def encode_batch(batch, tokenizer, max_seq_length, return_tokens=False):
    """
    Tokenize and encode the given input; stack it in a batch with fixed length
    according to the longest sequence in the batch.

    Return a tuple.  If asked to return tokens, return the encoded batch and a nested
    list of the tokenized texts.  Otherwise, just return the batch.
    """
    # Truncate input based on the sequence length
    # before stacking as a batch
    # Also return a boolean array indicating which tokens are masked
    encoded_texts = []

    for text in batch:
        encoded_text = torch.tensor(tokenizer.encode(text))[:max_seq_length]
        encoded_texts.append(encoded_text)

    encoded_batch = torch.nn.utils.rnn.pad_sequence(encoded_texts, batch_first=True)

    if return_tokens:
        tokens = [tokenizer.tokenize(text) for text in batch]
        return encoded_batch, tokens
    else:
        return (encoded_batch,)