chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,7 @@
|
||||
# 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.
|
||||
|
||||
from .hub_interface import * # noqa
|
||||
from .model import * # noqa
|
||||
@@ -0,0 +1,206 @@
|
||||
# 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 copy
|
||||
import logging
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.data import encoders
|
||||
from fairseq.hub_utils import GeneratorHubInterface
|
||||
from omegaconf import open_dict
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BARTHubInterface(GeneratorHubInterface):
|
||||
"""A simple PyTorch Hub interface to BART.
|
||||
|
||||
Usage: https://github.com/pytorch/fairseq/tree/master/examples/bart
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, task, model):
|
||||
super().__init__(cfg, task, [model])
|
||||
self.model = self.models[0]
|
||||
|
||||
def encode(
|
||||
self, sentence: str, *addl_sentences, no_separator=True
|
||||
) -> torch.LongTensor:
|
||||
"""
|
||||
BPE-encode a sentence (or multiple sentences).
|
||||
|
||||
Every sequence begins with a beginning-of-sentence (`<s>`) symbol.
|
||||
Every sentence ends with an end-of-sentence (`</s>`).
|
||||
|
||||
Example (single sentence): `<s> a b c </s>`
|
||||
Example (sentence pair): `<s> d e f </s> 1 2 3 </s>`
|
||||
|
||||
The BPE encoding follows GPT-2. One subtle detail is that the GPT-2 BPE
|
||||
requires leading spaces. For example::
|
||||
|
||||
>>> bart.encode('Hello world').tolist()
|
||||
[0, 31414, 232, 2]
|
||||
>>> bart.encode(' world').tolist()
|
||||
[0, 232, 2]
|
||||
>>> bart.encode('world').tolist()
|
||||
[0, 8331, 2]
|
||||
"""
|
||||
tokens = self.bpe.encode(sentence)
|
||||
if len(tokens.split(" ")) > min(self.max_positions) - 2:
|
||||
tokens = " ".join(tokens.split(" ")[: min(self.max_positions) - 2])
|
||||
bpe_sentence = "<s> " + tokens + " </s>"
|
||||
for s in addl_sentences:
|
||||
bpe_sentence += " </s>" if not no_separator else ""
|
||||
bpe_sentence += " " + self.bpe.encode(s) + " </s>"
|
||||
tokens = self.task.source_dictionary.encode_line(bpe_sentence, append_eos=False)
|
||||
return tokens.long()
|
||||
|
||||
def decode(self, tokens: torch.LongTensor):
|
||||
assert tokens.dim() == 1
|
||||
tokens = tokens.cpu().numpy()
|
||||
if tokens[0] == self.task.source_dictionary.bos():
|
||||
tokens = tokens[1:] # remove <s>
|
||||
eos_mask = tokens == self.task.source_dictionary.eos()
|
||||
doc_mask = eos_mask[1:] & eos_mask[:-1]
|
||||
sentences = np.split(tokens, doc_mask.nonzero()[0] + 1)
|
||||
sentences = [
|
||||
self.bpe.decode(self.task.source_dictionary.string(s)) for s in sentences
|
||||
]
|
||||
if len(sentences) == 1:
|
||||
return sentences[0]
|
||||
return sentences
|
||||
|
||||
def _build_sample(self, src_tokens: List[torch.LongTensor]):
|
||||
# assert torch.is_tensor(src_tokens)
|
||||
dataset = self.task.build_dataset_for_inference(
|
||||
src_tokens,
|
||||
[x.numel() for x in src_tokens],
|
||||
)
|
||||
sample = dataset.collater(dataset)
|
||||
sample = utils.apply_to_sample(lambda tensor: tensor.to(self.device), sample)
|
||||
return sample
|
||||
|
||||
def generate(
|
||||
self,
|
||||
tokenized_sentences: List[torch.LongTensor],
|
||||
*args,
|
||||
inference_step_args=None,
|
||||
skip_invalid_size_inputs=False,
|
||||
**kwargs
|
||||
) -> List[List[Dict[str, torch.Tensor]]]:
|
||||
inference_step_args = inference_step_args or {}
|
||||
if "prefix_tokens" in inference_step_args:
|
||||
raise NotImplementedError("prefix generation not implemented for BART")
|
||||
res = []
|
||||
for batch in self._build_batches(tokenized_sentences, skip_invalid_size_inputs):
|
||||
src_tokens = batch['net_input']['src_tokens']
|
||||
inference_step_args["prefix_tokens"] =src_tokens.new_full(
|
||||
(src_tokens.size(0), 1), fill_value=self.task.source_dictionary.bos()
|
||||
).to(device=self.device)
|
||||
results = super().generate(
|
||||
src_tokens,
|
||||
*args,
|
||||
inference_step_args=inference_step_args,
|
||||
skip_invalid_size_inputs=skip_invalid_size_inputs,
|
||||
**kwargs
|
||||
)
|
||||
res.extend(results)
|
||||
return res
|
||||
|
||||
def extract_features(
|
||||
self, tokens: torch.LongTensor, return_all_hiddens: bool = False
|
||||
) -> torch.Tensor:
|
||||
if tokens.dim() == 1:
|
||||
tokens = tokens.unsqueeze(0)
|
||||
if tokens.size(-1) > min(self.model.max_positions()):
|
||||
raise ValueError(
|
||||
"tokens exceeds maximum length: {} > {}".format(
|
||||
tokens.size(-1), self.model.max_positions()
|
||||
)
|
||||
)
|
||||
tokens.to(device=self.device),
|
||||
prev_output_tokens = tokens.clone()
|
||||
|
||||
prev_output_tokens[:, 0] = tokens.gather(
|
||||
1,
|
||||
(tokens.ne(self.task.source_dictionary.pad()).sum(dim=1) - 1).unsqueeze(-1),
|
||||
).squeeze()
|
||||
|
||||
prev_output_tokens[:, 1:] = tokens[:, :-1]
|
||||
features, extra = self.model(
|
||||
src_tokens=tokens,
|
||||
src_lengths=None,
|
||||
prev_output_tokens=prev_output_tokens,
|
||||
features_only=True,
|
||||
return_all_hiddens=return_all_hiddens,
|
||||
)
|
||||
if return_all_hiddens:
|
||||
# convert from T x B x C -> B x T x C
|
||||
inner_states = extra["inner_states"]
|
||||
return [inner_state.transpose(0, 1) for inner_state in inner_states]
|
||||
else:
|
||||
return features # just the last layer's features
|
||||
|
||||
def register_classification_head(
|
||||
self, name: str, num_classes: int = None, embedding_size: int = None, **kwargs
|
||||
):
|
||||
self.model.register_classification_head(
|
||||
name, num_classes=num_classes, embedding_size=embedding_size, **kwargs
|
||||
)
|
||||
|
||||
def predict(self, head: str, tokens: torch.LongTensor, return_logits: bool = False):
|
||||
if tokens.dim() == 1:
|
||||
tokens = tokens.unsqueeze(0)
|
||||
features = self.extract_features(tokens.to(device=self.device))
|
||||
sentence_representation = features[
|
||||
tokens.eq(self.task.source_dictionary.eos()), :
|
||||
].view(features.size(0), -1, features.size(-1))[:, -1, :]
|
||||
|
||||
logits = self.model.classification_heads[head](sentence_representation)
|
||||
if return_logits:
|
||||
return logits
|
||||
return F.log_softmax(logits, dim=-1)
|
||||
|
||||
def fill_mask(
|
||||
self,
|
||||
masked_inputs: List[str],
|
||||
topk: int = 5,
|
||||
match_source_len: bool = True,
|
||||
**generate_kwargs
|
||||
):
|
||||
masked_token = '<mask>'
|
||||
batch_tokens = []
|
||||
for masked_input in masked_inputs:
|
||||
assert masked_token in masked_input, \
|
||||
"please add one {} token for the input".format(masked_token)
|
||||
|
||||
text_spans = masked_input.split(masked_token)
|
||||
text_spans_bpe = (' {0} '.format(masked_token)).join(
|
||||
[self.bpe.encode(text_span.rstrip()) for text_span in text_spans]
|
||||
).strip()
|
||||
tokens = self.task.source_dictionary.encode_line(
|
||||
'<s> ' + text_spans_bpe + ' </s>',
|
||||
append_eos=False,
|
||||
add_if_not_exist=False,
|
||||
).long()
|
||||
batch_tokens.append(tokens)
|
||||
|
||||
# ensure beam size is at least as big as topk
|
||||
generate_kwargs['beam'] = max(
|
||||
topk,
|
||||
generate_kwargs.get('beam', -1),
|
||||
)
|
||||
generate_kwargs['match_source_len'] = match_source_len
|
||||
batch_hypos = self.generate(batch_tokens, **generate_kwargs)
|
||||
|
||||
return [
|
||||
[(self.decode(hypo['tokens']), hypo['score']) for hypo in hypos[:topk]]
|
||||
for hypos in batch_hypos
|
||||
]
|
||||
@@ -0,0 +1,384 @@
|
||||
# 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.
|
||||
"""
|
||||
BART: Denoising Sequence-to-Sequence Pre-training for
|
||||
Natural Language Generation, Translation, and Comprehension
|
||||
"""
|
||||
from typing import Optional
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from fairseq import utils
|
||||
from fairseq.models import register_model, register_model_architecture
|
||||
from fairseq.models.transformer import TransformerModel
|
||||
from fairseq.modules.transformer_sentence_encoder import init_bert_params
|
||||
|
||||
from .hub_interface import BARTHubInterface
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_model("bart")
|
||||
class BARTModel(TransformerModel):
|
||||
__jit_unused_properties__ = ["supported_targets"]
|
||||
|
||||
@classmethod
|
||||
def hub_models(cls):
|
||||
return {
|
||||
"bart.base": "http://dl.fbaipublicfiles.com/fairseq/models/bart.base.tar.gz",
|
||||
"bart.large": "http://dl.fbaipublicfiles.com/fairseq/models/bart.large.tar.gz",
|
||||
"bart.large.mnli": "http://dl.fbaipublicfiles.com/fairseq/models/bart.large.mnli.tar.gz",
|
||||
"bart.large.cnn": "http://dl.fbaipublicfiles.com/fairseq/models/bart.large.cnn.tar.gz",
|
||||
"bart.large.xsum": "http://dl.fbaipublicfiles.com/fairseq/models/bart.large.xsum.tar.gz",
|
||||
}
|
||||
|
||||
def __init__(self, args, encoder, decoder):
|
||||
super().__init__(args, encoder, decoder)
|
||||
|
||||
# We follow BERT's random weight initialization
|
||||
self.apply(init_bert_params)
|
||||
|
||||
self.classification_heads = nn.ModuleDict()
|
||||
if hasattr(self.encoder, "dictionary"):
|
||||
self.eos: int = self.encoder.dictionary.eos()
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
super(BARTModel, BARTModel).add_args(parser)
|
||||
parser.add_argument(
|
||||
"--pooler-dropout",
|
||||
type=float,
|
||||
metavar="D",
|
||||
help="dropout probability in the masked_lm pooler layers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pooler-activation-fn",
|
||||
choices=utils.get_available_activation_fns(),
|
||||
help="activation function to use for pooler layer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spectral-norm-classification-head",
|
||||
action="store_true",
|
||||
help="Apply spectral normalization on the classification head",
|
||||
)
|
||||
|
||||
@property
|
||||
def supported_targets(self):
|
||||
return {"self"}
|
||||
|
||||
def forward(
|
||||
self,
|
||||
src_tokens,
|
||||
src_lengths,
|
||||
prev_output_tokens,
|
||||
features_only: bool = False,
|
||||
classification_head_name: Optional[str] = None,
|
||||
token_embeddings: Optional[torch.Tensor] = None,
|
||||
return_all_hiddens: bool = True,
|
||||
alignment_layer: Optional[int] = None,
|
||||
alignment_heads: Optional[int] = None,
|
||||
):
|
||||
if classification_head_name is not None:
|
||||
features_only = True
|
||||
|
||||
encoder_out = self.encoder(
|
||||
src_tokens,
|
||||
src_lengths=src_lengths,
|
||||
token_embeddings=token_embeddings,
|
||||
return_all_hiddens=return_all_hiddens
|
||||
)
|
||||
x, extra = self.decoder(
|
||||
prev_output_tokens,
|
||||
encoder_out=encoder_out,
|
||||
features_only=features_only,
|
||||
alignment_layer=alignment_layer,
|
||||
alignment_heads=alignment_heads,
|
||||
src_lengths=src_lengths,
|
||||
return_all_hiddens=return_all_hiddens,
|
||||
)
|
||||
eos: int = self.eos
|
||||
if classification_head_name is not None:
|
||||
sentence_representation = x[
|
||||
src_tokens.eq(eos), :
|
||||
].view(x.size(0), -1, x.size(-1))[:, -1, :]
|
||||
for k, head in self.classification_heads.items():
|
||||
# for torch script only supports iteration
|
||||
if k == classification_head_name:
|
||||
x = head(sentence_representation)
|
||||
break
|
||||
return x, extra
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
model_name_or_path,
|
||||
checkpoint_file="model.pt",
|
||||
data_name_or_path=".",
|
||||
bpe="gpt2",
|
||||
sample_break_mode="eos",
|
||||
**kwargs,
|
||||
):
|
||||
from fairseq import hub_utils
|
||||
|
||||
x = hub_utils.from_pretrained(
|
||||
model_name_or_path,
|
||||
checkpoint_file,
|
||||
data_name_or_path,
|
||||
archive_map=cls.hub_models(),
|
||||
bpe=bpe,
|
||||
load_checkpoint_heads=True,
|
||||
sample_break_mode=sample_break_mode,
|
||||
**kwargs,
|
||||
)
|
||||
return BARTHubInterface(x["args"], x["task"], x["models"][0])
|
||||
|
||||
def register_classification_head(
|
||||
self, name, num_classes=None, inner_dim=None, **kwargs
|
||||
):
|
||||
"""Register a classification head."""
|
||||
logger.info("Registering classification head: {0}".format(name))
|
||||
if name in self.classification_heads:
|
||||
prev_num_classes = self.classification_heads[name].out_proj.out_features
|
||||
prev_inner_dim = self.classification_heads[name].dense.out_features
|
||||
if num_classes != prev_num_classes or inner_dim != prev_inner_dim:
|
||||
logger.warning(
|
||||
're-registering head "{}" with num_classes {} (prev: {}) '
|
||||
"and inner_dim {} (prev: {})".format(
|
||||
name, num_classes, prev_num_classes, inner_dim, prev_inner_dim
|
||||
)
|
||||
)
|
||||
self.classification_heads[name] = BARTClassificationHead(
|
||||
input_dim=self.args.encoder_embed_dim,
|
||||
inner_dim=inner_dim or self.args.encoder_embed_dim,
|
||||
num_classes=num_classes,
|
||||
activation_fn=self.args.pooler_activation_fn,
|
||||
pooler_dropout=self.args.pooler_dropout,
|
||||
do_spectral_norm=getattr(
|
||||
self.args, "spectral_norm_classification_head", False
|
||||
),
|
||||
)
|
||||
|
||||
def upgrade_state_dict_named(self, state_dict, name):
|
||||
super().upgrade_state_dict_named(state_dict, name)
|
||||
|
||||
prefix = name + "." if name != "" else ""
|
||||
current_head_names = (
|
||||
[]
|
||||
if not hasattr(self, "classification_heads")
|
||||
else self.classification_heads.keys()
|
||||
)
|
||||
|
||||
# Handle new classification heads present in the state dict.
|
||||
keys_to_delete = []
|
||||
for k in state_dict.keys():
|
||||
if not k.startswith(prefix + "classification_heads."):
|
||||
continue
|
||||
|
||||
head_name = k[len(prefix + "classification_heads.") :].split(".")[0]
|
||||
num_classes = state_dict[
|
||||
prefix + "classification_heads." + head_name + ".out_proj.weight"
|
||||
].size(0)
|
||||
inner_dim = state_dict[
|
||||
prefix + "classification_heads." + head_name + ".dense.weight"
|
||||
].size(0)
|
||||
|
||||
if getattr(self.args, "load_checkpoint_heads", False):
|
||||
if head_name not in current_head_names:
|
||||
self.register_classification_head(head_name, num_classes, inner_dim)
|
||||
else:
|
||||
if head_name not in current_head_names:
|
||||
logger.warning(
|
||||
"deleting classification head ({}) from checkpoint "
|
||||
"not present in current model: {}".format(head_name, k)
|
||||
)
|
||||
keys_to_delete.append(k)
|
||||
elif (
|
||||
num_classes
|
||||
!= self.classification_heads[head_name].out_proj.out_features
|
||||
or inner_dim
|
||||
!= self.classification_heads[head_name].dense.out_features
|
||||
):
|
||||
logger.warning(
|
||||
"deleting classification head ({}) from checkpoint "
|
||||
"with different dimensions than current model: {}".format(
|
||||
head_name, k
|
||||
)
|
||||
)
|
||||
keys_to_delete.append(k)
|
||||
for k in keys_to_delete:
|
||||
del state_dict[k]
|
||||
|
||||
def truncate_emb(key):
|
||||
if key in state_dict:
|
||||
state_dict[key] = state_dict[key][:-1, :]
|
||||
|
||||
# When finetuning on translation task, remove last row of
|
||||
# embedding matrix that corresponds to mask_idx token.
|
||||
loaded_dict_size = state_dict["encoder.embed_tokens.weight"].size(0)
|
||||
if (
|
||||
loaded_dict_size == len(self.encoder.dictionary) + 1
|
||||
and "<mask>" not in self.encoder.dictionary
|
||||
):
|
||||
truncate_emb("encoder.embed_tokens.weight")
|
||||
truncate_emb("decoder.embed_tokens.weight")
|
||||
truncate_emb("encoder.output_projection.weight")
|
||||
truncate_emb("decoder.output_projection.weight")
|
||||
|
||||
# When continued pretraining on new set of languages for mbart,
|
||||
# add extra lang embeddings at the end of embed_tokens.
|
||||
# Note: newly added languages are assumed to have been added at the end.
|
||||
if self.args.task == "multilingual_denoising" and loaded_dict_size < len(
|
||||
self.encoder.dictionary
|
||||
):
|
||||
logger.info(
|
||||
"Adding extra language embeddings not found in pretrained model for "
|
||||
"continued pretraining of MBART on new set of languages."
|
||||
)
|
||||
loaded_mask_token_embedding = state_dict["encoder.embed_tokens.weight"][
|
||||
-1, :
|
||||
]
|
||||
|
||||
num_langids_to_add = len(self.encoder.dictionary) - loaded_dict_size
|
||||
embed_dim = state_dict["encoder.embed_tokens.weight"].size(1)
|
||||
|
||||
new_lang_embed_to_add = torch.zeros(num_langids_to_add, embed_dim)
|
||||
nn.init.normal_(new_lang_embed_to_add, mean=0, std=embed_dim ** -0.5)
|
||||
new_lang_embed_to_add = new_lang_embed_to_add.to(
|
||||
dtype=state_dict["encoder.embed_tokens.weight"].dtype,
|
||||
)
|
||||
|
||||
state_dict["encoder.embed_tokens.weight"] = torch.cat(
|
||||
[
|
||||
state_dict["encoder.embed_tokens.weight"][
|
||||
: loaded_dict_size - 1, :
|
||||
],
|
||||
new_lang_embed_to_add,
|
||||
loaded_mask_token_embedding.unsqueeze(0),
|
||||
]
|
||||
)
|
||||
state_dict["decoder.embed_tokens.weight"] = torch.cat(
|
||||
[
|
||||
state_dict["decoder.embed_tokens.weight"][
|
||||
: loaded_dict_size - 1, :
|
||||
],
|
||||
new_lang_embed_to_add,
|
||||
loaded_mask_token_embedding.unsqueeze(0),
|
||||
]
|
||||
)
|
||||
|
||||
# Copy any newly-added classification heads into the state dict
|
||||
# with their current weights.
|
||||
if hasattr(self, "classification_heads"):
|
||||
cur_state = self.classification_heads.state_dict()
|
||||
for k, v in cur_state.items():
|
||||
if prefix + "classification_heads." + k not in state_dict:
|
||||
logger.info("Overwriting " + prefix + "classification_heads." + k)
|
||||
state_dict[prefix + "classification_heads." + k] = v
|
||||
|
||||
|
||||
class BARTClassificationHead(nn.Module):
|
||||
"""Head for sentence-level classification tasks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim,
|
||||
inner_dim,
|
||||
num_classes,
|
||||
activation_fn,
|
||||
pooler_dropout,
|
||||
do_spectral_norm=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(input_dim, inner_dim)
|
||||
self.activation_fn = utils.get_activation_fn(activation_fn)
|
||||
self.dropout = nn.Dropout(p=pooler_dropout)
|
||||
self.out_proj = nn.Linear(inner_dim, num_classes)
|
||||
|
||||
if do_spectral_norm:
|
||||
self.out_proj = torch.nn.utils.spectral_norm(self.out_proj)
|
||||
|
||||
def forward(self, features, **kwargs):
|
||||
x = features
|
||||
x = self.dropout(x)
|
||||
x = self.dense(x)
|
||||
x = self.activation_fn(x)
|
||||
x = self.dropout(x)
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
@register_model_architecture("bart", "bart_large")
|
||||
def bart_large_architecture(args):
|
||||
args.encoder_embed_path = getattr(args, "encoder_embed_path", None)
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4 * 1024)
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 12)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
|
||||
args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
|
||||
args.encoder_learned_pos = getattr(args, "encoder_learned_pos", True)
|
||||
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", 12)
|
||||
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
|
||||
args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
|
||||
args.decoder_learned_pos = getattr(args, "decoder_learned_pos", True)
|
||||
args.attention_dropout = getattr(args, "attention_dropout", 0.0)
|
||||
args.relu_dropout = getattr(args, "relu_dropout", 0.0)
|
||||
args.dropout = getattr(args, "dropout", 0.1)
|
||||
args.max_target_positions = getattr(args, "max_target_positions", 1024)
|
||||
args.max_source_positions = getattr(args, "max_source_positions", 1024)
|
||||
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", True
|
||||
)
|
||||
args.share_all_embeddings = getattr(args, "share_all_embeddings", True)
|
||||
|
||||
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)
|
||||
|
||||
args.no_scale_embedding = getattr(args, "no_scale_embedding", True)
|
||||
args.layernorm_embedding = getattr(args, "layernorm_embedding", True)
|
||||
|
||||
args.activation_fn = getattr(args, "activation_fn", "gelu")
|
||||
args.pooler_activation_fn = getattr(args, "pooler_activation_fn", "tanh")
|
||||
args.pooler_dropout = getattr(args, "pooler_dropout", 0.0)
|
||||
|
||||
|
||||
@register_model_architecture("bart", "bart_base")
|
||||
def bart_base_architecture(args):
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4 * 768)
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 6)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 12)
|
||||
args.decoder_layers = getattr(args, "decoder_layers", 6)
|
||||
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 12)
|
||||
bart_large_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("bart", "mbart_large")
|
||||
def mbart_large_architecture(args):
|
||||
args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
|
||||
bart_large_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("bart", "mbart_base")
|
||||
def mbart_base_architecture(args):
|
||||
args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
|
||||
bart_base_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("bart", "mbart_base_wmt20")
|
||||
def mbart_base_wmt20_architecture(args):
|
||||
args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
|
||||
mbart_base_architecture(args)
|
||||
Reference in New Issue
Block a user