pyiqa.archs.q_align.modeling_attn_mask_utils

Module Contents

class pyiqa.archs.q_align.modeling_attn_mask_utils.AttentionMaskConverter(is_causal: bool, sliding_window: int | None = None)

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

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.

to_causal_4d(batch_size: int, query_length: int, key_value_length: int, dtype: torch.dtype = torch.float32, device: torch.device | str = 'cpu') → torch.Tensor

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).

to_4d(attention_mask_2d: torch.Tensor, query_length: int, key_value_length: int | None = None, dtype: torch.dtype = torch.float32) → torch.Tensor

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.