llm_with_confidence / self_check_gpt.py
Lihuchen's picture
Upload 5 files
cf63839
raw
history blame
567 Bytes
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 = [1-s for s in sent_scores_nli]
return sent_scores_nli