| 8 | |
| 9 | |
| 10 | def padding_seq_to_same_length(input_ids, max_pad_length, pad_token = 0): |
| 11 | padding_length = max_pad_length - len(input_ids) |
| 12 | padding_ids = [pad_token] * padding_length |
| 13 | attention_mask = [] |
| 14 | |
| 15 | if padding_length <= 0: |
| 16 | attention_mask = [1] * max_pad_length |
| 17 | input_ids = input_ids[:max_pad_length] |
| 18 | else: |
| 19 | attention_mask = [1] * len(input_ids) + [0] * padding_length |
| 20 | input_ids = input_ids + padding_ids |
| 21 | |
| 22 | assert len(input_ids) == max_pad_length |
| 23 | assert len(attention_mask) == max_pad_length |
| 24 | |
| 25 | return input_ids, attention_mask |
| 26 | |
| 27 | |
| 28 | class Retrieval_Dataset(Dataset): |