(batch)
| 74 | |
| 75 | |
| 76 | def custom_collate(batch): |
| 77 | input_ids_list, target_ids_list = zip(*batch) |
| 78 | pad_idx = 0 |
| 79 | |
| 80 | padded_input_ids = pad_sequence(input_ids_list, batch_first=True, padding_value=pad_idx) |
| 81 | target_tokens = torch.tensor([x[-1].item() for x in target_ids_list], dtype=torch.long) |
| 82 | attention_mask = (padded_input_ids != pad_idx).long() |
| 83 | input_lengths = torch.tensor([len(ids) for ids in input_ids_list], dtype=torch.long) |
| 84 | |
| 85 | return padded_input_ids, target_tokens, attention_mask, input_lengths |
| 86 | |
| 87 | |
| 88 | def custom_collate_test(batch): |
nothing calls this directly
no outgoing calls
no test coverage detected