Views
No views yet
1import torch
2my_tokens = torch.load("train_tokens.pt")
3l = my_tokens['attention_mask'].sum(dim=-1)
4mask = ((170 <= l) & (l <= 180))
5
6train_215k = {}
7train_215k['input_ids'] = my_tokens['input_ids'][mask]
8train_215k['attention_mask'] = my_tokens['attention_mask'][mask]
9
10torch.save(train_215k, "train_tokens_215k.pt")
11check = torch.load("train_tokens_215k_2.pt")
12
13gap = (check['input_ids'] == 50256).sum(dim=-1)
14gap2 = (check['attention_mask'] == 0).sum(dim=-1)
15
16assert torch.all(gap == gap2)
17
18