chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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()
|
||||
Reference in New Issue
Block a user