Method knns_to_probs

knn_utils/saveEmbedMulti.py:188–195

(self, knns, neg_dists)

        return output 

    def knns_to_probs(self, knns, neg_dists):
        probs = torch.nn.functional.softmax(neg_dists / self.knn_temperature, dim=-1).to(torch.float32)
        vals_at_knns = self.vals[knns].squeeze(-1).to(probs.device) # (nonpad batch * time, k)
        knn_probs = torch.full(size=(vals_at_knns.shape[:-1] + (self.vocab_size,)), fill_value=0.0).to(self.device) \
            .scatter_add(dim=-1, index=vals_at_knns, src=probs) # (nonpad_batch * time, vocab)
        knn_probs = F.normalize(knn_probs, p=1, dim=-1)
        
        return knn_probs

    def knns_to_log_prob(self, knns, neg_dists):
        probs = torch.nn.functional.softmax(neg_dists / self.knn_temperature, dim=-1)