OFA-OCR / tasks /nlg_tasks /gigaword.py
JustinLin610's picture
first commit
ee21b96
raw
history blame
10.7 kB
# 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
import string
from typing import Optional
from argparse import Namespace
from fairseq import metrics
from fairseq.tasks import register_task
from fairseq.data import encoders
from tasks.ofa_task import OFATask, OFAConfig
from data.nlg_data.summary_dataset import SummaryDataset
from data.file_dataset import FileDataset
from datasets import load_metric
logger = logging.getLogger(__name__)
_tok_dict = {"(": "-lrb-", ")": "-rrb-",
"[": "-lsb-", "]": "-rsb-",
"{": "-lcb-", "}": "-rcb-",
"[UNK]": "UNK", '&': '&amp;', '<': '&lt;', '>': '&gt;'}
def _is_digit(w):
for ch in w:
if not(ch.isdigit() or ch == ','):
return False
return True
def fix_tokenization(text):
input_tokens = text.split()
output_tokens = []
has_left_quote = False
has_left_single_quote = False
i = 0
prev_dash = False
while i < len(input_tokens):
tok = input_tokens[i]
flag_prev_dash = False
if tok in _tok_dict.keys():
output_tokens.append(_tok_dict[tok])
i += 1
elif tok == "\"":
if has_left_quote:
output_tokens.append("''")
else:
output_tokens.append("``")
has_left_quote = not has_left_quote
i += 1
elif tok == "'" and len(output_tokens) > 0 and output_tokens[-1].endswith("n") and i < len(input_tokens) - 1 and input_tokens[i + 1] == "t":
output_tokens[-1] = output_tokens[-1][:-1]
output_tokens.append("n't")
i += 2
elif tok == "'" and i < len(input_tokens) - 1 and input_tokens[i + 1] in ("s", "d", "ll"):
output_tokens.append("'"+input_tokens[i + 1])
i += 2
elif tok == "'":
if has_left_single_quote:
output_tokens.append("'")
else:
output_tokens.append("`")
has_left_single_quote = not has_left_single_quote
i += 1
elif tok == "." and i < len(input_tokens) - 2 and input_tokens[i + 1] == "." and input_tokens[i + 2] == ".":
output_tokens.append("...")
i += 3
elif tok == "," and len(output_tokens) > 0 and _is_digit(output_tokens[-1]) and i < len(input_tokens) - 1 and _is_digit(input_tokens[i + 1]):
# $ 3 , 000 -> $ 3,000
output_tokens[-1] += ','+input_tokens[i + 1]
i += 2
elif tok == "." and len(output_tokens) > 0 and output_tokens[-1].isdigit() and i < len(input_tokens) - 1 and input_tokens[i + 1].isdigit():
# 3 . 03 -> $ 3.03
output_tokens[-1] += '.'+input_tokens[i + 1]
i += 2
elif tok == "." and len(output_tokens) > 0 and len(output_tokens[-1]) == 1 and output_tokens[-1].isupper() and i < len(input_tokens) - 2 and len(input_tokens[i + 1]) == 1 and input_tokens[i + 1].isupper() and input_tokens[i + 2] == '.':
# U . N . -> U.N.
k = i+3
while k+2 < len(input_tokens):
if len(input_tokens[k + 1]) == 1 and input_tokens[k + 1].isupper() and input_tokens[k + 2] == '.':
k += 2
else:
break
output_tokens[-1] += ''.join(input_tokens[i:k])
i += 2
elif tok == "-":
if i < len(input_tokens) - 1 and input_tokens[i + 1] == "-":
output_tokens.append("--")
i += 2
elif i == len(input_tokens) - 1 or i == 0:
output_tokens.append("-")
i += 1
elif output_tokens[-1] not in string.punctuation and input_tokens[i + 1][0] not in string.punctuation:
output_tokens[-1] += "-"
i += 1
flag_prev_dash = True
else:
output_tokens.append("-")
i += 1
elif prev_dash and len(output_tokens) > 0 and tok[0] not in string.punctuation:
output_tokens[-1] += tok
i += 1
else:
output_tokens.append(tok)
i += 1
prev_dash = flag_prev_dash
return " ".join(output_tokens)
@dataclass
class GigawordConfig(OFAConfig):
# options for reporting Rouge during validation
eval_rouge: bool = field(
default=False, metadata={"help": "evaluation with rouge 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_detok: str = field(
default="space",
metadata={
"help": "detokenize before computing BLEU or CIDEr (e.g., 'moses'); "
"required if using --eval-bleu or --eval-cider; "
"use 'space' to disable detokenization; see fairseq.data.encoders for other options"
},
)
eval_detok_args: Optional[str] = field(
default="{}",
metadata={"help": "args for building the tokenizer, if needed, as JSON string"},
)
eval_print_samples: bool = field(
default=False, metadata={"help": "print sample generations during validation"}
)
noise_ratio: float = field(
default=0.0, metadata={"help": "noise ratio for prev output"}
)
@register_task("gigaword", dataclass=GigawordConfig)
class GigawordTask(OFATask):
def __init__(self, cfg: GigawordConfig, 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] = SummaryDataset(
split,
dataset,
self.bpe,
self.src_dict,
self.tgt_dict,
code_dict_size=self.cfg.code_dict_size,
num_bins=self.cfg.num_bins,
max_src_length=self.cfg.max_src_length,
max_tgt_length=self.cfg.max_tgt_length,
noise_ratio=self.cfg.noise_ratio
)
def build_model(self, cfg):
model = super().build_model(cfg)
if self.cfg.eval_rouge:
detok_args = json.loads(self.cfg.eval_detok_args)
self.tokenizer = encoders.build_tokenizer(
Namespace(tokenizer=self.cfg.eval_detok, **detok_args)
)
gen_args = json.loads(self.cfg.eval_args)
self.sequence_generator = self.build_generator(
[model], Namespace(**gen_args)
)
self.metric = load_metric('../../utils/rouge.py')
return model
def valid_step(self, sample, model, criterion):
loss, sample_size, logging_output = super().valid_step(sample, model, criterion)
if self.cfg.eval_rouge:
hyps, refs = self._inference(self.sequence_generator, sample, model)
result = self.metric.compute(predictions=hyps, references=refs, use_agregator=False, use_stemmer=True)
result_recall = {key: sum([item.recall for item in value]) * 100 for key, value in result.items()}
result_f1 = {key: sum([item.fmeasure for item in value]) * 100 for key, value in result.items()}
logging_output['_rouge1_recall_sum'] = result_recall['rouge1']
logging_output['_rouge2_recall_sum'] = result_recall['rouge2']
logging_output['_rougeL_recall_sum'] = result_recall['rougeL']
logging_output['_rouge1_f1_sum'] = result_f1['rouge1']
logging_output['_rouge2_f1_sum'] = result_f1['rouge2']
logging_output['_rougeL_f1_sum'] = result_f1['rougeL']
logging_output['_rouge_cnt'] = len(hyps)
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 sum_logs("_rouge_cnt") > 0:
metrics.log_scalar("_rouge1_recall_sum", sum_logs("_rouge1_recall_sum"))
metrics.log_scalar("_rouge2_recall_sum", sum_logs("_rouge2_recall_sum"))
metrics.log_scalar("_rougeL_recall_sum", sum_logs("_rougeL_recall_sum"))
metrics.log_scalar("_rouge1_f1_sum", sum_logs("_rouge1_f1_sum"))
metrics.log_scalar("_rouge2_f1_sum", sum_logs("_rouge2_f1_sum"))
metrics.log_scalar("_rougeL_f1_sum", sum_logs("_rougeL_f1_sum"))
metrics.log_scalar("_rouge_cnt", sum_logs("_rouge_cnt"))
metrics.log_derived("rouge1_recall", lambda x: x["_rouge1_recall_sum"].sum / x["_rouge_cnt"].sum)
metrics.log_derived("rouge2_recall", lambda x: x["_rouge2_recall_sum"].sum / x["_rouge_cnt"].sum)
metrics.log_derived("rougeL_recall", lambda x: x["_rougeL_recall_sum"].sum / x["_rouge_cnt"].sum)
metrics.log_derived("rouge1_f1", lambda x: x["_rouge1_f1_sum"].sum / x["_rouge_cnt"].sum)
metrics.log_derived("rouge2_f1", lambda x: x["_rouge2_f1_sum"].sum / x["_rouge_cnt"].sum)
metrics.log_derived("rougeL_f1", lambda x: x["_rougeL_f1_sum"].sum / x["_rouge_cnt"].sum)
def _inference(self, generator, sample, model):
def decode(toks):
s = self.tgt_dict.string(toks.int().cpu())
if self.bpe:
s = self.bpe.decode(s)
if self.tokenizer:
s = self.tokenizer.decode(s)
return s
gen_out = self.inference_step(generator, [model], sample)
hyps, refs = [], []
for i in range(len(gen_out)):
hyp = decode(gen_out[i][0]["tokens"]).lower().strip()
hyp = fix_tokenization(hyp).replace('<unk>', ' unk').replace('1', '#')
ref = sample["target_strs"][i]
hyps.append(hyp)
refs.append(ref)
if self.cfg.eval_print_samples:
logger.info("example hypothesis: " + hyps[0])
logger.info("example reference: " + refs[0])
return hyps, refs