PoliBa-DeepLearning/NanoSocrates
2
0
Fork 0
Code Issues 3 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

added method fot batched attention_mask

Browse Source
This commit is contained in:
GassiGiuseppe
2025-10-06 12:00:11 +02:00
parent 7e40a36701
commit 87409fecd5
2 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
Project_Model/Libs/Transformer/Utils/attention_mask.py
View File

@@ -1,4 +1,9 @@
import torch
def get_attention_mask(embedding_dimension: int) -> torch.Tensor:
return torch.triu(torch.ones(embedding_dimension, embedding_dimension, dtype=torch.bool), diagonal=1)
def get_causal_attention_mask(embedding_dimension: int) -> torch.Tensor:
return torch.triu(torch.ones(embedding_dimension, embedding_dimension, dtype=torch.bool), diagonal=1)
def get_causal_attention_mask_batched(embedding_dimension: int, batch_size: int ) -> torch.Tensor:
base_mask = get_causal_attention_mask(embedding_dimension)
return base_mask.unsqueeze(0).expand(batch_size, -1, -1) # add another dimension at the beginning, big as batch_size
# the result is that z,x,y where x,y are repeated along z