Method post_forward_hook

knn_utils/saveEmbedMulti.py:155–186

(self, module, input, output)

        return self.original_forward_func(input_ids=input_ids, labels=labels, attention_mask=attention_mask, **kwargs)

    def post_forward_hook(self, module, input, output):
        batch, time_dim, vocab_size = output.shape
        shift = 0 if self.is_encoder_decoder else 1
        lm_logits = output
        lm_logits = torch.nn.functional.log_softmax(lm_logits, dim=-1) # (batch, time, vocab)
        queries = self.activation_capturer.captured # (batch, time, dim)

        if self.labels is None:
            nonpad_mask = torch.cat([
                torch.zeros([batch, time_dim - 1], dtype=torch.bool),
                torch.ones([batch, 1], dtype=torch.bool),
            ], axis=-1).to(self.device)
        else:
            nonpad_mask = torch.cat([
                self.labels[:, shift:] != -100, 
                torch.zeros([self.labels.shape[0], shift], dtype=torch.bool).to(self.device)
            ], axis=-1)

        lm_logits = lm_logits[nonpad_mask]
        queries = queries[nonpad_mask] # (nonpad, dim)
        
        dists, knns = self.get_knns(queries) # (nonpad batch * time, k)
        
        # Compute knn probs
        neg_dists = -dists
        knn_log_probs = self.knns_to_log_prob(knns, neg_dists) # (nonpad b*t, vocab_size)
        
        # Interpolate
        interpolated_scores = KNNWrapperMulti.interpolate(dists, knn_log_probs, lm_logits, self.lmbda) # (nonpad b * t, vocab)
        output[nonpad_mask] = interpolated_scores.to(output.dtype)
        
        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)