tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix attention mask collation (#1603)

Browse Source
This commit is contained in:
Wing Lian
2024-05-14 08:17:30 -04:00
committed by GitHub
parent 5d97e65f95
commit 02982733ec
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/utils/collators.py
View File

@@ -229,9 +229,8 @@ class PretrainingBatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
if feature == "attention_mask": if feature == "attention_mask":
if self.multipack_attn: if self.multipack_attn:
arrays = [ arrays = [
(i + 1) * np.array(item[feature]) (i + 1) * np.array(item)
for i, item in enumerate(features[feature]) for i, item in enumerate(features[feature])
if feature in item
] ]
else: else:
arrays = [(1) * np.array(item) for item in features[feature]] arrays = [(1) * np.array(item) for item in features[feature]]