import torch


def predict(model, dataloader, labeldict):
    """
    Predict the labels of an unlabelled test set with a pretrained model.

    Args:
        model: The torch module which must be used to make predictions.
        dataloader: A DataLoader object to iterate over some dataset.
        labeldict: A dictionary associating labels to integer values.

    Returns:
        A dictionary associating pair ids to predicted labels.
    """
    # Switch the model to eval mode.
    model.eval()
    device = model.device

    # Revert the labeldict to associate integers to labels.
    labels = {index: label for label, index in labeldict.items()}
    predictions = {}

    # Deactivate autograd for evaluation.
    with torch.no_grad():
        for batch in dataloader:

            # Move input and output data to the GPU if one is used.
            ids = batch["id"]
            premises = batch['premise'].to(device)
            premises_lengths = batch['premise_length'].to(device)
            hypotheses = batch['hypothesis'].to(device)
            hypotheses_lengths = batch['hypothesis_length'].to(device)

            _, probs = model(premises,
                             premises_lengths,
                             hypotheses,
                             hypotheses_lengths)

            _, preds = probs.max(dim=1)

            for i, pair_id in enumerate(ids):
                predictions[pair_id] = labels[int(preds[i])]

    return predictions