chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
"""isort:skip_file"""
|
||||
|
||||
from .fairseq_nat_model import *
|
||||
from .nonautoregressive_transformer import *
|
||||
@@ -0,0 +1,170 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from fairseq.models.transformer import (
|
||||
TransformerDecoder,
|
||||
TransformerEncoder,
|
||||
TransformerModel,
|
||||
)
|
||||
from fairseq.modules.transformer_sentence_encoder import init_bert_params
|
||||
|
||||
|
||||
def ensemble_encoder(func):
|
||||
def wrapper(self, *args, **kwargs):
|
||||
if self.ensemble_models is None or len(self.ensemble_models) == 1:
|
||||
return func(self, *args, **kwargs)
|
||||
encoder_outs = [func(model, *args, **kwargs, return_all_hiddens=True) for model in self.ensemble_models]
|
||||
_encoder_out = encoder_outs[0].copy()
|
||||
|
||||
def stack(key):
|
||||
outs = [e[key][0] for e in encoder_outs]
|
||||
return [torch.stack(outs, -1) if outs[0] is not None else None]
|
||||
|
||||
_encoder_out["encoder_out"] = stack("encoder_out")
|
||||
_encoder_out["encoder_embedding"] = stack("encoder_embedding")
|
||||
|
||||
num_layers = len(_encoder_out["encoder_states"])
|
||||
if num_layers > 0:
|
||||
_encoder_out["encoder_states"] = [
|
||||
torch.stack([e["encoder_states"][i] for e in encoder_outs], -1)
|
||||
for i in range(num_layers)
|
||||
]
|
||||
return _encoder_out
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def ensemble_decoder(func):
|
||||
def wrapper(self, normalize=False, encoder_out=None, *args, **kwargs):
|
||||
if self.ensemble_models is None or len(self.ensemble_models) == 1:
|
||||
return func(
|
||||
self, normalize=normalize, encoder_out=encoder_out, *args, **kwargs
|
||||
)
|
||||
|
||||
def _replace(encoder_out, new_val):
|
||||
new_encoder_out = encoder_out.copy()
|
||||
new_encoder_out["encoder_out"] = [new_val]
|
||||
return new_encoder_out
|
||||
|
||||
action_outs = [
|
||||
func(
|
||||
model,
|
||||
normalize=normalize,
|
||||
encoder_out=_replace(
|
||||
encoder_out,
|
||||
encoder_out["encoder_out"][0][:, :, :, i]
|
||||
),
|
||||
*args,
|
||||
**kwargs
|
||||
)
|
||||
for i, model in enumerate(self.ensemble_models)
|
||||
]
|
||||
|
||||
if not isinstance(action_outs[0], tuple): # return multiple values
|
||||
action_outs = [[a] for a in action_outs]
|
||||
else:
|
||||
action_outs = [list(a) for a in action_outs]
|
||||
|
||||
ensembled_outs = []
|
||||
for i in range(len(action_outs[0])):
|
||||
if i == 0 and normalize:
|
||||
ensembled_outs += [
|
||||
torch.logsumexp(
|
||||
torch.stack([a[i] for a in action_outs], -1), dim=-1
|
||||
)
|
||||
- math.log(len(self.ensemble_models))
|
||||
]
|
||||
elif action_outs[0][i] is not None:
|
||||
ensembled_outs += [torch.stack([a[i] for a in action_outs], -1)]
|
||||
else:
|
||||
ensembled_outs += [None]
|
||||
|
||||
if len(ensembled_outs) == 1:
|
||||
return ensembled_outs[0]
|
||||
return tuple(ensembled_outs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class FairseqNATModel(TransformerModel):
|
||||
"""
|
||||
Abstract class for all nonautoregressive-based models
|
||||
"""
|
||||
|
||||
def __init__(self, args, encoder, decoder):
|
||||
super().__init__(args, encoder, decoder)
|
||||
self.tgt_dict = decoder.dictionary
|
||||
self.bos = decoder.dictionary.bos()
|
||||
self.eos = decoder.dictionary.eos()
|
||||
self.pad = decoder.dictionary.pad()
|
||||
self.unk = decoder.dictionary.unk()
|
||||
|
||||
self.ensemble_models = None
|
||||
|
||||
@property
|
||||
def allow_length_beam(self):
|
||||
return False
|
||||
|
||||
@property
|
||||
def allow_ensemble(self):
|
||||
return True
|
||||
|
||||
def enable_ensemble(self, models):
|
||||
self.encoder.ensemble_models = [m.encoder for m in models]
|
||||
self.decoder.ensemble_models = [m.decoder for m in models]
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
TransformerModel.add_args(parser)
|
||||
parser.add_argument(
|
||||
"--apply-bert-init",
|
||||
action="store_true",
|
||||
help="use custom param initialization for BERT",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_decoder(cls, args, tgt_dict, embed_tokens):
|
||||
decoder = FairseqNATDecoder(args, tgt_dict, embed_tokens)
|
||||
if getattr(args, "apply_bert_init", False):
|
||||
decoder.apply(init_bert_params)
|
||||
return decoder
|
||||
|
||||
@classmethod
|
||||
def build_encoder(cls, args, src_dict, embed_tokens):
|
||||
encoder = FairseqNATEncoder(args, src_dict, embed_tokens)
|
||||
if getattr(args, "apply_bert_init", False):
|
||||
encoder.apply(init_bert_params)
|
||||
return encoder
|
||||
|
||||
def forward_encoder(self, encoder_inputs):
|
||||
return self.encoder(*encoder_inputs)
|
||||
|
||||
def forward_decoder(self, *args, **kwargs):
|
||||
return NotImplementedError
|
||||
|
||||
def initialize_output_tokens(self, *args, **kwargs):
|
||||
return NotImplementedError
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
return NotImplementedError
|
||||
|
||||
|
||||
class FairseqNATEncoder(TransformerEncoder):
|
||||
def __init__(self, args, dictionary, embed_tokens):
|
||||
super().__init__(args, dictionary, embed_tokens)
|
||||
self.ensemble_models = None
|
||||
|
||||
@ensemble_encoder
|
||||
def forward(self, *args, **kwargs):
|
||||
return super().forward(*args, **kwargs)
|
||||
|
||||
|
||||
class FairseqNATDecoder(TransformerDecoder):
|
||||
def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
|
||||
super().__init__(args, dictionary, embed_tokens, no_encoder_attn)
|
||||
self.ensemble_models = None
|
||||
@@ -0,0 +1,253 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fairseq.models.nat import (
|
||||
_apply_del_words,
|
||||
_apply_ins_masks,
|
||||
_apply_ins_words,
|
||||
_fill,
|
||||
_skip,
|
||||
_skip_encoder_out,
|
||||
)
|
||||
|
||||
|
||||
class _EnsembleModelEncoder(object):
|
||||
def __init__(self, models):
|
||||
self.models = models
|
||||
|
||||
def reorder_encoder_out(self, encoder_outs, new_order):
|
||||
encoder_outs = [
|
||||
model.encoder.reorder_encoder_out(encoder_out, new_order)
|
||||
for model, encoder_out in zip(self.models, encoder_outs)
|
||||
]
|
||||
return encoder_outs
|
||||
|
||||
|
||||
class BasicEnsembleModel(torch.nn.Module):
|
||||
"""A wrapper around an ensemble of models."""
|
||||
|
||||
def __init__(self, models):
|
||||
super().__init__()
|
||||
self.models = torch.nn.ModuleList(models)
|
||||
self.bos = self.models[0].decoder.dictionary.bos()
|
||||
self.eos = self.models[0].decoder.dictionary.eos()
|
||||
self.pad = self.models[0].decoder.dictionary.pad()
|
||||
self.unk = self.models[0].decoder.dictionary.unk()
|
||||
self.encoder = _EnsembleModelEncoder(self.models)
|
||||
|
||||
def has_encoder(self):
|
||||
return hasattr(self.models[0], "encoder")
|
||||
|
||||
def max_decoder_positions(self):
|
||||
return min(m.max_decoder_positions() for m in self.models)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_encoder(self, encoder_input):
|
||||
if not self.has_encoder():
|
||||
return None
|
||||
return [model.forward_encoder(encoder_input) for model in self.models]
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_decoder(self, *inputs):
|
||||
raise NotImplementedError
|
||||
|
||||
def initialize_output_tokens(self, *inputs):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class EnsembleLevT(BasicEnsembleModel):
|
||||
"""A wrapper around an ensemble of models."""
|
||||
|
||||
def __init__(self, models):
|
||||
super().__init__(models)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_decoder(
|
||||
self, decoder_out, encoder_outs, eos_penalty=0.0, max_ratio=None, **kwargs
|
||||
):
|
||||
# LevT ensembling
|
||||
# A pipeline of three steps: deletion, placeholder, and word insertion.
|
||||
# We need to average scores in each step in a pipeline way because of dependence.
|
||||
# deletion
|
||||
output_tokens = decoder_out.output_tokens
|
||||
output_scores = decoder_out.output_scores
|
||||
attn = decoder_out.attn
|
||||
|
||||
bsz = output_tokens.size(0)
|
||||
if max_ratio is None:
|
||||
max_lens = output_tokens.new().fill_(255)
|
||||
else:
|
||||
if not encoder_outs[0]["encoder_padding_mask"]:
|
||||
src_lens = (
|
||||
encoder_outs[0]["encoder_out"][0].new(bsz)
|
||||
.fill_(encoder_outs[0]["encoder_out"][0].size(1))
|
||||
)
|
||||
else:
|
||||
src_lens = (~encoder_outs[0]["encoder_padding_mask"][0]).sum(1)
|
||||
max_lens = (src_lens * max_ratio).clamp(min=10).long()
|
||||
|
||||
# delete words
|
||||
# do not delete tokens if it is <s> </s>
|
||||
can_del_word = output_tokens.ne(self.pad).sum(1) > 2
|
||||
if can_del_word.sum() != 0: # we cannot delete, skip
|
||||
output_tokens, output_scores, attn = self.forward_word_del(
|
||||
encoder_outs,
|
||||
output_tokens,
|
||||
output_scores,
|
||||
attn,
|
||||
can_del_word,
|
||||
)
|
||||
|
||||
# insert placeholders
|
||||
can_ins_mask = output_tokens.ne(self.pad).sum(1) < max_lens
|
||||
if can_ins_mask.sum() != 0:
|
||||
output_tokens, output_scores = self.forward_mask_ins(
|
||||
encoder_outs,
|
||||
output_tokens,
|
||||
output_scores,
|
||||
can_ins_mask,
|
||||
eos_penalty,
|
||||
max_lens,
|
||||
)
|
||||
|
||||
# insert words
|
||||
can_ins_word = output_tokens.eq(self.unk).sum(1) > 0
|
||||
if can_ins_word.sum() != 0:
|
||||
output_tokens, output_scores, attn = self.forward_word_ins(
|
||||
encoder_outs,
|
||||
output_tokens,
|
||||
output_scores,
|
||||
attn,
|
||||
can_ins_word,
|
||||
)
|
||||
|
||||
# delete some unnecessary paddings
|
||||
cut_off = output_tokens.ne(self.pad).sum(1).max()
|
||||
output_tokens = output_tokens[:, :cut_off]
|
||||
output_scores = output_scores[:, :cut_off]
|
||||
attn = None if attn is None else attn[:, :cut_off, :]
|
||||
return decoder_out._replace(
|
||||
output_tokens=output_tokens,
|
||||
output_scores=output_scores,
|
||||
attn=attn,
|
||||
history=None,
|
||||
)
|
||||
|
||||
def forward_word_del(
|
||||
self, encoder_outs, output_tokens, output_scores, attn, can_del_word
|
||||
):
|
||||
word_del_score_avg = []
|
||||
word_del_attn_avg = []
|
||||
for model, encoder_out in zip(self.models, encoder_outs):
|
||||
word_del_out, word_del_attn = model.decoder.forward_word_del(
|
||||
_skip(output_tokens, can_del_word),
|
||||
_skip_encoder_out(model.encoder, encoder_out, can_del_word),
|
||||
)
|
||||
word_del_score = F.log_softmax(word_del_out, 2)
|
||||
word_del_score_avg.append(word_del_score)
|
||||
word_del_attn_avg.append(word_del_attn)
|
||||
word_del_score_avg = torch.logsumexp(
|
||||
torch.stack(word_del_score_avg, dim=0), dim=0
|
||||
) - math.log(len(self.models))
|
||||
word_del_pred = word_del_score_avg.max(-1)[1].bool()
|
||||
if word_del_attn_avg[0] is not None:
|
||||
word_del_attn_avg = torch.stack(word_del_attn_avg, dim=0) / len(self.models)
|
||||
else:
|
||||
word_del_attn_avg = None
|
||||
|
||||
_tokens, _scores, _attn = _apply_del_words(
|
||||
output_tokens[can_del_word],
|
||||
output_scores[can_del_word],
|
||||
word_del_attn_avg,
|
||||
word_del_pred,
|
||||
self.pad,
|
||||
self.bos,
|
||||
self.eos,
|
||||
)
|
||||
output_tokens = _fill(output_tokens, can_del_word, _tokens, self.pad)
|
||||
output_scores = _fill(output_scores, can_del_word, _scores, 0)
|
||||
attn = _fill(attn, can_del_word, _attn, 0.0)
|
||||
return output_tokens, output_scores, attn
|
||||
|
||||
def forward_mask_ins(
|
||||
self,
|
||||
encoder_outs,
|
||||
output_tokens,
|
||||
output_scores,
|
||||
can_ins_mask,
|
||||
eos_penalty,
|
||||
max_lens,
|
||||
):
|
||||
mask_ins_score_avg = []
|
||||
for model, encoder_out in zip(self.models, encoder_outs):
|
||||
mask_ins_out, _ = model.decoder.forward_mask_ins(
|
||||
_skip(output_tokens, can_ins_mask),
|
||||
_skip_encoder_out(model.encoder, encoder_out, can_ins_mask),
|
||||
)
|
||||
mask_ins_score = F.log_softmax(mask_ins_out, 2)
|
||||
if eos_penalty > 0.0:
|
||||
mask_ins_score[:, :, 0] -= eos_penalty
|
||||
mask_ins_score_avg.append(mask_ins_score)
|
||||
mask_ins_score_avg = torch.logsumexp(
|
||||
torch.stack(mask_ins_score_avg, dim=0), dim=0
|
||||
) - math.log(len(self.models))
|
||||
mask_ins_pred = mask_ins_score_avg.max(-1)[1]
|
||||
mask_ins_pred = torch.min(
|
||||
mask_ins_pred, max_lens[can_ins_mask, None].expand_as(mask_ins_pred)
|
||||
)
|
||||
_tokens, _scores = _apply_ins_masks(
|
||||
output_tokens[can_ins_mask],
|
||||
output_scores[can_ins_mask],
|
||||
mask_ins_pred,
|
||||
self.pad,
|
||||
self.unk,
|
||||
self.eos,
|
||||
)
|
||||
output_tokens = _fill(output_tokens, can_ins_mask, _tokens, self.pad)
|
||||
output_scores = _fill(output_scores, can_ins_mask, _scores, 0)
|
||||
return output_tokens, output_scores
|
||||
|
||||
def forward_word_ins(
|
||||
self, encoder_outs, output_tokens, output_scores, attn, can_ins_word
|
||||
):
|
||||
word_ins_score_avg = []
|
||||
word_ins_attn_avg = []
|
||||
for model, encoder_out in zip(self.models, encoder_outs):
|
||||
word_ins_out, word_ins_attn = model.decoder.forward_word_ins(
|
||||
_skip(output_tokens, can_ins_word),
|
||||
_skip_encoder_out(model.encoder, encoder_out, can_ins_word),
|
||||
)
|
||||
word_ins_score = F.log_softmax(word_ins_out, 2)
|
||||
word_ins_score_avg.append(word_ins_score)
|
||||
word_ins_attn_avg.append(word_ins_attn)
|
||||
word_ins_score_avg = torch.logsumexp(
|
||||
torch.stack(word_ins_score_avg, dim=0), dim=0
|
||||
) - math.log(len(self.models))
|
||||
if word_ins_attn_avg[0] is not None:
|
||||
word_ins_attn_avg = torch.stack(word_ins_attn_avg, dim=0) / len(self.models)
|
||||
else:
|
||||
word_ins_attn_avg = None
|
||||
word_ins_score_max, word_ins_pred = word_ins_score_avg.max(-1)
|
||||
|
||||
_tokens, _scores = _apply_ins_words(
|
||||
output_tokens[can_ins_word],
|
||||
output_scores[can_ins_word],
|
||||
word_ins_pred,
|
||||
word_ins_score_max,
|
||||
self.unk,
|
||||
)
|
||||
|
||||
output_tokens = _fill(output_tokens, can_ins_word, _tokens, self.pad)
|
||||
output_scores = _fill(output_scores, can_ins_word, _scores, 0)
|
||||
attn = _fill(attn, can_ins_word, word_ins_attn, 0.0)
|
||||
return output_tokens, output_scores, attn
|
||||
|
||||
def initialize_output_tokens(self, encoder_outs, src_tokens):
|
||||
# LevT doesn't do length prediction.
|
||||
return self.models[0].initialize_output_tokens(encoder_outs[0], src_tokens)
|
||||
@@ -0,0 +1,450 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.iterative_refinement_generator import DecoderOut
|
||||
from fairseq.models import register_model, register_model_architecture
|
||||
from fairseq.models.nat import FairseqNATDecoder, FairseqNATEncoder, FairseqNATModel, ensemble_decoder, ensemble_encoder
|
||||
from fairseq.models.transformer import Embedding
|
||||
from fairseq.modules.transformer_sentence_encoder import init_bert_params
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from torch import Tensor
|
||||
|
||||
def _mean_pooling(enc_feats, src_masks):
|
||||
# enc_feats: T x B x C
|
||||
# src_masks: B x T or None
|
||||
if src_masks is None:
|
||||
enc_feats = enc_feats.mean(0)
|
||||
else:
|
||||
src_masks = (~src_masks).transpose(0, 1).type_as(enc_feats)
|
||||
enc_feats = (
|
||||
(enc_feats / src_masks.sum(0)[None, :, None]) * src_masks[:, :, None]
|
||||
).sum(0)
|
||||
return enc_feats
|
||||
|
||||
|
||||
def _argmax(x, dim):
|
||||
return (x == x.max(dim, keepdim=True)[0]).type_as(x)
|
||||
|
||||
|
||||
def _uniform_assignment(src_lens, trg_lens):
|
||||
max_trg_len = trg_lens.max()
|
||||
steps = (src_lens.float() - 1) / (trg_lens.float() - 1) # step-size
|
||||
# max_trg_len
|
||||
index_t = utils.new_arange(trg_lens, max_trg_len).float()
|
||||
index_t = steps[:, None] * index_t[None, :] # batch_size X max_trg_len
|
||||
index_t = torch.round(index_t).long().detach()
|
||||
return index_t
|
||||
|
||||
|
||||
@register_model("nonautoregressive_transformer")
|
||||
class NATransformerModel(FairseqNATModel):
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
FairseqNATModel.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
"--src-embedding-copy",
|
||||
action="store_true",
|
||||
help="copy encoder word embeddings as the initial input of the decoder",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_encoder(cls, args, tgt_dict, embed_tokens):
|
||||
encoder = NATransformerEncoder(args, tgt_dict, embed_tokens)
|
||||
if getattr(args, "apply_bert_init", False):
|
||||
encoder.apply(init_bert_params)
|
||||
return encoder
|
||||
|
||||
@classmethod
|
||||
def build_decoder(cls, args, tgt_dict, embed_tokens):
|
||||
decoder = NATransformerDecoder(args, tgt_dict, embed_tokens)
|
||||
if getattr(args, "apply_bert_init", False):
|
||||
decoder.apply(init_bert_params)
|
||||
return decoder
|
||||
|
||||
def forward(
|
||||
self, src_tokens, src_lengths, prev_output_tokens, tgt_tokens, **kwargs
|
||||
):
|
||||
# encoding
|
||||
encoder_out = self.encoder(src_tokens, src_lengths=src_lengths, **kwargs)
|
||||
|
||||
# decoding
|
||||
word_ins_out = self.decoder(
|
||||
normalize=False,
|
||||
prev_output_tokens=prev_output_tokens,
|
||||
encoder_out=encoder_out,
|
||||
)
|
||||
|
||||
return {
|
||||
"word_ins": {
|
||||
"out": word_ins_out,
|
||||
"tgt": tgt_tokens,
|
||||
"mask": tgt_tokens.ne(self.pad),
|
||||
"ls": self.args.label_smoothing,
|
||||
"nll_loss": True,
|
||||
}
|
||||
}
|
||||
|
||||
def forward_decoder(self, decoder_out, encoder_out, decoding_format=None, **kwargs):
|
||||
step = decoder_out.step
|
||||
output_tokens = decoder_out.output_tokens
|
||||
output_scores = decoder_out.output_scores
|
||||
history = decoder_out.history
|
||||
|
||||
# execute the decoder
|
||||
output_masks = output_tokens.ne(self.pad)
|
||||
_scores, _tokens = self.decoder(
|
||||
normalize=True,
|
||||
prev_output_tokens=output_tokens,
|
||||
encoder_out=encoder_out,
|
||||
step=step,
|
||||
).max(-1)
|
||||
|
||||
output_tokens.masked_scatter_(output_masks, _tokens[output_masks])
|
||||
output_scores.masked_scatter_(output_masks, _scores[output_masks])
|
||||
if history is not None:
|
||||
history.append(output_tokens.clone())
|
||||
|
||||
return decoder_out._replace(
|
||||
output_tokens=output_tokens,
|
||||
output_scores=output_scores,
|
||||
attn=None,
|
||||
history=history,
|
||||
)
|
||||
|
||||
|
||||
class NATransformerDecoder(FairseqNATDecoder):
|
||||
def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
|
||||
super().__init__(
|
||||
args, dictionary, embed_tokens, no_encoder_attn=no_encoder_attn
|
||||
)
|
||||
self.dictionary = dictionary
|
||||
self.bos = dictionary.bos()
|
||||
self.unk = dictionary.unk()
|
||||
self.eos = dictionary.eos()
|
||||
|
||||
self.encoder_embed_dim = args.encoder_embed_dim
|
||||
self.src_embedding_copy = getattr(args, "src_embedding_copy", False)
|
||||
if self.src_embedding_copy:
|
||||
self.copy_attn = torch.nn.Linear(self.embed_dim, self.embed_dim, bias=False)
|
||||
|
||||
@ensemble_decoder
|
||||
def forward(self, normalize, encoder_out, prev_output_tokens, step=0, **unused):
|
||||
features, _ = self.extract_features(
|
||||
prev_output_tokens,
|
||||
encoder_out=encoder_out,
|
||||
embedding_copy=(step == 0) & self.src_embedding_copy,
|
||||
)
|
||||
decoder_out = self.output_layer(features)
|
||||
return F.log_softmax(decoder_out, -1) if normalize else decoder_out
|
||||
|
||||
def extract_features(
|
||||
self,
|
||||
prev_output_tokens,
|
||||
encoder_out=None,
|
||||
early_exit=None,
|
||||
embedding_copy=False,
|
||||
**unused
|
||||
):
|
||||
"""
|
||||
Similar to *forward* but only return features.
|
||||
|
||||
Inputs:
|
||||
prev_output_tokens: Tensor(B, T)
|
||||
encoder_out: a dictionary of hidden states and masks
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- the decoder's features of shape `(batch, tgt_len, embed_dim)`
|
||||
- a dictionary with any model-specific outputs
|
||||
the LevenshteinTransformer decoder has full-attention to all generated tokens
|
||||
"""
|
||||
positions = (
|
||||
self.embed_positions(prev_output_tokens)
|
||||
if self.embed_positions is not None
|
||||
else None
|
||||
)
|
||||
# embedding
|
||||
if embedding_copy:
|
||||
src_embd = encoder_out["encoder_embedding"][0]
|
||||
if len(encoder_out["encoder_padding_mask"]) > 0:
|
||||
src_mask = encoder_out["encoder_padding_mask"][0]
|
||||
else:
|
||||
src_mask = None
|
||||
|
||||
bsz, seq_len = prev_output_tokens.size()
|
||||
attn_score = torch.bmm(self.copy_attn(positions),
|
||||
(src_embd + encoder_out['encoder_pos'][0]).transpose(1, 2))
|
||||
if src_mask is not None:
|
||||
attn_score = attn_score.masked_fill(src_mask.unsqueeze(1).expand(-1, seq_len, -1), float('-inf'))
|
||||
attn_weight = F.softmax(attn_score, dim=-1)
|
||||
x = torch.bmm(attn_weight, src_embd)
|
||||
mask_target_x, decoder_padding_mask = self.forward_embedding(prev_output_tokens)
|
||||
output_mask = prev_output_tokens.eq(self.unk)
|
||||
cat_x = torch.cat([mask_target_x.unsqueeze(2), x.unsqueeze(2)], dim=2).view(-1, x.size(2))
|
||||
# torch.arange(bsz * seq_len).cuda()
|
||||
x = cat_x.index_select(dim=0, index=torch.arange(bsz * seq_len).cuda() * 2 +
|
||||
output_mask.view(-1).long()).reshape(bsz, seq_len, x.size(2))
|
||||
else:
|
||||
|
||||
x, decoder_padding_mask = self.forward_embedding(prev_output_tokens)
|
||||
|
||||
# B x T x C -> T x B x C
|
||||
x = x.transpose(0, 1)
|
||||
positions = positions.transpose(0, 1)
|
||||
attn = None
|
||||
inner_states = [x]
|
||||
|
||||
# decoder layers
|
||||
for i, layer in enumerate(self.layers):
|
||||
|
||||
# early exit from the decoder.
|
||||
if (early_exit is not None) and (i >= early_exit):
|
||||
break
|
||||
|
||||
if positions is not None:
|
||||
x += positions
|
||||
x = self.dropout_module(x)
|
||||
|
||||
x, attn, _ = layer(
|
||||
x,
|
||||
encoder_out["encoder_out"][0]
|
||||
if (encoder_out is not None and len(encoder_out["encoder_out"]) > 0)
|
||||
else None,
|
||||
encoder_out["encoder_padding_mask"][0]
|
||||
if (
|
||||
encoder_out is not None
|
||||
and len(encoder_out["encoder_padding_mask"]) > 0
|
||||
)
|
||||
else None,
|
||||
self_attn_mask=None,
|
||||
self_attn_padding_mask=decoder_padding_mask,
|
||||
)
|
||||
inner_states.append(x)
|
||||
|
||||
if self.layer_norm:
|
||||
x = self.layer_norm(x)
|
||||
|
||||
# T x B x C -> B x T x C
|
||||
x = x.transpose(0, 1)
|
||||
|
||||
if self.project_out_dim is not None:
|
||||
x = self.project_out_dim(x)
|
||||
|
||||
return x, {"attn": attn, "inner_states": inner_states}
|
||||
|
||||
def forward_embedding(self, prev_output_tokens, states=None):
|
||||
# embed tokens
|
||||
if states is None:
|
||||
x = self.embed_tokens(prev_output_tokens)
|
||||
if self.project_in_dim is not None:
|
||||
x = self.project_in_dim(x)
|
||||
else:
|
||||
x = states
|
||||
|
||||
decoder_padding_mask = prev_output_tokens.eq(self.padding_idx)
|
||||
return x, decoder_padding_mask
|
||||
|
||||
def forward_copying_source(self, src_embeds, src_masks, tgt_masks):
|
||||
length_sources = src_masks.sum(1)
|
||||
length_targets = tgt_masks.sum(1)
|
||||
mapped_inputs = _uniform_assignment(length_sources, length_targets).masked_fill(
|
||||
~tgt_masks, 0
|
||||
)
|
||||
copied_embedding = torch.gather(
|
||||
src_embeds,
|
||||
1,
|
||||
mapped_inputs.unsqueeze(-1).expand(
|
||||
*mapped_inputs.size(), src_embeds.size(-1)
|
||||
),
|
||||
)
|
||||
return copied_embedding
|
||||
|
||||
|
||||
class NATransformerEncoder(FairseqNATEncoder):
|
||||
def __init__(self, args, dictionary, embed_tokens):
|
||||
super().__init__(args, dictionary, embed_tokens)
|
||||
|
||||
@ensemble_encoder
|
||||
def forward(
|
||||
self,
|
||||
src_tokens,
|
||||
src_lengths: Optional[torch.Tensor] = None,
|
||||
return_all_hiddens: bool = False,
|
||||
token_embeddings: Optional[torch.Tensor] = None,
|
||||
):
|
||||
encoder_padding_mask = src_tokens.eq(self.padding_idx)
|
||||
has_pads = (src_tokens.device.type == "xla" or encoder_padding_mask.any())
|
||||
|
||||
x, encoder_embedding = self.forward_embedding(src_tokens, token_embeddings)
|
||||
|
||||
encoder_pos = self.embed_positions(src_tokens)
|
||||
# account for padding while computing the representation
|
||||
if encoder_padding_mask is not None:
|
||||
x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
|
||||
|
||||
# B x T x C -> T x B x C
|
||||
x = x.transpose(0, 1)
|
||||
|
||||
encoder_states = []
|
||||
|
||||
if return_all_hiddens:
|
||||
encoder_states.append(x)
|
||||
# encoder layers
|
||||
for layer in self.layers:
|
||||
x = layer(
|
||||
x, encoder_padding_mask=encoder_padding_mask if has_pads else None
|
||||
)
|
||||
if return_all_hiddens:
|
||||
assert encoder_states is not None
|
||||
encoder_states.append(x)
|
||||
|
||||
if self.layer_norm is not None:
|
||||
x = self.layer_norm(x)
|
||||
|
||||
# The Pytorch Mobile lite interpreter does not supports returning NamedTuple in
|
||||
# `forward` so we use a dictionary instead.
|
||||
# TorchScript does not support mixed values so the values are all lists.
|
||||
# The empty list is equivalent to None.
|
||||
return {
|
||||
"encoder_out": [x], # T x B x C
|
||||
"encoder_padding_mask": [encoder_padding_mask], # B x T
|
||||
"encoder_embedding": [encoder_embedding], # B x T x C
|
||||
"encoder_pos": [encoder_pos],
|
||||
"encoder_states": encoder_states, # List[T x B x C]
|
||||
"src_tokens": [],
|
||||
"src_lengths": [],
|
||||
}
|
||||
|
||||
def forward_embedding(
|
||||
self, src_tokens, token_embedding: Optional[torch.Tensor] = None
|
||||
):
|
||||
# embed tokens and positions
|
||||
if token_embedding is None:
|
||||
token_embedding = self.embed_tokens(src_tokens)
|
||||
x = embed = token_embedding
|
||||
if self.embed_positions is not None:
|
||||
x = embed + self.embed_positions(src_tokens)
|
||||
if self.layernorm_embedding is not None:
|
||||
x = self.layernorm_embedding(x)
|
||||
x = self.dropout_module(x)
|
||||
if self.quant_noise is not None:
|
||||
x = self.quant_noise(x)
|
||||
return x, embed
|
||||
|
||||
@torch.jit.export
|
||||
def reorder_encoder_out(self, encoder_out: Dict[str, List[Tensor]], new_order):
|
||||
"""
|
||||
Reorder encoder output according to *new_order*.
|
||||
|
||||
Args:
|
||||
encoder_out: output from the ``forward()`` method
|
||||
new_order (LongTensor): desired order
|
||||
|
||||
Returns:
|
||||
*encoder_out* rearranged according to *new_order*
|
||||
"""
|
||||
if len(encoder_out["encoder_out"]) == 0:
|
||||
new_encoder_out = []
|
||||
else:
|
||||
new_encoder_out = [encoder_out["encoder_out"][0].index_select(1, new_order)]
|
||||
if len(encoder_out["encoder_padding_mask"]) == 0:
|
||||
new_encoder_padding_mask = []
|
||||
else:
|
||||
new_encoder_padding_mask = [
|
||||
encoder_out["encoder_padding_mask"][0].index_select(0, new_order)
|
||||
]
|
||||
if len(encoder_out["encoder_embedding"]) == 0:
|
||||
new_encoder_embedding = []
|
||||
else:
|
||||
new_encoder_embedding = [
|
||||
encoder_out["encoder_embedding"][0].index_select(0, new_order)
|
||||
]
|
||||
if len(encoder_out["encoder_pos"]) == 0:
|
||||
new_encoder_pos = []
|
||||
else:
|
||||
new_encoder_pos = [
|
||||
encoder_out["encoder_pos"][0].index_select(0, new_order)
|
||||
]
|
||||
|
||||
if len(encoder_out["src_tokens"]) == 0:
|
||||
src_tokens = []
|
||||
else:
|
||||
src_tokens = [(encoder_out["src_tokens"][0]).index_select(0, new_order)]
|
||||
|
||||
if len(encoder_out["src_lengths"]) == 0:
|
||||
src_lengths = []
|
||||
else:
|
||||
src_lengths = [(encoder_out["src_lengths"][0]).index_select(0, new_order)]
|
||||
|
||||
encoder_states = encoder_out["encoder_states"]
|
||||
if len(encoder_states) > 0:
|
||||
for idx, state in enumerate(encoder_states):
|
||||
encoder_states[idx] = state.index_select(1, new_order)
|
||||
|
||||
return {
|
||||
"encoder_out": new_encoder_out, # T x B x C
|
||||
"encoder_padding_mask": new_encoder_padding_mask, # B x T
|
||||
"encoder_embedding": new_encoder_embedding, # B x T x C
|
||||
"encoder_pos": new_encoder_pos,
|
||||
"encoder_states": encoder_states, # List[T x B x C]
|
||||
"src_tokens": src_tokens, # B x T
|
||||
"src_lengths": src_lengths, # B x 1
|
||||
}
|
||||
|
||||
@register_model_architecture(
|
||||
"nonautoregressive_transformer", "nonautoregressive_transformer"
|
||||
)
|
||||
def base_architecture(args):
|
||||
args.encoder_embed_path = getattr(args, "encoder_embed_path", None)
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 6)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
|
||||
args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
|
||||
args.encoder_learned_pos = getattr(args, "encoder_learned_pos", False)
|
||||
args.decoder_embed_path = getattr(args, "decoder_embed_path", None)
|
||||
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
|
||||
args.decoder_ffn_embed_dim = getattr(
|
||||
args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
|
||||
)
|
||||
args.decoder_layers = getattr(args, "decoder_layers", 6)
|
||||
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
|
||||
args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
|
||||
args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
|
||||
args.attention_dropout = getattr(args, "attention_dropout", 0.0)
|
||||
args.activation_dropout = getattr(args, "activation_dropout", 0.0)
|
||||
args.activation_fn = getattr(args, "activation_fn", "relu")
|
||||
args.dropout = getattr(args, "dropout", 0.1)
|
||||
args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
|
||||
args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
|
||||
args.share_decoder_input_output_embed = getattr(
|
||||
args, "share_decoder_input_output_embed", False
|
||||
)
|
||||
args.share_all_embeddings = getattr(args, "share_all_embeddings", False)
|
||||
args.no_token_positional_embeddings = getattr(
|
||||
args, "no_token_positional_embeddings", False
|
||||
)
|
||||
args.adaptive_input = getattr(args, "adaptive_input", False)
|
||||
args.apply_bert_init = getattr(args, "apply_bert_init", False)
|
||||
|
||||
args.decoder_output_dim = getattr(
|
||||
args, "decoder_output_dim", args.decoder_embed_dim
|
||||
)
|
||||
args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
|
||||
|
||||
# --- special arguments ---
|
||||
args.src_embedding_copy = getattr(args, "src_embedding_copy", False)
|
||||
|
||||
|
||||
@register_model_architecture(
|
||||
"nonautoregressive_transformer", "nonautoregressive_transformer_wmt_en_de"
|
||||
)
|
||||
def nonautoregressive_transformer_wmt_en_de(args):
|
||||
base_architecture(args)
|
||||
Reference in New Issue
Block a user