import torch


def create_batch(sentences, params, dico):
    """ Convert a list of tokenized sentences into a Pytorch batch

    args:
        sentences: list of sentences
        params: attribute params of the loaded model
        dico: dictionary

    returns:
        word_ids: indices of the tokens
        lengths: lengths of each sentence in the batch
    """
    bs = len(sentences)
    slen = max([len(sent) for sent in sentences])

    word_ids = torch.LongTensor(slen, bs).fill_(params.pad_index)
    for i in range(len(sentences)):
        sent = torch.LongTensor([dico.index(w) for w in sentences[i]])
        word_ids[:len(sent), i] = sent

    lengths = torch.LongTensor([len(sent) for sent in sentences])
    return word_ids, lengths