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