import torch
from selfcheckgpt.modeling_selfcheck import SelfCheckNLI
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


selfcheck_nli = SelfCheckNLI(device=device) # set device to 'cuda' if GPU is available


def nli_confidence(proofs, sentences):

    sent_scores_nli = selfcheck_nli.predict(
        sentences = sentences,                          # list of sentences
        sampled_passages = proofs, # list of sampled passages
    )
    #print(sent_scores_nli)
    sent_scores_nli = [round(1-s, 4) for s in sent_scores_nli]
    return sent_scores_nli