# Copyright 2022 The OFA-Sys Team. # All rights reserved. # This source code is licensed under the Apache 2.0 license # found in the LICENSE file in the root directory. from dataclasses import dataclass, field import json import logging from typing import Optional from argparse import Namespace from itertools import zip_longest from collections import OrderedDict import numpy as np import sacrebleu import string from fairseq import metrics, utils from fairseq.tasks import register_task from tasks.ofa_task import OFATask, OFAConfig from data.mm_data.caption_dataset import CaptionDataset from data.file_dataset import FileDataset from utils.cider.pyciderevalcap.ciderD.ciderD import CiderD EVAL_BLEU_ORDER = 4 logger = logging.getLogger(__name__) @dataclass class CaptionConfig(OFAConfig): eval_bleu: bool = field( default=False, metadata={"help": "evaluation with BLEU scores"} ) eval_cider: bool = field( default=False, metadata={"help": "evaluation with CIDEr scores"} ) eval_args: Optional[str] = field( default='{}', metadata={ "help": 'generation args for BLUE or CIDEr scoring, e.g., \'{"beam": 4, "lenpen": 0.6}\', as JSON string' }, ) eval_print_samples: bool = field( default=False, metadata={"help": "print sample generations during validation"} ) eval_cider_cached_tokens: Optional[str] = field( default=None, metadata={"help": "path to cached cPickle file used to calculate CIDEr scores"}, ) scst: bool = field( default=False, metadata={"help": "Self-critical sequence training"} ) scst_args: str = field( default='{}', metadata={ "help": 'generation args for Self-critical sequence training, as JSON string' }, ) @register_task("caption", dataclass=CaptionConfig) class CaptionTask(OFATask): def __init__(self, cfg: CaptionConfig, src_dict, tgt_dict): super().__init__(cfg, src_dict, tgt_dict) def load_dataset(self, split, epoch=1, combine=False, **kwargs): paths = self.cfg.data.split(',') assert len(paths) > 0 if split == 'train': file_path = paths[(epoch - 1) % (len(paths) - 1)] else: file_path = paths[-1] dataset = FileDataset(file_path, self.cfg.selected_cols) self.datasets[split] = CaptionDataset( split, dataset, self.bpe, self.src_dict, self.tgt_dict, max_src_length=self.cfg.max_src_length, max_tgt_length=self.cfg.max_tgt_length, patch_image_size=self.cfg.patch_image_size, imagenet_default_mean_and_std=self.cfg.imagenet_default_mean_and_std, scst=getattr(self.cfg, 'scst', False) ) def build_model(self, cfg): model = super().build_model(cfg) if self.cfg.eval_bleu or self.cfg.eval_cider: gen_args = json.loads(self.cfg.eval_args) self.sequence_generator = self.build_generator( [model], Namespace(**gen_args) ) if self.cfg.eval_cider: self.CiderD_scorer = CiderD(df=self.cfg.eval_cider_cached_tokens) if self.cfg.scst: scst_args = json.loads(self.cfg.scst_args) self.scst_generator = self.build_generator( [model], Namespace(**scst_args) ) return model def _calculate_cider_scores(self, gen_res, gt_res): ''' gen_res: generated captions, list of str gt_idx: list of int, of the same length as gen_res gt_res: ground truth captions, list of list of str. gen_res[i] corresponds to gt_res[gt_idx[i]] Each image can have multiple ground truth captions ''' gen_res_size = len(gen_res) res = OrderedDict() for i in range(gen_res_size): res[i] = [gen_res[i].strip()] gts = OrderedDict() gt_res_ = [ [gt_res[i][j].strip() for j in range(len(gt_res[i]))] for i in range(len(gt_res)) ] for i in range(gen_res_size): gts[i] = gt_res_[i] res_ = [{'image_id': i, 'caption': res[i]} for i in range(len(res))] _, scores = self.CiderD_scorer.compute_score(gts, res_) return scores def valid_step(self, sample, model, criterion): loss, sample_size, logging_output = criterion(model, sample) model.eval() if self.cfg.eval_bleu or self.cfg.eval_cider: hyps, refs = self._inference(self.sequence_generator, sample, model) if self.cfg.eval_bleu: if self.cfg.eval_tokenized_bleu: bleu = sacrebleu.corpus_bleu(hyps, list(zip_longest(*refs)), tokenize="none") else: bleu = sacrebleu.corpus_bleu(hyps, list(zip_longest(*refs))) logging_output["_bleu_sys_len"] = bleu.sys_len logging_output["_bleu_ref_len"] = bleu.ref_len # we split counts into separate entries so that they can be # summed efficiently across workers using fast-stat-sync assert len(bleu.counts) == EVAL_BLEU_ORDER for i in range(EVAL_BLEU_ORDER): logging_output["_bleu_counts_" + str(i)] = bleu.counts[i] logging_output["_bleu_totals_" + str(i)] = bleu.totals[i] if self.cfg.eval_cider: scores = self._calculate_cider_scores(hyps, refs) logging_output["_cider_score_sum"] = scores.sum() logging_output["_cider_cnt"] = scores.size return loss, sample_size, logging_output def reduce_metrics(self, logging_outputs, criterion): super().reduce_metrics(logging_outputs, criterion) def sum_logs(key): import torch result = sum(log.get(key, 0) for log in logging_outputs) if torch.is_tensor(result): result = result.cpu() return result if self.cfg.eval_bleu: counts, totals = [], [] for i in range(EVAL_BLEU_ORDER): counts.append(sum_logs("_bleu_counts_" + str(i))) totals.append(sum_logs("_bleu_totals_" + str(i))) if max(totals) > 0: # log counts as numpy arrays -- log_scalar will sum them correctly metrics.log_scalar("_bleu_counts", np.array(counts)) metrics.log_scalar("_bleu_totals", np.array(totals)) metrics.log_scalar("_bleu_sys_len", sum_logs("_bleu_sys_len")) metrics.log_scalar("_bleu_ref_len", sum_logs("_bleu_ref_len")) def compute_bleu(meters): import inspect import sacrebleu fn_sig = inspect.getfullargspec(sacrebleu.compute_bleu)[0] if "smooth_method" in fn_sig: smooth = {"smooth_method": "exp"} else: smooth = {"smooth": "exp"} bleu = sacrebleu.compute_bleu( correct=meters["_bleu_counts"].sum, total=meters["_bleu_totals"].sum, sys_len=meters["_bleu_sys_len"].sum, ref_len=meters["_bleu_ref_len"].sum, **smooth ) return round(bleu.score, 2) metrics.log_derived("bleu", compute_bleu) if self.cfg.eval_cider: def compute_cider(meters): cider = meters["_cider_score_sum"].sum / meters["_cider_cnt"].sum cider = cider if isinstance(cider, float) else cider.item() return round(cider, 3) if sum_logs("_cider_cnt") > 0: metrics.log_scalar("_cider_score_sum", sum_logs("_cider_score_sum")) metrics.log_scalar("_cider_cnt", sum_logs("_cider_cnt")) metrics.log_derived("cider", compute_cider) def _inference(self, generator, sample, model): def decode(toks, escape_unk=False): s = self.tgt_dict.string( toks.int().cpu(), # The default unknown string in fairseq is ``, but # this is tokenized by sacrebleu as `< unk >`, inflating # BLEU scores. Instead, we use a somewhat more verbose # alternative that is unlikely to appear in the real # reference, but doesn't get split into multiple tokens. unk_string=("UNKNOWNTOKENINREF" if escape_unk else "UNKNOWNTOKENINHYP"), ) if self.bpe: s = self.bpe.decode(s) return s gen_out = self.inference_step(generator, [model], sample) hyps, refs = [], [] transtab = str.maketrans({key: None for key in string.punctuation}) for i in range(len(gen_out)): decode_tokens = decode(gen_out[i][0]["tokens"]) hyps.append(decode_tokens.translate(transtab).strip()) refs.append( [ sent.translate(transtab).strip() for sent in decode( utils.strip_pad(sample["target"][i], self.tgt_dict.pad()), escape_unk=True, # don't count as matches to the hypo ).split('&&') ] ) if self.cfg.eval_print_samples: logger.info("example hypothesis: " + hyps[0]) logger.info("example reference: " + ' && '.join(refs[0])) return hyps, refs