set position ids and use block diagonal attn mask

This commit is contained in:
Wing Lian
2023-07-17 01:56:32 -04:00
parent 36b0e30a9d
commit 4964b0d345
4 changed files with 64 additions and 3 deletions

View File

@@ -55,10 +55,14 @@ class TestPacking(unittest.TestCase):
# first example doesn't have mask reset
assert example["input_ids"][0] == self.tokenizer.bos_token_id
assert example["attention_mask"][0] == 1
assert example["position_ids"][0] == 0
assert example["position_ids"][1] == 1
# but subsequent one does
assert example["input_ids"][next_bos_index] == self.tokenizer.bos_token_id
assert example["attention_mask"][next_bos_index] == 2
assert example["position_ids"][next_bos_index] == 0
assert example["position_ids"][next_bos_index + 1] == 1
if __name__ == "__main__":