added method fot batched attention_mask

2025-10-06 12:00:11 +02:00
parent 7e40a36701
commit 87409fecd5
2 changed files with 9 additions and 4 deletions
--- a/Project_Model/Libs/Transformer/Utils/init.py
+++ b/Project_Model/Libs/Transformer/Utils/init.py
@@ -1,4 +1,4 @@
-from .attention_mask import get_attention_mask
+from .attention_mask import get_causal_attention_mask,get_causal_attention_mask_batched
 from .task_type import TaskType

-__all__ = ["get_attention_mask", "TaskType"]
+__all__ = ["get_causal_attention_mask", "TaskType", "get_causal_attention_mask_batched"]