ICML2022/OFA

1import string2import math3import json4from itertools import chain5import os6 7import torch8import torch.distributed as dist9 10from data import data_utils11 12 13def get_symbols_to_strip_from_output(generator):14    if hasattr(generator, "symbols_to_strip_from_output"):15        return generator.symbols_to_strip_from_output16    else:17        return {generator.bos, generator.eos}18 19 20def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None):21    x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator))22    if bpe is not None:23        x = bpe.decode(x)24    if tokenizer is not None:25        x = tokenizer.decode(x)26    return x27 28 29def eval_caption(task, generator, models, sample, **kwargs):30    transtab = str.maketrans({key: None for key in string.punctuation})31    hypos = task.inference_step(generator, models, sample)32    results = []33    for i, sample_id in enumerate(sample["id"].tolist()):34        detok_hypo_str = decode_fn(hypos[i][0]["tokens"], task.tgt_dict, task.bpe, generator)35        results.append({"image_id": str(sample_id), "caption": detok_hypo_str.translate(transtab).strip()})36    return results, None37 38 39def eval_vqa_gen(task, generator, models, sample, **kwargs):40    if kwargs['beam_search_vqa_eval']:41        hypos = task.inference_step(generator, models, sample, prefix_tokens=sample['prefix_tokens'])42        results = []43        for i, sample_id in enumerate(sample["id"].tolist()):44            prefix_len = sample['prefix_tokens'][i].ne(1).sum().item()45            detok_hypo_str = decode_fn(hypos[i][0]["tokens"][prefix_len:], task.tgt_dict, task.bpe, generator)46            results.append({"question_id": int(sample_id), "answer": detok_hypo_str.strip()})47        scores = [ref_dict.get(result['answer'], 0) for ref_dict, result in zip(sample['ref_dict'], results)]48        return results, scores49 50    encoder_out = models[0].encoder(51        sample["net_input"]["src_tokens"],52        src_lengths=sample["net_input"]["src_lengths"],53        patch_images=sample["net_input"]["patch_images"],54        patch_masks=sample["net_input"]["patch_masks"]55    )56    device = sample["net_input"]["src_tokens"].device57    eos_item = torch.tensor([task.src_dict.eos()])58    pad = task.src_dict.pad()59    valid_result = []60    for valid_answers, valid_constraint_masks in zip(task.valid_answers_list, task.valid_constraint_masks_list):61        valid_size = len(valid_answers)62        valid_tgt_items = [63            torch.cat([torch.tensor(decoder_prompt[1:]), valid_answer, eos_item])64            for decoder_prompt in sample["decoder_prompts"] for valid_answer in valid_answers65        ]66        valid_prev_items = [67            torch.cat([torch.tensor(decoder_prompt), valid_answer])68            for decoder_prompt in sample["decoder_prompts"] for valid_answer in valid_answers69        ]70        valid_constraint_mask_items = [71            torch.cat(72                [torch.zeros(len(decoder_prompt) - 1, valid_constraint_mask.size(1)).bool(), valid_constraint_mask],73                dim=074            )75            for decoder_prompt in sample["decoder_prompts"] for valid_constraint_mask in valid_constraint_masks76        ]77        valid_tgt = data_utils.collate_tokens(valid_tgt_items, pad_idx=pad).to(device)78        valid_prev_output = data_utils.collate_tokens(valid_prev_items, pad_idx=pad).to(device)79        valid_constraint_masks = data_utils.collate_tokens(valid_constraint_mask_items, pad_idx=pad).to(device)80 81        new_encoder_out = {}82        new_encoder_out["encoder_out"] = [83            encoder_out["encoder_out"][0].repeat_interleave(valid_size, dim=1)84        ]85        new_encoder_out["encoder_padding_mask"] = [86            encoder_out["encoder_padding_mask"][0].repeat_interleave(valid_size, dim=0)87        ]88        new_encoder_out["position_embeddings"] = [89            encoder_out["position_embeddings"][0].repeat_interleave(valid_size, dim=0)90        ]91 92        decoder_out = models[0].decoder(valid_prev_output, encoder_out=new_encoder_out)93        decoder_out[0].masked_fill_(~valid_constraint_masks, -math.inf)94        lprobs = models[0].get_normalized_probs(decoder_out, log_probs=True)95        scores = lprobs.gather(dim=-1, index=valid_tgt.unsqueeze(-1)).squeeze(-1)96        scores = scores.masked_fill(valid_tgt.eq(task.tgt_dict.pad()), 0)97        scores = scores.masked_fill((~valid_constraint_masks).all(2), 0)98        scores = scores.sum(1)99        scores = scores.view(-1, valid_size)100        valid_result.append(scores)101    valid_result = torch.cat(valid_result, dim=-1)102    predicts = valid_result.argmax(1).tolist()103    hyps = [task.index2ans[predict_index] for predict_index in predicts]104    results = [{"question_id": int(id), "answer": hyp} for id, hyp in zip(sample["id"].tolist(), hyps)]105    scores = [ref_dict.get(hyp, 0) for ref_dict, hyp in zip(sample['ref_dict'], hyps)]106    return results, scores107 108 109def eval_refcoco(task, generator, models, sample, **kwargs):110    def _calculate_ap_score(hyps, refs, thresh=0.5):111        interacts = torch.cat(112            [torch.where(hyps[:, :2] < refs[:, :2], refs[:, :2], hyps[:, :2]),113             torch.where(hyps[:, 2:] < refs[:, 2:], hyps[:, 2:], refs[:, 2:])],114            dim=1115        )116        area_predictions = (hyps[:, 2] - hyps[:, 0]) * (hyps[:, 3] - hyps[:, 1])117        area_targets = (refs[:, 2] - refs[:, 0]) * (refs[:, 3] - refs[:, 1])118        interacts_w = interacts[:, 2] - interacts[:, 0]119        interacts_h = interacts[:, 3] - interacts[:, 1]120        area_interacts = interacts_w * interacts_h121        ious = area_interacts / (area_predictions + area_targets - area_interacts + 1e-6)122        return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float()123 124    gen_out = task.inference_step(generator, models, sample)125    hyps = []126    for i in range(len(gen_out)):127        hyps.append(gen_out[i][0]["tokens"][:-1] - len(task.src_dict) + task.cfg.num_bins)128    hyps = torch.stack(hyps, dim=0)129    hyps = hyps / (task.cfg.num_bins - 1) * task.cfg.max_image_size130    hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)131    hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)132 133    results = [134        {"uniq_id": sample_id,135         "box": [hyps[i][0].item(), hyps[i][1].item(), hyps[i][2].item(), hyps[i][3].item()]}136        for i, sample_id in enumerate(sample["id"].tolist())137    ]138    scores = _calculate_ap_score(hyps, sample['region_coords'].float())139    return results, scores140 141 142def eval_snli_ve(task, generator, models, sample, **kwargs):143    encoder_out = models[0].encoder(144        sample["net_input"]["src_tokens"],145        src_lengths=sample["net_input"]["src_lengths"],146        patch_images=sample["net_input"]["patch_images"],147        patch_masks=sample["net_input"]["patch_masks"]148    )149    device = sample["net_input"]["src_tokens"].device150    eos_item = torch.tensor([task.src_dict.eos()])151    pad = task.src_dict.pad()152    valid_result = []153    for valid_answers, valid_constraint_masks in zip(task.valid_answers_list, task.valid_constraint_masks_list):154        valid_size = len(valid_answers)155        valid_tgt_items = [156            torch.cat([torch.tensor(decoder_prompt[1:]), valid_answer, eos_item])157            for decoder_prompt in sample["decoder_prompts"] for valid_answer in valid_answers158        ]159        valid_prev_items = [160            torch.cat([torch.tensor(decoder_prompt), valid_answer])161            for decoder_prompt in sample["decoder_prompts"] for valid_answer in valid_answers162        ]163        valid_constraint_mask_items = [164            torch.cat(165                [torch.zeros(len(decoder_prompt) - 1, valid_constraint_mask.size(1)).bool(), valid_constraint_mask],166                dim=0167            )168            for decoder_prompt in sample["decoder_prompts"] for valid_constraint_mask in valid_constraint_masks169        ]170        valid_tgt = data_utils.collate_tokens(valid_tgt_items, pad_idx=pad).to(device)171        valid_prev_output = data_utils.collate_tokens(valid_prev_items, pad_idx=pad).to(device)172        valid_constraint_masks = data_utils.collate_tokens(valid_constraint_mask_items, pad_idx=pad).to(device)173 174        new_encoder_out = {}175        new_encoder_out["encoder_out"] = [176            encoder_out["encoder_out"][0].repeat_interleave(valid_size, dim=1)177        ]178        new_encoder_out["encoder_padding_mask"] = [179            encoder_out["encoder_padding_mask"][0].repeat_interleave(valid_size, dim=0)180        ]181        new_encoder_out["position_embeddings"] = [182            encoder_out["position_embeddings"][0].repeat_interleave(valid_size, dim=0)183        ]184 185        decoder_out = models[0].decoder(valid_prev_output, encoder_out=new_encoder_out)186        decoder_out[0].masked_fill_(~valid_constraint_masks, -math.inf)187        lprobs = models[0].get_normalized_probs(decoder_out, log_probs=True)188        scores = lprobs.gather(dim=-1, index=valid_tgt.unsqueeze(-1)).squeeze(-1)189        scores = scores.masked_fill(valid_tgt.eq(task.tgt_dict.pad()), 0)190        scores = scores.masked_fill((~valid_constraint_masks).all(2), 0)191        scores = scores.sum(1)192        scores = scores.view(-1, valid_size)193        valid_result.append(scores)194    valid_result = torch.cat(valid_result, dim=-1)195    predicts = valid_result.argmax(1).tolist()196    hyps = [task.index2ans[predict_index] for predict_index in predicts]197    results = [{"uniq_id": id, "answer": hyp} for id, hyp in zip(sample["id"].tolist(), hyps)]198    scores = [ref_dict.get(hyp, 0) for ref_dict, hyp in zip(sample['ref_dict'], hyps)]199    return results, scores200 201 202def eval_image_gen(task, generator, models, sample, **kwargs):203    hypos, _ = task.inference_image(generator, sample, models)204    tokens = sample['net_input']['src_tokens'][0].view(-1).tolist()205    caption = task.bpe.decode(task.tgt_dict.string([token for token in tokens if token >= 4]))[206              38:].replace('/', '')207 208    text_similarity_score, indices = task.compute_text_similarity(hypos, caption,209                                                                  sample['net_input']['src_tokens'].device)210    results = []211    for i, indice in enumerate(indices):212        results.append({"sample_id": str(sample["id"][0]), "score": text_similarity_score[i], "image": hypos[indice]})213 214    scores = [max(text_similarity_score).item()]215    return results, scores216 217 218def eval_glue(task, generator, models, sample, **kwargs):219    net_output = models[0](**sample["net_input"])220    net_output[0].masked_fill_(~sample["constraint_masks"], -math.inf)221    last_token_ids = sample["net_input"]["prev_output_tokens"].ne(task.src_dict.pad()).sum(1, keepdim=True) - 1222    logits = net_output[0].gather(1, last_token_ids.unsqueeze(2).expand(-1, -1, net_output[0].size(2)))223    logits = logits.squeeze(1)224    predicts = logits.argmax(1).tolist()225    hyps = [task.bpe.decode(task.src_dict[predict]).strip() for predict in predicts]226    results = [{"hyp": hyp, "ref": ref_dict.keys()[0]} for hyp, ref_dict in zip(hyps, sample['ref_dict'])]227    return results, None228 229 230def eval_step(task, generator, models, sample, **kwargs):231    if task.cfg._name == 'caption':232        return eval_caption(task, generator, models, sample, **kwargs)233    elif task.cfg._name == 'vqa_gen':234        return eval_vqa_gen(task, generator, models, sample, **kwargs)235    elif task.cfg._name == 'refcoco':236        return eval_refcoco(task, generator, models, sample, **kwargs)237    elif task.cfg._name == 'snli_ve':238        return eval_snli_ve(task, generator, models, sample, **kwargs)239    elif task.cfg._name == 'image_gen':240        return eval_image_gen(task, generator, models, sample, **kwargs)241    elif task.cfg._name in {'cola', 'mnli', 'mrpc', 'qnli', 'qqp', 'rte', 'sst2'}:242        return eval_glue(task, generator, models, sample, **kwargs)243    else:244        raise NotImplementedError245 246 247def merge_results(task, cfg, logger, score_cnt, score_sum, results):248    if task.cfg._name == 'image_gen':249        if cfg.distributed_training.distributed_world_size > 1:250            dist.all_reduce(score_sum.data)251            dist.all_reduce(score_cnt.data)252        if score_cnt.item() > 0:253            logger.info("score_sum: {}, score_cnt: {}, score: {}".format(254                score_sum, score_cnt, round(score_sum.item() / score_cnt.item(), 4)255            ))256    else:257        gather_results = None258        if cfg.distributed_training.distributed_world_size > 1:259            gather_results = [None for _ in range(dist.get_world_size())]260            dist.all_gather_object(gather_results, results)261            dist.all_reduce(score_sum.data)262            dist.all_reduce(score_cnt.data)263        if score_cnt.item() > 0:264            logger.info("score_sum: {}, score_cnt: {}, score: {}".format(265                score_sum, score_cnt, round(score_sum.item() / score_cnt.item(), 4)266            ))267 268        if cfg.distributed_training.distributed_world_size == 1 or dist.get_rank() == 0:269            os.makedirs(cfg.common_eval.results_path, exist_ok=True)270            output_path = os.path.join(cfg.common_eval.results_path, "{}_predict.json".format(cfg.dataset.gen_subset))271            gather_results = list(chain(*gather_results)) if gather_results is not None else results272            with open(output_path, 'w') as fw:273                json.dump(gather_results, fw)274