chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,347 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""
|
||||
Evaluate the perplexity of a trained language model.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from argparse import Namespace
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
import torch
|
||||
import fairseq
|
||||
from fairseq import checkpoint_utils, distributed_utils, options, tasks, utils
|
||||
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
||||
from fairseq.logging import progress_bar
|
||||
from fairseq.logging.meters import StopwatchMeter
|
||||
from fairseq.sequence_scorer import SequenceScorer
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=sys.stdout,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.eval_lm")
|
||||
|
||||
|
||||
def eval_lm(
|
||||
models: List[fairseq.models.FairseqModel],
|
||||
source_dictionary: fairseq.data.Dictionary,
|
||||
batch_iterator: Iterable,
|
||||
post_process: Optional[str] = None,
|
||||
output_word_probs: bool = False,
|
||||
output_word_stats: bool = False,
|
||||
target_dictionary: Optional[fairseq.data.Dictionary] = None,
|
||||
softmax_batch: int = False,
|
||||
remove_bos_token: bool = False,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
models (List[~fairseq.models.FairseqModel]): list of models to
|
||||
evaluate. Models are essentially `nn.Module` instances, but
|
||||
must be compatible with fairseq's `SequenceScorer`.
|
||||
source_dictionary (~fairseq.data.Dictionary): dictionary for
|
||||
applying any relevant post processing or outputing word
|
||||
probs/stats.
|
||||
batch_iterator (Iterable): yield batches of data
|
||||
post_process (Optional[str]): post-process text by removing BPE,
|
||||
letter segmentation, etc. Valid options can be found in
|
||||
fairseq.data.utils.post_process, although not all options
|
||||
are implemented here.
|
||||
output_word_probs (Optional[bool]): output words and their
|
||||
predicted log probabilities
|
||||
output_word_stats (Optional[bool]): output word statistics such
|
||||
as word count and average probability
|
||||
target_dictionary (Optional[~fairseq.data.Dictionary]): output
|
||||
dictionary (defaults to *source_dictionary*)
|
||||
softmax_batch (Optional[bool]): if BxT is more than this, will
|
||||
batch the softmax over vocab to this amount of tokens, in
|
||||
order to fit into GPU memory
|
||||
remove_bos_token (Optional[bool]): if True, confirm that the
|
||||
first token is the beginning-of-sentence symbol (according
|
||||
to the relevant dictionary) and remove it from the output
|
||||
device (Optional[torch.device]): device to use for evaluation
|
||||
(defaults to device of first model parameter)
|
||||
"""
|
||||
if target_dictionary is None:
|
||||
target_dictionary = source_dictionary
|
||||
if device is None:
|
||||
device = next(models[0].parameters()).device
|
||||
|
||||
gen_timer = StopwatchMeter()
|
||||
scorer = SequenceScorer(target_dictionary, softmax_batch)
|
||||
|
||||
score_sum = 0.0
|
||||
count = 0
|
||||
|
||||
if post_process is not None:
|
||||
if post_process in {"subword_nmt", "@@ "}:
|
||||
bpe_cont = post_process.rstrip()
|
||||
bpe_toks = {
|
||||
i
|
||||
for i in range(len(source_dictionary))
|
||||
if source_dictionary[i].endswith(bpe_cont)
|
||||
}
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"--post-process={post_process} is not implemented"
|
||||
)
|
||||
bpe_len = len(bpe_cont)
|
||||
else:
|
||||
bpe_toks = None
|
||||
bpe_len = 0
|
||||
|
||||
word_stats = dict()
|
||||
|
||||
for sample in batch_iterator:
|
||||
if "net_input" not in sample:
|
||||
continue
|
||||
|
||||
sample = utils.move_to_cuda(sample, device=device)
|
||||
|
||||
gen_timer.start()
|
||||
hypos = scorer.generate(models, sample)
|
||||
gen_timer.stop(sample["ntokens"])
|
||||
|
||||
for i, hypos_i in enumerate(hypos):
|
||||
hypo = hypos_i[0]
|
||||
sample_id = sample["id"][i]
|
||||
|
||||
tokens = hypo["tokens"]
|
||||
tgt_len = tokens.numel()
|
||||
pos_scores = hypo["positional_scores"].float()
|
||||
|
||||
if remove_bos_token:
|
||||
assert hypo["tokens"][0].item() == target_dictionary.bos()
|
||||
tokens = tokens[1:]
|
||||
pos_scores = pos_scores[1:]
|
||||
|
||||
skipped_toks = 0
|
||||
if bpe_toks is not None:
|
||||
for i in range(tgt_len - 1):
|
||||
if tokens[i].item() in bpe_toks:
|
||||
skipped_toks += 1
|
||||
pos_scores[i + 1] += pos_scores[i]
|
||||
pos_scores[i] = 0
|
||||
|
||||
inf_scores = pos_scores.eq(float("inf")) | pos_scores.eq(float("-inf"))
|
||||
if inf_scores.any():
|
||||
logger.info(
|
||||
"skipping tokens with inf scores:",
|
||||
target_dictionary.string(tokens[inf_scores.nonzero()]),
|
||||
)
|
||||
pos_scores = pos_scores[(~inf_scores).nonzero()]
|
||||
score_sum += pos_scores.sum().cpu()
|
||||
count += pos_scores.numel() - skipped_toks
|
||||
|
||||
if output_word_probs or output_word_stats:
|
||||
w = ""
|
||||
word_prob = []
|
||||
is_bpe = False
|
||||
for i in range(len(tokens)):
|
||||
w_ind = tokens[i].item()
|
||||
w += source_dictionary[w_ind]
|
||||
if bpe_toks is not None and w_ind in bpe_toks:
|
||||
w = w[:-bpe_len]
|
||||
is_bpe = True
|
||||
else:
|
||||
word_prob.append((w, pos_scores[i].item()))
|
||||
|
||||
next_prob = None
|
||||
ind = i + 1
|
||||
while ind < len(tokens):
|
||||
if pos_scores[ind].item() != 0:
|
||||
next_prob = pos_scores[ind]
|
||||
break
|
||||
ind += 1
|
||||
|
||||
word_stats.setdefault(w, WordStat(w, is_bpe)).add(
|
||||
pos_scores[i].item(), next_prob
|
||||
)
|
||||
is_bpe = False
|
||||
w = ""
|
||||
if output_word_probs:
|
||||
logger.info(
|
||||
str(int(sample_id))
|
||||
+ " "
|
||||
+ (
|
||||
"\t".join(
|
||||
"{} [{:2f}]".format(x[0], x[1]) for x in word_prob
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
avg_nll_loss = (
|
||||
-score_sum / count / math.log(2) if count > 0 else 0
|
||||
) # convert to base 2
|
||||
logger.info(
|
||||
"Evaluated {:,} tokens in {:.1f}s ({:.2f} tokens/s)".format(
|
||||
gen_timer.n, gen_timer.sum, 1.0 / gen_timer.avg if gen_timer.avg > 0 else 0
|
||||
)
|
||||
)
|
||||
|
||||
if output_word_stats:
|
||||
for ws in sorted(word_stats.values(), key=lambda x: x.count, reverse=True):
|
||||
logger.info(ws)
|
||||
|
||||
return {
|
||||
"loss": avg_nll_loss,
|
||||
"perplexity": 2 ** avg_nll_loss,
|
||||
}
|
||||
|
||||
|
||||
class WordStat(object):
|
||||
def __init__(self, word, is_bpe):
|
||||
self.word = word
|
||||
self.is_bpe = is_bpe
|
||||
self.log_prob = 0
|
||||
self.next_word_prob = 0
|
||||
self.count = 0
|
||||
self.missing_next_words = 0
|
||||
|
||||
def add(self, log_prob, next_word_prob):
|
||||
"""increments counters for the sum of log probs of current word and next
|
||||
word (given context ending at current word). Since the next word might be at the end of the example,
|
||||
or it might be not counted because it is not an ending subword unit,
|
||||
also keeps track of how many of those we have seen"""
|
||||
if next_word_prob is not None:
|
||||
self.next_word_prob += next_word_prob
|
||||
else:
|
||||
self.missing_next_words += 1
|
||||
self.log_prob += log_prob
|
||||
self.count += 1
|
||||
|
||||
def __str__(self):
|
||||
return "{}\t{}\t{}\t{}\t{}\t{}".format(
|
||||
self.word,
|
||||
self.count,
|
||||
self.log_prob,
|
||||
self.is_bpe,
|
||||
self.next_word_prob,
|
||||
self.count - self.missing_next_words,
|
||||
)
|
||||
|
||||
|
||||
def main(cfg: DictConfig, **unused_kwargs):
|
||||
if isinstance(cfg, Namespace):
|
||||
cfg = convert_namespace_to_omegaconf(cfg)
|
||||
|
||||
utils.import_user_module(cfg.common)
|
||||
|
||||
logger.info(cfg)
|
||||
|
||||
if cfg.eval_lm.context_window > 0:
|
||||
# reduce tokens per sample by the required context window size
|
||||
cfg.task.tokens_per_sample -= cfg.eval_lm.context_window
|
||||
|
||||
# Initialize the task using the current *cfg*
|
||||
task = tasks.setup_task(cfg.task)
|
||||
|
||||
# Load ensemble
|
||||
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
|
||||
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
|
||||
[cfg.common_eval.path],
|
||||
arg_overrides=eval(cfg.common_eval.model_overrides),
|
||||
suffix=cfg.checkpoint.checkpoint_suffix,
|
||||
strict=(cfg.checkpoint.checkpoint_shard_count == 1),
|
||||
num_shards=cfg.checkpoint.checkpoint_shard_count,
|
||||
task=task,
|
||||
)
|
||||
|
||||
use_fp16 = cfg.common.fp16
|
||||
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
|
||||
if use_cuda:
|
||||
torch.cuda.set_device(cfg.distributed_training.device_id)
|
||||
|
||||
# Optimize ensemble for generation and set the source and dest dicts on the model
|
||||
# (required by scorer)
|
||||
for model in models:
|
||||
if use_fp16:
|
||||
model.half()
|
||||
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
|
||||
model.cuda()
|
||||
model.prepare_for_inference_(cfg)
|
||||
|
||||
assert len(models) > 0
|
||||
|
||||
logger.info(
|
||||
"num. model params: {:,}".format(sum(p.numel() for p in models[0].parameters()))
|
||||
)
|
||||
|
||||
# Load dataset splits
|
||||
task.load_dataset(cfg.dataset.gen_subset)
|
||||
dataset = task.dataset(cfg.dataset.gen_subset)
|
||||
logger.info(
|
||||
"{} {} {:,} examples".format(
|
||||
cfg.task.data, cfg.dataset.gen_subset, len(dataset)
|
||||
)
|
||||
)
|
||||
|
||||
itr = task.eval_lm_dataloader(
|
||||
dataset=dataset,
|
||||
max_tokens=cfg.dataset.max_tokens or 36000,
|
||||
batch_size=cfg.dataset.batch_size,
|
||||
max_positions=utils.resolve_max_positions(
|
||||
*[model.max_positions() for model in models]
|
||||
),
|
||||
num_shards=max(
|
||||
cfg.dataset.num_shards,
|
||||
cfg.distributed_training.distributed_world_size,
|
||||
),
|
||||
shard_id=max(
|
||||
cfg.dataset.shard_id,
|
||||
cfg.distributed_training.distributed_rank,
|
||||
),
|
||||
num_workers=cfg.dataset.num_workers,
|
||||
data_buffer_size=cfg.dataset.data_buffer_size,
|
||||
context_window=cfg.eval_lm.context_window,
|
||||
)
|
||||
|
||||
itr = progress_bar.progress_bar(
|
||||
itr,
|
||||
log_format=cfg.common.log_format,
|
||||
log_interval=cfg.common.log_interval,
|
||||
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
|
||||
)
|
||||
|
||||
results = eval_lm(
|
||||
models=models,
|
||||
source_dictionary=task.source_dictionary,
|
||||
batch_iterator=itr,
|
||||
post_process=cfg.common_eval.post_process,
|
||||
output_word_probs=cfg.eval_lm.output_word_probs,
|
||||
output_word_stats=cfg.eval_lm.output_word_stats,
|
||||
target_dictionary=task.target_dictionary,
|
||||
softmax_batch=cfg.eval_lm.softmax_batch,
|
||||
remove_bos_token=getattr(cfg.task, "add_bos_token", False),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Loss (base 2): {:.4f}, Perplexity: {:.2f}".format(
|
||||
results["loss"], results["perplexity"]
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = options.get_eval_lm_parser()
|
||||
args = options.parse_args_and_arch(parser)
|
||||
|
||||
distributed_utils.call_main(convert_namespace_to_omegaconf(args), main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,408 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
Translate pre-processed data with a trained model.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from argparse import Namespace
|
||||
from itertools import chain
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq import checkpoint_utils, options, scoring, tasks, utils
|
||||
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
||||
from fairseq.logging import progress_bar
|
||||
from fairseq.logging.meters import StopwatchMeter, TimeMeter
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
def main(cfg: DictConfig):
|
||||
|
||||
if isinstance(cfg, Namespace):
|
||||
cfg = convert_namespace_to_omegaconf(cfg)
|
||||
|
||||
assert cfg.common_eval.path is not None, "--path required for generation!"
|
||||
assert (
|
||||
not cfg.generation.sampling or cfg.generation.nbest == cfg.generation.beam
|
||||
), "--sampling requires --nbest to be equal to --beam"
|
||||
assert (
|
||||
cfg.generation.replace_unk is None or cfg.dataset.dataset_impl == "raw"
|
||||
), "--replace-unk requires a raw text dataset (--dataset-impl=raw)"
|
||||
|
||||
if cfg.common_eval.results_path is not None:
|
||||
os.makedirs(cfg.common_eval.results_path, exist_ok=True)
|
||||
output_path = os.path.join(
|
||||
cfg.common_eval.results_path,
|
||||
"generate-{}.txt".format(cfg.dataset.gen_subset),
|
||||
)
|
||||
with open(output_path, "w", buffering=1, encoding="utf-8") as h:
|
||||
return _main(cfg, h)
|
||||
else:
|
||||
return _main(cfg, sys.stdout)
|
||||
|
||||
|
||||
def get_symbols_to_strip_from_output(generator):
|
||||
if hasattr(generator, "symbols_to_strip_from_output"):
|
||||
return generator.symbols_to_strip_from_output
|
||||
else:
|
||||
return {generator.eos}
|
||||
|
||||
|
||||
def _main(cfg: DictConfig, output_file):
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=output_file,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.generate")
|
||||
|
||||
utils.import_user_module(cfg.common)
|
||||
|
||||
if cfg.dataset.max_tokens is None and cfg.dataset.batch_size is None:
|
||||
cfg.dataset.max_tokens = 12000
|
||||
logger.info(cfg)
|
||||
|
||||
# Fix seed for stochastic decoding
|
||||
if cfg.common.seed is not None and not cfg.generation.no_seed_provided:
|
||||
np.random.seed(cfg.common.seed)
|
||||
utils.set_torch_seed(cfg.common.seed)
|
||||
|
||||
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
|
||||
|
||||
# Load dataset splits
|
||||
task = tasks.setup_task(cfg.task)
|
||||
|
||||
|
||||
# Set dictionaries
|
||||
try:
|
||||
src_dict = getattr(task, "source_dictionary", None)
|
||||
except NotImplementedError:
|
||||
src_dict = None
|
||||
tgt_dict = task.target_dictionary
|
||||
|
||||
overrides = ast.literal_eval(cfg.common_eval.model_overrides)
|
||||
|
||||
# Load ensemble
|
||||
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
|
||||
models, saved_cfg = checkpoint_utils.load_model_ensemble(
|
||||
utils.split_paths(cfg.common_eval.path),
|
||||
arg_overrides=overrides,
|
||||
task=task,
|
||||
suffix=cfg.checkpoint.checkpoint_suffix,
|
||||
strict=(cfg.checkpoint.checkpoint_shard_count == 1),
|
||||
num_shards=cfg.checkpoint.checkpoint_shard_count,
|
||||
)
|
||||
|
||||
# loading the dataset should happen after the checkpoint has been loaded so we can give it the saved task config
|
||||
task.load_dataset(cfg.dataset.gen_subset, task_cfg=saved_cfg.task)
|
||||
|
||||
if cfg.generation.lm_path is not None:
|
||||
overrides["data"] = cfg.task.data
|
||||
|
||||
try:
|
||||
lms, _ = checkpoint_utils.load_model_ensemble(
|
||||
[cfg.generation.lm_path], arg_overrides=overrides, task=None
|
||||
)
|
||||
except:
|
||||
logger.warning(
|
||||
f"Failed to load language model! Please make sure that the language model dict is the same "
|
||||
f"as target dict and is located in the data dir ({cfg.task.data})"
|
||||
)
|
||||
raise
|
||||
|
||||
assert len(lms) == 1
|
||||
else:
|
||||
lms = [None]
|
||||
|
||||
# Optimize ensemble for generation
|
||||
for model in chain(models, lms):
|
||||
if model is None:
|
||||
continue
|
||||
if cfg.common.fp16:
|
||||
model.half()
|
||||
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
|
||||
model.cuda()
|
||||
model.prepare_for_inference_(cfg)
|
||||
|
||||
# Load alignment dictionary for unknown word replacement
|
||||
# (None if no unknown word replacement, empty if no path to align dictionary)
|
||||
align_dict = utils.load_align_dict(cfg.generation.replace_unk)
|
||||
|
||||
# Load dataset (possibly sharded)
|
||||
itr = task.get_batch_iterator(
|
||||
dataset=task.dataset(cfg.dataset.gen_subset),
|
||||
max_tokens=cfg.dataset.max_tokens,
|
||||
max_sentences=cfg.dataset.batch_size,
|
||||
max_positions=utils.resolve_max_positions(
|
||||
task.max_positions(), *[m.max_positions() for m in models]
|
||||
),
|
||||
ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
|
||||
required_batch_size_multiple=cfg.dataset.required_batch_size_multiple,
|
||||
seed=cfg.common.seed,
|
||||
num_shards=cfg.distributed_training.distributed_world_size,
|
||||
shard_id=cfg.distributed_training.distributed_rank,
|
||||
num_workers=cfg.dataset.num_workers,
|
||||
data_buffer_size=cfg.dataset.data_buffer_size,
|
||||
).next_epoch_itr(shuffle=False)
|
||||
progress = progress_bar.progress_bar(
|
||||
itr,
|
||||
log_format=cfg.common.log_format,
|
||||
log_interval=cfg.common.log_interval,
|
||||
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
|
||||
)
|
||||
|
||||
# Initialize generator
|
||||
gen_timer = StopwatchMeter()
|
||||
|
||||
extra_gen_cls_kwargs = {"lm_model": lms[0], "lm_weight": cfg.generation.lm_weight}
|
||||
generator = task.build_generator(
|
||||
models, cfg.generation, extra_gen_cls_kwargs=extra_gen_cls_kwargs
|
||||
)
|
||||
|
||||
# Handle tokenization and BPE
|
||||
tokenizer = task.build_tokenizer(cfg.tokenizer)
|
||||
bpe = task.build_bpe(cfg.bpe)
|
||||
|
||||
def decode_fn(x):
|
||||
if bpe is not None:
|
||||
x = bpe.decode(x)
|
||||
if tokenizer is not None:
|
||||
x = tokenizer.decode(x)
|
||||
return x
|
||||
|
||||
scorer = scoring.build_scorer(cfg.scoring, tgt_dict)
|
||||
|
||||
num_sentences = 0
|
||||
has_target = True
|
||||
wps_meter = TimeMeter()
|
||||
for sample in progress:
|
||||
sample = utils.move_to_cuda(sample) if use_cuda else sample
|
||||
if "net_input" not in sample:
|
||||
continue
|
||||
|
||||
prefix_tokens = None
|
||||
if cfg.generation.prefix_size > 0:
|
||||
prefix_tokens = sample["target"][:, : cfg.generation.prefix_size]
|
||||
|
||||
constraints = None
|
||||
if "constraints" in sample:
|
||||
constraints = sample["constraints"]
|
||||
|
||||
gen_timer.start()
|
||||
hypos = task.inference_step(
|
||||
generator,
|
||||
models,
|
||||
sample,
|
||||
prefix_tokens=prefix_tokens,
|
||||
constraints=constraints,
|
||||
)
|
||||
num_generated_tokens = sum(len(h[0]["tokens"]) for h in hypos)
|
||||
gen_timer.stop(num_generated_tokens)
|
||||
|
||||
for i, sample_id in enumerate(sample["id"].tolist()):
|
||||
has_target = sample["target"] is not None
|
||||
|
||||
# Remove padding
|
||||
if "src_tokens" in sample["net_input"]:
|
||||
src_tokens = utils.strip_pad(
|
||||
sample["net_input"]["src_tokens"][i, :], tgt_dict.pad()
|
||||
)
|
||||
else:
|
||||
src_tokens = None
|
||||
|
||||
target_tokens = None
|
||||
if has_target:
|
||||
target_tokens = (
|
||||
utils.strip_pad(sample["target"][i, :], tgt_dict.pad()).int().cpu()
|
||||
)
|
||||
|
||||
# Either retrieve the original sentences or regenerate them from tokens.
|
||||
if align_dict is not None:
|
||||
src_str = task.dataset(cfg.dataset.gen_subset).src.get_original_text(
|
||||
sample_id
|
||||
)
|
||||
target_str = task.dataset(cfg.dataset.gen_subset).tgt.get_original_text(
|
||||
sample_id
|
||||
)
|
||||
else:
|
||||
if src_dict is not None:
|
||||
src_str = src_dict.string(src_tokens, cfg.common_eval.post_process)
|
||||
else:
|
||||
src_str = ""
|
||||
if has_target:
|
||||
target_str = tgt_dict.string(
|
||||
target_tokens,
|
||||
cfg.common_eval.post_process,
|
||||
escape_unk=True,
|
||||
extra_symbols_to_ignore=get_symbols_to_strip_from_output(
|
||||
generator
|
||||
),
|
||||
)
|
||||
|
||||
src_str = decode_fn(src_str)
|
||||
if has_target:
|
||||
target_str = decode_fn(target_str)
|
||||
|
||||
if not cfg.common_eval.quiet:
|
||||
if src_dict is not None:
|
||||
print("S-{}\t{}".format(sample_id, src_str), file=output_file)
|
||||
if has_target:
|
||||
print("T-{}\t{}".format(sample_id, target_str), file=output_file)
|
||||
|
||||
# Process top predictions
|
||||
for j, hypo in enumerate(hypos[i][: cfg.generation.nbest]):
|
||||
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
|
||||
hypo_tokens=hypo["tokens"].int().cpu(),
|
||||
src_str=src_str,
|
||||
alignment=hypo["alignment"],
|
||||
align_dict=align_dict,
|
||||
tgt_dict=tgt_dict,
|
||||
remove_bpe=cfg.common_eval.post_process,
|
||||
extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator),
|
||||
)
|
||||
detok_hypo_str = decode_fn(hypo_str)
|
||||
if not cfg.common_eval.quiet:
|
||||
score = hypo["score"] / math.log(2) # convert to base 2
|
||||
# original hypothesis (after tokenization and BPE)
|
||||
print(
|
||||
"H-{}\t{}\t{}".format(sample_id, score, hypo_str),
|
||||
file=output_file,
|
||||
)
|
||||
# detokenized hypothesis
|
||||
print(
|
||||
"D-{}\t{}\t{}".format(sample_id, score, detok_hypo_str),
|
||||
file=output_file,
|
||||
)
|
||||
print(
|
||||
"P-{}\t{}".format(
|
||||
sample_id,
|
||||
" ".join(
|
||||
map(
|
||||
lambda x: "{:.4f}".format(x),
|
||||
# convert from base e to base 2
|
||||
hypo["positional_scores"]
|
||||
.div_(math.log(2))
|
||||
.tolist(),
|
||||
)
|
||||
),
|
||||
),
|
||||
file=output_file,
|
||||
)
|
||||
|
||||
if cfg.generation.print_alignment == "hard":
|
||||
print(
|
||||
"A-{}\t{}".format(
|
||||
sample_id,
|
||||
" ".join(
|
||||
[
|
||||
"{}-{}".format(src_idx, tgt_idx)
|
||||
for src_idx, tgt_idx in alignment
|
||||
]
|
||||
),
|
||||
),
|
||||
file=output_file,
|
||||
)
|
||||
if cfg.generation.print_alignment == "soft":
|
||||
print(
|
||||
"A-{}\t{}".format(
|
||||
sample_id,
|
||||
" ".join(
|
||||
[
|
||||
",".join(src_probs)
|
||||
for src_probs in alignment
|
||||
]
|
||||
),
|
||||
),
|
||||
file=output_file,
|
||||
)
|
||||
|
||||
if cfg.generation.print_step:
|
||||
print(
|
||||
"I-{}\t{}".format(sample_id, hypo["steps"]),
|
||||
file=output_file,
|
||||
)
|
||||
|
||||
if cfg.generation.retain_iter_history:
|
||||
for step, h in enumerate(hypo["history"]):
|
||||
_, h_str, _ = utils.post_process_prediction(
|
||||
hypo_tokens=h["tokens"].int().cpu(),
|
||||
src_str=src_str,
|
||||
alignment=None,
|
||||
align_dict=None,
|
||||
tgt_dict=tgt_dict,
|
||||
remove_bpe=None,
|
||||
)
|
||||
print(
|
||||
"E-{}_{}\t{}".format(sample_id, step, h_str),
|
||||
file=output_file,
|
||||
)
|
||||
|
||||
# Score only the top hypothesis
|
||||
if has_target and j == 0:
|
||||
if align_dict is not None or cfg.common_eval.post_process is not None:
|
||||
# Convert back to tokens for evaluation with unk replacement and/or without BPE
|
||||
target_tokens = tgt_dict.encode_line(
|
||||
target_str, add_if_not_exist=True
|
||||
)
|
||||
hypo_tokens = tgt_dict.encode_line(
|
||||
detok_hypo_str, add_if_not_exist=True
|
||||
)
|
||||
if hasattr(scorer, "add_string"):
|
||||
scorer.add_string(target_str, detok_hypo_str)
|
||||
else:
|
||||
scorer.add(target_tokens, hypo_tokens)
|
||||
|
||||
wps_meter.update(num_generated_tokens)
|
||||
progress.log({"wps": round(wps_meter.avg)})
|
||||
num_sentences += (
|
||||
sample["nsentences"] if "nsentences" in sample else sample["id"].numel()
|
||||
)
|
||||
|
||||
logger.info("NOTE: hypothesis and token scores are output in base 2")
|
||||
logger.info(
|
||||
"Translated {:,} sentences ({:,} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)".format(
|
||||
num_sentences,
|
||||
gen_timer.n,
|
||||
gen_timer.sum,
|
||||
num_sentences / gen_timer.sum,
|
||||
1.0 / gen_timer.avg,
|
||||
)
|
||||
)
|
||||
if has_target:
|
||||
if cfg.bpe and not cfg.generation.sacrebleu:
|
||||
if cfg.common_eval.post_process:
|
||||
logger.warning(
|
||||
"BLEU score is being computed by splitting detokenized string on spaces, this is probably not what you want. Use --sacrebleu for standard 13a BLEU tokenization"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"If you are using BPE on the target side, the BLEU score is computed on BPE tokens, not on proper words. Use --sacrebleu for standard 13a BLEU tokenization"
|
||||
)
|
||||
# use print to be consistent with other main outputs: S-, H-, T-, D- and so on
|
||||
print(
|
||||
"Generate {} with beam={}: {}".format(
|
||||
cfg.dataset.gen_subset, cfg.generation.beam, scorer.result_string()
|
||||
),
|
||||
file=output_file,
|
||||
)
|
||||
|
||||
return scorer
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = options.get_generation_parser()
|
||||
args = options.parse_args_and_arch(parser)
|
||||
main(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,95 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
from fairseq.dataclass.initialize import add_defaults, hydra_init
|
||||
from fairseq_cli.train import main as pre_main
|
||||
from fairseq import distributed_utils, metrics
|
||||
from fairseq.dataclass.configs import FairseqConfig
|
||||
|
||||
import hydra
|
||||
from hydra.core.hydra_config import HydraConfig
|
||||
import torch
|
||||
from omegaconf import OmegaConf, open_dict
|
||||
|
||||
|
||||
logger = logging.getLogger("fairseq_cli.hydra_train")
|
||||
|
||||
|
||||
@hydra.main(config_path=os.path.join("..", "fairseq", "config"), config_name="config")
|
||||
def hydra_main(cfg: FairseqConfig) -> float:
|
||||
add_defaults(cfg)
|
||||
|
||||
if cfg.common.reset_logging:
|
||||
reset_logging() # Hydra hijacks logging, fix that
|
||||
else:
|
||||
with open_dict(cfg):
|
||||
# make hydra logging work with ddp (see # see https://github.com/facebookresearch/hydra/issues/1126)
|
||||
cfg.job_logging_cfg = OmegaConf.to_container(HydraConfig.get().job_logging, resolve=True)
|
||||
|
||||
cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=True, enum_to_str=True))
|
||||
OmegaConf.set_struct(cfg, True)
|
||||
|
||||
try:
|
||||
if cfg.common.profile:
|
||||
with torch.cuda.profiler.profile():
|
||||
with torch.autograd.profiler.emit_nvtx():
|
||||
distributed_utils.call_main(cfg, pre_main)
|
||||
else:
|
||||
distributed_utils.call_main(cfg, pre_main)
|
||||
except BaseException as e:
|
||||
if not cfg.common.suppress_crashes:
|
||||
raise
|
||||
else:
|
||||
logger.error("Crashed! " + str(e))
|
||||
|
||||
# get best val and return - useful for sweepers
|
||||
try:
|
||||
best_val = metrics.get_smoothed_value(
|
||||
"valid", cfg.checkpoint.best_checkpoint_metric
|
||||
)
|
||||
except:
|
||||
best_val = None
|
||||
|
||||
if best_val is None:
|
||||
best_val = float("inf")
|
||||
|
||||
return best_val
|
||||
|
||||
|
||||
def reset_logging():
|
||||
root = logging.getLogger()
|
||||
for handler in root.handlers:
|
||||
root.removeHandler(handler)
|
||||
root.setLevel(os.environ.get("LOGLEVEL", "INFO").upper())
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setFormatter(
|
||||
logging.Formatter(
|
||||
fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
)
|
||||
root.addHandler(handler)
|
||||
|
||||
|
||||
def cli_main():
|
||||
try:
|
||||
from hydra._internal.utils import get_args
|
||||
|
||||
cfg_name = get_args().config_name or "config"
|
||||
except:
|
||||
logger.warning("Failed to get config name from hydra args")
|
||||
cfg_name = "config"
|
||||
|
||||
hydra_init(cfg_name)
|
||||
hydra_main()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,316 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
Translate raw text with a trained model. Batches data on-the-fly.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import fileinput
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from argparse import Namespace
|
||||
from collections import namedtuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq import checkpoint_utils, distributed_utils, options, tasks, utils
|
||||
from fairseq.dataclass.configs import FairseqConfig
|
||||
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
||||
from fairseq.token_generation_constraints import pack_constraints, unpack_constraints
|
||||
from fairseq_cli.generate import get_symbols_to_strip_from_output
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=sys.stdout,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.interactive")
|
||||
|
||||
|
||||
Batch = namedtuple("Batch", "ids src_tokens src_lengths constraints")
|
||||
Translation = namedtuple("Translation", "src_str hypos pos_scores alignments")
|
||||
|
||||
|
||||
def buffered_read(input, buffer_size):
|
||||
buffer = []
|
||||
with fileinput.input(files=[input], openhook=fileinput.hook_encoded("utf-8")) as h:
|
||||
for src_str in h:
|
||||
buffer.append(src_str.strip())
|
||||
if len(buffer) >= buffer_size:
|
||||
yield buffer
|
||||
buffer = []
|
||||
|
||||
if len(buffer) > 0:
|
||||
yield buffer
|
||||
|
||||
|
||||
def make_batches(lines, cfg, task, max_positions, encode_fn):
|
||||
def encode_fn_target(x):
|
||||
return encode_fn(x)
|
||||
|
||||
if cfg.generation.constraints:
|
||||
# Strip (tab-delimited) contraints, if present, from input lines,
|
||||
# store them in batch_constraints
|
||||
batch_constraints = [list() for _ in lines]
|
||||
for i, line in enumerate(lines):
|
||||
if "\t" in line:
|
||||
lines[i], *batch_constraints[i] = line.split("\t")
|
||||
|
||||
# Convert each List[str] to List[Tensor]
|
||||
for i, constraint_list in enumerate(batch_constraints):
|
||||
batch_constraints[i] = [
|
||||
task.target_dictionary.encode_line(
|
||||
encode_fn_target(constraint),
|
||||
append_eos=False,
|
||||
add_if_not_exist=False,
|
||||
)
|
||||
for constraint in constraint_list
|
||||
]
|
||||
|
||||
if cfg.generation.constraints:
|
||||
constraints_tensor = pack_constraints(batch_constraints)
|
||||
else:
|
||||
constraints_tensor = None
|
||||
|
||||
tokens, lengths = task.get_interactive_tokens_and_lengths(lines, encode_fn)
|
||||
|
||||
itr = task.get_batch_iterator(
|
||||
dataset=task.build_dataset_for_inference(
|
||||
tokens, lengths, constraints=constraints_tensor
|
||||
),
|
||||
max_tokens=cfg.dataset.max_tokens,
|
||||
max_sentences=cfg.dataset.batch_size,
|
||||
max_positions=max_positions,
|
||||
ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
|
||||
).next_epoch_itr(shuffle=False)
|
||||
for batch in itr:
|
||||
ids = batch["id"]
|
||||
src_tokens = batch["net_input"]["src_tokens"]
|
||||
src_lengths = batch["net_input"]["src_lengths"]
|
||||
constraints = batch.get("constraints", None)
|
||||
|
||||
yield Batch(
|
||||
ids=ids,
|
||||
src_tokens=src_tokens,
|
||||
src_lengths=src_lengths,
|
||||
constraints=constraints,
|
||||
)
|
||||
|
||||
|
||||
def main(cfg: FairseqConfig):
|
||||
if isinstance(cfg, Namespace):
|
||||
cfg = convert_namespace_to_omegaconf(cfg)
|
||||
|
||||
start_time = time.time()
|
||||
total_translate_time = 0
|
||||
|
||||
utils.import_user_module(cfg.common)
|
||||
|
||||
if cfg.interactive.buffer_size < 1:
|
||||
cfg.interactive.buffer_size = 1
|
||||
if cfg.dataset.max_tokens is None and cfg.dataset.batch_size is None:
|
||||
cfg.dataset.batch_size = 1
|
||||
|
||||
assert (
|
||||
not cfg.generation.sampling or cfg.generation.nbest == cfg.generation.beam
|
||||
), "--sampling requires --nbest to be equal to --beam"
|
||||
assert (
|
||||
not cfg.dataset.batch_size
|
||||
or cfg.dataset.batch_size <= cfg.interactive.buffer_size
|
||||
), "--batch-size cannot be larger than --buffer-size"
|
||||
|
||||
logger.info(cfg)
|
||||
|
||||
# Fix seed for stochastic decoding
|
||||
if cfg.common.seed is not None and not cfg.generation.no_seed_provided:
|
||||
np.random.seed(cfg.common.seed)
|
||||
utils.set_torch_seed(cfg.common.seed)
|
||||
|
||||
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
|
||||
|
||||
# Setup task, e.g., translation
|
||||
task = tasks.setup_task(cfg.task)
|
||||
|
||||
# Load ensemble
|
||||
overrides = ast.literal_eval(cfg.common_eval.model_overrides)
|
||||
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
|
||||
models, _model_args = checkpoint_utils.load_model_ensemble(
|
||||
utils.split_paths(cfg.common_eval.path),
|
||||
arg_overrides=overrides,
|
||||
task=task,
|
||||
suffix=cfg.checkpoint.checkpoint_suffix,
|
||||
strict=(cfg.checkpoint.checkpoint_shard_count == 1),
|
||||
num_shards=cfg.checkpoint.checkpoint_shard_count,
|
||||
)
|
||||
|
||||
# Set dictionaries
|
||||
src_dict = task.source_dictionary
|
||||
tgt_dict = task.target_dictionary
|
||||
|
||||
# Optimize ensemble for generation
|
||||
for model in models:
|
||||
if model is None:
|
||||
continue
|
||||
if cfg.common.fp16:
|
||||
model.half()
|
||||
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
|
||||
model.cuda()
|
||||
model.prepare_for_inference_(cfg)
|
||||
|
||||
# Initialize generator
|
||||
generator = task.build_generator(models, cfg.generation)
|
||||
|
||||
# Handle tokenization and BPE
|
||||
tokenizer = task.build_tokenizer(cfg.tokenizer)
|
||||
bpe = task.build_bpe(cfg.bpe)
|
||||
|
||||
def encode_fn(x):
|
||||
if tokenizer is not None:
|
||||
x = tokenizer.encode(x)
|
||||
if bpe is not None:
|
||||
x = bpe.encode(x)
|
||||
return x
|
||||
|
||||
def decode_fn(x):
|
||||
if bpe is not None:
|
||||
x = bpe.decode(x)
|
||||
if tokenizer is not None:
|
||||
x = tokenizer.decode(x)
|
||||
return x
|
||||
|
||||
# Load alignment dictionary for unknown word replacement
|
||||
# (None if no unknown word replacement, empty if no path to align dictionary)
|
||||
align_dict = utils.load_align_dict(cfg.generation.replace_unk)
|
||||
|
||||
max_positions = utils.resolve_max_positions(
|
||||
task.max_positions(), *[model.max_positions() for model in models]
|
||||
)
|
||||
|
||||
if cfg.generation.constraints:
|
||||
logger.warning(
|
||||
"NOTE: Constrained decoding currently assumes a shared subword vocabulary."
|
||||
)
|
||||
|
||||
if cfg.interactive.buffer_size > 1:
|
||||
logger.info("Sentence buffer size: %s", cfg.interactive.buffer_size)
|
||||
logger.info("NOTE: hypothesis and token scores are output in base 2")
|
||||
logger.info("Type the input sentence and press return:")
|
||||
start_id = 0
|
||||
for inputs in buffered_read(cfg.interactive.input, cfg.interactive.buffer_size):
|
||||
results = []
|
||||
for batch in make_batches(inputs, cfg, task, max_positions, encode_fn):
|
||||
bsz = batch.src_tokens.size(0)
|
||||
src_tokens = batch.src_tokens
|
||||
src_lengths = batch.src_lengths
|
||||
constraints = batch.constraints
|
||||
if use_cuda:
|
||||
src_tokens = src_tokens.cuda()
|
||||
src_lengths = src_lengths.cuda()
|
||||
if constraints is not None:
|
||||
constraints = constraints.cuda()
|
||||
|
||||
sample = {
|
||||
"net_input": {
|
||||
"src_tokens": src_tokens,
|
||||
"src_lengths": src_lengths,
|
||||
},
|
||||
}
|
||||
translate_start_time = time.time()
|
||||
translations = task.inference_step(
|
||||
generator, models, sample, constraints=constraints
|
||||
)
|
||||
translate_time = time.time() - translate_start_time
|
||||
total_translate_time += translate_time
|
||||
list_constraints = [[] for _ in range(bsz)]
|
||||
if cfg.generation.constraints:
|
||||
list_constraints = [unpack_constraints(c) for c in constraints]
|
||||
for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
|
||||
src_tokens_i = utils.strip_pad(src_tokens[i], tgt_dict.pad())
|
||||
constraints = list_constraints[i]
|
||||
results.append(
|
||||
(
|
||||
start_id + id,
|
||||
src_tokens_i,
|
||||
hypos,
|
||||
{
|
||||
"constraints": constraints,
|
||||
"time": translate_time / len(translations),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# sort output to match input order
|
||||
for id_, src_tokens, hypos, info in sorted(results, key=lambda x: x[0]):
|
||||
src_str = ''
|
||||
if src_dict is not None:
|
||||
src_str = src_dict.string(src_tokens, cfg.common_eval.post_process)
|
||||
print("S-{}\t{}".format(id_, src_str))
|
||||
print("W-{}\t{:.3f}\tseconds".format(id_, info["time"]))
|
||||
for constraint in info["constraints"]:
|
||||
print(
|
||||
"C-{}\t{}".format(
|
||||
id_, tgt_dict.string(constraint, cfg.common_eval.post_process)
|
||||
)
|
||||
)
|
||||
|
||||
# Process top predictions
|
||||
for hypo in hypos[: min(len(hypos), cfg.generation.nbest)]:
|
||||
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
|
||||
hypo_tokens=hypo["tokens"].int().cpu(),
|
||||
src_str=src_str,
|
||||
alignment=hypo["alignment"],
|
||||
align_dict=align_dict,
|
||||
tgt_dict=tgt_dict,
|
||||
remove_bpe=cfg.common_eval.post_process,
|
||||
extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator),
|
||||
)
|
||||
detok_hypo_str = decode_fn(hypo_str)
|
||||
score = hypo["score"] / math.log(2) # convert to base 2
|
||||
# original hypothesis (after tokenization and BPE)
|
||||
print("H-{}\t{}\t{}".format(id_, score, hypo_str))
|
||||
# detokenized hypothesis
|
||||
print("D-{}\t{}\t{}".format(id_, score, detok_hypo_str))
|
||||
print(
|
||||
"P-{}\t{}".format(
|
||||
id_,
|
||||
" ".join(
|
||||
map(
|
||||
lambda x: "{:.4f}".format(x),
|
||||
# convert from base e to base 2
|
||||
hypo["positional_scores"].div_(math.log(2)).tolist(),
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
if cfg.generation.print_alignment:
|
||||
alignment_str = " ".join(
|
||||
["{}-{}".format(src, tgt) for src, tgt in alignment]
|
||||
)
|
||||
print("A-{}\t{}".format(id_, alignment_str))
|
||||
|
||||
# update running id_ counter
|
||||
start_id += len(inputs)
|
||||
|
||||
logger.info(
|
||||
"Total time: {:.3f} seconds; translation time: {:.3f}".format(
|
||||
time.time() - start_time, total_translate_time
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = options.get_interactive_generation_parser()
|
||||
args = options.parse_args_and_arch(parser)
|
||||
distributed_utils.call_main(convert_namespace_to_omegaconf(args), main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,398 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
Data pre-processing: build vocabularies and binarize training data.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from collections import Counter
|
||||
from itertools import zip_longest
|
||||
from multiprocessing import Pool
|
||||
|
||||
from fairseq import options, tasks, utils
|
||||
from fairseq.binarizer import Binarizer
|
||||
from fairseq.data import indexed_dataset
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=sys.stdout,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.preprocess")
|
||||
|
||||
|
||||
def main(args):
|
||||
utils.import_user_module(args)
|
||||
|
||||
os.makedirs(args.destdir, exist_ok=True)
|
||||
|
||||
logger.addHandler(
|
||||
logging.FileHandler(
|
||||
filename=os.path.join(args.destdir, "preprocess.log"),
|
||||
)
|
||||
)
|
||||
logger.info(args)
|
||||
|
||||
task = tasks.get_task(args.task)
|
||||
|
||||
def train_path(lang):
|
||||
return "{}{}".format(args.trainpref, ("." + lang) if lang else "")
|
||||
|
||||
def file_name(prefix, lang):
|
||||
fname = prefix
|
||||
if lang is not None:
|
||||
fname += ".{lang}".format(lang=lang)
|
||||
return fname
|
||||
|
||||
def dest_path(prefix, lang):
|
||||
return os.path.join(args.destdir, file_name(prefix, lang))
|
||||
|
||||
def dict_path(lang):
|
||||
return dest_path("dict", lang) + ".txt"
|
||||
|
||||
def build_dictionary(filenames, src=False, tgt=False):
|
||||
assert src ^ tgt
|
||||
return task.build_dictionary(
|
||||
filenames,
|
||||
workers=args.workers,
|
||||
threshold=args.thresholdsrc if src else args.thresholdtgt,
|
||||
nwords=args.nwordssrc if src else args.nwordstgt,
|
||||
padding_factor=args.padding_factor,
|
||||
)
|
||||
|
||||
target = not args.only_source
|
||||
|
||||
if not args.srcdict and os.path.exists(dict_path(args.source_lang)):
|
||||
raise FileExistsError(dict_path(args.source_lang))
|
||||
if target and not args.tgtdict and os.path.exists(dict_path(args.target_lang)):
|
||||
raise FileExistsError(dict_path(args.target_lang))
|
||||
|
||||
if args.joined_dictionary:
|
||||
assert (
|
||||
not args.srcdict or not args.tgtdict
|
||||
), "cannot use both --srcdict and --tgtdict with --joined-dictionary"
|
||||
|
||||
if args.srcdict:
|
||||
src_dict = task.load_dictionary(args.srcdict)
|
||||
elif args.tgtdict:
|
||||
src_dict = task.load_dictionary(args.tgtdict)
|
||||
else:
|
||||
assert (
|
||||
args.trainpref
|
||||
), "--trainpref must be set if --srcdict is not specified"
|
||||
src_dict = build_dictionary(
|
||||
{train_path(lang) for lang in [args.source_lang, args.target_lang]},
|
||||
src=True,
|
||||
)
|
||||
tgt_dict = src_dict
|
||||
else:
|
||||
if args.srcdict:
|
||||
src_dict = task.load_dictionary(args.srcdict)
|
||||
else:
|
||||
assert (
|
||||
args.trainpref
|
||||
), "--trainpref must be set if --srcdict is not specified"
|
||||
src_dict = build_dictionary([train_path(args.source_lang)], src=True)
|
||||
|
||||
if target:
|
||||
if args.tgtdict:
|
||||
tgt_dict = task.load_dictionary(args.tgtdict)
|
||||
else:
|
||||
assert (
|
||||
args.trainpref
|
||||
), "--trainpref must be set if --tgtdict is not specified"
|
||||
tgt_dict = build_dictionary([train_path(args.target_lang)], tgt=True)
|
||||
else:
|
||||
tgt_dict = None
|
||||
|
||||
src_dict.save(dict_path(args.source_lang))
|
||||
if target and tgt_dict is not None:
|
||||
tgt_dict.save(dict_path(args.target_lang))
|
||||
|
||||
def make_binary_dataset(vocab, input_prefix, output_prefix, lang, num_workers):
|
||||
logger.info("[{}] Dictionary: {} types".format(lang, len(vocab)))
|
||||
n_seq_tok = [0, 0]
|
||||
replaced = Counter()
|
||||
|
||||
def merge_result(worker_result):
|
||||
replaced.update(worker_result["replaced"])
|
||||
n_seq_tok[0] += worker_result["nseq"]
|
||||
n_seq_tok[1] += worker_result["ntok"]
|
||||
|
||||
input_file = "{}{}".format(
|
||||
input_prefix, ("." + lang) if lang is not None else ""
|
||||
)
|
||||
offsets = Binarizer.find_offsets(input_file, num_workers)
|
||||
pool = None
|
||||
if num_workers > 1:
|
||||
pool = Pool(processes=num_workers - 1)
|
||||
for worker_id in range(1, num_workers):
|
||||
prefix = "{}{}".format(output_prefix, worker_id)
|
||||
pool.apply_async(
|
||||
binarize,
|
||||
(
|
||||
args,
|
||||
input_file,
|
||||
vocab,
|
||||
prefix,
|
||||
lang,
|
||||
offsets[worker_id],
|
||||
offsets[worker_id + 1],
|
||||
),
|
||||
callback=merge_result,
|
||||
)
|
||||
pool.close()
|
||||
|
||||
ds = indexed_dataset.make_builder(
|
||||
dataset_dest_file(args, output_prefix, lang, "bin"),
|
||||
impl=args.dataset_impl,
|
||||
vocab_size=len(vocab),
|
||||
)
|
||||
merge_result(
|
||||
Binarizer.binarize(
|
||||
input_file, vocab, lambda t: ds.add_item(t), offset=0, end=offsets[1]
|
||||
)
|
||||
)
|
||||
if num_workers > 1:
|
||||
pool.join()
|
||||
for worker_id in range(1, num_workers):
|
||||
prefix = "{}{}".format(output_prefix, worker_id)
|
||||
temp_file_path = dataset_dest_prefix(args, prefix, lang)
|
||||
ds.merge_file_(temp_file_path)
|
||||
os.remove(indexed_dataset.data_file_path(temp_file_path))
|
||||
os.remove(indexed_dataset.index_file_path(temp_file_path))
|
||||
|
||||
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
|
||||
|
||||
logger.info(
|
||||
"[{}] {}: {} sents, {} tokens, {:.3}% replaced by {}".format(
|
||||
lang,
|
||||
input_file,
|
||||
n_seq_tok[0],
|
||||
n_seq_tok[1],
|
||||
100 * sum(replaced.values()) / n_seq_tok[1],
|
||||
vocab.unk_word,
|
||||
)
|
||||
)
|
||||
|
||||
def make_binary_alignment_dataset(input_prefix, output_prefix, num_workers):
|
||||
nseq = [0]
|
||||
|
||||
def merge_result(worker_result):
|
||||
nseq[0] += worker_result["nseq"]
|
||||
|
||||
input_file = input_prefix
|
||||
offsets = Binarizer.find_offsets(input_file, num_workers)
|
||||
pool = None
|
||||
if num_workers > 1:
|
||||
pool = Pool(processes=num_workers - 1)
|
||||
for worker_id in range(1, num_workers):
|
||||
prefix = "{}{}".format(output_prefix, worker_id)
|
||||
pool.apply_async(
|
||||
binarize_alignments,
|
||||
(
|
||||
args,
|
||||
input_file,
|
||||
utils.parse_alignment,
|
||||
prefix,
|
||||
offsets[worker_id],
|
||||
offsets[worker_id + 1],
|
||||
),
|
||||
callback=merge_result,
|
||||
)
|
||||
pool.close()
|
||||
|
||||
ds = indexed_dataset.make_builder(
|
||||
dataset_dest_file(args, output_prefix, None, "bin"), impl=args.dataset_impl
|
||||
)
|
||||
|
||||
merge_result(
|
||||
Binarizer.binarize_alignments(
|
||||
input_file,
|
||||
utils.parse_alignment,
|
||||
lambda t: ds.add_item(t),
|
||||
offset=0,
|
||||
end=offsets[1],
|
||||
)
|
||||
)
|
||||
if num_workers > 1:
|
||||
pool.join()
|
||||
for worker_id in range(1, num_workers):
|
||||
prefix = "{}{}".format(output_prefix, worker_id)
|
||||
temp_file_path = dataset_dest_prefix(args, prefix, None)
|
||||
ds.merge_file_(temp_file_path)
|
||||
os.remove(indexed_dataset.data_file_path(temp_file_path))
|
||||
os.remove(indexed_dataset.index_file_path(temp_file_path))
|
||||
|
||||
ds.finalize(dataset_dest_file(args, output_prefix, None, "idx"))
|
||||
|
||||
logger.info("[alignments] {}: parsed {} alignments".format(input_file, nseq[0]))
|
||||
|
||||
def make_dataset(vocab, input_prefix, output_prefix, lang, num_workers=1):
|
||||
if args.dataset_impl == "raw":
|
||||
# Copy original text file to destination folder
|
||||
output_text_file = dest_path(
|
||||
output_prefix + ".{}-{}".format(args.source_lang, args.target_lang),
|
||||
lang,
|
||||
)
|
||||
shutil.copyfile(file_name(input_prefix, lang), output_text_file)
|
||||
else:
|
||||
make_binary_dataset(vocab, input_prefix, output_prefix, lang, num_workers)
|
||||
|
||||
def make_all(lang, vocab):
|
||||
if args.trainpref:
|
||||
make_dataset(vocab, args.trainpref, "train", lang, num_workers=args.workers)
|
||||
if args.validpref:
|
||||
for k, validpref in enumerate(args.validpref.split(",")):
|
||||
outprefix = "valid{}".format(k) if k > 0 else "valid"
|
||||
make_dataset(
|
||||
vocab, validpref, outprefix, lang, num_workers=args.workers
|
||||
)
|
||||
if args.testpref:
|
||||
for k, testpref in enumerate(args.testpref.split(",")):
|
||||
outprefix = "test{}".format(k) if k > 0 else "test"
|
||||
make_dataset(vocab, testpref, outprefix, lang, num_workers=args.workers)
|
||||
|
||||
def make_all_alignments():
|
||||
if args.trainpref and os.path.exists(args.trainpref + "." + args.align_suffix):
|
||||
make_binary_alignment_dataset(
|
||||
args.trainpref + "." + args.align_suffix,
|
||||
"train.align",
|
||||
num_workers=args.workers,
|
||||
)
|
||||
if args.validpref and os.path.exists(args.validpref + "." + args.align_suffix):
|
||||
make_binary_alignment_dataset(
|
||||
args.validpref + "." + args.align_suffix,
|
||||
"valid.align",
|
||||
num_workers=args.workers,
|
||||
)
|
||||
if args.testpref and os.path.exists(args.testpref + "." + args.align_suffix):
|
||||
make_binary_alignment_dataset(
|
||||
args.testpref + "." + args.align_suffix,
|
||||
"test.align",
|
||||
num_workers=args.workers,
|
||||
)
|
||||
|
||||
make_all(args.source_lang, src_dict)
|
||||
if target:
|
||||
make_all(args.target_lang, tgt_dict)
|
||||
if args.align_suffix:
|
||||
make_all_alignments()
|
||||
|
||||
logger.info("Wrote preprocessed data to {}".format(args.destdir))
|
||||
|
||||
if args.alignfile:
|
||||
assert args.trainpref, "--trainpref must be set if --alignfile is specified"
|
||||
src_file_name = train_path(args.source_lang)
|
||||
tgt_file_name = train_path(args.target_lang)
|
||||
freq_map = {}
|
||||
with open(args.alignfile, "r", encoding="utf-8") as align_file:
|
||||
with open(src_file_name, "r", encoding="utf-8") as src_file:
|
||||
with open(tgt_file_name, "r", encoding="utf-8") as tgt_file:
|
||||
for a, s, t in zip_longest(align_file, src_file, tgt_file):
|
||||
si = src_dict.encode_line(s, add_if_not_exist=False)
|
||||
ti = tgt_dict.encode_line(t, add_if_not_exist=False)
|
||||
ai = list(map(lambda x: tuple(x.split("-")), a.split()))
|
||||
for sai, tai in ai:
|
||||
srcidx = si[int(sai)]
|
||||
tgtidx = ti[int(tai)]
|
||||
if srcidx != src_dict.unk() and tgtidx != tgt_dict.unk():
|
||||
assert srcidx != src_dict.pad()
|
||||
assert srcidx != src_dict.eos()
|
||||
assert tgtidx != tgt_dict.pad()
|
||||
assert tgtidx != tgt_dict.eos()
|
||||
|
||||
if srcidx not in freq_map:
|
||||
freq_map[srcidx] = {}
|
||||
if tgtidx not in freq_map[srcidx]:
|
||||
freq_map[srcidx][tgtidx] = 1
|
||||
else:
|
||||
freq_map[srcidx][tgtidx] += 1
|
||||
|
||||
align_dict = {}
|
||||
for srcidx in freq_map.keys():
|
||||
align_dict[srcidx] = max(freq_map[srcidx], key=freq_map[srcidx].get)
|
||||
|
||||
with open(
|
||||
os.path.join(
|
||||
args.destdir,
|
||||
"alignment.{}-{}.txt".format(args.source_lang, args.target_lang),
|
||||
),
|
||||
"w",
|
||||
encoding="utf-8",
|
||||
) as f:
|
||||
for k, v in align_dict.items():
|
||||
print("{} {}".format(src_dict[k], tgt_dict[v]), file=f)
|
||||
|
||||
|
||||
def binarize(args, filename, vocab, output_prefix, lang, offset, end, append_eos=True):
|
||||
ds = indexed_dataset.make_builder(
|
||||
dataset_dest_file(args, output_prefix, lang, "bin"),
|
||||
impl=args.dataset_impl,
|
||||
vocab_size=len(vocab),
|
||||
)
|
||||
|
||||
def consumer(tensor):
|
||||
ds.add_item(tensor)
|
||||
|
||||
res = Binarizer.binarize(
|
||||
filename, vocab, consumer, append_eos=append_eos, offset=offset, end=end
|
||||
)
|
||||
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
|
||||
return res
|
||||
|
||||
|
||||
def binarize_alignments(args, filename, parse_alignment, output_prefix, offset, end):
|
||||
ds = indexed_dataset.make_builder(
|
||||
dataset_dest_file(args, output_prefix, None, "bin"),
|
||||
impl=args.dataset_impl,
|
||||
vocab_size=None,
|
||||
)
|
||||
|
||||
def consumer(tensor):
|
||||
ds.add_item(tensor)
|
||||
|
||||
res = Binarizer.binarize_alignments(
|
||||
filename, parse_alignment, consumer, offset=offset, end=end
|
||||
)
|
||||
ds.finalize(dataset_dest_file(args, output_prefix, None, "idx"))
|
||||
return res
|
||||
|
||||
|
||||
def dataset_dest_prefix(args, output_prefix, lang):
|
||||
base = "{}/{}".format(args.destdir, output_prefix)
|
||||
if lang is not None:
|
||||
lang_part = ".{}-{}.{}".format(args.source_lang, args.target_lang, lang)
|
||||
elif args.only_source:
|
||||
lang_part = ""
|
||||
else:
|
||||
lang_part = ".{}-{}".format(args.source_lang, args.target_lang)
|
||||
|
||||
return "{}{}".format(base, lang_part)
|
||||
|
||||
|
||||
def dataset_dest_file(args, output_prefix, lang, extension):
|
||||
base = dataset_dest_prefix(args, output_prefix, lang)
|
||||
return "{}.{}".format(base, extension)
|
||||
|
||||
|
||||
def get_offsets(input_file, num_workers):
|
||||
return Binarizer.find_offsets(input_file, num_workers)
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = options.get_preprocessing_parser()
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,102 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
BLEU scoring of generated translations against reference translations.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
from fairseq.data import dictionary
|
||||
from fairseq.scoring import bleu
|
||||
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Command-line script for BLEU scoring."
|
||||
)
|
||||
# fmt: off
|
||||
parser.add_argument('-s', '--sys', default='-', help='system output')
|
||||
parser.add_argument('-r', '--ref', required=True, help='references')
|
||||
parser.add_argument('-o', '--order', default=4, metavar='N',
|
||||
type=int, help='consider ngrams up to this order')
|
||||
parser.add_argument('--ignore-case', action='store_true',
|
||||
help='case-insensitive scoring')
|
||||
parser.add_argument('--sacrebleu', action='store_true',
|
||||
help='score with sacrebleu')
|
||||
parser.add_argument('--sentence-bleu', action='store_true',
|
||||
help='report sentence-level BLEUs (i.e., with +1 smoothing)')
|
||||
# fmt: on
|
||||
return parser
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = get_parser()
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
|
||||
assert args.sys == "-" or os.path.exists(
|
||||
args.sys
|
||||
), "System output file {} does not exist".format(args.sys)
|
||||
assert os.path.exists(args.ref), "Reference file {} does not exist".format(args.ref)
|
||||
|
||||
dict = dictionary.Dictionary()
|
||||
|
||||
def readlines(fd):
|
||||
for line in fd.readlines():
|
||||
if args.ignore_case:
|
||||
yield line.lower()
|
||||
else:
|
||||
yield line
|
||||
|
||||
if args.sacrebleu:
|
||||
import sacrebleu
|
||||
|
||||
def score(fdsys):
|
||||
with open(args.ref) as fdref:
|
||||
print(sacrebleu.corpus_bleu(fdsys, [fdref]).format())
|
||||
|
||||
elif args.sentence_bleu:
|
||||
|
||||
def score(fdsys):
|
||||
with open(args.ref) as fdref:
|
||||
scorer = bleu.Scorer(dict.pad(), dict.eos(), dict.unk())
|
||||
for i, (sys_tok, ref_tok) in enumerate(
|
||||
zip(readlines(fdsys), readlines(fdref))
|
||||
):
|
||||
scorer.reset(one_init=True)
|
||||
sys_tok = dict.encode_line(sys_tok)
|
||||
ref_tok = dict.encode_line(ref_tok)
|
||||
scorer.add(ref_tok, sys_tok)
|
||||
print(i, scorer.result_string(args.order))
|
||||
|
||||
else:
|
||||
|
||||
def score(fdsys):
|
||||
with open(args.ref) as fdref:
|
||||
scorer = bleu.Scorer(
|
||||
bleu.BleuConfig(
|
||||
pad=dict.pad(),
|
||||
eos=dict.eos(),
|
||||
unk=dict.unk(),
|
||||
)
|
||||
)
|
||||
for sys_tok, ref_tok in zip(readlines(fdsys), readlines(fdref)):
|
||||
sys_tok = dict.encode_line(sys_tok)
|
||||
ref_tok = dict.encode_line(ref_tok)
|
||||
scorer.add(ref_tok, sys_tok)
|
||||
print(scorer.result_string(args.order))
|
||||
|
||||
if args.sys == "-":
|
||||
score(sys.stdin)
|
||||
else:
|
||||
with open(args.sys, "r") as f:
|
||||
score(f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,454 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""
|
||||
Train a new model on one or across multiple GPUs.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict, Optional, Any, List, Tuple, Callable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq import (
|
||||
checkpoint_utils,
|
||||
distributed_utils,
|
||||
options,
|
||||
quantization_utils,
|
||||
tasks,
|
||||
utils,
|
||||
)
|
||||
from fairseq.data import iterators
|
||||
from fairseq.dataclass.configs import FairseqConfig
|
||||
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
||||
from fairseq.distributed_utils import is_master
|
||||
from fairseq.logging import meters, metrics, progress_bar
|
||||
from fairseq.model_parallel.megatron_trainer import MegatronTrainer
|
||||
from fairseq.trainer import Trainer
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=sys.stdout,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.train")
|
||||
|
||||
|
||||
def main(cfg: FairseqConfig) -> None:
|
||||
if isinstance(cfg, argparse.Namespace):
|
||||
cfg = convert_namespace_to_omegaconf(cfg)
|
||||
|
||||
utils.import_user_module(cfg.common)
|
||||
|
||||
if is_master(cfg.distributed_training) and "job_logging_cfg" in cfg:
|
||||
# make hydra logging work with ddp (see # see https://github.com/facebookresearch/hydra/issues/1126)
|
||||
logging.config.dictConfig(OmegaConf.to_container(cfg.job_logging_cfg))
|
||||
|
||||
assert (
|
||||
cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None
|
||||
), "Must specify batch size either with --max-tokens or --batch-size"
|
||||
metrics.reset()
|
||||
|
||||
np.random.seed(cfg.common.seed)
|
||||
utils.set_torch_seed(cfg.common.seed)
|
||||
|
||||
if distributed_utils.is_master(cfg.distributed_training):
|
||||
checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
|
||||
|
||||
# Print args
|
||||
logger.info(cfg)
|
||||
|
||||
# Setup task, e.g., translation, language modeling, etc.
|
||||
task = tasks.setup_task(cfg.task)
|
||||
# Load valid dataset (we load training data below, based on the latest checkpoint)
|
||||
for valid_sub_split in cfg.dataset.valid_subset.split(","):
|
||||
task.load_dataset(valid_sub_split, combine=False, epoch=1)
|
||||
|
||||
assert cfg.criterion, "Please specify criterion to train a model"
|
||||
|
||||
# Build model and criterion
|
||||
model = task.build_model(cfg.model)
|
||||
criterion = task.build_criterion(cfg.criterion)
|
||||
logger.info(model)
|
||||
logger.info("task: {}".format(task.__class__.__name__))
|
||||
logger.info("model: {}".format(model.__class__.__name__))
|
||||
logger.info("criterion: {}".format(criterion.__class__.__name__))
|
||||
logger.info(
|
||||
"num. model params: {:,} (num. trained: {:,})".format(
|
||||
sum(p.numel() for p in model.parameters()),
|
||||
sum(p.numel() for p in model.parameters() if p.requires_grad),
|
||||
)
|
||||
)
|
||||
|
||||
# (optionally) Configure quantization
|
||||
if cfg.common.quantization_config_path is not None:
|
||||
quantizer = quantization_utils.Quantizer(
|
||||
config_path=cfg.common.quantization_config_path,
|
||||
max_epoch=cfg.optimization.max_epoch,
|
||||
max_update=cfg.optimization.max_update,
|
||||
)
|
||||
else:
|
||||
quantizer = None
|
||||
|
||||
# Build trainer
|
||||
if cfg.common.model_parallel_size == 1:
|
||||
trainer = Trainer(cfg, task, model, criterion, quantizer)
|
||||
else:
|
||||
trainer = MegatronTrainer(cfg, task, model, criterion)
|
||||
|
||||
logger.info(
|
||||
"training on {} devices (GPUs/TPUs)".format(
|
||||
cfg.distributed_training.distributed_world_size
|
||||
)
|
||||
)
|
||||
logger.info(
|
||||
"max tokens per GPU = {} and batch size per GPU = {}".format(
|
||||
cfg.dataset.max_tokens,
|
||||
cfg.dataset.batch_size,
|
||||
)
|
||||
)
|
||||
|
||||
# Load the latest checkpoint if one is available and restore the
|
||||
# corresponding train iterator
|
||||
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
|
||||
cfg.checkpoint,
|
||||
trainer,
|
||||
# don't cache epoch iterators for sharded datasets
|
||||
disable_iterator_cache=task.has_sharded_data("train"),
|
||||
)
|
||||
|
||||
max_epoch = cfg.optimization.max_epoch or math.inf
|
||||
lr = trainer.get_lr()
|
||||
train_meter = meters.StopwatchMeter()
|
||||
train_meter.start()
|
||||
while epoch_itr.next_epoch_idx <= max_epoch:
|
||||
if lr <= cfg.optimization.stop_min_lr:
|
||||
logger.info(
|
||||
f"stopping training because current learning rate ({lr}) is smaller "
|
||||
"than or equal to minimum learning rate "
|
||||
f"(--stop-min-lr={cfg.optimization.stop_min_lr})"
|
||||
)
|
||||
break
|
||||
|
||||
# train for one epoch
|
||||
valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
# only use first validation loss to update the learning rate
|
||||
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
|
||||
|
||||
epoch_itr = trainer.get_train_iterator(
|
||||
epoch_itr.next_epoch_idx,
|
||||
# sharded data: get train iterator for next epoch
|
||||
load_dataset=task.has_sharded_data("train"),
|
||||
# don't cache epoch iterators for sharded datasets
|
||||
disable_iterator_cache=task.has_sharded_data("train"),
|
||||
)
|
||||
train_meter.stop()
|
||||
logger.info("done training in {:.1f} seconds".format(train_meter.sum))
|
||||
|
||||
|
||||
def should_stop_early(cfg: DictConfig, valid_loss: float) -> bool:
|
||||
# skip check if no validation was done in the current epoch
|
||||
if valid_loss is None:
|
||||
return False
|
||||
if cfg.checkpoint.patience <= 0:
|
||||
return False
|
||||
|
||||
def is_better(a, b):
|
||||
return a > b if cfg.checkpoint.maximize_best_checkpoint_metric else a < b
|
||||
|
||||
prev_best = getattr(should_stop_early, "best", None)
|
||||
if prev_best is None or is_better(valid_loss, prev_best):
|
||||
should_stop_early.best = valid_loss
|
||||
should_stop_early.num_runs = 0
|
||||
return False
|
||||
else:
|
||||
should_stop_early.num_runs += 1
|
||||
if should_stop_early.num_runs >= cfg.checkpoint.patience:
|
||||
logger.info(
|
||||
"early stop since valid performance hasn't improved for last {} runs".format(
|
||||
cfg.checkpoint.patience
|
||||
)
|
||||
)
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
@metrics.aggregate("train")
|
||||
def train(
|
||||
cfg: DictConfig, trainer: Trainer, task: tasks.FairseqTask, epoch_itr
|
||||
) -> Tuple[List[Optional[float]], bool]:
|
||||
"""Train the model for one epoch and return validation losses."""
|
||||
# Initialize data iterator
|
||||
itr = epoch_itr.next_epoch_itr(
|
||||
fix_batches_to_gpus=cfg.distributed_training.fix_batches_to_gpus,
|
||||
shuffle=(epoch_itr.next_epoch_idx > cfg.dataset.curriculum),
|
||||
)
|
||||
update_freq = (
|
||||
cfg.optimization.update_freq[epoch_itr.epoch - 1]
|
||||
if epoch_itr.epoch <= len(cfg.optimization.update_freq)
|
||||
else cfg.optimization.update_freq[-1]
|
||||
)
|
||||
itr = iterators.GroupedIterator(itr, update_freq)
|
||||
if cfg.common.tpu:
|
||||
itr = utils.tpu_data_loader(itr)
|
||||
progress = progress_bar.progress_bar(
|
||||
itr,
|
||||
log_format=cfg.common.log_format,
|
||||
log_interval=cfg.common.log_interval,
|
||||
epoch=epoch_itr.epoch,
|
||||
tensorboard_logdir=(
|
||||
cfg.common.tensorboard_logdir
|
||||
if distributed_utils.is_master(cfg.distributed_training)
|
||||
else None
|
||||
),
|
||||
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
|
||||
wandb_project=(
|
||||
cfg.common.wandb_project
|
||||
if distributed_utils.is_master(cfg.distributed_training)
|
||||
else None
|
||||
),
|
||||
wandb_run_name=os.environ.get(
|
||||
"WANDB_NAME", os.path.basename(cfg.checkpoint.save_dir)
|
||||
),
|
||||
azureml_logging=(
|
||||
cfg.common.azureml_logging
|
||||
if distributed_utils.is_master(cfg.distributed_training)
|
||||
else False
|
||||
),
|
||||
)
|
||||
progress.update_config(_flatten_config(cfg))
|
||||
|
||||
trainer.begin_epoch(epoch_itr.epoch)
|
||||
|
||||
valid_subsets = cfg.dataset.valid_subset.split(",")
|
||||
should_stop = False
|
||||
num_updates = trainer.get_num_updates()
|
||||
logger.info("Start iterating over samples")
|
||||
for i, samples in enumerate(progress):
|
||||
with metrics.aggregate("train_inner"), torch.autograd.profiler.record_function(
|
||||
"train_step-%d" % i
|
||||
):
|
||||
log_output = trainer.train_step(samples)
|
||||
|
||||
if log_output is not None: # not OOM, overflow, ...
|
||||
# log mid-epoch stats
|
||||
num_updates = trainer.get_num_updates()
|
||||
if num_updates % cfg.common.log_interval == 0:
|
||||
stats = get_training_stats(metrics.get_smoothed_values("train_inner"))
|
||||
progress.log(stats, tag="train_inner", step=num_updates)
|
||||
|
||||
# reset mid-epoch stats after each log interval
|
||||
# the end-of-epoch stats will still be preserved
|
||||
metrics.reset_meters("train_inner")
|
||||
|
||||
end_of_epoch = not itr.has_next()
|
||||
valid_losses, should_stop = validate_and_save(
|
||||
cfg, trainer, task, epoch_itr, valid_subsets, end_of_epoch
|
||||
)
|
||||
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
# log end-of-epoch stats
|
||||
logger.info("end of epoch {} (average epoch stats below)".format(epoch_itr.epoch))
|
||||
stats = get_training_stats(metrics.get_smoothed_values("train"))
|
||||
progress.print(stats, tag="train", step=num_updates)
|
||||
|
||||
# reset epoch-level meters
|
||||
metrics.reset_meters("train")
|
||||
return valid_losses, should_stop
|
||||
|
||||
|
||||
def _flatten_config(cfg: DictConfig):
|
||||
config = OmegaConf.to_container(cfg)
|
||||
# remove any legacy Namespaces and replace with a single "args"
|
||||
namespace = None
|
||||
for k, v in list(config.items()):
|
||||
if isinstance(v, argparse.Namespace):
|
||||
namespace = v
|
||||
del config[k]
|
||||
if namespace is not None:
|
||||
config["args"] = vars(namespace)
|
||||
return config
|
||||
|
||||
|
||||
def validate_and_save(
|
||||
cfg: DictConfig,
|
||||
trainer: Trainer,
|
||||
task: tasks.FairseqTask,
|
||||
epoch_itr,
|
||||
valid_subsets: List[str],
|
||||
end_of_epoch: bool,
|
||||
) -> Tuple[List[Optional[float]], bool]:
|
||||
num_updates = trainer.get_num_updates()
|
||||
max_update = cfg.optimization.max_update or math.inf
|
||||
|
||||
# Stopping conditions (and an additional one based on validation loss later
|
||||
# on)
|
||||
should_stop = False
|
||||
if num_updates >= max_update:
|
||||
should_stop = True
|
||||
logger.info(
|
||||
f"Stopping training due to "
|
||||
f"num_updates: {num_updates} >= max_update: {max_update}"
|
||||
)
|
||||
|
||||
training_time_hours = trainer.cumulative_training_time() / (60 * 60)
|
||||
if (
|
||||
cfg.optimization.stop_time_hours > 0
|
||||
and training_time_hours > cfg.optimization.stop_time_hours
|
||||
):
|
||||
should_stop = True
|
||||
logger.info(
|
||||
f"Stopping training due to "
|
||||
f"cumulative_training_time: {training_time_hours} > "
|
||||
f"stop_time_hours: {cfg.optimization.stop_time_hours} hour(s)"
|
||||
)
|
||||
|
||||
do_save = (
|
||||
(end_of_epoch and epoch_itr.epoch % cfg.checkpoint.save_interval == 0)
|
||||
or should_stop
|
||||
or (
|
||||
cfg.checkpoint.save_interval_updates > 0
|
||||
and num_updates > 0
|
||||
and num_updates % cfg.checkpoint.save_interval_updates == 0
|
||||
and num_updates >= cfg.dataset.validate_after_updates
|
||||
)
|
||||
)
|
||||
do_validate = (
|
||||
(not end_of_epoch and do_save) # validate during mid-epoch saves
|
||||
or (end_of_epoch and epoch_itr.epoch % cfg.dataset.validate_interval == 0)
|
||||
or should_stop
|
||||
or (
|
||||
cfg.dataset.validate_interval_updates > 0
|
||||
and num_updates > 0
|
||||
and num_updates % cfg.dataset.validate_interval_updates == 0
|
||||
)
|
||||
) and not cfg.dataset.disable_validation
|
||||
|
||||
# Validate
|
||||
valid_losses = [None]
|
||||
if do_validate:
|
||||
valid_losses = validate(cfg, trainer, task, epoch_itr, valid_subsets)
|
||||
|
||||
should_stop |= should_stop_early(cfg, valid_losses[0])
|
||||
|
||||
# Save checkpoint
|
||||
if do_save or should_stop:
|
||||
checkpoint_utils.save_checkpoint(
|
||||
cfg.checkpoint, trainer, epoch_itr, valid_losses[0]
|
||||
)
|
||||
|
||||
return valid_losses, should_stop
|
||||
|
||||
|
||||
def get_training_stats(stats: Dict[str, Any]) -> Dict[str, Any]:
|
||||
stats["wall"] = round(metrics.get_meter("default", "wall").elapsed_time, 0)
|
||||
return stats
|
||||
|
||||
|
||||
def validate(
|
||||
cfg: DictConfig,
|
||||
trainer: Trainer,
|
||||
task: tasks.FairseqTask,
|
||||
epoch_itr,
|
||||
subsets: List[str],
|
||||
) -> List[Optional[float]]:
|
||||
"""Evaluate the model on the validation set(s) and return the losses."""
|
||||
|
||||
if cfg.dataset.fixed_validation_seed is not None:
|
||||
# set fixed seed for every validation
|
||||
utils.set_torch_seed(cfg.dataset.fixed_validation_seed)
|
||||
|
||||
trainer.begin_valid_epoch(epoch_itr.epoch)
|
||||
valid_losses = []
|
||||
for subset in subsets:
|
||||
logger.info('begin validation on "{}" subset'.format(subset))
|
||||
|
||||
# Initialize data iterator
|
||||
itr = trainer.get_valid_iterator(subset).next_epoch_itr(
|
||||
shuffle=False, set_dataset_epoch=False # use a fixed valid set
|
||||
)
|
||||
if cfg.common.tpu:
|
||||
itr = utils.tpu_data_loader(itr)
|
||||
progress = progress_bar.progress_bar(
|
||||
itr,
|
||||
log_format=cfg.common.log_format,
|
||||
log_interval=cfg.common.log_interval,
|
||||
epoch=epoch_itr.epoch,
|
||||
prefix=f"valid on '{subset}' subset",
|
||||
tensorboard_logdir=(
|
||||
cfg.common.tensorboard_logdir
|
||||
if distributed_utils.is_master(cfg.distributed_training)
|
||||
else None
|
||||
),
|
||||
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
|
||||
wandb_project=(
|
||||
cfg.common.wandb_project
|
||||
if distributed_utils.is_master(cfg.distributed_training)
|
||||
else None
|
||||
),
|
||||
wandb_run_name=os.environ.get(
|
||||
"WANDB_NAME", os.path.basename(cfg.checkpoint.save_dir)
|
||||
),
|
||||
)
|
||||
|
||||
# create a new root metrics aggregator so validation metrics
|
||||
# don't pollute other aggregators (e.g., train meters)
|
||||
with metrics.aggregate(new_root=True) as agg:
|
||||
for sample in progress:
|
||||
trainer.valid_step(sample)
|
||||
|
||||
# log validation stats
|
||||
stats = get_valid_stats(cfg, trainer, agg.get_smoothed_values())
|
||||
progress.print(stats, tag=subset, step=trainer.get_num_updates())
|
||||
|
||||
valid_losses.append(stats[cfg.checkpoint.best_checkpoint_metric])
|
||||
return valid_losses
|
||||
|
||||
|
||||
def get_valid_stats(
|
||||
cfg: DictConfig, trainer: Trainer, stats: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
stats["num_updates"] = trainer.get_num_updates()
|
||||
if hasattr(checkpoint_utils.save_checkpoint, "best"):
|
||||
key = "best_{0}".format(cfg.checkpoint.best_checkpoint_metric)
|
||||
best_function = max if cfg.checkpoint.maximize_best_checkpoint_metric else min
|
||||
stats[key] = best_function(
|
||||
checkpoint_utils.save_checkpoint.best,
|
||||
stats[cfg.checkpoint.best_checkpoint_metric],
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
def cli_main(
|
||||
modify_parser: Optional[Callable[[argparse.ArgumentParser], None]] = None
|
||||
) -> None:
|
||||
parser = options.get_training_parser()
|
||||
args = options.parse_args_and_arch(parser, modify_parser=modify_parser)
|
||||
|
||||
cfg = convert_namespace_to_omegaconf(args)
|
||||
|
||||
if args.profile:
|
||||
with torch.cuda.profiler.profile():
|
||||
with torch.autograd.profiler.emit_nvtx():
|
||||
distributed_utils.call_main(cfg, main)
|
||||
else:
|
||||
distributed_utils.call_main(cfg, main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
@@ -0,0 +1,151 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from argparse import Namespace
|
||||
from itertools import chain
|
||||
|
||||
import torch
|
||||
from fairseq import checkpoint_utils, distributed_utils, options, utils
|
||||
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
||||
from fairseq.logging import metrics, progress_bar
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
||||
stream=sys.stdout,
|
||||
)
|
||||
logger = logging.getLogger("fairseq_cli.validate")
|
||||
|
||||
|
||||
def main(cfg: DictConfig, override_args=None):
|
||||
if isinstance(cfg, Namespace):
|
||||
cfg = convert_namespace_to_omegaconf(cfg)
|
||||
|
||||
utils.import_user_module(cfg.common)
|
||||
|
||||
assert (
|
||||
cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None
|
||||
), "Must specify batch size either with --max-tokens or --batch-size"
|
||||
|
||||
use_fp16 = cfg.common.fp16
|
||||
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
|
||||
|
||||
if use_cuda:
|
||||
torch.cuda.set_device(cfg.distributed_training.device_id)
|
||||
|
||||
if cfg.distributed_training.distributed_world_size > 1:
|
||||
data_parallel_world_size = distributed_utils.get_data_parallel_world_size()
|
||||
data_parallel_rank = distributed_utils.get_data_parallel_rank()
|
||||
else:
|
||||
data_parallel_world_size = 1
|
||||
data_parallel_rank = 0
|
||||
|
||||
if override_args is not None:
|
||||
overrides = vars(override_args)
|
||||
overrides.update(eval(getattr(override_args, "model_overrides", "{}")))
|
||||
else:
|
||||
overrides = None
|
||||
|
||||
# Load ensemble
|
||||
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
|
||||
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
|
||||
[cfg.common_eval.path],
|
||||
arg_overrides=overrides,
|
||||
suffix=cfg.checkpoint.checkpoint_suffix,
|
||||
)
|
||||
model = models[0]
|
||||
|
||||
# Move models to GPU
|
||||
for model in models:
|
||||
if use_fp16:
|
||||
model.half()
|
||||
if use_cuda:
|
||||
model.cuda()
|
||||
|
||||
# Print args
|
||||
logger.info(saved_cfg)
|
||||
|
||||
# Build criterion
|
||||
criterion = task.build_criterion(saved_cfg.criterion)
|
||||
criterion.eval()
|
||||
|
||||
for subset in cfg.dataset.valid_subset.split(","):
|
||||
try:
|
||||
task.load_dataset(subset, combine=False, epoch=1, task_cfg=saved_cfg.task)
|
||||
dataset = task.dataset(subset)
|
||||
except KeyError:
|
||||
raise Exception("Cannot find dataset: " + subset)
|
||||
|
||||
# Initialize data iterator
|
||||
itr = task.get_batch_iterator(
|
||||
dataset=dataset,
|
||||
max_tokens=cfg.dataset.max_tokens,
|
||||
max_sentences=cfg.dataset.batch_size,
|
||||
max_positions=utils.resolve_max_positions(
|
||||
task.max_positions(),
|
||||
*[m.max_positions() for m in models],
|
||||
),
|
||||
ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
|
||||
required_batch_size_multiple=cfg.dataset.required_batch_size_multiple,
|
||||
seed=cfg.common.seed,
|
||||
num_shards=data_parallel_world_size,
|
||||
shard_id=data_parallel_rank,
|
||||
num_workers=cfg.dataset.num_workers,
|
||||
data_buffer_size=cfg.dataset.data_buffer_size,
|
||||
).next_epoch_itr(shuffle=False)
|
||||
progress = progress_bar.progress_bar(
|
||||
itr,
|
||||
log_format=cfg.common.log_format,
|
||||
log_interval=cfg.common.log_interval,
|
||||
prefix=f"valid on '{subset}' subset",
|
||||
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
|
||||
)
|
||||
|
||||
log_outputs = []
|
||||
for i, sample in enumerate(progress):
|
||||
sample = utils.move_to_cuda(sample) if use_cuda else sample
|
||||
_loss, _sample_size, log_output = task.valid_step(sample, model, criterion)
|
||||
progress.log(log_output, step=i)
|
||||
log_outputs.append(log_output)
|
||||
|
||||
if data_parallel_world_size > 1:
|
||||
log_outputs = distributed_utils.all_gather_list(
|
||||
log_outputs,
|
||||
max_size=cfg.common.all_gather_list_size,
|
||||
group=distributed_utils.get_data_parallel_group(),
|
||||
)
|
||||
log_outputs = list(chain.from_iterable(log_outputs))
|
||||
|
||||
with metrics.aggregate() as agg:
|
||||
task.reduce_metrics(log_outputs, criterion)
|
||||
log_output = agg.get_smoothed_values()
|
||||
|
||||
progress.print(log_output, tag=subset, step=i)
|
||||
|
||||
|
||||
def cli_main():
|
||||
parser = options.get_validation_parser()
|
||||
args = options.parse_args_and_arch(parser)
|
||||
|
||||
# only override args that are explicitly given on the command line
|
||||
override_parser = options.get_validation_parser()
|
||||
override_args = options.parse_args_and_arch(
|
||||
override_parser, suppress_defaults=True
|
||||
)
|
||||
|
||||
distributed_utils.call_main(
|
||||
convert_namespace_to_omegaconf(args), main, override_args=override_args
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_main()
|
||||
Reference in New Issue
Block a user