Pipeline fix and added a util to decode
This commit is contained in:
@@ -41,18 +41,19 @@ class Decoder(nn.Module):
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor
|
||||
]
|
||||
): # -> list[torch.Tensor]: # k_x = v_x . While x_q = x
|
||||
# WARNING: args is needed to have sequential
|
||||
x, k_x, v_x, padding_mask = args
|
||||
x, k_x, v_x, src_padding_mask, tgt_padding_mask = args
|
||||
|
||||
# build of attention mask
|
||||
attention_mask = get_causal_attention_mask(x.size(1))
|
||||
|
||||
# 1) Masked Attention
|
||||
MASKED_ATTENTION = self.__masked_attention(
|
||||
x, x, x, key_padding_mask=padding_mask, attention_mask=attention_mask
|
||||
x, x, x, key_padding_mask=tgt_padding_mask, attention_mask=attention_mask
|
||||
)
|
||||
|
||||
# 2) Dropout
|
||||
@@ -68,7 +69,7 @@ class Decoder(nn.Module):
|
||||
|
||||
# 5) Encoder–decoder (cross) attention
|
||||
CROSS_ATTENTION = self.__cross_attention(
|
||||
x, k_x, v_x, key_padding_mask=padding_mask
|
||||
x, k_x, v_x, key_padding_mask=src_padding_mask
|
||||
)
|
||||
|
||||
# 6) Dropout
|
||||
@@ -96,7 +97,7 @@ class Decoder(nn.Module):
|
||||
# 12) Layer Normalization
|
||||
x = self.__layer_norm_3(x)
|
||||
|
||||
return (x, k_x, v_x, padding_mask)
|
||||
return (x, k_x, v_x, src_padding_mask, tgt_padding_mask)
|
||||
|
||||
|
||||
# use eval to disable dropout ecc
|
||||
|
||||
Reference in New Issue
Block a user