import torch

def precision_at_k(documents: list, k=3):
    """
    documents: list, each item is [logit, label]
    """
    documents = sorted(documents, key=lambda x: x[0], reverse=True)
    right_number = len([item for item in documents[:k] if item[1] == 1])
    return right_number / k
