scoringmethod.py

import torch
from torch import nn
torch.cuda.empty_cache()

class ScoringMethod(nn.Module):
    def __init__(self, label):
        super(ScoringMethod, self).__init__()
        self.label = label


class PllScoringMethod(ScoringMethod):
    def __init__(self, label):
        super(PllScoringMethod, self).__init__(label)

    def forward(self, probs, origids, return_all=False, **kwargs):
        mask = origids >= 0
        origids[~mask] = 0
        slen = len(probs) - 1
        dia = torch.diag(probs[1:].gather(-1, origids.unsqueeze(0).repeat(slen, 1).unsqueeze(-1)).squeeze(-1), diagonal=0)[mask]
        dia_list = dia.tolist()
        prob = torch.mean(torch.log(dia), dim=-1).detach().item()
        if return_all:
            return prob, dia_list
        return prob


class ComparativeScoringMethod(ScoringMethod):
    def __init__(self, label):
        super(ComparativeScoringMethod, self).__init__(label)

    def forward(self, probs, return_all=False, **kwargs):
        slen = len(probs) - 1
        dia = self.calc(probs[0, :slen], probs[torch.arange(1, slen + 1), torch.arange(slen)])
        dia_list = dia.tolist()
        prob = torch.mean(torch.log(dia), dim=-1).detach().item()
        if return_all:
            return prob, dia_list
        return prob

    def calc(self, p: torch.tensor, q: torch.tensor):
        raise NotImplementedError


class JSD(ComparativeScoringMethod):
    def __init__(self):
        super(JSD, self).__init__("jsd")
        self.kl = nn.KLDivLoss(reduction='none', log_target=True)

    def calc(self, p: torch.tensor, q: torch.tensor):
        m = torch.log((0.5 * (p + q)))
        return 1 - (0.5 * (torch.sum(self.kl(m, p.log()), dim=-1) + torch.sum(self.kl(m, q.log()), dim=-1)))


class PLL(PllScoringMethod):
    def __init__(self):
        super(PLL, self).__init__("pll")


class CSD(ComparativeScoringMethod):
    def __init__(self):
        super(CSD, self).__init__("csd")
        self.csd = torch.nn.CosineSimilarity(dim=1)

    def calc(self, p: torch.tensor, q: torch.tensor):
        return self.csd(p, q)


class ESD(ComparativeScoringMethod):
    def __init__(self):
        super(ESD, self).__init__("esd")
        self.pwd = torch.nn.PairwiseDistance()
        self.sqrt = torch.sqrt(torch.tensor(2, requires_grad=False))

    def norm(self, dist):
        return (torch.relu(self.sqrt - dist) + 0.000001) / self.sqrt

    def calc(self, p: torch.tensor, q: torch.tensor):
        return self.norm(self.pwd(p, q))


class MSD(ComparativeScoringMethod):
    def __init__(self):
        super(MSD, self).__init__("msd")
        self.mse = torch.nn.MSELoss(reduction="none")

    def calc(self, p: torch.tensor, q: torch.tensor):
        return self.mse(p, q).mean(axis=-1)


class HSD(ComparativeScoringMethod):
    def __init__(self):
        super(HSD, self).__init__("hsd")
        self.sqrt = torch.sqrt(torch.tensor(2, requires_grad=False))

    def calc(self, p: torch.Tensor, q: torch.Tensor):
        p = p.clone()
        q = q.clone()
        return 1 - torch.sqrt_(torch.sum(torch.pow(torch.sqrt_(p) - torch.sqrt_(q), 2), dim=-1)) / self.sqrt


KNOWN_METHODS = [CSD(), ESD(), JSD(), MSD(), HSD(), PLL()]
KNOWN_METHODS = {m.label: m for m in KNOWN_METHODS}