2023 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license a utility attention mask class that allows one to create a causal 4d mask create a causal 4d mask with slided window convert a 2d attention mask batchsize querylength to a 4d attention mask batchsize 1 querylength keyvaluelength that can be multiplied with attention scores examples python import torch from transformers modelingattnmaskutils import attentionmaskconverter converter attentionmaskconvertertrue converter to4dtorch tensor0 0 0 1 1 5 5 tensor3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 3 4028e38 0 0000e00 3 4028e38 3 4028e38 3 4028e38 3 4028e38 0 0000e00 0 0000e00 parameters iscausal bool whether the attention mask should be a unidirectional causal or bidirectional mask slidingwindow int optional optionally the sliding window masks can be created if slidingwindow is defined to a positive integer creates a causal 4d mask of bsz headdim1 querylength keyvaluelength shape and adds large negative bias to upper right hand triangular matrix causal mask if shape is not cached create a new causal mask and cache it create causal mask bsz seqlen bsz 1 tgtseqlen srcseqlen converts 2d attention mask to 4d attention mask by expanding mask to bsz headdim1 querylength keyvaluelength shape and by adding a large negative bias to notattended positions if attentionmask is causal a causal mask will be added create causal mask bsz seqlen bsz 1 tgtseqlen srcseqlen bsz seqlen bsz 1 tgtseqlen srcseqlen expandedattnmask causal4dmask can cause some overflow make causal mask used for bidirectional selfattention add lower triangular sliding window mask if necessary expands attentionmask from bsz seqlen to bsz 1 tgtseqlen srcseqlen creates a causal 4d mask of shape batchsize 1 querylength keyvaluelength from a 2d mask of shape batchsize keyvaluelength args attentionmask torch tensor or none a 2d attention mask of shape batchsize keyvaluelength inputshape tupleint or listint or torch size the input shape should be a tuple that defines batchsize querylength inputsembeds torch tensor the embedded inputs as a torch tensor pastkeyvalueslength int the length of the key value cache slidingwindow int optional if the model uses windowed attention a sliding window should be passed 4d mask is passed through the layers creates a noncausal 4d mask of shape batchsize 1 querylength keyvaluelength from a 2d mask of shape batchsize keyvaluelength args mask torch tensor or none a 2d attention mask of shape batchsize keyvaluelength dtype torch dtype the torch dtype the created mask shall have tgtlen int the target length or query length the created mask shall have creates a causal 4d mask of shape batchsize 1 querylength keyvaluelength args inputshape tupleint or listint or torch size the input shape should be a tuple that defines batchsize querylength dtype torch dtype the torch dtype the created mask shall have device int the torch device the created mask shall have slidingwindow int optional if the model uses windowed attention a sliding window should be passed 2023 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license a utility attention mask class that allows one to create a causal 4d mask create a causal 4d mask with slided window convert a 2d attention mask batch_size query_length to a 4d attention mask batch_size 1 query_length key_value_length that can be multiplied with attention scores examples python import torch from transformers modeling_attn_mask_utils import attentionmaskconverter converter attentionmaskconverter true converter to_4d torch tensor 0 0 0 1 1 5 5 tensor 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 0 0000e 00 3 4028e 38 3 4028e 38 3 4028e 38 3 4028e 38 0 0000e 00 0 0000e 00 parameters is_causal bool whether the attention mask should be a uni directional causal or bi directional mask sliding_window int optional optionally the sliding window masks can be created if sliding_window is defined to a positive integer creates a causal 4d mask of bsz head_dim 1 query_length key_value_length shape and adds large negative bias to upper right hand triangular matrix causal mask if shape is not cached create a new causal mask and cache it create causal mask bsz seq_len bsz 1 tgt_seq_len src_seq_len converts 2d attention mask to 4d attention mask by expanding mask to bsz head_dim 1 query_length key_value_length shape and by adding a large negative bias to not attended positions if attention_mask is causal a causal mask will be added create causal mask bsz seq_len bsz 1 tgt_seq_len src_seq_len bsz seq_len bsz 1 tgt_seq_len src_seq_len expanded_attn_mask causal_4d_mask can cause some overflow make causal mask used for bi directional self attention add lower triangular sliding window mask if necessary expands attention_mask from bsz seq_len to bsz 1 tgt_seq_len src_seq_len creates a causal 4d mask of shape batch_size 1 query_length key_value_length from a 2d mask of shape batch_size key_value_length args attention_mask torch tensor or none a 2d attention mask of shape batch_size key_value_length input_shape tuple int or list int or torch size the input shape should be a tuple that defines batch_size query_length inputs_embeds torch tensor the embedded inputs as a torch tensor past_key_values_length int the length of the key value cache sliding_window int optional if the model uses windowed attention a sliding window should be passed 4d mask is passed through the layers creates a non causal 4d mask of shape batch_size 1 query_length key_value_length from a 2d mask of shape batch_size key_value_length args mask torch tensor or none a 2d attention mask of shape batch_size key_value_length dtype torch dtype the torch dtype the created mask shall have tgt_len int the target length or query length the created mask shall have creates a causal 4d mask of shape batch_size 1 query_length key_value_length args input_shape tuple int or list int or torch size the input shape should be a tuple that defines batch_size query_length dtype torch dtype the torch dtype the created mask shall have device int the torch device the created mask shall have sliding_window int optional if the model uses windowed attention a sliding window should be passed
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch


@dataclass
class AttentionMaskConverter:
    

    is_causal: bool
    sliding_window: int

    def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
        self.is_causal = is_causal
        self.sliding_window = sliding_window

        if self.sliding_window is not None and self.sliding_window <= 0:
            raise ValueError(
                f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
            )

    def to_causal_4d(
        self,
        batch_size: int,
        query_length: int,
        key_value_length: int,
        dtype: torch.dtype = torch.float32,
        device: Union[torch.device, "str"] = "cpu",
    ) -> torch.Tensor:
        
        if not self.is_causal:
            raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")

        
        input_shape = (batch_size, query_length)
        past_key_values_length = key_value_length - query_length

        
        
        causal_4d_mask = None
        if input_shape[-1] > 1 or self.sliding_window is not None:
            causal_4d_mask = self._make_causal_mask(
                input_shape,
                dtype,
                device=device,
                past_key_values_length=past_key_values_length,
                sliding_window=self.sliding_window,
            )

        return causal_4d_mask

    def to_4d(
        self,
        attention_mask_2d: torch.Tensor,
        query_length: int,
        key_value_length: Optional[int] = None,
        dtype: torch.dtype = torch.float32,
    ) -> torch.Tensor:
        
        input_shape = (attention_mask_2d.shape[0], query_length)

        
        
        causal_4d_mask = None
        if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
            if key_value_length is None:
                raise ValueError(
                    "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
                )

            past_key_values_length = key_value_length - query_length
            causal_4d_mask = self._make_causal_mask(
                input_shape,
                dtype,
                device=attention_mask_2d.device,
                past_key_values_length=past_key_values_length,
                sliding_window=self.sliding_window,
            )
        elif self.sliding_window is not None:
            raise NotImplementedError("Sliding window is currently only implemented for causal masking")

        
        expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
            attention_mask_2d.device
        )
        if causal_4d_mask is not None:
            expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)

        
        expanded_4d_mask = expanded_attn_mask

        return expanded_4d_mask

    @staticmethod
    def _make_causal_mask(
        input_ids_shape: torch.Size,
        dtype: torch.dtype,
        device: torch.device,
        past_key_values_length: int = 0,
        sliding_window: Optional[int] = None,
    ):
        
        bsz, tgt_len = input_ids_shape
        mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
        mask_cond = torch.arange(mask.size(-1), device=device)
        mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)

        mask = mask.to(dtype)

        if past_key_values_length > 0:
            mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)

        
        if sliding_window is not None:
            diagonal = past_key_values_length - sliding_window + 1

            context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
            mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)

        return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)

    @staticmethod
    def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
        
        bsz, src_len = mask.size()
        tgt_len = tgt_len if tgt_len is not None else src_len

        expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

        inverted_mask = 1.0 - expanded_mask

        return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)


def _prepare_4d_causal_attention_mask(
    attention_mask: Optional[torch.Tensor],
    input_shape: Union[torch.Size, Tuple, List],
    inputs_embeds: torch.Tensor,
    past_key_values_length: int,
    sliding_window: Optional[int] = None,
):
    
    attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)

    key_value_length = input_shape[-1] + past_key_values_length

    
    if attention_mask is not None:
        attention_mask = attn_mask_converter.to_4d(
            attention_mask, input_shape[-1], key_value_length, dtype=inputs_embeds.dtype
        )
    else:
        attention_mask = attn_mask_converter.to_causal_4d(
            input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
        )

    return attention_mask


def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    
    return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)


def _create_4d_causal_attention_mask(
    input_shape: Union[torch.Size, Tuple, List],
    dtype: torch.dtype,
    device: torch.device,
    past_key_values_length: int = 0,
    sliding_window: Optional[int] = None,
):
    
    attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)

    key_value_length = past_key_values_length + input_shape[-1]
    attention_mask = attn_mask_converter.to_causal_4d(
        input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
    )

    return attention_mask