def __init__(self, args, dictionary, embed_tokens, proj_to_decoder=True):
super().__init__(dictionary)
self.transformer_embedding = TransformerEmbedding(
args=args, embed_tokens=embed_tokens)
self.transformer_encoder_given_embeddings = TransformerEncoderGivenEmbeddings(
args=args, proj_to_decoder=proj_to_decoder)
# Variable tracker
self.tracker = VariableTracker()
# Initialize adversarial mode
self.set_gradient_tracking_mode(False)
self.set_embed_noising_mode(False)
def __init__(
self,
args,
dictionary,
embed_tokens,
num_chars=50,
embed_dim=32,
char_cnn_params="[(128, 3), (128, 5)]",
char_cnn_nonlinear_fn="tanh",
char_cnn_pool_type="max",
char_cnn_num_highway_layers=0,
char_cnn_output_dim=-1,
use_pretrained_weights=False,
finetune_pretrained_weights=False,
weights_file=None,
):
super().__init__(dictionary)
convolutions_params = literal_eval(char_cnn_params)
self.char_cnn_encoder = char_encoder.CharCNNModel(
dictionary,
num_chars,
embed_dim,
convolutions_params,
char_cnn_nonlinear_fn,
char_cnn_pool_type,
char_cnn_num_highway_layers,
char_cnn_output_dim,
use_pretrained_weights,
finetune_pretrained_weights,
weights_file,
)
self.embed_tokens = embed_tokens
token_embed_dim = embed_tokens.embedding_dim
self.word_layer_norm = nn.LayerNorm(token_embed_dim)
char_embed_dim = (
char_cnn_output_dim
if char_cnn_output_dim != -1
else sum(out_dim for (out_dim, _) in convolutions_params)
)
self.char_layer_norm = nn.LayerNorm(char_embed_dim)
self.word_dim = char_embed_dim + token_embed_dim
self.char_scale = math.sqrt(char_embed_dim / self.word_dim)
self.word_scale = math.sqrt(token_embed_dim / self.word_dim)
if self.word_dim != args.encoder_embed_dim:
self.word_to_transformer_embed = fairseq_transformer.Linear(
self.word_dim, args.encoder_embed_dim
)
self.dropout = args.dropout
self.padding_idx = dictionary.pad()
self.embed_positions = fairseq_transformer.PositionalEmbedding(
1024,
args.encoder_embed_dim,
self.padding_idx,
learned=args.encoder_learned_pos,
)
self.transformer_encoder_given_embeddings = TransformerEncoderGivenEmbeddings(
args=args, proj_to_decoder=True
)
# Variable tracker
self.tracker = VariableTracker()
# Initialize adversarial mode
self.set_gradient_tracking_mode(False)
self.set_embed_noising_mode(False)
# disables sorting and word-length thresholding if True
# (enables ONNX tracing of length-sorted input with batch_size = 1)
self.onnx_export_model = False
class TransformerEncoder(FairseqEncoder):
"""Transformer encoder."""
def __init__(self, args, dictionary, embed_tokens, proj_to_decoder=True):
super().__init__(dictionary)
self.transformer_embedding = TransformerEmbedding(
args=args, embed_tokens=embed_tokens)
self.transformer_encoder_given_embeddings = TransformerEncoderGivenEmbeddings(
args=args, proj_to_decoder=proj_to_decoder)
# Variable tracker
self.tracker = VariableTracker()
# Initialize adversarial mode
self.set_gradient_tracking_mode(False)
self.set_embed_noising_mode(False)
def forward(self, src_tokens, src_lengths):
# Initialize the tracker to keep track of internal variables
self.tracker.reset()
x, encoder_padding_mask, positions = self.transformer_embedding(
src_tokens=src_tokens, src_lengths=src_lengths)
# Track token embeddings
self.tracker.track(x,
"token_embeddings",
retain_grad=self.track_gradients)
x = self.transformer_encoder_given_embeddings(
x=x,
positions=positions,
encoder_padding_mask=encoder_padding_mask)
# TODO(jamesreed): this is kinda a hack because we can't annotate an
# Optional[Tensor] output for encoder_padding_mask
if encoder_padding_mask is None:
encoder_padding_mask = torch.empty([])
return x, src_tokens, encoder_padding_mask
def reorder_encoder_out(self, encoder_out, new_order):
(x, src_tokens, encoder_padding_mask) = encoder_out
src_tokens_tensor = pytorch_translate_utils.get_source_tokens_tensor(
src_tokens)
if x is not None:
x = x.index_select(1, new_order)
if src_tokens_tensor is not None:
src_tokens_tensor = src_tokens_tensor.index_select(0, new_order)
if encoder_padding_mask.shape == torch.Size([]):
encoder_padding_mask = None
if encoder_padding_mask is not None:
encoder_padding_mask = encoder_padding_mask.index_select(
0, new_order)
return (x, src_tokens_tensor, encoder_padding_mask)
def max_positions(self):
"""Maximum input length supported by the encoder."""
return self.transformer_embedding.embed_positions.max_positions()
def upgrade_state_dict_named(self, state_dict, name):
if isinstance(self.transformer_embedding.embed_positions,
SinusoidalPositionalEmbedding):
if f"{name}.transformer_embedding.embed_positions.weights" in state_dict:
del state_dict[
f"{name}.transformer_embedding.embed_positions.weights"]
state_dict[
f"{name}.transformer_embedding.embed_positions._float_tensor"] = torch.FloatTensor(
1)
self.transformer_encoder_given_embeddings.upgrade_state_dict_named(
state_dict, f"{name}.transformer_encoder_given_embeddings")
return state_dict
def set_gradient_tracking_mode(self, mode=True):
self.tracker.reset()
self.track_gradients = mode
def set_embed_noising_mode(self, mode=True):
"""This allows adversarial trainer to turn on and off embedding noising
layers. In regular training, this mode is off, and it is not included
in forward pass.
"""
self.embed_noising_mode = mode