def main():
parser = options.get_parser('Generation')
parser.add_argument('--path', metavar='FILE', required=True, action='append',
help='path(s) to model file(s)')
options.add_dataset_args(parser)
options.add_generation_args(parser)
args = parser.parse_args()
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, model_args = utils.load_ensemble_for_inference(args.path, data_dir=args.data)
src_dict, dst_dict = models[0].src_dict, models[0].dst_dict
print('| [{}] dictionary: {} types'.format(model_args.source_lang, len(src_dict)))
print('| [{}] dictionary: {} types'.format(model_args.target_lang, len(dst_dict)))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
translator = SequenceGenerator(
models, beam_size=args.beam, stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
print('| Type the input sentence and press return:')
for src_str in sys.stdin:
src_str = src_str.strip()
src_tokens = tokenizer.Tokenizer.tokenize(src_str, src_dict, add_if_not_exist=False).long()
if use_cuda:
src_tokens = src_tokens.cuda()
translations = translator.generate(Variable(src_tokens.view(1, -1)))
hypos = translations[0]
print('O\t{}'.format(src_str))
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dst_dict,
remove_bpe=args.remove_bpe)
print('H\t{}\t{}'.format(hypo['score'], hypo_str))
print('A\t{}'.format(' '.join(map(str, alignment))))
def _generate(self, opt, src_tokens):
translator = SequenceGenerator(
[self.trainer.get_model()],
self.fairseq_dict,
beam_size=opt.beam,
stop_early=(not opt.no_early_stop),
normalize_scores=(not opt.unnormalized),
len_penalty=opt.lenpen)
translator.cuda()
tokens = src_tokens
translations = translator.generate(
Variable(tokens), Variable(self._positions_for_tokens(tokens)))
results = [t[0] for t in translations]
output_lines = [[] for _ in range(len(results))]
for i in range(len(results)):
output_lines[i] = ' '.join(self.fairseq_dict[idx]
for idx in results[i]['tokens'][:-1])
return output_lines
class FairseqAgent(TorchAgent):
"""Generic wrapper around fairseq for use in ParlAI"""
metrics = {}
@classmethod
def add_cmdline_args(cls, argparser):
"""Add command-line arguments specifically for this agent."""
# first we need to add the general torch agent operations
super(FairseqAgent, cls).add_cmdline_args(argparser)
# let's store any defaults that were overridden
old_defaults = argparser._defaults
if 'clip_norm' not in old_defaults:
# fairseq has a few awful defaults
old_defaults['clip_norm'] = 1.0
if 'optimizer' not in old_defaults:
old_defaults['optimizer'] = 'adam'
old_defaults['adam_betas'] = '(0.9,0.98)'
agent = argparser.add_argument_group('Fairseq Arguments')
agent.add_argument('--fp16',
default=False,
type='bool',
help='Use fp16 training')
agent.add_argument(
'--fp16-init-scale',
default=2**7,
type=int,
help='default FP16 loss scale',
)
agent.add_argument(
'--seed',
default=1,
type=int,
metavar='N',
help='pseudo random number generator seed',
)
agent.add_argument(
'--skip-generation',
default=False,
type='bool',
metavar='BOOL',
help=
'Skips test time beam search. Much faster if you only need PPL',
)
# Check subargs for generation, optimizers, criterions, archs, etc
options.add_generation_args(argparser)
options.add_optimization_args(argparser)
options.add_checkpoint_args(argparser)
# restore any user set defaults that fairseq possibly overrode
argparser.set_defaults(**old_defaults)
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "optimizer"):
optimizer = known_args.optimizer
opt_group = argparser.add_argument_group(
'{} optimizer arguments'.format(optimizer))
optim.OPTIMIZER_REGISTRY[optimizer].add_args(opt_group)
if hasattr(known_args, "lr_scheduler"):
lr_scheduler = known_args.lr_scheduler
lr_group = argparser.add_argument_group(
'{} scheduler arguments'.format(lr_scheduler))
optim.lr_scheduler.LR_SCHEDULER_REGISTRY[lr_scheduler].add_args(
lr_group)
# We need to find out the fairseq model-specific options, so grab the
# architecture stuff and look up its options
arch_group = options.add_model_args(argparser)
# Fairseq marks the arch flag as required, but it may be specified
# by a saved model cache, so we do some weird stuff to undo that
for a in arch_group._actions:
if a.dest == "arch":
a.required = False
a.default = None
break
# once again restore any user-set defaults
argparser.set_defaults(**old_defaults)
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "arch") and known_args.arch is not None:
arch = known_args.arch
arch_group = argparser.add_argument_group(
"{} architecture arguments".format(arch))
models.ARCH_MODEL_REGISTRY[arch].add_args(arch_group)
if hasattr(known_args, "criterion"):
crit_group = argparser.add_argument_group(
'{} criterion arguments'.format(known_args.criterion))
criterions.CRITERION_REGISTRY[known_args.criterion].add_args(
crit_group)
# one last time, restore any user set defaults
argparser.set_defaults(**old_defaults)
@staticmethod
def dictionary_class():
# Force use of the Fairseq Dictionary
return _FairseqDictionary
def __init__(self, opt, shared=None):
# In general use a basic TorchAgent wherever possible
super().__init__(opt, shared)
if not shared:
# this is not a shared instance of this class, so do full initialization
# check early if we're going to be loading the model from a checkpoint
model_file_exists = self.opt.get('model_file') and os.path.isfile(
self.opt['model_file'])
# fairseq expects options to be in argparse format, instead of a dict
# We also need to do some argument postprocessing and whatnot
# We'll skip pretrained embeddings if we're going to override them with
# a model checkpoint anyway
self.args, self.opt = _fairseq_opt_wrapper(opt, model_file_exists)
# seed the RNG
torch.manual_seed(self.args.seed)
# Just some identifying info
self.id = "fairseq:{}".format(self.args.arch)
# We need a placeholder task for fairseq
self.task = _ParlaiTask(self.dict)
# meters for keeping track of loss, ppl, etc.
self.meters = defaultdict(AverageMeter)
# actually construct the model and generator
self.model = self.build_model()
# Construct the generator and scorer
self.generator = SequenceGenerator(
[self.model],
tgt_dict=self.dict,
beam_size=self.args.beam,
stop_early=(not self.args.no_early_stop),
normalize_scores=(not self.args.unnormalized),
len_penalty=self.args.lenpen,
unk_penalty=self.args.unkpen,
sampling=self.args.sampling,
sampling_topk=self.args.sampling_topk,
sampling_temperature=self.args.sampling_temperature,
)
self.scorer = SequenceScorer([self.model], self.dict)
# set up the grader and the trainer
self.criterion = criterions.build_criterion(self.args, self.task)
# TODO: we might choose to add a --no-fp16 opt in the future to
# explicitly disable fp16 instead
if not self.args.fp16 and torch.cuda.get_device_capability(
0)[0] >= 7:
print("Heads up: using --fp16 could be a lot faster!")
if self.use_cuda:
self.trainer = trainer.Trainer(self.args, self.task,
self.model, self.criterion,
None)
self.trainer._build_optimizer()
else:
self.trainer = None
# if the model already existed, let's preload it and the trainer
if model_file_exists:
print('Loading existing model params from ' +
self.opt['model_file'])
self.load(self.opt.get('model_file'))
# move things to the GPU if possible
if self.use_cuda:
self.model = self.model.cuda()
self.generator = self.generator.cuda()
else:
self.model = shared['model']
self.trainer = shared['trainer']
self.generator = shared['generator']
self.dict = shared['dict']
self.args = shared['args']
self.meters = shared['meters']
# Start things off clean
self.reset()
def _check_opts_unchanged(self, saved_opts, current_opts):
"""Verify that critical options do not differ in command line vs saved model"""
for k in NON_OVERRIDABLE_ARGS:
if k not in saved_opts or k not in current_opts:
# if it's not an option needed by this fairseq model, don't stress
continue
if saved_opts[k] != current_opts[k]:
raise ValueError(
'{} cannot be overridden when --model-file is specified'.
format(k))
def build_model(self):
"""
Construct the actual Fairseq model. Default implementation is to use
Fairseq's arch builder, but this method may be overridden to build custom
models.
"""
model_class = models.ARCH_MODEL_REGISTRY[self.args.arch]
model = model_class.build_model(self.args, self.task)
if self.args.embedding_type != 'random':
self._copy_embeddings(model.encoder.embed_tokens.weight,
self.args.embedding_type)
return model
def share(self):
shared = super().share()
shared['model'] = self.model
shared['trainer'] = self.trainer
shared['generator'] = self.generator
shared['dict'] = self.dict
shared['args'] = self.args
shared['meters'] = self.meters
return shared
def save(self, path):
"""Save using fairseq's checkpointing."""
if not path:
return
self.trainer.save_checkpoint(path, {'opt': self.opt, 'epoch': 0})
# Parlai expects options to also be saved
with open(path + '.opt', 'w') as handle:
# overridden options shouldn't be stored, only the main ones
if 'override' in self.opt:
del self.opt['override']
json.dump(self.opt, handle)
# force save the dict
self.dict.save(path + '.dict', sort=False)
def load(self, path):
"""Load using fairseq's checkpointing."""
if self.trainer:
old_options = self.trainer.load_checkpoint(
path, self.args.reset_optimizer)
self._check_opts_unchanged(old_options, self.opt)
else:
load_model_state(path, self.model)
def shutdown(self):
if not hasattr(self, 'trainer'):
# looks like this is a "fake" model that isn't actually used for batch_act.
# we don't need to save this one.
return
super().shutdown()
def reset(self):
"""Reset observation and episode_done."""
super().reset()
self.reset_metrics()
def is_valid(self, obs):
"""Override from TorchAgent.
Check if an observation has no tokens in it."""
return len(obs.get('text_vec', [])) > 0
def batchify(self, obs_batch):
"""
Override parent batchify to set requirements for fairseq.
Fairseq depends on sorted batch inputs for a call to rnn.pad_packed_sequence.
Fairseq models cannot handle zero length sentences
"""
return super().batchify(obs_batch, sort=True)
def _update_metrics(self, metrics, sample):
if metrics is None:
# probably got an overflow in fp16 mode. don't count this sample
return
bsz = len(sample['target'])
ntok = sample['ntokens']
ssize = metrics['sample_size']
for k, v in metrics.items():
if k in {'ntokens', 'nsentences', 'sample_size'}:
# don't need these
continue
elif k == "nll_loss":
# nll loss is always normalized by ntokens
self.meters[k].update(v, ntok)
elif k == "loss":
# loss is explicitly normalized by passed up sample size
self.meters[k].update(v, ssize)
else:
# assume everything else it's averaged over bsz
self.meters[k].update(v, bsz)
def train_step(self, batch):
"""Process batch of inputs and targets and train on them.
:param batch: parlai.core.torch_agent.Batch, contains tensorized
version of observations.
"""
if batch.text_vec is None:
return
self.is_training = True
sample = self._make_sample(batch)
self.model.train()
metrics = self.trainer.train_step([sample])
self._update_metrics(metrics, sample)
def eval_step(self, batch):
"""Process batch of inputs.
If the batch includes labels, calculate validation metrics as well.
If --skip-generation is not set, return a prediction for each input.
:param batch: parlai.core.torch_agent.Batch, contains tensorized
version of observations.
"""
if batch.text_vec is None:
return
self.is_training = False
samples = self._make_sample(batch)
self.model.eval()
if batch.label_vec is not None and self.trainer is not None:
# Interactive mode won't have a gold label
metrics = self.trainer.valid_step(samples)
self._update_metrics(metrics, samples)
# Output placeholders
reranked_cands = None
generated_output = None
# Grade each of the candidate sequences
if batch.candidate_vecs is not None:
bsz = len(batch.text_vec)
reranked_cands = []
# score the candidates for each item in the batch separately, so that
# we can support variable number of candidates
for i in range(bsz):
cands = batch.candidate_vecs[i]
if not cands:
reranked_cands.append(None)
continue
ncand = len(cands)
# repeat the input many times
xs = batch.text_vec[i].unsqueeze(0).expand(ncand, -1)
# some models crash if there's leading padding on every example
xs = xs[:, :batch.text_lengths[i]]
# and appropriately pack the outputs
ys, _ = padded_tensor(cands, self.NULL_IDX, self.use_cuda)
s = self._make_sample(xs=xs, ys=ys)
# perform the actual grading, extract the scores
scored = list(
self.scorer.score_batched_itr([s], cuda=self.use_cuda))
scores = [s[3][0]['score'].item() for s in scored]
# intentional hanging comma here; argsort returns a list
ranked, = argsort(scores, batch.candidates[i], descending=True)
reranked_cands.append(ranked)
# Next generate freely to create our response
if not self.args.skip_generation:
generated_output = self._generate(samples)
elif reranked_cands:
# we're skiping generation, but we're also grading candidates
# so output the highest ranked candidate
# In the case of zero candidates, we don't have something to rank,
# so we may need to pass on that None
generated_output = [
ranked and ranked[0] or None for ranked in reranked_cands
]
else:
# no output at all
pass
return Output(generated_output, reranked_cands)
def _generate(self, samples):
no_prev_token = {
k: v
for k, v in samples['net_input'].items()
if k != 'prev_output_tokens'
}
gens = self.generator.generate(no_prev_token, maxlen=64)
bsz = samples['net_input']['src_tokens'].size(0)
responses = []
for i in range(bsz):
beams = gens[i]
selected = max(beams, key=lambda x: x["score"])
tokens = selected["tokens"]
start = 0
end = -1
for i, t in enumerate(tokens):
t = t.item()
if t == self.dict.bos_index:
# don't include <s> token
start = i + 1
continue
if t == self.dict.eos_index:
# stop (and don't include) </s> token
end = i
break
responses.append(self.dict.vec2txt(tokens[start:end]))
return responses
def report(self):
"""Return metrics calculated by the model."""
# if we haven't initialized yet, just return a dummy object
if not hasattr(self, "trainer"):
return {}
output = {k: v.avg for k, v in self.meters.items()}
if "nll_loss" in self.meters:
# special case, we used sentence averaging so ppl comes from nll_loss
output["ppl"] = np.exp2(self.meters["nll_loss"].avg)
else:
# normal case, just use loss
output["ppl"] = np.exp2(self.meters["loss"].avg)
# Fairseq trainer metrics we'll pass up the way
trainer_metrics = {"ups", "wps", "gnorm", "clip"}
if self.is_training:
for k in trainer_metrics:
output[k] = self.trainer.meters[k].avg
# for display purposes
output = {k: round_sigfigs(v, 4) for k, v in output.items()}
return output
def reset_metrics(self):
"""Reset metrics calculated by the model back to zero."""
if not hasattr(self, "trainer"):
# We haven't set up the trainer yet, so we don't have any metrics
return
# We need to reset everything
self.meters.clear()
if self.trainer:
for k in self.trainer.meters:
self.trainer.meters[k].reset()
def receive_metrics(self, metrics_dict):
"""Update lr scheduler with validation loss."""
# TODO: this should be smarter
self.trainer.lr_step(-1, metrics_dict["loss"])
# Helper functions
def _seq_length(self, xs):
"""Compute length of the sequence (non-padded size)."""
return xs.ne(self.dict.pad_index).long().sum(dim=-1)
def _right_shifted_ys(self, ys):
"""Replace first token with EOS and shift remaining tokens right 1."""
result = torch.LongTensor(ys.size())
result[:, 0] = self.dict.eos_index
result[:, 1:] = ys[:, :-1]
return result
def _make_sample(self, batch=None, xs=None, ys=None):
"""Generate a sample object that Fairseq expects."""
# add extra info to samples
if batch is None and xs is None:
raise ValueError("Must supply either batch or xs")
if batch is None and ys is None:
raise ValueError("Must supply either batch or ys")
if xs is None:
xs = batch.text_vec
if ys is None:
ys = batch.label_vec
repadded = convert_padding_direction(xs,
self.dict.pad(),
right_to_left=True)
sample = {}
sample["id"] = torch.arange(len(xs) - 1)
sample["net_input"] = {
"src_tokens": repadded,
"src_lengths": self._seq_length(xs),
}
if ys is not None:
sample["target"] = ys
sample["ntokens"] = sum(self._seq_length(ys)).item()
sample["net_input"]["prev_output_tokens"] = self._right_shifted_ys(
ys)
return sample
def main():
parser = options.get_parser('Generation')
parser.add_argument('--path', metavar='FILE', required=True, action='append',
help='path(s) to model file(s)')
dataset_args = options.add_dataset_args(parser)
dataset_args.add_argument('--batch-size', default=32, type=int, metavar='N',
help='batch size')
dataset_args.add_argument('--gen-subset', default='test', metavar='SPLIT',
help='data subset to generate (train, valid, test)')
options.add_generation_args(parser)
args = parser.parse_args()
if args.no_progress_bar and args.log_format is None:
args.log_format = 'none'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset
if args.replace_unk is None:
dataset = data.load_dataset(args.data, [args.gen_subset], args.source_lang, args.target_lang)
else:
dataset = data.load_raw_text_dataset(args.data, [args.gen_subset], args.source_lang, args.target_lang)
if args.source_lang is None or args.target_lang is None:
# record inferred languages in args
args.source_lang, args.target_lang = dataset.src, dataset.dst
# Load ensemble
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict, dataset.dst_dict)
print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
translator = SequenceGenerator(
models, beam_size=args.beam, stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
# Generate and compute BLEU score
scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
max_positions = min(model.max_encoder_positions() for model in models)
itr = dataset.eval_dataloader(
args.gen_subset, max_sentences=args.batch_size, max_positions=max_positions,
skip_invalid_size_inputs_valid_test=args.skip_invalid_size_inputs_valid_test)
num_sentences = 0
with utils.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
cuda_device=0 if use_cuda else None, timer=gen_timer)
for sample_id, src_tokens, target_tokens, hypos in translations:
# Process input and ground truth
target_tokens = target_tokens.int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = dataset.splits[args.gen_subset].src.get_original_text(sample_id)
target_str = dataset.splits[args.gen_subset].dst.get_original_text(sample_id)
else:
src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)
target_str = dataset.dst_dict.string(target_tokens, args.remove_bpe, escape_unk=True)
if not args.quiet:
print('S-{}\t{}'.format(sample_id, src_str))
print('T-{}\t{}'.format(sample_id, target_str))
# Process top predictions
for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dataset.dst_dict,
remove_bpe=args.remove_bpe)
if not args.quiet:
print('H-{}\t{}\t{}'.format(sample_id, hypo['score'], hypo_str))
print('A-{}\t{}'.format(sample_id, ' '.join(map(str, alignment))))
# Score only the top hypothesis
if i == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tokenizer.Tokenizer.tokenize(target_str,
dataset.dst_dict,
add_if_not_exist=True)
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(src_tokens.size(0))
t.log({'wps': round(wps_meter.avg)})
num_sentences += 1
print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} tokens/s)'.format(
num_sentences, gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))
print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(model_paths, task)
# Set dictionaries
src_dict = task.source_dictionary
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
minlen=args.min_len,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str='O\t{}'.format(src_str),
hypos=[],
alignments=[],
)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
result.hypos.append('H\t{}\t{}'.format(hypo['score'], hypo_str))
result.alignments.append('A\t{}'.format(' '.join(
map(lambda x: str(utils.item(x)), alignment))))
return result
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
translations = translator.generate(
tokens,
lengths,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
return [
make_result(batch.srcs[i], t) for i, t in enumerate(translations)
]
if args.buffer_size > 1:
print('| Sentence buffer size:', args.buffer_size)
print('| Type the input sentence and press return:')
for inputs in buffered_read(args.buffer_size):
indices = []
results = []
for batch, batch_indices in make_batches(inputs, args, src_dict,
models[0].max_positions()):
indices.extend(batch_indices)
results += process_batch(batch)
for i in np.argsort(indices):
result = results[i]
print(result.src_str)
for hypo, align in zip(result.hypos, result.alignments):
print(hypo)
print(align)
def model_fn(model_dir):
model_name = 'checkpoint_best.pt'
model_path = os.path.join(model_dir, model_name)
logger.info('Loading the model')
with open(model_path, 'rb') as f:
model_info = torch.load(f, map_location=torch.device('cpu'))
# Will be overidden by the model_info['args'] - need to keep for pre-trained models
parser = options.get_generation_parser(interactive=True)
# get args for FairSeq by converting the hyperparameters as if they were command-line arguments
argv_copy = copy.deepcopy(sys.argv)
# remove the modifications we did in the command-line arguments
sys.argv[1:] = ['--path', model_path, model_dir]
args = options.parse_args_and_arch(parser)
# restore previous command-line args
sys.argv = argv_copy
saved_args = model_info['args']
for key, value in vars(saved_args).items():
setattr(args, key, value)
args.data = [model_dir]
print(args)
# Setup task, e.g., translation
task = tasks.setup_task(args)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info('Current device: {}'.format(device))
model_paths = [os.path.join(model_dir, model_name)]
models, model_args = utils.load_ensemble_for_inference(
model_paths, task, model_arg_overrides={})
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if device.type == 'cuda':
translator.cuda()
# 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(args.replace_unk)
align_dict = utils.load_align_dict(None)
max_positions = utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models])
return dict(
translator=translator,
task=task,
max_positions=max_positions,
align_dict=align_dict,
tgt_dict=tgt_dict,
args=args,
device=device,
)
def main():
parser = options.get_parser('Generation')
parser.add_argument('--path',
metavar='FILE',
required=True,
action='append',
help='path(s) to model file(s)')
dataset_args = options.add_dataset_args(parser)
dataset_args.add_argument('--batch-size',
default=32,
type=int,
metavar='N',
help='batch size')
dataset_args.add_argument(
'--gen-subset',
default='test',
metavar='SPLIT',
help='data subset to generate (train, valid, test)')
dataset_args.add_argument('--num-shards',
default=1,
type=int,
metavar='N',
help='shard generation over N shards')
dataset_args.add_argument(
'--shard-id',
default=0,
type=int,
metavar='ID',
help='id of the shard to generate (id < num_shards)')
options.add_generation_args(parser)
args = parser.parse_args()
if args.no_progress_bar and args.log_format is None:
args.log_format = 'none'
# print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
if hasattr(torch, 'set_grad_enabled'):
torch.set_grad_enabled(False)
# Load dataset
if args.replace_unk is None:
dataset = data.load_dataset(args.data, [args.gen_subset],
args.source_lang, args.target_lang)
else:
dataset = data.load_raw_text_dataset(args.data, [args.gen_subset],
args.source_lang,
args.target_lang)
if args.source_lang is None or args.target_lang is None:
# record inferred languages in args
args.source_lang, args.target_lang = dataset.src, dataset.dst
# Load ensemble
# print('| loading model(s) from {}'.format(', '.join(args.path)))
models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict,
dataset.dst_dict)
# print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
# print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
# print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
translator = SequenceGenerator(models,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
# Generate and compute BLEU score
#scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
max_positions = min(model.max_encoder_positions() for model in models)
itr = dataset.eval_dataloader(args.gen_subset,
max_sentences=args.batch_size,
max_positions=max_positions,
skip_invalid_size_inputs_valid_test=args.
skip_invalid_size_inputs_valid_test)
if args.num_shards > 1:
if args.shard_id < 0 or args.shard_id >= args.num_shards:
raise ValueError('--shard-id must be between 0 and num_shards')
itr = data.sharded_iterator(itr, args.num_shards, args.shard_id)
num_sentences = 0
with utils.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t,
maxlen_a=args.max_len_a,
maxlen_b=args.max_len_b,
cuda_device=0 if use_cuda else None,
timer=gen_timer)
correct = 0
total = 0
for sample_id, src_tokens, target_tokens, hypos in translations:
# Process input and ground truth
target_tokens = target_tokens.int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = dataset.splits[
args.gen_subset].src.get_original_text(sample_id)
target_str = dataset.splits[
args.gen_subset].dst.get_original_text(sample_id)
else:
src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)
target_str = dataset.dst_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
# if not args.quiet:
# print('S-{}\t{}'.format(sample_id, src_str))
# print('T-{}\t{}'.format(sample_id, target_str))
total += 1
# Process top predictions
for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dataset.dst_dict,
remove_bpe=args.remove_bpe)
#if src_str == 'walk around right thrice after jump opposite left twice':
# import pdb; pdb.set_trace()
# if not args.quiet:
# print('H-{}\t{}\t{}'.format(sample_id, hypo['score'], hypo_str))
# print('A-{}\t{}'.format(sample_id, ' '.join(map(str, alignment))))
# Score only the top hypothesis
if i == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tokenizer.Tokenizer.tokenize(
target_str,
dataset.dst_dict,
add_if_not_exist=True)
#scorer.add(target_tokens, hypo_tokens)
mat = ''
for row in hypo['attention']:
for column in row:
mat += str(column) + '\t'
mat += '\n'
tar = '/' + target_str
tra = '=' + str(target_str == hypo_str)
to_write.write(mat)
to_write.write(src_str)
to_write.write('\n')
to_write.write(hypo_str)
to_write.write('\n')
to_write.write(tar)
to_write.write('\n')
to_write.write(tra)
to_write.write('\n')
to_write.write('-----------')
to_write.write('\n')
if hypo_str == target_str:
correct += 1
wps_meter.update(src_tokens.size(0))
t.log({'wps': round(wps_meter.avg)})
num_sentences += 1
print('| Correct : {} - Total: {}. Accuracy: {:.5f}'.format(
correct, total, correct / total))
def main():
parser = options.get_parser('Generation')
parser.add_argument('--path', metavar='FILE', required=True, action='append',
help='path(s) to model file(s)')
dataset_args = options.add_dataset_args(parser)
dataset_args.add_argument('-i', '--interactive', action='store_true',
help='generate translations in interactive mode')
dataset_args.add_argument('--batch-size', default=32, type=int, metavar='N',
help='batch size')
dataset_args.add_argument('--gen-subset', default='test', metavar='SPLIT',
help='data subset to generate (train, valid, test)')
options.add_generation_args(parser)
args = parser.parse_args()
print(args)
if args.no_progress_bar:
progress_bar.enabled = False
use_cuda = torch.cuda.is_available() and not args.cpu
# Load model and dataset
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, dataset = utils.load_ensemble_for_inference(args.path, args.data)
print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
if not args.interactive:
print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
# Optimize model for generation
for model in models:
model.make_generation_fast_(not args.no_beamable_mm)
# Initialize generator
translator = SequenceGenerator(models, dataset.dst_dict, beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen)
align_dict = {}
if args.unk_replace_dict != '':
assert args.interactive, "Unkown words replacing requires access to original source and is only" \
"supported in interactive mode"
with open(args.unk_replace_dict, 'r') as f:
for line in f:
l = line.split()
align_dict[l[0]] = l[1]
def replace_unk(hypo_str, align_str, src, unk):
hypo_tokens = hypo_str.split()
src_tokens = tokenizer.tokenize_line(src)
align_idx = [int(i) for i in align_str.split()]
for i, ht in enumerate(hypo_tokens):
if ht == unk:
src_token = src_tokens[align_idx[i]]
if src_token in align_dict:
hypo_tokens[i] = align_dict[src_token]
else:
hypo_tokens[i] = src_token
return ' '.join(hypo_tokens)
if use_cuda:
translator.cuda()
bpe_symbol = '@@ ' if args.remove_bpe else None
def display_hypotheses(id, src, orig, ref, hypos):
id_str = '' if id is None else '-{}'.format(id)
src_str = to_sentence(dataset.src_dict, src, bpe_symbol)
print('S{}\t{}'.format(id_str, src_str))
if orig is not None:
print('O{}\t{}'.format(id_str, orig.strip()))
if ref is not None:
print('T{}\t{}'.format(id_str, to_sentence(dataset.dst_dict, ref, bpe_symbol, ref_unk=True)))
for hypo in hypos:
hypo_str = to_sentence(dataset.dst_dict, hypo['tokens'], bpe_symbol)
align_str = ' '.join(map(str, hypo['alignment']))
if args.unk_replace_dict != '':
hypo_str = replace_unk(hypo_str, align_str, orig, unk_symbol(dataset.dst_dict))
print('H{}\t{}\t{}'.format(
id_str, hypo['score'], hypo_str))
print('A{}\t{}'.format(id_str, align_str))
if args.interactive:
for line in sys.stdin:
tokens = tokenizer.Tokenizer.tokenize(line, dataset.src_dict, add_if_not_exist=False).long()
start = dataset.src_dict.pad() + 1
positions = torch.arange(start, start + len(tokens)).type_as(tokens)
if use_cuda:
positions = positions.cuda()
tokens = tokens.cuda()
translations = translator.generate(Variable(tokens.view(1, -1)), Variable(positions.view(1, -1)))
hypos = translations[0]
display_hypotheses(None, tokens, line, None, hypos[:min(len(hypos), args.nbest)])
else:
def maybe_remove_bpe(tokens):
"""Helper for removing BPE symbols from a hypothesis."""
if not args.remove_bpe:
return tokens
assert (tokens == dataset.dst_dict.pad()).sum() == 0
hypo_minus_bpe = to_sentence(dataset.dst_dict, tokens, bpe_symbol)
return tokenizer.Tokenizer.tokenize(hypo_minus_bpe, dataset.dst_dict, add_if_not_exist=True)
# Generate and compute BLEU score
scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
itr = dataset.dataloader(args.gen_subset, batch_size=args.batch_size, max_positions=args.max_positions)
num_sentences = 0
with progress_bar(itr, smoothing=0, leave=False) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
cuda_device=0 if use_cuda else None, timer=gen_timer)
for id, src, ref, hypos in translations:
ref = ref.int().cpu()
top_hypo = hypos[0]['tokens'].int().cpu()
scorer.add(maybe_remove_bpe(ref), maybe_remove_bpe(top_hypo))
display_hypotheses(id, src, None, ref, hypos[:min(len(hypos), args.nbest)])
wps_meter.update(src.size(0))
t.set_postfix(wps='{:5d}'.format(round(wps_meter.avg)))
num_sentences += 1
print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} tokens/s)'.format(
num_sentences, gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))
print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
def main():
parser = argparse.ArgumentParser(description='Batch translate')
#parser.add_argument('--no-progress-bar', action='store_true', help='disable progress bar')
parser.add_argument('--model', metavar='FILE', required=True, action='append',
help='path(s) to model file(s)')
parser.add_argument('--dictdir', metavar='DIR', required=True, help='directory of dictionary files')
parser.add_argument('--batch-size', default=32, type=int, metavar='N',
help='batch size')
parser.add_argument('--beam', default=5, type=int, metavar='N',
help='beam size (default: 5)')
#parser.add_argument('--nbest', default=1, type=int, metavar='N',
# help='number of hypotheses to output')
#parser.add_argument('--remove-bpe', nargs='?', const='@@ ', default=None,
# help='remove BPE tokens before scoring')
parser.add_argument('--no-early-stop', action='store_true',
help=('continue searching even after finalizing k=beam '
'hypotheses; this is more correct, but increases '
'generation time by 50%%'))
#parser.add_argument('--unnormalized', action='store_true',
# help='compare unnormalized hypothesis scores')
parser.add_argument('--cpu', action='store_true', help='generate on CPU')
parser.add_argument('--no-beamable-mm', action='store_true',
help='don\'t use BeamableMM in attention layers')
parser.add_argument('--lenpen', default=1, type=float,
help='length penalty: <1.0 favors shorter, >1.0 favors longer sentences')
parser.add_argument('--unkpen', default=0, type=float,
help='unknown word penalty: <0 produces more unks, >0 produces fewer')
#parser.add_argument('--replace-unk', nargs='?', const=True, default=None,
# help='perform unknown replacement (optionally with alignment dictionary)')
#parser.add_argument('--quiet', action='store_true',
# help='Only print final scores')
parser.add_argument('input', metavar='INPUT', help='Input file')
args = parser.parse_args()
# required by progress bar
args.log_format = None
USE_CUDA = not args.cpu and torch.cuda.is_available()
print('Loading model...', file=sys.stderr)
models, _ = utils.load_ensemble_for_inference(args.model, data_dir=args.dictdir)
src_dic = models[0].src_dict
dst_dic = models[0].dst_dict
for model in models:
model.make_generation_fast_(beamable_mm_beam_size=args.beam)
translator = SequenceGenerator(
models, beam_size=args.beam, stop_early=(not args.no_early_stop),
len_penalty=args.lenpen, unk_penalty=args.unkpen)
if USE_CUDA:
translator.cuda()
max_positions = min(model.max_encoder_positions() for model in models)
print('Loading input data...', file=sys.stderr)
raw_dataset = indexed_dataset.IndexedRawTextDataset(args.input, src_dic)
dataset = fairseq.data.LanguageDatasets(SRC_LANG, TGT_LANG, src_dic, dst_dic)
dataset.splits['test'] = fairseq.data.LanguagePairDataset(
raw_dataset, raw_dataset, pad_idx=dataset.src_dict.pad(),
eos_idx=dataset.src_dict.eos())
# itr = dataset.eval_dataloader(
# 'test', max_sentences=args.batch_size, max_positions=max_positions)
itr = dataset.eval_dataloader('test', max_sentences=args.batch_size)
itr = utils.build_progress_bar(args, itr)
#out = []
for sample_id, src_tokens, _, hypos in translator.generate_batched_itr(
itr, cuda_device=0 if USE_CUDA else None):
src_str = dataset.src_dict.string(src_tokens, '@@ ')
#print(src_str)
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypos[0]['tokens'].int().cpu(),
src_str=src_str,
alignment=hypos[0]['alignment'].int().cpu(),
align_dict=None,
dst_dict=dataset.dst_dict,
remove_bpe='@@ ')
#out.append((sample_id, hypo_str))
print('{}\t{}'.format(sample_id, hypo_str), flush=True)
class FairseqAgent(Agent):
"""Agent which takes an input sequence and produces an output sequence.
For more information, see Convolutional Sequence to Sequence Learning
`(Gehring et al. 2017) <https://arxiv.org/abs/1705.03122>`_.
"""
@staticmethod
def add_cmdline_args(argparser):
"""Add command-line arguments specifically for this agent."""
DictionaryAgent.add_cmdline_args(argparser)
agent = argparser.add_argument_group('Fairseq Arguments')
agent.add_argument(
'-tr', '--truncate',
type=int, default=-1,
help='truncate input & output lengths to speed up training (may '
'reduce accuracy). This fixes all input and output to have a '
'maximum length. This reduces the total amount of padding in '
'the batches.')
agent.add_argument(
'--max-positions',
default=1024,
type=int,
metavar='N',
help='max number of tokens in the sequence')
agent.add_argument(
'--seed',
default=1,
type=int,
metavar='N',
help='pseudo random number generator seed')
options.add_optimization_args(argparser)
options.add_generation_args(argparser)
options.add_model_args(argparser)
def __init__(self, opt, shared=None):
# initialize defaults first
super().__init__(opt, shared)
if not shared:
# this is not a shared instance of this class, so do full
# initialization. if shared is set, only set up shared members.
saved_state = None
if opt.get('model_file') and os.path.isfile(opt['model_file']):
# load model parameters if available
print('Loading existing model params from ' +
opt['model_file'])
new_opt, saved_state = self.load(opt['model_file'])
# override options with stored ones
opt = self._override_opt(new_opt)
self.args = OptWrapper(opt)
self.fairseq_dict = _make_fairseq_dict(DictionaryAgent(opt))
self.id = 'Fairseq'
self.truncate = opt['truncate'] if opt['truncate'] > 0 else None
self.EOS = self.fairseq_dict[self.fairseq_dict.eos()]
self.EOS_TENSOR = (torch.LongTensor(1, 1)
.fill_(self.fairseq_dict.eos()))
self.NULL_IDX = self.fairseq_dict.pad()
encoder = fconv.Encoder(
self.fairseq_dict,
embed_dim=self.args.encoder_embed_dim,
convolutions=eval(self.args.encoder_layers),
dropout=self.args.dropout,
max_positions=self.args.max_positions)
decoder = fconv.Decoder(
self.fairseq_dict,
embed_dim=self.args.decoder_embed_dim,
convolutions=eval(self.args.decoder_layers),
out_embed_dim=self.args.decoder_out_embed_dim,
attention=eval(self.args.decoder_attention),
dropout=self.args.dropout,
max_positions=self.args.max_positions)
self.model = fconv.FConvModel(encoder, decoder)
# from fairseq's build_criterion()
if self.args.label_smoothing > 0:
self.criterion = criterions.LabelSmoothedCrossEntropyCriterion(
self.args.label_smoothing, self.NULL_IDX)
else:
self.criterion = criterions.CrossEntropyCriterion(
self.NULL_IDX)
self.trainer = MultiprocessingTrainer(self.args, self.model, self.criterion)
if saved_state is not None:
self.set_states(saved_state)
self.reset()
def _override_opt(self, new_opt):
"""Set overridable opts from loaded opt file.
Print out each added key and each overriden key.
Only override args specific to the model.
"""
model_args = {
'arch',
'encoder-embed-dim',
'encoder-layers',
'decoder-embed-dim',
'decoder-layers',
'decoder-out-embed-dim',
'decoder-attention',
}
for k, v in new_opt.items():
if k not in model_args:
# skip non-model args
continue
if k not in self.opt:
print('Adding new option [ {k}: {v} ]'.format(k=k, v=v))
elif self.opt[k] != v:
print('Overriding option [ {k}: {old} => {v}]'.format(
k=k, old=self.opt[k], v=v))
self.opt[k] = v
return self.opt
def reset(self):
"""Reset observation and episode_done."""
self.observation = None
self.episode_done = True
def observe(self, observation):
# shallow copy observation (deep copy can be expensive)
observation = observation.copy()
if not self.episode_done:
# if the last example wasn't the end of an episode, then we need to
# recall what was said in that example
prev_dialogue = self.observation['text']
observation['text'] = prev_dialogue + '\n' + observation['text']
self.observation = observation
self.episode_done = observation['episode_done']
return observation
def act(self):
# call batch_act with this batch of one
return self.batch_act([self.observation])[0]
def batch_act(self, observations):
bsz = len(observations)
# initialize a table of replies with this agent's id
batch_reply = [{'id': self.getID()} for _ in range(bsz)]
# convert the observations into batches of inputs and targets
# valid_inds tells us the indices of all valid examples
# e.g. for input [{}, {'text': 'hello'}, {}, {}], valid_inds is [1]
# since the other three elements had no 'text' field
# also, split observations into sub-batches based on number of gpus
obs_split = np.array_split(observations, self.trainer.num_replicas)
samples = [self.batchify(obs) for obs in obs_split]
samples = [s for s in samples if s[0] is not None]
any_valid = any(len(s[0]) > 0 for s in samples)
if not any_valid:
# no valid examples, just return the empty responses we set up
return batch_reply
# produce predictions if testing; otherwise, train
has_targets = any(s[1] is not None for s in samples)
if not has_targets:
offset = 0
for s in samples:
xs = s[0]
valid_inds = s[2]
predictions = self._generate(self.args, xs)
for i in range(len(predictions)):
# map the predictions back to non-empty examples in the batch
batch_reply[valid_inds[i] + offset]['text'] = predictions[i]
if i == 0:
print('prediction:', predictions[i])
offset += len(valid_inds)
else:
loss = self._train(samples)
batch_reply[0]['metrics'] = {}
for k, v in loss.items():
batch_reply[0]['metrics'][k] = v * bsz
if k == 'loss':
batch_reply[0]['metrics']['perplexity'] = 2 ** v * bsz
return batch_reply
def parse(self, string):
return [self.fairseq_dict.index(word) for word in string.split(' ')]
def batchify(self, observations):
"""Convert a list of observations into input & target tensors."""
# valid examples
exs = [ex for ex in observations if 'text' in ex]
# the indices of the valid (non-empty) tensors
valid_inds = [i for i, ex in enumerate(observations) if 'text' in ex]
# set up the input tensors
batchsize = len(exs)
if batchsize == 0:
return None, None, None
# tokenize the text
parsed_x = [deque(maxlen=self.truncate) for _ in exs]
for dq, ex in zip(parsed_x, exs):
dq += self.parse(ex['text'])
# parsed = [self.parse(ex['text']) for ex in exs]
max_x_len = max((len(x) for x in parsed_x))
for x in parsed_x:
# left pad with zeros
x.extendleft([self.fairseq_dict.pad()] * (max_x_len - len(x)))
xs = torch.LongTensor(parsed_x)
# set up the target tensors
ys = None
if 'labels' in exs[0]:
# randomly select one of the labels to update on, if multiple
labels = [random.choice(ex.get('labels', [''])) for ex in exs]
parsed_y = [deque(maxlen=self.truncate) for _ in labels]
for dq, y in zip(parsed_y, labels):
dq.extendleft(reversed(self.parse(y)))
for y in parsed_y:
y.append(self.fairseq_dict.eos())
# append EOS to each label
max_y_len = max(len(y) for y in parsed_y)
for y in parsed_y:
y += [self.fairseq_dict.pad()] * (max_y_len - len(y))
ys = torch.LongTensor(parsed_y)
return xs, ys, valid_inds
def _positions_for_tokens(self, tokens):
size = tokens.size()
not_pad = tokens.ne(self.fairseq_dict.pad()).long()
new_pos = tokens.new(size).fill_(self.fairseq_dict.pad())
new_pos += not_pad
for i in range(1, size[1]):
new_pos[:, i] += new_pos[:, i-1] - 1
return new_pos
def _right_shifted_ys(self, ys):
result = torch.LongTensor(ys.size())
result[:, 0] = self.fairseq_dict.index(self.EOS)
result[:, 1:] = ys[:, :-1]
return result
def _generate(self, opt, src_tokens):
if not hasattr(self, 'translator'):
self.translator = SequenceGenerator(
[self.trainer.get_model()],
beam_size=opt.beam,
stop_early=(not opt.no_early_stop),
normalize_scores=(not opt.unnormalized),
len_penalty=opt.lenpen)
self.translator.cuda()
tokens = src_tokens.cuda(async=True)
token_pos = Variable(self._positions_for_tokens(tokens).cuda())
translations = self.translator.generate(Variable(tokens), token_pos)
results = [t[0] for t in translations]
output_lines = [[] for _ in range(len(results))]
for i in range(len(results)):
output_lines[i] = ' '.join(self.fairseq_dict[idx]
for idx in results[i]['tokens'][:-1])
return output_lines
def _train(self, samples):
"""Update the model using the targets."""
for i, sample in enumerate(samples):
# add extra info to samples
sample = {
'src_tokens': sample[0],
'input_tokens': self._right_shifted_ys(sample[1]),
'target': sample[1],
'id': None
}
sample['ntokens'] = sum(len(t) for t in sample['target'])
sample['src_positions'] = self._positions_for_tokens(
sample['src_tokens'])
sample['input_positions'] = self._positions_for_tokens(
sample['input_tokens'])
samples[i] = sample
return self.trainer.train_step(samples)
def save(self, path=None):
path = self.opt.get('model_file', None) if path is None else path
if path and hasattr(self, 'trainer'):
model = {}
model['state_dict'] = self.trainer.get_model().state_dict()
model['opt'] = self.opt
with open(path, 'wb') as write:
torch.save(model, write)
def shutdown(self):
"""Save the state of the model when shutdown."""
path = self.opt.get('model_file', None)
if path is not None:
self.save(path + '.shutdown_state')
super().shutdown()
def load(self, path):
"""Return opt and model states."""
with open(path, 'rb') as read:
model = torch.load(read)
return model['opt'], model['state_dict']
def set_states(self, state_dict):
"""Set the state dict of the model from saved states."""
self.trainer.get_model().load_state_dict(state_dict)
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str='O\t{}'.format(src_str),
hypos=[],
pos_scores=[],
alignments=[],
)
# Process top predictions
def forward(self, model, sample, reduce=True):
# sample mode
#print('!!!RL loss.')
model.eval()
# src_dict = self.task.source_dictionary
tgt_dict = self.task.target_dictionary
eos_idx = self.task.target_dictionary.eos()
sample_beam = 1
translator = SequenceGenerator([model], tgt_dict=tgt_dict, sampling=self.args.multinomial_sample_train,
beam_size=sample_beam, minlen=1)
translator.cuda()
ct = 0
translations = []
s = utils.move_to_cuda(sample)
input = s['net_input']
max_len = 200
with torch.no_grad():
hypos = translator.generate(
input['src_tokens'],
input['src_lengths'],
beam_size=sample_beam,
maxlen=max_len,
)
for i, id in enumerate(s['id'].data):
src = input['src_tokens'].data[i, :]
# remove padding from ref
ref = utils.strip_pad(s['target'].data[i, :], tgt_dict.pad()) if s['target'] is not None else None
translations.append((id, src, ref, hypos[i]))
ct += 1
# print("sample batch size:", ct)
model.train()
# MLE loss
mle_net_output = model(**sample['net_input'])
mle_lprobs = model.get_normalized_probs(mle_net_output, log_probs=True)
mle_lprobs = mle_lprobs.view(-1, mle_lprobs.size(-1))
mle_target = model.get_targets(sample, mle_net_output).view(-1)
mle_loss = F.nll_loss(mle_lprobs, mle_target, size_average=False,
ignore_index=self.padding_idx, reduce=reduce)
mle_tokens = sample['ntokens']
avg_mle_loss = mle_loss / mle_tokens
print('avg_mle_loss:', avg_mle_loss)
# RL loss
batch_rl_loss = 0
batch_tokens = 0
id = 0
result = []
for sample_id, src_tokens, tgt_tokens, hypos in translations:
# calculate bleu
id += 1
hypo = hypos[0] # only extract the first hypo (beam1 or sample1)
trans_tokens = hypo['tokens']
if self.args.delta_reward:
reward = self.compute_sentence_bleu(tgt_tokens.cpu(), trans_tokens.cpu()).cuda()
else:
reward = self.compute_sentence_total_bleu(tgt_tokens.cpu(), trans_tokens.cpu()).cuda()
result.append((id, reward.item(), tgt_tokens.size(0), trans_tokens.size(0)))
# one_sample loss calculation
tgt_input_tokens = trans_tokens.new(trans_tokens.shape).fill_(0)
assert trans_tokens[-1] == eos_idx
tgt_input_tokens[0] = eos_idx
tgt_input_tokens[1:] = trans_tokens[:-1]
train_sample = {
'net_input': {
'src_tokens': src_tokens.view(1, -1),
'src_lengths': torch.LongTensor(src_tokens.numel()).view(1, -1),
'prev_output_tokens': tgt_input_tokens.view(1, -1),
},
'target': trans_tokens.view(1, -1)
}
train_sample = utils.move_to_cuda(train_sample)
net_output = model(**train_sample['net_input'])
lprobs = model.get_normalized_probs(net_output, log_probs=True)
lprobs = lprobs.view(-1, lprobs.size(-1))
target = model.get_targets(train_sample, net_output).view(-1, 1)
non_pad_mask = target.ne(tgt_dict.pad())
lprob = -lprobs.gather(dim=-1, index=target)[non_pad_mask]
rl_loss = torch.sum(lprob * reward) # one sample loss
ntokens = len(train_sample['target'])
batch_tokens += ntokens
batch_rl_loss += rl_loss
avg_rl_loss = batch_rl_loss / batch_tokens
with open('./results/reward/v0_m'+str(self.args.mle_weight)+'r'+str(self.args.rl_weight)+'_lr'+str(self.args.lr)+'_r.csv','a', newline='') as csv_file:
csv_writer = csv.writer(csv_file)
for r in result:
csv_writer.writerow(r)
print('avg_rl_loss:', avg_rl_loss)
if self.args.mle_weight:
assert self.args.rl_weight
total_loss = self.args.mle_weight * avg_mle_loss + self.args.rl_weight * avg_rl_loss
total_tokens = batch_tokens + mle_tokens
else:
total_loss = avg_rl_loss
total_tokens = batch_tokens
logging_output = {
'loss': utils.item(total_loss.data),
'ntokens': total_tokens,
'sample_size': total_tokens,
}
print('total: ',total_loss)
with open('./results/loss/v0_m'+str(self.args.mle_weight)+'r'+str(self.args.rl_weight)+'_lr'+str(self.args.lr)+'_l.csv','a', newline='') as csv_file:
csv_writer = csv.writer(csv_file)
csv_writer.writerow((avg_mle_loss.item(), avg_rl_loss.item(), total_loss.item(), total_tokens))
return total_loss, total_tokens, logging_output
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(
model_paths, task, model_arg_overrides=eval(args.model_overrides))
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str='O\t{}'.format(src_str),
hypos=[],
pos_scores=[],
alignments=[],
)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
result.hypos.append('H\t{}\t{}'.format(hypo['score'], hypo_str))
result.pos_scores.append('P\t{}'.format(' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))))
result.alignments.append('A\t{}'.format(' '.join(
map(lambda x: str(utils.item(x)), alignment))) if args.
print_alignment else None)
return result
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
encoder_input = {'src_tokens': tokens, 'src_lengths': lengths}
translations = translator.generate(
encoder_input,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
return [
make_result(batch.srcs[i], t) for i, t in enumerate(translations)
]
def translate_one(source, args, task, max_positions):
for batch, batch_indices in make_batches([source], args, task,
max_positions):
result = process_batch(batch)[0]
for hypo, pos_scores, align in zip(result.hypos, result.pos_scores,
result.alignments):
hypo = hypo.split("\t")[-1]
break
return hypo
def process_unrolled(steps, args, task, max_positions):
is_long = False
pairs = []
collect_outcomes = dict()
for source, target in steps:
for i, token in enumerate(source):
if "*" in token: source[i] = collect_outcomes[token]
source2 = []
for token in source:
if type(token) != list:
source2.append(token)
else:
source2.extend(token)
source = source2
source = " ".join(source)
predicted_target = translate_one(source, args, task, max_positions)
pairs.append((source, predicted_target))
if "*" in target[0]:
collect_outcomes[target[0]] = predicted_target
else:
return predicted_target, pairs
logging.error("Unrolling stopped prematurely.")
max_positions = utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models])
print(
translate_one("copy G17 P19 Z18 E1 S13 J15 A3 A3", args, task,
max_positions))
data = []
with open(args.src) as f:
sample = {"unrolled": [], "original": ("", "")}
for line in f:
[line_type, source, target] = line.split("\t")
source = source.strip().split()
target = target.strip().split()
if line_type == "unrolled":
sample[line_type].append((source, target))
else:
sample[line_type] = (source, target)
data.append(sample)
sample = {"unrolled": [], "original": ("", "")}
predictions_equal = []
scores_per_input_length = defaultdict(list)
scores_per_target_length = defaultdict(list)
all_pairs = []
random.shuffle(data)
for sample in tqdm(data[:1000]):
unrolled_predicted, is_long, pairs = process_unrolled(
sample["unrolled"], args, task, max_positions)
source, target = sample["original"]
target = " ".join(target)
source = " ".join(source)
original_predicted = translate_one(source, args, task, max_positions)
local_score = original_predicted == unrolled_predicted
all_pairs.append((pairs, local_score))
predictions_equal.append(local_score)
scores_per_input_length[len(source)].append(local_score)
scores_per_target_length[len(target)].append(local_score)
print(f"Localism {np.mean(predictions_equal)}")
with open("trace_localism.txt", 'w') as f:
for pairs, score in all_pairs:
f.write("------------------------------\n")
for s, t in pairs:
f.write("{} -> {}\n".format(s, t))
f.write("{}".format(score))
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(
model_paths, task, model_arg_overrides=eval(args.model_overrides))
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
# Initialize fluency scorer (and language model)
fluency_scorer = FluencyScorer(args.lang_model_path, args.lang_model_data)
def make_result(src_str, hypos, tgt_str='', iteration=0):
results = []
# compute fluency score for source string
# the source string itself is an entry
result0 = Correction()
result0.iteration = iteration
result0.src_str = result0.hypo_str = src_str
fluency_scores = fluency_scorer.score_sentence(src_str).item()
result0.fluency_scores = fluency_scores
result0.fluency_scores_str = "Fluency Score: {:0.4f}".format(
fluency_scores)
results.append(result0)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
result = Correction()
result.iteration = iteration + 1
result.src_str = src_str
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
# result.hypos.append('H\t{}\t{}'.format(hypo['score'], hypo_str))
result.hypo_str = hypo_str
result.hypo_score = result.hypo_score_str = hypo['score']
result.pos_scores_str = 'P\t{}'.format(' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
)))
result.alignments_str = ('A\t{}'.format(' '.join(
map(lambda x: str(utils.item(x)), alignment)))
if args.print_alignment else None)
# compute GLEU if target is provided
if tgt_str:
gleu_calculator = GLEU(args.n)
gleu_calculator.load_text_sources([src_str])
gleu_calculator.load_text_references([[tgt_str]])
gleu_scores = gleu_calculator.run_iterations(
num_iterations=args.iter,
hypothesis=[hypo_str],
per_sent=args.sent)
gleu_score = [g for g in gleu_scores][0][0] * 100
result.gleu_scores = gleu_score
result.gleu_scores_str = 'GLEU {:2.2f}'.format(gleu_score)
else:
result.gleu_scores_str = 'GLEU N/A (no target was provided. use format "source sentence|target setence" to provide a target/reference)'
# compute fluency score
fluency_scores = fluency_scorer.score_sentence(hypo_str).item()
result.fluency_scores = fluency_scores
result.fluency_scores_str = "Fluency Score: {:0.4f}".format(
fluency_scores)
results.append(result)
return results
def process_batch(batch, tgts, iteration):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
encoder_input = {'src_tokens': tokens, 'src_lengths': lengths}
translations = translator.generate(
encoder_input,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
return [
make_result(batch.srcs[i], t, tgts[i], iteration)
for i, t in enumerate(translations)
]
max_positions = utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models])
if not args.server:
listen_to_stdin(args, max_positions, task, process_batch)
else:
listen_to_web(args, max_positions, task, process_batch)
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(model_paths, task)
# Set dictionaries
src_dict = task.source_dictionary
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
minlen=args.min_len,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str='O\t{}'.format(src_str),
hypos=[],
alignments=[],
)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
result.hypos.append('H\t{}\t{}'.format(hypo['score'], hypo_str))
result.alignments.append('A\t{}'.format(' '.join(
map(lambda x: str(utils.item(x)), alignment))))
return result
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
translations = translator.generate(
Variable(tokens),
Variable(lengths),
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
# print("translations",translations)
# print(batch.srcs[0])
return [
make_result(batch.srcs[i], t) for i, t in enumerate(translations)
]
# indices.extend(batch_indices)
# results += process_batch(batch)
def forward(c):
batch, batch_indices = next(
make_batches([args.target], args, src_dict,
models[0].max_positions()))
translations = translator.generate(
Variable(batch.tokens),
Variable(batch.lengths),
maxlen=int(args.max_len_a * batch.tokens.size(1) + args.max_len_b),
prefix_tokens=c,
)
return translations, [
tgt_dict[x] for x in range(translations.shape[0])
]
# np.max(translations),tgt_dict.symbols[np.argsort(-translations)[0]]
batch_2, batch_indices_2 = next(
make_batches([args.distractor], args, src_dict,
models[0].max_positions()))
# return [tgt_dict.symbols[x] for x in np.argsort(-translations)[:10]]
translations_2 = translator.generate(
Variable(batch_2.tokens),
Variable(batch_2.lengths),
maxlen=int(args.max_len_a * batch_2.tokens.size(1) +
args.max_len_b),
prefix_tokens=c,
)
# l0 = translations - scipy.misc.logsumexp([translations, translations_2],axis=0)
s1_0 = translations - scipy.misc.logsumexp(
[translations, translations_2], axis=0)
# print([tgt_dict.symbols[x] for x in np.argsort(-translations)].index("ces"))
# print([tgt_dict.symbols[x] for x in np.argsort(-translations_2)].index("ces"))
# print([tgt_dict.symbols[x] for x in np.argsort(-s1_0)].index("ces"))
# raise Exception
# l0 = translations - (scipy.misc.logsumexp([translations, translations_2],axis=0))
# s1_1 = translations_2 - scipy.misc.logsumexp([translations, translations_2],axis=0)
# print([tgt_dict.symbols[x] for x in np.argsort(-translations)[:10]] )
# print([tgt_dict.symbols[x] for x in np.argsort(-translations_2)[:10]] )
# print(np.argsort(-s1_0)[:10])
# print([tgt_dict.symbols[x] for x in np.argsort(-(0.001*np.exp(s1_0)+0.999*np.exp(translations)))[:10]] )
# print([tgt_dict.symbols[x] for x in np.argsort(s1_1)[:10]] )
return tgt_dict.symbols[np.argsort(-(1.0 * np.exp(s1_0) +
0.0 * np.exp(translations)))[0]]
return forward(args.context_sentence)
for i in range(10):
print("RESULT:", args.context_sentence)
next_word = forward(args.context_sentence)
if next_word == '</s>':
break
args.context_sentence += [next_word]
return " ".join(args.context_sentence)
# print(l0)
# print()
print([tgt_dict.symbols[x] for x in np.argsort(-s1_0)[:10]], "s1_0")
print([tgt_dict.symbols[x] for x in np.argsort(-s1_1)[:10]], "s1_1")
raise Exception
if args.buffer_size > 1:
print('| Sentence buffer size:', args.buffer_size)
print('| Type the input sentence and press return:')
for inputs in buffered_read(args.buffer_size):
# print("inputs, interactive",inputs)
# for inputs in ["my name is John."]:
indices = []
results = []
for batch, batch_indices in make_batches(inputs, args, src_dict,
models[0].max_positions()):
print(batch.tokens, batch.lengths)
raise Exception
indices.extend(batch_indices)
results += process_batch(batch)
# raise Exception
for i in np.argsort(indices):
result = results[i]
print(result.src_str)
for hypo, align in zip(result.hypos, result.alignments):
print(hypo)
print(align)
class FairseqAgent(TorchAgent):
"""Generic wrapper around fairseq for use in ParlAI"""
metrics = {}
# TODO: merge with TorchAgent.add_cmdline_args
@staticmethod
def add_cmdline_args(argparser):
"""Add command-line arguments specifically for this agent."""
# first we need to add the general torch agent operations
TorchAgent.add_cmdline_args(argparser)
agent = argparser.add_argument_group('Fairseq Arguments')
agent.add_argument('--seed',
default=1,
type=int,
metavar='N',
help='pseudo random number generator seed')
agent.add_argument(
'--skip-generation',
default=False,
type=bool,
metavar='BOOL',
help=
'Skips test time beam search. Much faster if you only need PPL',
)
# Dictionary construction stuff. Using the subclass in case we end up
# needing any fairseq specific things
_FairseqDictionary.add_cmdline_args(argparser)
# Optimization and learning rate schedule specific arguments
options.add_optimization_args(argparser)
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "optimizer"):
optimizer = known_args.optimizer
opt_group = argparser.add_argument_group(
'{} optimizer arguments'.format(optimizer))
optim.OPTIMIZER_REGISTRY[optimizer].add_args(opt_group)
if hasattr(known_args, "lr_scheduler"):
lr_scheduler = known_args.lr_scheduler
lr_group = argparser.add_argument_group(
'{} scheduler arguments'.format(lr_scheduler))
optim.lr_scheduler.LR_SCHEDULER_REGISTRY[lr_scheduler].add_args(
lr_group)
# Generation arguments
options.add_generation_args(argparser)
# We need to find out the fairseq model-specific options, so grab the
# architecture stuff and look up its options
arch_group = options.add_model_args(argparser)
# Fairseq marks the arch flag as required, but it may be specified
# by a saved model cache, so we do some weird stuff to undo that
for a in arch_group._actions:
if a.dest == "arch":
a.required = False
a.default = None
break
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "arch") and known_args.arch is not None:
arch = known_args.arch
arch_group = argparser.add_argument_group(
"{} architecture arguments".format(arch))
models.ARCH_MODEL_REGISTRY[arch].add_args(arch_group)
# Override a few defaults from within fairseq to more sensible defaults
argparser.set_defaults(clip_norm=0.1, adam_betas="(0.9,0.98)")
def __init__(self, opt, shared=None):
# In general use a basic TorchAgent wherever possible
super().__init__(opt, shared)
if not shared:
# this is not a shared instance of this class, so do full initialization
# fairseq expects options to be in argparse format, instead of a dict
# We also need to do some argument postprocessing and whatnot
self.args, self.opt = _fairseq_opt_wrapper(opt)
# seed the RNG
torch.manual_seed(self.args.seed)
# Just some identifying info
self.id = "fairseq:{}".format(self.args.arch)
# construct dictionaries for parlai frontend and fairseq backend
self.dict = _FairseqDictionary(self.opt)
# We need a placeholder task for fairseq
self.task = _ParlaiTask(self.dict)
# actually construct the model and generator
model_class = models.ARCH_MODEL_REGISTRY[self.args.arch]
self.model = model_class.build_model(self.args, self.task)
self.generator = SequenceGenerator(
[self.model],
tgt_dict=self.dict,
beam_size=self.args.beam,
stop_early=(not self.args.no_early_stop),
normalize_scores=(not self.args.unnormalized),
len_penalty=self.args.lenpen,
)
# set up the grader and the trainer
# TODO: maybe support label smoothing here
self.criterion = CrossEntropyCriterion(self.args, self.task)
if self.args.fp16:
self.trainer = fp16_trainer.FP16Trainer(
self.args, self.task, self.model, self.criterion)
else:
# TODO: we might choose to add a --no-fp16 opt in the future to
# explicitly disable fp16 instead
if torch.cuda.get_device_capability(0)[0] >= 7:
print("Heads up: using --fp16 could be a lot faster!")
self.trainer = trainer.Trainer(self.args, self.task,
self.model, self.criterion)
# if the model already existed, let's preload it and the trainer
if self.opt.get('model_file') and os.path.isfile(
self.opt['model_file']):
print('Loading existing model params from ' +
self.opt['model_file'])
self.load(self.opt.get('model_file'))
# move things to the GPU if possible
if self.use_cuda:
self.model = self.model.cuda()
self.generator = self.generator.cuda()
# Start things off clean
self.reset()
def _check_opts_unchanged(self, saved_opts, current_opts):
"""Verify that critical options do not differ in command line vs saved model"""
for k in NON_OVERRIDABLE_ARGS:
if k not in saved_opts or k not in current_opts:
# if it's not an option needed by this fairseq model, don't stress
continue
if saved_opts[k] != current_opts[k]:
raise ValueError(
'{} cannot be overridden when --model-file is specified'.
format(k))
def save(self, path):
"""Save using fairseq's checkpointing."""
if not path:
return
self.trainer.save_checkpoint(path, {'opt': self.opt, 'epoch': 0})
# Parlai expects options to also be saved
with open(path + ".opt", 'wb') as handle:
# overridden options shouldn't be stored, only the main ones
if 'override' in self.opt:
del self.opt['override']
pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)
def load(self, path):
"""Load using fairseq's checkpointing."""
old_options = self.trainer.load_checkpoint(path)
self._check_opts_unchanged(old_options, self.opt)
def shutdown(self):
if not hasattr(self, 'trainer'):
# looks like this is a "fake" model that isn't actually used for batch_act.
# we don't need to save this one.
return
super().shutdown()
def reset(self):
"""Reset observation and episode_done."""
super().reset()
self.reset_metrics()
def batch_act(self, observations):
bsz = len(observations)
# initialize a table of replies with this agent's id
batch_reply = [{"id": self.getID()} for _ in range(bsz)]
# torchagent boilerplate
self.is_training = any(["labels" in obs for obs in observations])
vec_obs = [self.vectorize(obs) for obs in observations]
xs, _, ys, _, valid_inds = self.map_valid(vec_obs)
if xs is None:
return batch_reply
# here begins fairseq specific stuff
samples = self._make_sample(xs, ys)
if self.is_training:
self.model.train()
self.trainer.train_step(samples)
else:
# grade the evaluation label
self.model.eval()
if ys is not None:
# Interactive mode won't have a gold label
self.trainer.valid_step(samples)
# Grade each of the candidate sequences
# TODO: grade everything in observations[i]['label_candidates']
# Next generate freely to create our response
if self.args.skip_generation:
# skip the generation step
for i in valid_inds:
batch_reply[i]["text"] = ""
else:
# actually do the generation
for i, response in zip(valid_inds, self._generate(samples)):
batch_reply[i]["text"] = response
return batch_reply
def _generate(self, samples):
src_tokens = samples["net_input"]["src_tokens"]
src_lengths = samples["net_input"]["src_lengths"]
gens = self.generator.generate(src_tokens, src_lengths, maxlen=64)
responses = []
for i in range(len(src_tokens)):
beams = gens[i]
selected = max(beams, key=lambda x: x["score"])
response = []
for t in selected["tokens"]:
t = t.item()
if t == self.dict.bos_index:
# don't include <s> token
continue
if t == self.dict.eos_index:
# stop (and don't include) </s> token
break
response.append(self.dict[t])
responses.append(" ".join(response))
return responses
def report(self):
# if we haven't initialized yet, just return a dummy object
if not hasattr(self, "trainer"):
return {}
# These are the metrics we'll pass up the way, and their new names
train_metrics = {"train_loss", "ups", "wps", "gnorm", "clip"}
valid_metrics = {"valid_loss"}
metrics = train_metrics if self.is_training else valid_metrics
m = {k: self.trainer.meters[k].avg for k in metrics}
# additionally output perplexity. note that fairseq models use base 2
# in cross_entropy:
# github.com/pytorch/fairseq/blob/master/fairseq/criterions/cross_entropy.py#L55
if "train_loss" in m:
m["train_ppl"] = np.exp2(m["train_loss"])
if "valid_loss" in m:
m["ppl"] = np.exp2(m["valid_loss"])
for k, v in m.items():
# clean up: rounds to sigfigs and converts tensors to floats
m[k] = round_sigfigs(v, 4)
return m
def reset_metrics(self):
if not hasattr(self, "trainer"):
# We haven't initialized the trainer yet, so we don't have any metrics
return
# We need to reset everything
for k in self.trainer.meters:
self.trainer.meters[k].reset()
def receive_metrics(self, metrics_dict):
"""Used to update lr scheduler."""
self.trainer.lr_step(-1, metrics_dict["valid_loss"])
# Helper functions
def _seq_length(self, xs):
"""Computes length of the sequence (non-padded size)"""
return xs.ne(self.dict.pad_index).long().sum(dim=-1)
def _right_shifted_ys(self, ys):
"""Replaces first token with EOS and shifts the remaining tokens right one."""
result = torch.LongTensor(ys.size())
result[:, 0] = self.dict.eos_index
result[:, 1:] = ys[:, :-1]
return result
def _make_sample(self, xs, ys):
"""Generates a sample object that Fairseq expects."""
# add extra info to samples
# TODO: should the right/left padding thing be in torch agent?
repadded = convert_padding_direction(xs,
self.dict.pad(),
right_to_left=True)
sample = {}
sample["net_input"] = {
"src_tokens": repadded,
"src_lengths": self._seq_length(xs),
}
if ys is not None:
sample["target"] = ys
sample["ntokens"] = sum(self._seq_length(ys)).item()
sample["net_input"]["prev_output_tokens"] = self._right_shifted_ys(
ys)
return sample
class FairseqAgent(TorchAgent):
"""Generic wrapper around fairseq for use in ParlAI"""
DEFAULT_OPTIONS = {
"adam_betas": "(0.9,0.98)",
"optimizer": "adam",
"clip_norm": 0.1,
}
metrics = {}
@classmethod
def add_cmdline_args(cls, argparser):
"""Add command-line arguments specifically for this agent."""
# first we need to add the general torch agent operations
TorchAgent.add_cmdline_args(argparser)
agent = argparser.add_argument_group('Fairseq Arguments')
agent.add_argument('--fp16',
default=False,
type=bool,
help='Use fp16 training')
agent.add_argument('--seed',
default=1,
type=int,
metavar='N',
help='pseudo random number generator seed')
agent.add_argument(
'--skip-generation',
default=False,
type=bool,
metavar='BOOL',
help=
'Skips test time beam search. Much faster if you only need PPL',
)
# Dictionary construction stuff. Using the subclass in case we end up
# needing any fairseq specific things
cls.dictionary_class().add_cmdline_args(argparser)
# Check subargs for generation, optimizers, criterions, archs, etc
options.add_generation_args(argparser)
options.add_optimization_args(argparser)
# make sure we set defaults according to the model before parsing
argparser.set_defaults(**cls.DEFAULT_OPTIONS)
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "optimizer"):
optimizer = known_args.optimizer
opt_group = argparser.add_argument_group(
'{} optimizer arguments'.format(optimizer))
optim.OPTIMIZER_REGISTRY[optimizer].add_args(opt_group)
if hasattr(known_args, "lr_scheduler"):
lr_scheduler = known_args.lr_scheduler
lr_group = argparser.add_argument_group(
'{} scheduler arguments'.format(lr_scheduler))
optim.lr_scheduler.LR_SCHEDULER_REGISTRY[lr_scheduler].add_args(
lr_group)
# We need to find out the fairseq model-specific options, so grab the
# architecture stuff and look up its options
arch_group = options.add_model_args(argparser)
# Fairseq marks the arch flag as required, but it may be specified
# by a saved model cache, so we do some weird stuff to undo that
for a in arch_group._actions:
if a.dest == "arch":
a.required = False
a.default = None
break
# make sure we set defaults according to parlai model before parsing
argparser.set_defaults(**cls.DEFAULT_OPTIONS)
known_args = argparser.parse_known_args(nohelp=True)[0]
if hasattr(known_args, "arch") and known_args.arch is not None:
arch = known_args.arch
arch_group = argparser.add_argument_group(
"{} architecture arguments".format(arch))
models.ARCH_MODEL_REGISTRY[arch].add_args(arch_group)
if hasattr(known_args, "criterion"):
crit_group = argparser.add_argument_group(
'{} criterion arguments'.format(known_args.criterion))
criterions.CRITERION_REGISTRY[known_args.criterion].add_args(
crit_group)
# As one final check, let's make sure we set defaults correctly
argparser.set_defaults(**cls.DEFAULT_OPTIONS)
@staticmethod
def dictionary_class():
# Force use of the Fairseq Dictionary
return _FairseqDictionary
def __init__(self, opt, shared=None):
# In general use a basic TorchAgent wherever possible
super().__init__(opt, shared)
if not shared:
# this is not a shared instance of this class, so do full initialization
# check early if we're going to be loading the model from a checkpoint
model_file_exists = (self.opt.get('model_file')
and os.path.isfile(self.opt['model_file']))
# fairseq expects options to be in argparse format, instead of a dict
# We also need to do some argument postprocessing and whatnot
# We'll skip pretrained embeddings if we're going to override them with
# a model checkpoint anyway
self.args, self.opt = _fairseq_opt_wrapper(opt, model_file_exists)
# seed the RNG
torch.manual_seed(self.args.seed)
# Just some identifying info
self.id = "fairseq:{}".format(self.args.arch)
# We need a placeholder task for fairseq
self.task = _ParlaiTask(self.dict)
# actually construct the model and generator
self.model = self.build_model()
# Construct the generator and scorer
self.generator = SequenceGenerator(
[self.model],
tgt_dict=self.dict,
beam_size=self.args.beam,
stop_early=(not self.args.no_early_stop),
normalize_scores=(not self.args.unnormalized),
len_penalty=self.args.lenpen,
unk_penalty=self.args.unkpen,
sampling=self.args.sampling,
sampling_topk=self.args.sampling_topk,
sampling_temperature=self.args.sampling_temperature,
)
self.scorer = SequenceScorer([self.model], self.dict)
# set up the grader and the trainer
self.criterion = criterions.build_criterion(self.args, self.task)
if getattr(self.args, 'fp16', None):
self.trainer = fp16_trainer.FP16Trainer(
self.args, self.task, self.model, self.criterion)
else:
# TODO: we might choose to add a --no-fp16 opt in the future to
# explicitly disable fp16 instead
if torch.cuda.get_device_capability(0)[0] >= 7:
print("Heads up: using --fp16 could be a lot faster!")
self.trainer = trainer.Trainer(self.args, self.task,
self.model, self.criterion)
# if the model already existed, let's preload it and the trainer
if model_file_exists:
print('Loading existing model params from ' +
self.opt['model_file'])
self.load(self.opt.get('model_file'))
# move things to the GPU if possible
if self.use_cuda:
self.model = self.model.cuda()
self.generator = self.generator.cuda()
else:
self.model = shared['model']
self.trainer = shared['trainer']
self.generator = shared['generator']
self.dict = shared['dict']
self.args = shared['args']
# Start things off clean
self.reset()
def _check_opts_unchanged(self, saved_opts, current_opts):
"""Verify that critical options do not differ in command line vs saved model"""
for k in NON_OVERRIDABLE_ARGS:
if k not in saved_opts or k not in current_opts:
# if it's not an option needed by this fairseq model, don't stress
continue
if saved_opts[k] != current_opts[k]:
raise ValueError(
'{} cannot be overridden when --model-file is specified'.
format(k))
def build_model(self):
"""
Construct the actual Fairseq model. Default implementation is to use
Fairseq's arch builder, but this method may be overridden to build custom
models.
"""
model_class = models.ARCH_MODEL_REGISTRY[self.args.arch]
return model_class.build_model(self.args, self.task)
def share(self):
shared = super().share()
shared['model'] = self.model
shared['trainer'] = self.trainer
shared['generator'] = self.generator
shared['dict'] = self.dict
shared['args'] = self.args
return shared
def save(self, path):
"""Save using fairseq's checkpointing."""
if not path:
return
self.trainer.save_checkpoint(path, {'opt': self.opt, 'epoch': 0})
# Parlai expects options to also be saved
with open(path + ".opt", 'wb') as handle:
# overridden options shouldn't be stored, only the main ones
if 'override' in self.opt:
del self.opt['override']
pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)
def load(self, path):
"""Load using fairseq's checkpointing."""
old_options = self.trainer.load_checkpoint(path)
self._check_opts_unchanged(old_options, self.opt)
def shutdown(self):
if not hasattr(self, 'trainer'):
# looks like this is a "fake" model that isn't actually used for batch_act.
# we don't need to save this one.
return
super().shutdown()
def reset(self):
"""Reset observation and episode_done."""
super().reset()
self.reset_metrics()
def batchify(self, obs_batch):
"""
Override parent batchify to set requirements for fairseq.
Fairseq depends on sorted batch inputs for a call to rnn.pad_packed_sequence.
Fairseq models cannot handle zero length sentences
"""
return super().batchify(obs_batch,
sort=True,
is_valid=_is_nonempty_observation)
def train_step(self, batch):
"""Process batch of inputs and targets and train on them.
:param batch: parlai.core.torch_agent.Batch, contains tensorized
version of observations.
"""
if batch.text_vec is None:
return
self.is_training = True
samples = self._make_sample(batch.text_vec, batch.label_vec)
self.model.train()
self.trainer.train_step(samples)
def eval_step(self, batch):
"""Process batch of inputs.
If the batch includes labels, calculate validation metrics as well.
If --skip-generation is not set, return a prediction for each input.
:param batch: parlai.core.torch_agent.Batch, contains tensorized
version of observations.
"""
if batch.text_vec is None:
return
self.is_training = False
samples = self._make_sample(batch.text_vec, batch.label_vec)
self.model.eval()
if batch.label_vec is not None:
# Interactive mode won't have a gold label
self.trainer.valid_step(samples)
# Output placeholders
reranked_cands = None
generated_output = None
# Grade each of the candidate sequences
if batch.candidate_vecs is not None:
bsz = len(batch.text_vec)
reranked_cands = []
# score the candidates for each item in the batch separately, so that
# we can support variable number of candidates
for i in range(bsz):
cands = batch.candidate_vecs[i]
if not cands:
reranked_cands.append(None)
continue
ncand = len(cands)
# repeat the input many times
xs = batch.text_vec[i].unsqueeze(0).expand(ncand, -1)
# some models crash if there's leading padding on every example
xs = xs[:, :batch.text_lengths[i]]
# and appropriately pack the outputs
ys, _ = padded_tensor(cands, self.NULL_IDX, self.use_cuda)
s = self._make_sample(xs, ys)
# perform the actual grading, extract the scores
scored = list(
self.scorer.score_batched_itr([s], cuda=self.use_cuda))
scores = [s[3][0]['score'].item() for s in scored]
# intentional hanging comma here; argsort returns a list
ranked, = argsort(scores, batch.candidates[i], descending=True)
reranked_cands.append(ranked)
# Next generate freely to create our response
if not self.args.skip_generation:
generated_output = self._generate(samples)
elif reranked_cands:
# we're skiping generation, but we're also grading candidates
# so output the highest ranked candidate
# In the case of zero candidates, we don't have something to rank,
# so we may need to pass on that None
generated_output = [
ranked and ranked[0] or None for ranked in reranked_cands
]
else:
# no output at all
pass
return Output(generated_output, reranked_cands)
def _generate(self, samples):
src_tokens = samples["net_input"]["src_tokens"]
src_lengths = samples["net_input"]["src_lengths"]
gens = self.generator.generate(src_tokens, src_lengths, maxlen=64)
responses = []
for i in range(len(src_tokens)):
beams = gens[i]
selected = max(beams, key=lambda x: x["score"])
tokens = selected["tokens"]
start = 0
end = -1
for i, t in enumerate(tokens):
t = t.item()
if t == self.dict.bos_index:
# don't include <s> token
start = i + 1
continue
if t == self.dict.eos_index:
# stop (and don't include) </s> token
end = i
break
responses.append(self.dict.vec2txt(tokens[start:end]))
return responses
def report(self):
"""Return metrics calculated by the model."""
# if we haven't initialized yet, just return a dummy object
if not hasattr(self, "trainer"):
return {}
# These are the metrics we'll pass up the way, and their new names
train_metrics = {"train_loss", "ups", "wps", "gnorm", "clip"}
valid_metrics = {"valid_loss"}
metrics = train_metrics if self.is_training else valid_metrics
m = {k: self.trainer.meters[k].avg for k in metrics}
# additionally output perplexity. note that fairseq models use base 2
# in cross_entropy:
# github.com/pytorch/fairseq/blob/master/fairseq/criterions/cross_entropy.py#L55
if "train_loss" in m:
m["train_ppl"] = np.exp2(m["train_loss"])
if "valid_loss" in m:
m["ppl"] = np.exp2(m["valid_loss"])
for k, v in m.items():
# clean up: rounds to sigfigs and converts tensors to floats
m[k] = round_sigfigs(v, 4)
return m
def reset_metrics(self):
"""Reset metrics calculated by the model back to zero."""
if not hasattr(self, "trainer"):
# We haven't set up the trainer yet, so we don't have any metrics
return
# We need to reset everything
for k in self.trainer.meters:
self.trainer.meters[k].reset()
def receive_metrics(self, metrics_dict):
"""Update lr scheduler with validation loss."""
self.trainer.lr_step(-1, metrics_dict["valid_loss"])
# Helper functions
def _seq_length(self, xs):
"""Compute length of the sequence (non-padded size)."""
return xs.ne(self.dict.pad_index).long().sum(dim=-1)
def _right_shifted_ys(self, ys):
"""Replace first token with EOS and shift remaining tokens right 1."""
result = torch.LongTensor(ys.size())
result[:, 0] = self.dict.eos_index
result[:, 1:] = ys[:, :-1]
return result
def _make_sample(self, xs, ys):
"""Generate a sample object that Fairseq expects."""
# add extra info to samples
# TODO: should the right/left padding thing be in torch agent?
sample = {}
sample["id"] = torch.arange(len(xs) - 1)
sample["net_input"] = {
"src_tokens": xs,
"src_lengths": self._seq_length(xs),
}
if ys is not None:
sample["target"] = ys
sample["ntokens"] = sum(self._seq_length(ys)).item()
sample["net_input"]["prev_output_tokens"] = self._right_shifted_ys(
ys)
return sample
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(
model_paths, task, model_arg_overrides=eval(args.model_overrides))
# Set dictionaries
tgt_dict = copy.deepcopy(task.target_dictionary)
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
if len(args.transformer_mask_heads) > 0:
# Determine which head to prune
to_prune = parse_head_pruning_descriptors(
args.transformer_mask_heads,
reverse_descriptors=args.transformer_mask_all_but_one_head,
n_heads=model.encoder.layers[0].self_attn.num_heads)
print(to_prune)
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
translate_corpus(
translator,
task,
buffer_size=args.buffer_size,
replace_unk=args.replace_unk,
use_cuda=use_cuda,
print_directly=True,
nbest=args.nbest,
remove_bpe=args.remove_bpe,
print_alignment=args.print_alignment,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
)
def main():
parser = options.get_parser('Generation')
parser.add_argument('--path',
metavar='FILE',
required=True,
action='append',
help='path(s) to model file(s)')
options.add_dataset_args(parser)
options.add_generation_args(parser)
args = parser.parse_args()
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, model_args = utils.load_ensemble_for_inference(args.path,
data_dir=args.data)
src_dict, dst_dict = models[0].src_dict, models[0].dst_dict
print('| [{}] dictionary: {} types'.format(model_args.source_lang,
len(src_dict)))
print('| [{}] dictionary: {} types'.format(model_args.target_lang,
len(dst_dict)))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
translator = SequenceGenerator(models,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
print('| Type the input sentence and press return:')
for src_str in sys.stdin:
src_str = src_str.strip()
src_tokens = tokenizer.Tokenizer.tokenize(
src_str, src_dict, add_if_not_exist=False).long()
if use_cuda:
src_tokens = src_tokens.cuda()
translations = translator.generate(Variable(src_tokens.view(1, -1)))
hypos = translations[0]
print('O\t{}'.format(src_str))
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dst_dict,
remove_bpe=args.remove_bpe)
print('H\t{}\t{}'.format(hypo['score'], hypo_str))
print('A\t{}'.format(' '.join(map(str, alignment))))
def main(args):
print(args)
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences, \
'--max-sentences/--batch-size is not supported in interactive mode'
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, model_args = utils.load_ensemble_for_inference(args.path, data_dir=args.data)
src_dict, dst_dict = models[0].src_dict, models[0].dst_dict
print('| [{}] dictionary: {} types'.format(model_args.source_lang, len(src_dict)))
print('| [{}] dictionary: {} types'.format(model_args.target_lang, len(dst_dict)))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
)
# Initialize generator
translator = SequenceGenerator(
models, beam_size=args.beam, stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
print('| Type the input sentence and press return:')
for src_str in sys.stdin:
src_str = src_str.strip()
src_tokens = tokenizer.Tokenizer.tokenize(src_str, src_dict, add_if_not_exist=False).long()
if use_cuda:
src_tokens = src_tokens.cuda()
src_lengths = src_tokens.new([src_tokens.numel()])
translations = translator.generate(
Variable(src_tokens.view(1, -1)),
Variable(src_lengths.view(-1)),
)
hypos = translations[0]
print('O\t{}'.format(src_str))
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dst_dict,
remove_bpe=args.remove_bpe,
)
print('H\t{}\t{}'.format(hypo['score'], hypo_str))
print('A\t{}'.format(' '.join(map(lambda x: str(utils.item(x)), alignment))))
def forward(self, model, sample, reduce=True):
# sample mode
#print('!!!RL loss.')
model.eval()
# src_dict = self.task.source_dictionary
tgt_dict = self.task.target_dictionary
eos_idx = self.task.target_dictionary.eos()
sample_beam = self.args.sample_beam
translator = SequenceGenerator(
[model],
tgt_dict=tgt_dict,
sampling=self.args.multinomial_sample_train,
beam_size=sample_beam,
minlen=1)
translator.cuda()
ct = 0
translations = []
s = utils.move_to_cuda(sample)
input = s['net_input']
max_len = 200
with torch.no_grad():
hypos = translator.generate(
input['src_tokens'],
input['src_lengths'],
beam_size=sample_beam,
maxlen=max_len,
)
for i, id in enumerate(s['id'].data):
src = input['src_tokens'].data[i, :]
# remove padding from ref
ref = utils.strip_pad(
s['target'].data[i, :],
tgt_dict.pad()) if s['target'] is not None else None
translations.append((id, src, ref, hypos[i]))
ct += 1
# print("sample batch size:", ct)
model.train()
# MLE loss
mle_net_output = model(**sample['net_input'])
mle_lprobs = model.get_normalized_probs(mle_net_output, log_probs=True)
mle_lprobs = mle_lprobs.view(-1, mle_lprobs.size(-1))
mle_target = model.get_targets(sample, mle_net_output).view(-1)
mle_loss = F.nll_loss(mle_lprobs,
mle_target,
size_average=False,
ignore_index=self.padding_idx,
reduce=reduce)
mle_tokens = sample['ntokens']
avg_mle_loss = mle_loss / mle_tokens
print('avg_mle_loss:', avg_mle_loss)
# RL loss
batch_rl_loss = 0
batch_tokens = 0
sample_ind = 0
for sample_id, src_tokens, tgt_tokens, hypos in translations:
# calculate bleu
sample_ind += 1
rewards = torch.Tensor(sample_beam).float().cuda()
logprobs = torch.Tensor(sample_beam).float().cuda()
for i in range(sample_beam):
hypo = hypos[i]
trans_tokens = hypo['tokens']
rewards[i] = self.compute_gleu(tgt_tokens.cpu(),
trans_tokens.cpu(),
max_order=self.args.max_order,
gram=self.args.gram).cuda()
# one_sample loss calculation
tgt_input_tokens = trans_tokens.new(
trans_tokens.shape).fill_(0)
assert trans_tokens[-1] == eos_idx
tgt_input_tokens[0] = eos_idx
tgt_input_tokens[1:] = trans_tokens[:-1]
train_sample = {
'net_input': {
'src_tokens':
src_tokens.view(1, -1),
'src_lengths':
torch.LongTensor(src_tokens.numel()).view(1, -1),
'prev_output_tokens':
tgt_input_tokens.view(1, -1),
},
'target': trans_tokens.view(1, -1)
}
train_sample = utils.move_to_cuda(train_sample)
net_output = model(**train_sample['net_input'])
lprobs = model.get_normalized_probs(net_output, log_probs=True)
lprobs = lprobs.view(-1, lprobs.size(-1))
target = model.get_targets(train_sample,
net_output).view(-1, 1)
non_pad_mask = target.ne(tgt_dict.pad())
lprob = -lprobs.gather(dim=-1, index=target)[non_pad_mask]
logprobs[i] = torch.sum(lprob)
ntokens = len(train_sample['target'])
batch_tokens += ntokens
rl_loss = torch.sum(logprobs *
(rewards - rewards.mean())) # one sample loss
batch_rl_loss += rl_loss
avg_rl_loss = batch_rl_loss / batch_tokens
print('avg_rl_loss:', avg_rl_loss)
if self.args.mle_weight:
assert self.args.rl_weight
total_loss = self.args.mle_weight * avg_mle_loss + self.args.rl_weight * avg_rl_loss
total_tokens = batch_tokens + mle_tokens
else:
total_loss = avg_rl_loss
total_tokens = batch_tokens
logging_output = {
'loss': utils.item(total_loss.data),
'ntokens': total_tokens,
'sample_size': total_tokens,
}
print('total: ', total_loss)
return total_loss, total_tokens, logging_output
def main(args):
setup_logger(args)
args.interactive = sys.stdin.isatty(
) and not args.file # Just make the code more understendable
if args.file:
data_descriptor = open(args.file, 'r')
else:
data_descriptor = sys.stdin
if args.interactive:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
if args.buffer_size > 50000:
print(
"WARNING: To prevent memory exhaustion buffer size is set to 50000",
file=sys.stderr)
args.buffer_size = 50000
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args, file=sys.stderr)
use_cuda = torch.cuda.is_available() and not args.cpu
processing_start = time.time()
# Load ensemble
print('| loading model(s) from {}'.format(args.path), file=sys.stderr)
model_paths = args.path.split(':')
models, model_args, src_dict, tgt_dict = load_ensemble_for_inference(
model_paths)
if args.fp16:
for model in models:
model.half()
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(need_attn=args.print_alignment)
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict.get_metadata(),
maxlen=args.max_target_positions,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
minlen=args.min_len,
sampling_temperature=args.sampling_temperature)
if use_cuda:
translator.cuda()
# Load BPE codes file
if args.bpe_codes:
codes = open(args.bpe_codes, 'r')
bpe = BPE(codes)
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str=src_str,
hypos=[],
pos_scores=[],
alignments=[],
)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
hypo_str = tokenizer.Tokenizer.detokenize(hypo_str, 'de').strip()
result.hypos.append((hypo['score'], hypo_str))
result.pos_scores.append('P\t' + ' '.join(
f'{x:.4f}' for x in hypo['positional_scores'].tolist()))
result.alignments.append('A\t' + ' '.join(
str(utils.item(x))
for x in alignment) if args.print_alignment else None)
return result
gen_timer = StopwatchMeter()
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
translation_start = time.time()
gen_timer.start()
translations = translator.generate(
tokens,
lengths,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
gen_timer.stop(sum(len(h[0]['tokens']) for h in translations))
dllogger.log(step='infer',
data={'latency': time.time() - translation_start})
return [
make_result(batch.srcs[i], t) for i, t in enumerate(translations)
]
if args.interactive:
print('| Type the input sentence and press return:')
for inputs in buffered_read(args.buffer_size, data_descriptor):
indices = []
results = []
for batch, batch_indices in make_batches(inputs, args, src_dict,
args.max_positions, bpe):
indices.extend(batch_indices)
results += process_batch(batch)
for i in np.argsort(indices):
result = results[i]
print(result.src_str, file=sys.stderr)
for hypo, pos_scores, align in zip(result.hypos, result.pos_scores,
result.alignments):
print(f'Score {hypo[0]}', file=sys.stderr)
print(hypo[1])
print(pos_scores, file=sys.stderr)
if align is not None:
print(align, file=sys.stderr)
if args.file:
data_descriptor.close()
log_dict = {
'throughput': 1. / gen_timer.avg,
'latency_avg': sum(gen_timer.intervals) / len(gen_timer.intervals),
'latency_p90': gen_timer.p(90),
'latency_p95': gen_timer.p(95),
'latency_p99': gen_timer.p(99),
'total_infernece_time': gen_timer.sum,
'total_run_time': time.time() - processing_start,
}
print('Translation time: {} s'.format(log_dict['total_infernece_time']),
file=sys.stderr)
print('Model throughput (beam {}): {} tokens/s'.format(
args.beam, log_dict['throughput']),
file=sys.stderr)
print(
'Latency:\n\tAverage {:.3f}s\n\tp90 {:.3f}s\n\tp95 {:.3f}s\n\tp99 {:.3f}s'
.format(log_dict['latency_avg'], log_dict['latency_p90'],
log_dict['latency_p95'], log_dict['latency_p99']),
file=sys.stderr)
print('End to end time: {} s'.format(log_dict['total_run_time']),
file=sys.stderr)
dllogger.log(step=(), data=log_dict)
def generate(self, text: list):
"""
text: Takes in a list of sentences as input, don't call by ourself if you don't know what you are doing
"""
resultStorage: list = []
if self.args.buffer_size < 1: # set's buffer size to a min of 1
self.args.buffer_size = 1
# if not number of max tokens and max_sentences is given -> set max_sentences to
if self.args.max_tokens is None and self.args.max_sentences is None:
self.args.max_sentences = 1
assert not self.args.sampling or self.args.nbest == self.args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not self.args.max_sentences or self.args.max_sentences <= self.args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
# print(args)# print arguments
# checks if cuda can be used
use_cuda = torch.cuda.is_available() and not self.args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(self.args) # idk??
# Load ensemble
print('| loading model(s) from {}'.format(
self.args.path)) # useless info
# model_paths = self.args.path.split(':')#useless
model_paths = [self.chkpath]
models, model_args = utils.load_ensemble_for_inference(
model_paths,
task,
model_arg_overrides=eval(self.args.model_overrides)) # load models
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models: # iterate through models
model.make_generation_fast_( # uses make generateion fast method
# uses onliner to check if beamable_mm_beam_size is None or not
beamable_mm_beam_size=None
if self.args.no_beamable_mm else self.args.beam,
need_attn=self.args.
print_alignment, # boolean if a print alignment is needed
)
if self.args.fp16: # half the model????
model.half()
# Initialize generator
translator = SequenceGenerator(
tgt_dict,
beam_size=self.args.beam,
min_len=self.args.min_len,
stop_early=(not self.args.no_early_stop),
sampling=self.args.sampling,
normalize_scores=(not self.args.unnormalized),
len_penalty=self.args.lenpen,
unk_penalty=self.args.unkpen,
sampling_topk=self.args.
sampling_topk, #can't be used in fairseq 0.6.2
temperature=self.args.sampling_temperature,
diverse_beam_strength=self.args.diverse_beam_strength,
diverse_beam_groups=self.args.diverse_beam_groups,
)
if use_cuda:
translator.cuda() # if cuda can be used use it
# 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(self.args.replace_unk)
# Uses named tuple to init values for later and palce src sentence as src_strs
def make_result(src_str, hypos):
result = self.Translation( # create a result Tupple (named tupple result)
src_str='O\t{}'.format(src_str),
hypos=[],
pos_scores=[],
alignments=[],
)
#print(f"Hypos at beginning of make_result: {hypos}")
# Process top predictions
# iterates through top predictions?
for hypo in hypos[:min(len(hypos), self.args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction( # post processes the prediction
# pass tokens of current prediction and convert to int and then make it for cpu compatible?
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str, # gives the input string
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else
None, # checks if alignmemt is needed
align_dict=
align_dict, # passes dict with words for alignment
# gives the target dictionary (decode one???)
tgt_dict=tgt_dict,
remove_bpe=self.args.remove_bpe, # bool idk what for
)
# use the hypons list of the result to save formatted post_process_prediction score
#print(f"Hypo tokens and string in make result after post process prediction:{hypo_tokens}, {hypo_str} ")
result.hypos.append('H\t{}\t{}'.format(hypo['score'],
hypo_str))
result.pos_scores.append(
'P\t{}'.
format( # does the same for the positional score in list format
' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))))
result.alignments.append( # saves the formatted stuff in the alignements section of the result, if print_alignment is not false
'A\t{}'.format(' '.join(
map(lambda x: str(utils.item(x)), alignment)))
if self.args.print_alignment else None)
return result # returns result
def process_batch(batch): # takes in a batch
"""
Processes the batch, don't call directly except you know what you are doing
"""
tokens = batch.tokens # sets tokens to batch tokens
#print(f"Tokens in process batch: {tokens}")
lengths = batch.lengths # sets length to batch lenght
if use_cuda:
tokens = tokens.cuda() # loads tokens on cuda
lengths = lengths.cuda() # loads lengths on cuda
# prepare encoder input with tokens and and src_lengths
encoder_input = {
"net_input": {
'src_tokens': tokens,
'src_lengths': lengths
}
}
#print(f"\t Tokens: \t {tokens}")
# Problem should appear around here!
translations = translator.generate( # generate actual translation from encoder input and maxlen
models,
encoder_input,
maxlen=int(self.args.max_len_a * \
tokens.size(1) + self.args.max_len_b),
)
####
# print("----------------------")
# return a list of results
#print(f"\n > {translations}\n")
# print("----------------------")
return [
make_result(batch.srcs[i], t)
for i, t in enumerate(translations)
]
max_positions = utils.resolve_max_positions( # resolves max positions
task.max_positions(), # how???
*[model.max_positions() for model in models])
if self.args.buffer_size > 1: # checks buffer size
# prints current buffer size
print('| Sentence buffer size:', self.args.buffer_size)
#print('| Type the input sentence and press return:')
for inputs in text: # self.buffered_read(self.args.buffer_size):
# stores indicies of batches (for later structering the answer?)
# print(inputs)
indices = []
results = [] # stores results
# takes user input and generates batch and corresponding ID for iteration
for batch, batch_indices in self.make_batches(
inputs, self.args, task, max_positions):
# adds batch indecies to indicies list
indices.extend(batch_indices)
# results will be returned by process batch
results += process_batch(batch)
for i in np.argsort(indices): # iterates through sorted arrays?
# takes result corresponding to the input batch id
result = results[i]
#print(f"result: {result}")
# print(result.src_str) # print the input string (i.e. the question)
resultStorage.append(result.hypos) # stores result
# print(result.hypos)
# prints other stuff not needed for us
# for hypo, pos_scores, align in zip(result.hypos, result.pos_scores, result.alignments):
# print(hypo)
# print(pos_scores)
# if align is not None:
# print(align)
return resultStorage
def _backtranslation_dataset_helper(
self, remove_eos_from_input_src, remove_eos_from_output_src,
):
tgt_dataset = LanguagePairDataset(
src=self.tgt_dataset,
src_sizes=self.tgt_dataset.sizes,
src_dict=self.tgt_dict,
tgt=None,
tgt_sizes=None,
tgt_dict=None,
)
generator = SequenceGenerator(
models=[self.model],
tgt_dict=self.tgt_dict,
beam_size=2,
unk_penalty=0,
sampling=False,
)
if self.cuda:
generator.cuda()
backtranslation_dataset = BacktranslationDataset(
tgt_dataset=TransformEosDataset(
dataset=tgt_dataset,
eos=self.tgt_dict.eos(),
# remove eos from the input src
remove_eos_from_src=remove_eos_from_input_src,
),
backtranslation_fn=generator.generate,
max_len_a=0,
max_len_b=200,
output_collater=TransformEosDataset(
dataset=tgt_dataset,
eos=self.tgt_dict.eos(),
# if we remove eos from the input src, then we need to add it
# back to the output tgt
append_eos_to_tgt=remove_eos_from_input_src,
remove_eos_from_src=remove_eos_from_output_src,
).collater,
cuda=self.cuda,
)
dataloader = torch.utils.data.DataLoader(
backtranslation_dataset,
batch_size=2,
collate_fn=backtranslation_dataset.collater,
)
backtranslation_batch_result = next(iter(dataloader))
eos, pad, w1, w2 = self.tgt_dict.eos(), self.tgt_dict.pad(), self.w1, self.w2
# Note that we sort by src_lengths and add left padding, so actually
# ids will look like: [1, 0]
expected_src = torch.LongTensor([[w1, w2, w1, eos], [pad, pad, w1, eos]])
if remove_eos_from_output_src:
expected_src = expected_src[:, :-1]
expected_tgt = torch.LongTensor([[w1, w2, eos], [w1, w2, eos]])
generated_src = backtranslation_batch_result["net_input"]["src_tokens"]
tgt_tokens = backtranslation_batch_result["target"]
self.assertTensorEqual(expected_src, generated_src)
self.assertTensorEqual(expected_tgt, tgt_tokens)
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(model_paths, task, model_arg_overrides=eval(args.model_overrides))
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models, tgt_dict, beam_size=args.beam, minlen=args.min_len,
stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen, unk_penalty=args.unkpen,
sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
result = Translation(
src_str='O\t{}'.format(src_str),
hypos=[],
pos_scores=[],
alignments=[],
)
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu() if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
result.hypos.append('H\t{}\t{}'.format(hypo['score'], hypo_str))
result.pos_scores.append('P\t{}'.format(
' '.join(map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))
))
result.alignments.append(
'A\t{}'.format(' '.join(map(lambda x: str(utils.item(x)), alignment)))
if args.print_alignment else None
)
return result
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
translations = translator.generate(
tokens,
lengths,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
return [make_result(batch.srcs[i], t) for i, t in enumerate(translations)]
max_positions = utils.resolve_max_positions(
task.max_positions(),
*[model.max_positions() for model in models]
)
if args.buffer_size > 1:
print('| Sentence buffer size:', args.buffer_size)
print('| Type the input sentence and press return:')
for inputs in buffered_read(args.buffer_size):
indices = []
results = []
for batch, batch_indices in make_batches(inputs, args, task, max_positions):
indices.extend(batch_indices)
results += process_batch(batch)
for i in np.argsort(indices):
result = results[i]
print(result.src_str)
for hypo, pos_scores, align in zip(result.hypos, result.pos_scores, result.alignments):
print(hypo)
print(pos_scores)
if align is not None:
print(align)
def main(args):
if args.buffer_size < 1:
args.buffer_size = 1
if args.max_tokens is None and args.max_sentences is None:
args.max_sentences = 1
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert not args.max_sentences or args.max_sentences <= args.buffer_size, \
'--max-sentences/--batch-size cannot be larger than --buffer-size'
# print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Setup task, e.g., translation
task = tasks.setup_task(args)
# Load ensemble
# print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, model_args = utils.load_ensemble_for_inference(
model_paths, task, model_arg_overrides=eval(args.model_overrides))
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models,
tgt_dict,
beam_size=args.beam,
minlen=args.min_len,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
sampling=args.sampling,
sampling_topk=args.sampling_topk,
sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups,
diverse_beam_strength=args.diverse_beam_strength,
)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
def make_result(src_str, hypos):
hypo_strs = []
# Process top predictions
for hypo in hypos[:min(len(hypos), args.nbest)]:
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
hypo_strs.append(hypo_str)
return hypo_strs
def process_batch(batch):
tokens = batch.tokens
lengths = batch.lengths
if use_cuda:
tokens = tokens.cuda()
lengths = lengths.cuda()
encoder_input = {'src_tokens': tokens, 'src_lengths': lengths}
translations = translator.generate(
encoder_input,
maxlen=int(args.max_len_a * tokens.size(1) + args.max_len_b),
)
return [
make_result(batch.srcs[i], t) for i, t in enumerate(translations)
]
max_positions = utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models])
bichig_vocab = " ᠠᠢᠡᠨᠭᠤᠦᠬᠷᠳᠯᠪᠶᠰᠮᠣᠴᠲᠵᠥᠩ᠋ᠸᠱᠫᠧᠺᠽᠹ᠍ᠼᠾ=ᡀ"
def prepare(inputs):
cleaned_inputs = []
for line in inputs:
line = "".join([c for c in line if c in bichig_vocab])
if len(line) > 0:
cleaned_inputs.append(" ".join(line.replace(' ', '_')))
return cleaned_inputs
def compress(line):
return line.replace(' ', '').replace('_', ' ')
for inputs in buffered_read(args.buffer_size):
prepared_inputs = prepare(inputs)
indices = []
results = []
for batch, batch_indices in make_batches(prepared_inputs, args, task,
max_positions):
indices.extend(batch_indices)
results += process_batch(batch)
for i in np.argsort(indices):
for result in results[i]:
print(compress(result))
def model_fn(model_dir):
model_name = 'checkpoint_best.pt'
model_path = os.path.join(model_dir, model_name)
logger.info('Loading the model')
with open(model_path, 'rb') as f:
model_info = torch.load(f, map_location=torch.device('cpu'))
# Will be overidden by the model_info['args'] - need to keep for pre-trained models
parser = options.get_generation_parser(interactive=True)
# get args for FairSeq by converting the hyperparameters as if they were command-line arguments
argv_copy = copy.deepcopy(sys.argv)
# remove the modifications we did in the command-line arguments
sys.argv[1:] = ['--path', model_path, model_dir]
args = options.parse_args_and_arch(parser)
# restore previous command-line args
sys.argv = argv_copy
saved_args = model_info['args']
for key, value in vars(saved_args).items():
setattr(args, key, value)
args.data = [model_dir]
print(args)
# Setup task, e.g., translation
task = tasks.setup_task(args)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info('Current device: {}'.format(device))
model_paths = [os.path.join(model_dir, model_name)]
models, model_args = utils.load_ensemble_for_inference(model_paths, task, model_arg_overrides={})
# Set dictionaries
tgt_dict = task.target_dictionary
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Initialize generator
translator = SequenceGenerator(
models, tgt_dict, beam_size=args.beam, minlen=args.min_len,
stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen, unk_penalty=args.unkpen,
sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
)
if device.type == 'cuda':
translator.cuda()
# 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(args.replace_unk)
align_dict = utils.load_align_dict(None)
max_positions = utils.resolve_max_positions(
task.max_positions(),
*[model.max_positions() for model in models]
)
return dict(
translator=translator,
task=task,
max_positions=max_positions,
align_dict=align_dict,
tgt_dict=tgt_dict,
args=args,
device=device,
)
def main():
parser = options.get_parser('Generation')
parser.add_argument('--path',
metavar='FILE',
required=True,
action='append',
help='path(s) to model file(s)')
dataset_args = options.add_dataset_args(parser)
dataset_args.add_argument('--batch-size',
default=32,
type=int,
metavar='N',
help='batch size')
dataset_args.add_argument(
'--gen-subset',
default='test',
metavar='SPLIT',
help='data subset to generate (train, valid, test)')
options.add_generation_args(parser)
args = parser.parse_args()
if args.no_progress_bar and args.log_format is None:
args.log_format = 'none'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset
if args.replace_unk is None:
dataset = data.load_dataset(args.data, [args.gen_subset],
args.source_lang, args.target_lang)
else:
dataset = data.load_raw_text_dataset(args.data, [args.gen_subset],
args.source_lang,
args.target_lang)
if args.source_lang is None or args.target_lang is None:
# record inferred languages in args
args.source_lang, args.target_lang = dataset.src, dataset.dst
# Load ensemble
print('| loading model(s) from {}'.format(', '.join(args.path)))
models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict,
dataset.dst_dict)
print('| [{}] dictionary: {} types'.format(dataset.src,
len(dataset.src_dict)))
print('| [{}] dictionary: {} types'.format(dataset.dst,
len(dataset.dst_dict)))
print('| {} {} {} examples'.format(args.data, args.gen_subset,
len(dataset.splits[args.gen_subset])))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
translator = SequenceGenerator(models,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen)
if use_cuda:
translator.cuda()
# 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(args.replace_unk)
# Generate and compute BLEU score
scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(),
dataset.dst_dict.unk())
max_positions = min(model.max_encoder_positions() for model in models)
itr = dataset.eval_dataloader(args.gen_subset,
max_sentences=args.batch_size,
max_positions=max_positions,
skip_invalid_size_inputs_valid_test=args.
skip_invalid_size_inputs_valid_test)
num_sentences = 0
with utils.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t,
maxlen_a=args.max_len_a,
maxlen_b=args.max_len_b,
cuda_device=0 if use_cuda else None,
timer=gen_timer)
for sample_id, src_tokens, target_tokens, hypos in translations:
# Process input and ground truth
target_tokens = target_tokens.int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = dataset.splits[
args.gen_subset].src.get_original_text(sample_id)
target_str = dataset.splits[
args.gen_subset].dst.get_original_text(sample_id)
else:
src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)
target_str = dataset.dst_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
if not args.quiet:
print('S-{}\t{}'.format(sample_id, src_str))
print('T-{}\t{}'.format(sample_id, target_str))
# Process top predictions
for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu(),
align_dict=align_dict,
dst_dict=dataset.dst_dict,
remove_bpe=args.remove_bpe)
if not args.quiet:
print('H-{}\t{}\t{}'.format(sample_id, hypo['score'],
hypo_str))
print('A-{}\t{}'.format(sample_id,
' '.join(map(str, alignment))))
# Score only the top hypothesis
if i == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tokenizer.Tokenizer.tokenize(
target_str,
dataset.dst_dict,
add_if_not_exist=True)
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(src_tokens.size(0))
t.log({'wps': round(wps_meter.avg)})
num_sentences += 1
print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} tokens/s)'.
format(num_sentences, gen_timer.n, gen_timer.sum,
1. / gen_timer.avg))
print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam,
scorer.result_string()))
