(batch, max_len=25)
| 86 | |
| 87 | |
| 88 | def custom_collate(batch, max_len=25): |
| 89 | input_ids_list, target_ids_list = zip(*batch) |
| 90 | pad_idx = 0 |
| 91 | |
| 92 | padded_input_ids = [] |
| 93 | input_lengths = [] |
| 94 | for ids in input_ids_list: |
| 95 | length = len(ids) |
| 96 | input_lengths.append(min(length, max_len)) |
| 97 | if length < max_len: |
| 98 | padded = torch.cat([ids, torch.full((max_len - length,), pad_idx, dtype=torch.long)]) |
| 99 | else: |
| 100 | padded = ids[:max_len] |
| 101 | padded_input_ids.append(padded) |
| 102 | |
| 103 | padded_input_ids = torch.stack(padded_input_ids) |
| 104 | |
| 105 | # padded_input_ids = pad_sequence(input_ids_list, batch_first=True, padding_value=pad_idx) |
| 106 | target_tokens = torch.tensor([x[-1].item() for x in target_ids_list], dtype=torch.long) |
| 107 | attention_mask = (padded_input_ids != pad_idx).long() |
| 108 | input_lengths = torch.tensor([len(ids) for ids in input_ids_list], dtype=torch.long) |
| 109 | |
| 110 | return padded_input_ids, target_tokens, attention_mask, input_lengths |
| 111 | |
| 112 | |
| 113 | def custom_collate_test(batch, max_len=25): |
nothing calls this directly
no outgoing calls
no test coverage detected