dataset.py (23802B)
1 from typing import List 2 3 import json 4 import os 5 import pathlib 6 import pandas 7 import random 8 import collections 9 10 import torch 11 import PIL 12 import torchvision 13 import mmsrl.utils 14 15 16 class MemeDataset(torch.utils.data.IterableDataset): 17 """ 18 Creating a custom dataset for reading the dataset and 19 loading it into the dataloader to pass it to the 20 neural network for finetuning the model 21 22 """ 23 24 def __init__(self, config, split, is_eval, tokenizer=None, ofa_task=None, image_processor=None, image_features_loader=None, vbert_tokenizer = None): 25 """ 26 Initializes a Dataset class 27 28 Args: 29 dataname in ['uspolitics','covid19', 'all'] 30 split in ['train', 'val', 'test', 'unseen_test'] 31 tokenizer (transformers.tokenizer): Transformers tokenizer 32 """ 33 super().__init__() 34 self.config = config 35 self.tokenizer = tokenizer 36 self.split = split 37 self.is_eval = is_eval 38 self.ofa_task = ofa_task 39 self.has_labels = ('unseen' not in split) 40 41 if os.environ.get("DATA_PATH", "") == "": 42 raise RuntimeError("Environment variable DATA_PATH is not set. Run the following command: export DATA_PATH=\"$PWD/../data\"") 43 data_path = pathlib.Path(os.environ['DATA_PATH']) 44 annot_file = data_path / "mmsrl" / self.config.dataname / 'annotations' / f"{split}.jsonl" 45 with annot_file.open('r') as json_file: 46 self.dataset = list(map(json.loads, json_file)) 47 48 self.images_path = data_path / "mmsrl" / self.config.dataname / 'images' 49 if self.config.get('use_caption') and self.config.dataname == "all": 50 caption_path = data_path / "mmsrl" / self.config.dataname / "captions.csv" 51 self.captions_df = pandas.read_csv(caption_path, sep='\t', header=None) 52 53 mean = [0.5, 0.5, 0.5] 54 std = [0.5, 0.5, 0.5] 55 self.image_transform = image_processor 56 if image_features_loader is not None: 57 self.image_features_loader = image_features_loader 58 if vbert_tokenizer is not None: 59 self.vbert_tokenizer = vbert_tokenizer 60 elif self.config.model in ["ofa", "mmf"]: 61 self.image_transform = torchvision.transforms.Compose([ 62 lambda image: image.convert("RGB"), 63 self.resize_image, 64 torchvision.transforms.ToTensor(), 65 torchvision.transforms.Normalize(mean=mean, std=std), 66 ]) 67 68 if self.config.model == "ofa" and not self.config.get("group_image_entities"): 69 self.dataset_len = 0 70 for sample in self.dataset: 71 if "entity_list" in sample: 72 self.dataset_len += len(set(sample["entity_list"])) 73 else: 74 self.dataset_len += sum(len(sample[label]) for label in mmsrl.utils.LABELS) 75 # If we ever remove entities appearing twice in samples: 76 #entities = set() 77 #for label in mmsrl.utils.LABELS: 78 # entities.update(sample[label]) 79 #self.dataset_len += len(entities) 80 81 if self.ofa_task: 82 self.ofa_bos = torch.LongTensor([self.ofa_task.src_dict.bos()]) 83 self.ofa_eos = torch.LongTensor([self.ofa_task.src_dict.eos()]) 84 85 if self.config.get("maybe_is", "") in ["yes", "no"]: 86 self.snli_answers = ["no", "maybe", "yes"] 87 else: 88 self.snli_answers = ["no", "yes"] 89 90 if self.config.get("constrain_output"): 91 assert(self.config.model == "ofa") 92 assert(not self.config.get("only_score_answer")) 93 assert(not self.config.get("only_train_answer")) 94 assert(not self.config.get("ofa_classification_head")) 95 assert(self.config.get("decoder_input") is not None and self.config.decoder_input == "{label}") 96 if self.config.ofa_task == "vqa_gen": 97 self.build_trie(mmsrl.utils.LABELS) 98 else: 99 self.build_trie(self.snli_answers) 100 101 self.seed(0) 102 self.compute_subsample(epoch=0) 103 104 def seed(self, id: int): 105 self.common_rng = random.Random(0) 106 self.rng = random.Random(id) 107 108 def __len__(self): 109 """returns the length of dataframe""" 110 if not self.is_eval and self.config.get("batch_per_epoch"): 111 return self.config.batch_per_epoch * self.config.batch_size 112 elif self.config.model == "ofa" and not self.config.get("group_image_entities"): 113 return self.dataset_len 114 else: 115 # Approximation due to sub-sampling entities deleting empty samples. 116 return len(self.dataset) 117 118 def get_label_weights(self): 119 if self.config.get("class_weights"): 120 label_weights = torch.zeros(4, dtype=torch.float32, requires_grad=False) 121 for sample in self.dataset: 122 for i, label in enumerate(mmsrl.utils.LABELS): 123 label_weights[i] += len(sample[label]) 124 label_weights = 1/label_weights 125 label_weights /= label_weights.mean() # normalise 126 else: 127 label_weights = torch.ones(4, dtype=torch.float32, requires_grad=False) 128 return label_weights 129 130 def compute_subsample(self, epoch): 131 if self.is_eval: 132 return 133 if self.config.get("cyclic_subsample") is not None: 134 config = mmsrl.utils.dotdict(self.config.cyclic_subsample[epoch % len(self.config.cyclic_subsample)]) 135 elif epoch>0 and not self.config.get("subsample_labels", "").startswith("interpolate"): 136 return 137 else: 138 config = self.config 139 140 frequencies = self.config.label_frequencies 141 if config.get("subsample_labels", "micro") == "micro": 142 target_frequencies = frequencies 143 elif config.subsample_labels == "macro": 144 target_frequencies = [0.25]*4 145 elif config.subsample_labels == "interpolate_micro_to_macro": 146 initial_frequencies = frequencies 147 final_frequencies = [0.25]*4 148 factor = epoch / (self.config.max_epoch - 1) 149 target_frequencies = [ 150 factor * final + (1 - factor) * initial 151 for initial, final in zip(initial_frequencies, final_frequencies)] 152 elif config.subsample_labels == "affine": 153 target_frequencies = [frequency * config.frequency_factor + config.frequency_bias for frequency in frequencies] 154 elif config.subsample_labels == "target": 155 target_frequencies = config.frequency_target 156 else: 157 raise RuntimeError(f"Unknown value for config.subsample_labels: {config.subsample_labels}") 158 target_frequencies[3] *= config.get("subsample_other", 1) 159 target_frequencies = [frequency / sum(target_frequencies) for frequency in target_frequencies] 160 sampling_frequencies = [frequency / target for frequency, target in zip(frequencies, target_frequencies)] 161 sampling_frequencies = [frequency / sum(sampling_frequencies) for frequency in sampling_frequencies] 162 self.subsample_probs = [min(sampling_frequencies) / frequency for frequency in sampling_frequencies] 163 164 repr_frequencies = ", ".join(f"{f:5.3f}" for f in target_frequencies) 165 repr_subsample = ", ".join(f"{p:5.3f}" for p in self.subsample_probs) 166 if self.split == "train": 167 print(f"\033[33mSubsampling factors are [{repr_subsample}]\033[0m (target frequencies: [{repr_frequencies}])") 168 169 def resize_image(self, image): 170 if max(image.size) > self.config.image_max_edge: 171 factor: float = self.config.image_max_edge / max(image.size) 172 height: int = round(image.size[0] * factor) 173 width: int = round(image.size[1] * factor) 174 image = torchvision.transforms.functional.resize(image, (height, width), interpolation=torchvision.transforms.InterpolationMode.BICUBIC) 175 return image 176 177 def process_image(self, image_path: pathlib.Path): 178 if self.config.model in ["baseline", "none"]: 179 return {} 180 image = PIL.Image.open(image_path) 181 if self.config.model in ['clip_momenta', 'clip', 'generic']: 182 image = self.image_transform(image) 183 if self.config.features_path: 184 feature = self.image_features_loader(image_path) 185 return {'image': image, 'image_feature': feature} 186 else: 187 return {'image': image} 188 elif self.config.model == 'visualbert': 189 image = self.image_transform(image_path) 190 elif self.image_transform is not None: 191 image = self.image_transform(image) 192 return {'image': image} 193 194 def process_text(self, text): 195 if self.config.model == "ofa": 196 if isinstance(text, str): 197 text = text.lower().replace('-', ' ').replace('/', ' ') 198 if not text: 199 return torch.tensor([], dtype=torch.int64) 200 encoded = self.ofa_task.tgt_dict.encode_line( 201 line=self.ofa_task.bpe.encode(f" {text}"), 202 add_if_not_exist=False, 203 append_eos=False 204 ).long() 205 return encoded[:self.config.max_text_len] 206 elif isinstance(text, tuple): 207 prefix, infix, suffix = tuple(map(self.process_text, text)) 208 total_length = prefix.shape[0] + infix.shape[0] + suffix.shape[0] + 2 209 if total_length > self.config.max_text_len: 210 infix = infix[:max(total_length - self.config.max_text_len, 0)] 211 encoded = torch.cat([self.ofa_bos, prefix, infix, suffix, self.ofa_eos]) 212 return encoded[:self.config.max_text_len] 213 else: 214 raise RuntimeError("text is of wrong type {type(text)}") 215 elif self.config.model in ["clip", "clip_momenta", "generic"]: 216 return self.tokenizer(text, truncate=True).squeeze(0)[:self.config.max_text_len] 217 elif self.config.model in ["baseline", 'visualbert']: 218 return self.tokenizer.encode_plus( 219 text, 220 max_length=self.config.max_text_len, 221 truncation=True, 222 return_tensors="pt", 223 )["input_ids"].squeeze(0).to(dtype=torch.long) 224 else: 225 return None 226 227 def create_affirmation(self, entity: str, label: str) -> str: 228 if label != "other": 229 return f"{entity} is a {label}." 230 else: 231 assert(self.config.get("neither_for_other")) # For another commit 232 classes = mmsrl.utils.LABELS.copy() 233 classes.remove("other") 234 self.rng.shuffle(classes) 235 class_string = ", ".join(classes[:-1]) + f" nor {classes[-1]}" 236 return f"{entity} is neither a {class_string}." 237 238 def create_question(self, entity, ocr): 239 classes = mmsrl.utils.LABELS.copy() 240 self.rng.shuffle(classes) 241 class_string = ", ".join(classes[:-1]) + f" or {classes[-1]}" 242 #question = f"{entity} belongs to one of the following categorie: {class_string}. Which one is it ?" 243 question = f"What is the category of {entity} between {class_string} ?" 244 245 if self.config.vqa_input == "question ocr": 246 return (question, ocr, "") 247 elif self.config.vqa_input == "ocr question": 248 return ("", ocr, question) 249 elif self.config.vqa_input == "question": 250 return ("", question, "") 251 else: 252 raise RuntimeError(f"Unknown value for config.vqa_input: {self.config.vqa_input}") 253 254 def create_premises(self, entity, ocr): 255 if self.config.get("ofa_classification_head"): 256 labels = ["other"] 257 else: 258 labels = mmsrl.utils.LABELS 259 260 premises = [] 261 for label in labels: 262 affirmation = self.create_affirmation(entity, label) 263 premises.append((('can image and text1 "', ocr, '" imply text2 " {affirmation} "?'), label)) 264 return premises 265 266 def prepare_ofa_io(self, ocr, entity, true_label): 267 test_answers: bool = self.is_eval and not self.config.get("ofa_classification_head") 268 if self.config.get("ofa_classification_head"): 269 true_label = "other" 270 271 encoder_input = [] 272 decoder_input = [] 273 if self.config.ofa_task == "vqa_gen": 274 source = self.create_question(entity, ocr) 275 encoder_input.append(source) 276 if test_answers: 277 for label in mmsrl.utils.LABELS: 278 if self.config.get("decoder_input") is not None: 279 decoder_input.append((self.config.decoder_input.format(label=label), "", "")) 280 else: 281 decoder_input.append((source[0], source[1], f"{source[2]} {label}")) 282 else: 283 if self.config.get("decoder_input") is not None: 284 decoder_input.append((self.config.decoder_input.format(label=true_label), "", "")) 285 else: 286 decoder_input.append((source[0], source[1], f"{source[2]} {true_label}")) 287 elif self.config.ofa_task == "snli_ve": 288 for source, label in self.create_premises(entity, ocr): 289 encoder_input.append(source) 290 if test_answers: 291 for answer in self.snli_answers: 292 if self.config.get("decoder_input") is not None: 293 decoder_input.append((self.config.decoder_input.format(label=answer), "", "")) 294 else: 295 decoder_input.append((source[0], source[1], f"{source[2]} {answer}")) 296 else: 297 answer = "yes" if label == true_label else "no" 298 if self.config.get("decoder_input") is not None: 299 decoder_input.append((self.config.decoder_input.format(label=answer), "", "")) 300 else: 301 decoder_input.append((source[0], source[1], f"{source[2]} {answer}")) 302 else: 303 raise RuntimeError(f"Unknown value for self.config.ofa_task: {self.config.ofa_task}") 304 return encoder_input, decoder_input 305 306 def subsample_labels(self, entities, labels): 307 if self.subsample_probs == [1, 1, 1, 1]: 308 return entities, labels 309 310 delete_me: List[int] = [] 311 for i, label in enumerate(labels): 312 if self.config.get("multilabel"): 313 for j, val in enumerate(label): 314 if val and torch.rand(1) > self.subsample_probs[j]: 315 label[j] = 0 316 if label[j] == [0,0,0,0]: 317 delete_me.append(i) 318 else: 319 if torch.rand(1) > self.subsample_probs[label]: 320 delete_me.append(i) 321 322 # We good a good score while doing this is old code… 323 if self.config.get("keep_all_batches"): 324 if len(delete_me) == len(entities): 325 self.rng.shuffle(delete_me) 326 delete_me = delete_me[1:] 327 328 delete_me.sort(reverse=True) 329 for index in delete_me: 330 labels.pop(index) 331 entities.pop(index) 332 333 return entities, labels 334 335 def generate_mask(self, input_ids, text): 336 mask = torch.zeros(input_ids.shape[0]-1, dtype=torch.bool) 337 mask[-2] = True 338 if text[-1].endswith("victim") or text[-1].endswith("villain"): 339 mask[-3] = True 340 return mask 341 342 def build_trie(self, answers: List[str]): 343 self.trie = collections.defaultdict(lambda: torch.zeros(len(self.ofa_task.tgt_dict), dtype=torch.bool)) 344 for answer in answers: 345 tokenized = torch.cat([self.process_text(answer), self.ofa_eos]) 346 for i in range(len(tokenized)): 347 self.trie[tuple(tokenized[:i].tolist())][tokenized[i]] = True 348 349 def get_constraint(self, target: torch.Tensor) -> torch.Tensor: 350 constraint = torch.zeros((target.shape[0], len(self.ofa_task.tgt_dict)), dtype=torch.bool) 351 for i in range(target.shape[0]): 352 constraint[i] = self.trie[tuple(target[:i].tolist())] 353 return constraint 354 355 def index_to_sample(self, index): 356 """return the input ids, images and target ids""" 357 sample = {} 358 359 sample["image_name"] = self.dataset[index]["image"] 360 image_data = self.process_image(self.images_path / sample["image_name"]) 361 if image_data is None: 362 return 363 if self.config.model == "visualbert": 364 sample['image_feature'] = image_data['image'] 365 else: 366 sample.update(image_data) 367 368 sample["raw_text"] = self.dataset[index]["OCR"] 369 if self.config.get('use_caption'): 370 caption = self.captions_df[self.captions_df[0]==sample["image_name"]][1].item() + '. ' 371 sample["caption"] = caption 372 sample["caption_input_ids"] = self.process_text(sample["caption"]) 373 374 if self.has_labels: 375 targets = {l: self.dataset[index][l] for l in mmsrl.utils.LABELS} 376 entities = [t for l in targets for t in targets[l]] 377 378 if self.config.get("multilabel"): 379 labels = [] 380 for ent in entities: 381 ent_lab = [mmsrl.utils.LABELS.index(l) for l in targets for t in targets[l] if t==ent] 382 ent_lab_distrib = [0]*len(mmsrl.utils.LABELS) 383 for l in ent_lab: ent_lab_distrib[l] = 1 384 labels.append(ent_lab_distrib) 385 else: 386 labels = [mmsrl.utils.LABELS.index(l) for l in targets for t in targets[l]] 387 388 if not self.is_eval: 389 entities, labels = self.subsample_labels(entities, labels) 390 if not entities: 391 return 392 393 ids = list(range(len(entities))) 394 self.rng.shuffle(ids) 395 ids = ids[:self.config.get("max_nb_ent", len(ids))] 396 entities = [entities[i] for i in ids] 397 labels = [labels[i] for i in ids] 398 399 sample["labels"] = torch.tensor(labels, dtype=torch.int64) 400 else: 401 entities = self.dataset[index]["entity_list"] 402 403 sample["entities_name"] = entities 404 405 if self.config.model == "ofa": 406 if self.config.get("group_image_entities"): 407 assert(self.config.get("ofa_classification_head")) 408 if self.has_labels: 409 sample["labels"] = torch.tensor(labels, dtype=torch.int64) 410 else: 411 None 412 sample["text"] = [] 413 sample["decoder_text"] = [] 414 for i, entity in enumerate(entities): 415 label = labels[i] if self.has_labels else None 416 label_text = mmsrl.utils.LABELS[label] if label is not None else None 417 encoder_input, decoder_input = self.prepare_ofa_io(sample["raw_text"], entity, label_text) 418 assert(len(encoder_input) == 1) 419 assert(len(decoder_input) == 1) 420 sample["text"].append(encoder_input[0]) 421 sample["decoder_text"].append(decoder_input[0]) 422 sample["input_ids"] = list(map(self.process_text, sample["text"])) 423 decoder_text = list(map(self.process_text, sample["decoder_text"])) 424 if not self.config.get("ofa_classification_head") and not self.config.get("keep_eof"): 425 decoder_text = [text[:-1] for text in decoder_text] 426 sample["decoder_input_ids"] = [text[:-1] for text in decoder_text] 427 sample["decoder_target"] = [text[1:] for text in decoder_text] 428 if self.config.get("constrain_output"): 429 sample["decoder_constraint"] = [self.get_constraint(target) for target in sample["decoder_target"]] 430 yield sample 431 else: 432 for i, entity in enumerate(entities): 433 current_sample = {} 434 current_sample["raw_text"] = sample["raw_text"] 435 current_sample["image_name"] = sample["image_name"] 436 current_sample["image"] = sample["image"] 437 438 if self.has_labels: 439 label = labels[i] 440 current_sample["labels"] = torch.tensor([label], dtype=torch.int64) 441 else: 442 label = None 443 444 label_text = mmsrl.utils.LABELS[label] if label is not None else None 445 encoder_input, decoder_input = self.prepare_ofa_io(sample["raw_text"], entity, label_text) 446 current_sample["text"] = encoder_input 447 current_sample["decoder_text"] = decoder_input 448 449 current_sample["input_ids"] = list(map(self.process_text, encoder_input)) 450 decoder_text = list(map(self.process_text, decoder_input)) 451 if not self.config.get("ofa_classification_head") and not self.config.get("keep_eof"): 452 decoder_text = [text[:-1] for text in decoder_text] 453 current_sample["decoder_input_ids"] = [text[:-1] for text in decoder_text] 454 current_sample["decoder_target"] = [text[1:] for text in decoder_text] 455 if self.config.get("only_score_answer") or self.config.get("only_train_answer"): 456 current_sample["decoder_answer_mask"] = [self.generate_mask(input_ids, text) for input_ids, text in zip(decoder_text, decoder_input)] 457 if self.config.get("constrain_output"): 458 current_sample["decoder_constraint"] = [self.get_constraint(target) for target in current_sample["decoder_target"]] 459 460 current_sample["entities_name"] = [entity] 461 yield current_sample 462 else: 463 sample["text"] = sample["raw_text"] 464 sample["input_ids"] = self.process_text(sample["text"]) 465 if self.config.get('use_visualbert'): 466 sample["vbert_input_ids"] = self.vbert_tokenizer.encode_plus( 467 sample["text"], 468 max_length=self.config.max_text_len, 469 truncation=True, 470 return_tensors="pt", 471 )["input_ids"].squeeze(0).to(dtype=torch.long) 472 sample["entities_input_ids"] = [self.process_text(entity) for entity in entities] 473 yield sample 474 475 def __iter__(self): 476 worker_info = torch.utils.data.get_worker_info() 477 if worker_info is None: 478 worker_modulo: int = 1 479 worker_residue: int = 0 480 else: 481 worker_modulo: int = worker_info.num_workers 482 worker_residue: int = worker_info.id 483 484 if not self.is_eval and self.config.get("batch_per_epoch"): 485 sample_per_worker = len(self) // worker_modulo + (1 if worker_residue < len(self) % worker_modulo else 0) 486 sample_generated = 0 487 while sample_generated < sample_per_worker: 488 samples = list(self.index_to_sample(self.rng.randint(0, len(self.dataset)-1))) 489 yield from samples[:sample_per_worker - sample_generated] 490 sample_generated += len(samples) 491 else: 492 self.order = list(range(len(self.dataset))) 493 if not self.is_eval: 494 self.common_rng.shuffle(self.order) 495 496 for i, index in enumerate(self.order): 497 if i % worker_modulo == worker_residue: 498 yield from self.index_to_sample(index)