def __call__()

in summarize_from_feedback/models/attention.py [0:0]

• 19 lines of code
• 1 McCabe index (conditional complexity)

    def __call__(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
        query = self.split_heads(query)
        key = self.split_heads(key, k=True)
        value = self.split_heads(value)

        # query: [batch, head, n_q, d_model]
        # key: [batch, head, d_model, n_k]
        # value: [batch, head, n_k, d_model]

        # Pre-divide by fp16_stability_scale to prevent fp16 overflow
        softmax_scale = 1.0 / np.sqrt(np.sqrt(query.size(-1)))
        query = query * softmax_scale
        key = key * softmax_scale

        w = torch.matmul(query, key)
        wtype = w.dtype
        w = w.float()

        # Dense attn with autoregressive mask
        n_q = w.size(-2)
        n_k = w.size(-1)

        # NOTE: Could use apex prefix softmax to speed this up

        mask = torch.ones(n_q, n_k, device=w.device).tril(diagonal=n_k - n_q).view(1, 1, n_q, n_k)

        # We make all values where the mask==0 into -inf so that they get
        # ignored when we do our softmax
        w = w * mask + -1e9 * (1 - mask)

        w = nn.Softmax(dim=-1)(w).type(wtype)

        w = self.attn_dropout_module(w)
        a = torch.matmul(w, value)
        # a: [batch, head, n_q, d_model]
        a = self.merge_heads(a)
        return a