forked from mindspore/mindspore
27 lines
953 B
Python
27 lines
953 B
Python
import mindspore as ms
|
|
import mindspore.ops as ops
|
|
|
|
class TriangularCausalMask:
|
|
def __init__(self, B, L, device="cpu"):
|
|
mask_shape = (B, 1, L, L)
|
|
self._mask = ms.Tensor(ops.triu(ops.ones(mask_shape, ms.bool_), diagonal=1), ms.bool_).to(device)
|
|
|
|
@property
|
|
def mask(self):
|
|
return self._mask
|
|
|
|
class ProbMask:
|
|
def __init__(self, B, H, L, index, scores, device="cpu"):
|
|
_mask = ops.triu(ops.ones((L, scores.shape[-1]), ms.bool_), diagonal=1).to(device)
|
|
_mask_ex = _mask[None, None, :].expand((B, H, L, scores.shape[-1]))
|
|
|
|
# Using advanced indexing to create the mask
|
|
indicator = _mask_ex[ms.Tensor(np.arange(B)[:, None, None]),
|
|
ms.Tensor(np.arange(H)[None, :, None]),
|
|
index, :].to(device)
|
|
|
|
self._mask = indicator.view(scores.shape).to(device)
|
|
|
|
@property
|
|
def mask(self):
|
|
return self._mask |