r""" convert the sample to batch tensor.
(samples, pad_idx, eos_idx)
| 7 | |
| 8 | |
| 9 | def collate_fn(samples, pad_idx, eos_idx): |
| 10 | r""" |
| 11 | convert the sample to batch tensor. |
| 12 | """ |
| 13 | if len(samples) == 0: |
| 14 | return {} |
| 15 | |
| 16 | def merge(key): |
| 17 | return collate_tokens([s[key] for s in samples], |
| 18 | pad_idx, |
| 19 | eos_idx=eos_idx) |
| 20 | |
| 21 | batch = { |
| 22 | 'nsentences': len(samples), |
| 23 | 'net_input': {}, |
| 24 | } |
| 25 | if samples[0].get('source', None) is not None: |
| 26 | batch['net_input']['input_ids'] = merge('source') |
| 27 | if samples[0].get('id', None) is not None: |
| 28 | batch['id'] = np.array([s.get('id') for s in samples]) |
| 29 | if samples[0].get('target', None) is not None: |
| 30 | batch['target'] = merge('target') |
| 31 | tgt_lengths = torch.LongTensor( |
| 32 | [s['target'].ne(pad_idx).long().sum() for s in samples]) |
| 33 | ntokens = tgt_lengths.sum().item() |
| 34 | batch['ntokens'] = ntokens |
| 35 | if samples[0].get('prev_output_tokens', None) is not None: |
| 36 | batch['net_input']['decoder_input_ids'] = merge('prev_output_tokens') |
| 37 | if samples[0].get('patch_image', None) is not None: |
| 38 | batch['net_input']['patch_images'] = torch.stack( |
| 39 | [sample['patch_image'] for sample in samples], dim=0) |
| 40 | if samples[0].get('patch_mask', None) is not None: |
| 41 | batch['net_input']['patch_masks'] = torch.cat( |
| 42 | [sample['patch_mask'] for sample in samples]) |
| 43 | # image generation |
| 44 | if samples[0].get('code_mask', None) is not None: |
| 45 | batch['net_input']['code_masks'] = torch.cat( |
| 46 | [sample['code_mask'] for sample in samples]) |
| 47 | if samples[0].get('code_image', None) is not None: |
| 48 | batch['code_images'] = torch.cat( |
| 49 | [sample['code_image'] for sample in samples]) |
| 50 | # For classification tasks (i.e., VQA, SNLI-VE, GLUE) |
| 51 | if samples[0].get('conf', None) is not None: |
| 52 | batch['conf'] = torch.cat([s['conf'] for s in samples], dim=0) |
| 53 | if samples[0].get('ref_dict', None) is not None: |
| 54 | batch['ref_dict'] = np.array([s['ref_dict'] for s in samples]) |
| 55 | if samples[0].get('label', None) is not None: |
| 56 | batch['labels'] = np.array([s['label'] for s in samples]).tolist() |
| 57 | if samples[0].get('constraint_mask', None) is not None: |
| 58 | batch['constraint_masks'] = merge('constraint_mask') |
| 59 | if samples[0].get('decoder_prompt', None) is not None: |
| 60 | batch['decoder_prompts'] = np.array( |
| 61 | [s['decoder_prompt'].tolist() for s in samples]) |
| 62 | if samples[0].get('prefix_token', None) is not None: |
| 63 | batch['prefix_tokens'] = merge('prefix_token') |
| 64 | # For detection and visual grounding |
| 65 | if samples[0].get('w_resize_ratio', None) is not None: |
| 66 | batch['w_resize_ratios'] = torch.stack( |