File size: 2,975 Bytes
2e7f2ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from typing import Optional
import torch
def _make_causal_mask(
attention_mask: torch.Tensor, dtype: torch.dtype, device: torch.device
):
"""
Make causal mask used for bi-directional self-attention.
"""
bsz, tgt_len = attention_mask.shape
mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), 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)
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len)
def _make_2dvison_mask(column_mask, dtype: torch.dtype, device: torch.device):
"""
"""
bsz, seq_length = column_mask.shape
cross_mask = torch.zeros((bsz, 1, seq_length, seq_length), dtype=dtype, device=device)
# 找到连续的 1 的区间
start = None
for bsz_idx in range(bsz):
for i in range(seq_length):
if column_mask[bsz_idx, i] == 1:
if start is None:
start = i
else:
if start is not None:
# 填充区间
cross_mask[bsz_idx, 0, start:i, start:i] = 1
start = None
# 处理最后一个区间
if start is not None:
cross_mask[bsz_idx, 0, start:seq_length, start:seq_length] = 1
return cross_mask
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
"""
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
"""
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 make_mask(attention_mask: torch.Tensor, dtype: torch.dtype=None, device: torch.device=None, mode: str="default", vision_mask: torch.Tensor=None, ):
if dtype is None:
dtype = attention_mask.dtype
if device is None:
device = attention_mask.device
expanded_attn_mask = _expand_mask(attention_mask, dtype).to(device)
causal_mask = _make_causal_mask(attention_mask, dtype, device).to(device)
if mode == "default":
return attention_mask
else:
assert vision_mask is not None, "vision_mask is None"
vision_mask = vision_mask.to(device)
bsz, seq_length = attention_mask.shape
vision_mask_bg = vision_mask[:, None, :, None]
vision_mask_2d = _make_2dvison_mask(vision_mask, dtype, device)
if mode == "bidirectional":
mask = expanded_attn_mask + causal_mask
mask = mask.clone().masked_fill_(vision_mask_2d.to(torch.bool), 0)
return mask
else:
raise NotImplementedError(f"mode {mode} is not implemented")
|