batcher.py (8514B)
1 from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple 2 import collections 3 4 import torch 5 6 import mmsrl.utils 7 8 9 class Batcher: 10 """ 11 Batch a group of sample together. 12 13 Two new features are derived from the "text": its length and a mask. 14 """ 15 16 def __init__(self, config: mmsrl.utils.dotdict, text_pad: int, label_pad: int, image_pad: int) -> None: 17 """ Initialize a Batcher, using the provided value to pad text. """ 18 self.config: mmsrl.utils.dotdict = config 19 self.text_pad: int = text_pad 20 self.label_pad: int = label_pad 21 self.image_pad: int = image_pad 22 23 def process_text(self, batch: Dict[str, Any], target: str, prefix: str, pad: int) -> None: 24 """ Build mask and text batch by padding sentences. """ 25 if target not in batch: 26 return 27 batch_size: int = len(batch[target]) 28 if isinstance(batch[target][0], torch.Tensor): 29 length: torch.Tensor = torch.tensor([len(text) for text in batch[target]], dtype=torch.int64) 30 max_seq_len: int = max(length.max().item(), self.config.get("min_text_len", 0)) 31 32 text: torch.Tensor = torch.empty((batch_size, max_seq_len), dtype=torch.int64) 33 mask: torch.Tensor = torch.empty((batch_size, max_seq_len), dtype=torch.bool) 34 for b, sentence in enumerate(batch[target]): 35 text[b, :sentence.shape[0]] = sentence 36 text[b, sentence.shape[0]:] = pad 37 mask[b, :sentence.shape[0]] = True 38 mask[b, sentence.shape[0]:] = False 39 40 batch[target] = text 41 batch[f"{prefix}length"] = length 42 batch[f"{prefix}attention_mask"] = mask 43 else: 44 assert(isinstance(batch[target][0], list) and isinstance(batch[target][0][0], torch.Tensor)) 45 generate_mask2: bool = "decoder_answer_mask" in batch and target == "decoder_target" 46 generate_mask3: bool = "decoder_constraint" in batch and target == "decoder_target" 47 48 sup_length: int = len(batch[target][0]) 49 length: torch.Tensor = torch.tensor([[len(text) for text in sup] for sup in batch[target]], dtype=torch.int64) 50 max_seq_len: int = max(length.max().item(), self.config.get("min_text_len", 0)) 51 52 text: torch.Tensor = torch.empty((batch_size, sup_length, max_seq_len), dtype=torch.int64) 53 mask: torch.Tensor = torch.empty((batch_size, sup_length, max_seq_len), dtype=torch.bool) 54 if generate_mask2: 55 mask2: torch.Tensor = torch.empty((batch_size, sup_length, max_seq_len), dtype=torch.bool) 56 if generate_mask3: 57 vocabulary_size: int = batch["decoder_constraint"][0][0].shape[1] 58 mask3: torch.Tensor = torch.empty((batch_size, sup_length, max_seq_len, vocabulary_size), dtype=torch.bool) 59 for b, sup in enumerate(batch[target]): 60 for i, sentence in enumerate(sup): 61 text[b, i, :sentence.shape[0]] = sentence 62 text[b, i, sentence.shape[0]:] = pad 63 mask[b, i, :sentence.shape[0]] = True 64 mask[b, i, sentence.shape[0]:] = False 65 if generate_mask2: 66 mask2[b, i, :sentence.shape[0]] = batch["decoder_answer_mask"][b][i] 67 mask2[b, i, sentence.shape[0]:] = False 68 if generate_mask3: 69 mask3[b, i, :sentence.shape[0]] = batch["decoder_constraint"][b][i] 70 mask3[b, i, sentence.shape[0]:] = False 71 72 batch[target] = text 73 batch[f"{prefix}length"] = length 74 batch[f"{prefix}attention_mask"] = mask 75 if generate_mask2: 76 batch["decoder_answer_mask"] = mask2 77 if generate_mask3: 78 batch["decoder_constraint"] = mask3 79 80 def process_entities(self, batch: Dict[str, Any]) -> None: 81 """ Build mask and entities text batch by padding list of entities and their input ids. """ 82 batch_size: int = len(batch["entities_name"]) 83 max_num_entities: int = max(len(entities) for entities in batch["entities_name"]) 84 max_num_entities = max(max_num_entities, self.config.get("min_nb_ent", 0)) 85 86 if "labels" in batch: 87 if batch["labels"][0].ndim == 2: 88 assert(batch["labels"][0].shape[1] == len(mmsrl.utils.LABELS)) 89 labels: torch.Tensor = torch.empty((batch_size, max_num_entities, len(mmsrl.utils.LABELS)), dtype=torch.bool) 90 else: 91 labels: torch.Tensor = torch.empty((batch_size, max_num_entities), dtype=torch.int64) 92 93 for b, entities in enumerate(batch["entities_name"]): 94 labels[b, :len(entities)] = batch["labels"][b] 95 labels[b, len(entities):] = self.label_pad 96 batch["labels"] = labels 97 98 mask: torch.Tensor = torch.empty((batch_size, max_num_entities), dtype=torch.bool) 99 for b, entities in enumerate(batch["entities_name"]): 100 mask[b, :len(entities)] = True 101 mask[b, len(entities):] = False 102 batch["entities_mask"] = mask 103 104 if "entities_input_ids" in batch: 105 lengths: torch.Tensor = torch.zeros((batch_size, max_num_entities), dtype=torch.int64) 106 for b, entities in enumerate(batch["entities_input_ids"]): 107 for i, entity in enumerate(entities): 108 lengths[b, i] = len(entity) 109 max_seq_len: int = max(lengths.max().item(), self.config.get("min_ent_len", 0)) 110 111 text: torch.Tensor = torch.empty((batch_size, max_num_entities, max_seq_len), dtype=torch.int64) 112 text_mask: torch.Tensor = torch.empty((batch_size, max_num_entities, max_seq_len), dtype=torch.bool) 113 for b, entities in enumerate(batch["entities_input_ids"]): 114 text[b, len(entities):] = self.text_pad 115 text_mask[b, len(entities):] = False 116 for i, entity in enumerate(entities): 117 text[b, i, :len(entity)] = entity 118 text[b, i, len(entity):] = self.text_pad 119 text_mask[b, i, :len(entity)] = True 120 text_mask[b, i, len(entity):] = False 121 122 batch["entities_input_ids"] = text 123 batch["entities_length"] = lengths 124 batch["entities_attention_mask"] = text_mask 125 126 def process_image(self, batch: Dict[str, Any]) -> None: 127 """ Build width, height and mask by padding images. """ 128 batch_size: int = len(batch["image"]) 129 heights: torch.Tensor = torch.tensor([image.shape[1] for image in batch["image"]], dtype=torch.int64) 130 widths: torch.Tensor = torch.tensor([image.shape[2] for image in batch["image"]], dtype=torch.int64) 131 height: int = max(heights.max().item(), self.config.get("height_min", 0)) 132 width: int = max(widths.max().item(), self.config.get("width_min", 0)) 133 images: torch.Tensor = torch.full((batch_size, 3, height, width), self.image_pad, dtype=torch.float32) 134 mask: torch.Tensor = torch.zeros((batch_size, height, width), dtype=torch.bool) 135 136 for b, image in enumerate(batch["image"]): 137 mask[b, :image.shape[1], :image.shape[2]] = True 138 images[b, :, :image.shape[1], :image.shape[2]] = image 139 140 batch["image"] = images 141 batch["image_mask"] = mask 142 batch["height"] = heights 143 batch["width"] = widths 144 145 def __call__(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: 146 """ Batch the provided samples """ 147 batch = collections.defaultdict(list) 148 for sample in samples: 149 for key, value in sample.items(): 150 batch[key].append(value) 151 152 self.process_text(batch, "input_ids", "", self.text_pad) 153 self.process_text(batch, "vbert_input_ids", "vbert_", self.text_pad) 154 self.process_text(batch, "caption_input_ids", "caption_", self.text_pad) 155 self.process_text(batch, "decoder_input_ids", "decoder_", self.text_pad) 156 self.process_text(batch, "decoder_target", "decoder_target_", self.label_pad) 157 self.process_entities(batch) 158 if batch.get("image"): 159 self.process_image(batch) 160 for feature in list(batch.keys()): 161 if isinstance(batch[feature], list) and isinstance(batch[feature][0], torch.Tensor): 162 batch[feature] = torch.stack(batch[feature]) 163 164 return batch