(batch_input, tokenizer)
| 82 | |
| 83 | @staticmethod |
| 84 | def collate_fn(batch_input, tokenizer): |
| 85 | input_ids = pad_sequence([torch.tensor( |
| 86 | [tokenizer.cls_token_id] + d['source'] + d['target'] + [tokenizer.sep_token_id] |
| 87 | ) for d in batch_input], batch_first=True) |
| 88 | |
| 89 | attention_mask = torch.ones_like(input_ids) |
| 90 | |
| 91 | target_mask = torch.stack([torch.cat([ |
| 92 | torch.zeros(len(d['source']) + 1), torch.ones(input_ids.size(1) - len(d['source']) - 1) |
| 93 | ]) for d in batch_input]) |
| 94 | |
| 95 | assert input_ids.size() == attention_mask.size() == target_mask.size() |
| 96 | return { |
| 97 | 'input_ids': input_ids, |
| 98 | 'attention_mask': attention_mask, |
| 99 | 'target_mask': target_mask, |
| 100 | } |
| 101 | |
| 102 | class QQPLoader(ConditionalLoader): |
| 103 | def __init__(self, tokenizer, return_source_length=False): |
nothing calls this directly
no outgoing calls
no test coverage detected