uncertainty_utils.py

import numpy as np

def temperature_scaling(logits, temperature):
    logits = np.array(logits)
    logits /= temperature

    # apply softmax
    logits -= logits.max()
    logits = logits - np.log(np.sum(np.exp(logits)))
    smx = np.exp(logits)
    return smx

def cal_uncertainty_entropy(token_logprobs, temperature):
    logits_softmax = temperature_scaling(token_logprobs, temperature)
    uncertainty = np.sum(-np.log(logits_softmax)*logits_softmax)
    # uncertainty = entropy(logits_softmax, base=2)
    return uncertainty

def cal_uncertainty_log_sum(token_logprobs, temperature):
    logits_softmax = temperature_scaling(token_logprobs, temperature)
    uncertainty = np.sum(-np.log(logits_softmax))
    return uncertainty

def cal_uncertainty_min(token_logprobs, temperature):
    logits_softmax = temperature_scaling(token_logprobs, temperature)
    uncertainty = -np.log(np.min(logits_softmax))
    return uncertainty

def cal_uncertainty_mean(token_logprobs, temperature):
    logits_softmax = temperature_scaling(token_logprobs, temperature)
    uncertainty = -np.log(np.mean(logits_softmax))
    return uncertainty

def cal_uncertainty_norm(token_logprobs, temperature):
    def normalize(v):
        norm = np.linalg.norm(v)
        if norm == 0: 
            return v
        return v / norm
    logits_softmax = temperature_scaling(token_logprobs, temperature)
    uncertainty = np.sum(-np.log(normalize(logits_softmax)))
    return uncertainty

def cal_uncertainty_single_token(token_logprobs):
    if len(token_logprobs) == 1:
        uncertainty = np.abs(token_logprobs[0])
    else:
        uncertainty = np.abs(np.min(token_logprobs))
    return uncertainty