import torch


def embed_passage_batch(passages, tokenizer, embedding_model, max_length, device="cuda:0"):
    """
    This is a refactoring of the embed_passages_for_retrieval function in the blog code.
    :param passages: (N element list) the batch of passages to be embedded. Each element in the list is a string
    :param tokenizer: the retribert tokenizer from https://huggingface.co/yjernite/retribert-base-uncased
    :param embedding_model: the retribert model from https://huggingface.co/yjernite/retribert-base-uncased
    :param max_length: (int) the maximum number of tokens for each embedded passage. The code either truncates or pads to this length
    :param device: the device (cpu or cuda:x) that the embeddings will be calculated on

    :return: N*max_length dimension numpy array containing the batch embeddings
    """

    # the tokenization base class has (I think) been updated since the blog post.
    # The blog code throws warnings. This should work.
    tokenized_passages = tokenizer(passages,
                                   max_length=max_length,
                                   padding="max_length",
                                   truncation=True,
                                   return_tensors='pt')

    # now make the embeddings
    # note we are moving this to the gpu
    # we don't want to calculate gradients here because we are only doing inference
    with torch.no_grad():
        embedded_passages = embedding_model.embed_answers(tokenized_passages["input_ids"].to(device),
                                                          tokenized_passages["attention_mask"].to(device))

    # now change these to numpy and return (because we will save them as a np.memmap)
    # note if we hadn't used the torch.no_grad() we would have had to detach the gradients.
    return embedded_passages.cpu().numpy()