chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,10 @@
|
||||
# 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
|
||||
from .model_camembert import * # noqa
|
||||
from .model_gottbert import * # noqa
|
||||
from .model_xlmr import * # noqa
|
||||
@@ -0,0 +1,118 @@
|
||||
# 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 collections import Counter
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def align_bpe_to_words(roberta, bpe_tokens: torch.LongTensor, other_tokens: List[str]):
|
||||
"""
|
||||
Helper to align GPT-2 BPE to other tokenization formats (e.g., spaCy).
|
||||
|
||||
Args:
|
||||
roberta (RobertaHubInterface): RoBERTa instance
|
||||
bpe_tokens (torch.LongTensor): GPT-2 BPE tokens of shape `(T_bpe)`
|
||||
other_tokens (List[str]): other tokens of shape `(T_words)`
|
||||
|
||||
Returns:
|
||||
List[str]: mapping from *other_tokens* to corresponding *bpe_tokens*.
|
||||
"""
|
||||
assert bpe_tokens.dim() == 1
|
||||
assert bpe_tokens[0] == 0
|
||||
|
||||
def clean(text):
|
||||
return text.strip()
|
||||
|
||||
# remove whitespaces to simplify alignment
|
||||
bpe_tokens = [roberta.task.source_dictionary.string([x]) for x in bpe_tokens]
|
||||
bpe_tokens = [
|
||||
clean(roberta.bpe.decode(x) if x not in {"<s>", ""} else x) for x in bpe_tokens
|
||||
]
|
||||
other_tokens = [clean(str(o)) for o in other_tokens]
|
||||
|
||||
# strip leading <s>
|
||||
bpe_tokens = bpe_tokens[1:]
|
||||
assert "".join(bpe_tokens) == "".join(other_tokens)
|
||||
|
||||
# create alignment from every word to a list of BPE tokens
|
||||
alignment = []
|
||||
bpe_toks = filter(lambda item: item[1] != "", enumerate(bpe_tokens, start=1))
|
||||
j, bpe_tok = next(bpe_toks)
|
||||
for other_tok in other_tokens:
|
||||
bpe_indices = []
|
||||
while True:
|
||||
if other_tok.startswith(bpe_tok):
|
||||
bpe_indices.append(j)
|
||||
other_tok = other_tok[len(bpe_tok) :]
|
||||
try:
|
||||
j, bpe_tok = next(bpe_toks)
|
||||
except StopIteration:
|
||||
j, bpe_tok = None, None
|
||||
elif bpe_tok.startswith(other_tok):
|
||||
# other_tok spans multiple BPE tokens
|
||||
bpe_indices.append(j)
|
||||
bpe_tok = bpe_tok[len(other_tok) :]
|
||||
other_tok = ""
|
||||
else:
|
||||
raise Exception('Cannot align "{}" and "{}"'.format(other_tok, bpe_tok))
|
||||
if other_tok == "":
|
||||
break
|
||||
assert len(bpe_indices) > 0
|
||||
alignment.append(bpe_indices)
|
||||
assert len(alignment) == len(other_tokens)
|
||||
|
||||
return alignment
|
||||
|
||||
|
||||
def align_features_to_words(roberta, features, alignment):
|
||||
"""
|
||||
Align given features to words.
|
||||
|
||||
Args:
|
||||
roberta (RobertaHubInterface): RoBERTa instance
|
||||
features (torch.Tensor): features to align of shape `(T_bpe x C)`
|
||||
alignment: alignment between BPE tokens and words returned by
|
||||
func:`align_bpe_to_words`.
|
||||
"""
|
||||
assert features.dim() == 2
|
||||
|
||||
bpe_counts = Counter(j for bpe_indices in alignment for j in bpe_indices)
|
||||
assert bpe_counts[0] == 0 # <s> shouldn't be aligned
|
||||
denom = features.new([bpe_counts.get(j, 1) for j in range(len(features))])
|
||||
weighted_features = features / denom.unsqueeze(-1)
|
||||
|
||||
output = [weighted_features[0]]
|
||||
largest_j = -1
|
||||
for bpe_indices in alignment:
|
||||
output.append(weighted_features[bpe_indices].sum(dim=0))
|
||||
largest_j = max(largest_j, *bpe_indices)
|
||||
for j in range(largest_j + 1, len(features)):
|
||||
output.append(weighted_features[j])
|
||||
output = torch.stack(output)
|
||||
assert torch.all(torch.abs(output.sum(dim=0) - features.sum(dim=0)) < 1e-4)
|
||||
return output
|
||||
|
||||
|
||||
def spacy_nlp():
|
||||
if getattr(spacy_nlp, "_nlp", None) is None:
|
||||
try:
|
||||
from spacy.lang.en import English
|
||||
|
||||
spacy_nlp._nlp = English()
|
||||
except ImportError:
|
||||
raise ImportError("Please install spacy with: pip install spacy")
|
||||
return spacy_nlp._nlp
|
||||
|
||||
|
||||
def spacy_tokenizer():
|
||||
if getattr(spacy_tokenizer, "_tokenizer", None) is None:
|
||||
try:
|
||||
nlp = spacy_nlp()
|
||||
spacy_tokenizer._tokenizer = nlp.Defaults.create_tokenizer(nlp)
|
||||
except ImportError:
|
||||
raise ImportError("Please install spacy with: pip install spacy")
|
||||
return spacy_tokenizer._tokenizer
|
||||
@@ -0,0 +1,235 @@
|
||||
# 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 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
|
||||
|
||||
|
||||
class RobertaHubInterface(nn.Module):
|
||||
"""A simple PyTorch Hub interface to RoBERTa.
|
||||
|
||||
Usage: https://github.com/pytorch/fairseq/tree/master/examples/roberta
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, task, model):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.task = task
|
||||
self.model = model
|
||||
|
||||
self.bpe = encoders.build_bpe(cfg.bpe)
|
||||
|
||||
# this is useful for determining the device
|
||||
self.register_buffer("_float_tensor", torch.tensor([0], dtype=torch.float))
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self._float_tensor.device
|
||||
|
||||
def encode(
|
||||
self, sentence: str, *addl_sentences, no_separator=False
|
||||
) -> 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>`) and we use an
|
||||
extra end-of-sentence (`</s>`) as a separator.
|
||||
|
||||
Example (single sentence): `<s> a b c </s>`
|
||||
Example (sentence pair): `<s> d e f </s> </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::
|
||||
|
||||
>>> roberta.encode('Hello world').tolist()
|
||||
[0, 31414, 232, 2]
|
||||
>>> roberta.encode(' world').tolist()
|
||||
[0, 232, 2]
|
||||
>>> roberta.encode('world').tolist()
|
||||
[0, 8331, 2]
|
||||
"""
|
||||
bpe_sentence = "<s> " + self.bpe.encode(sentence) + " </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, add_if_not_exist=False
|
||||
)
|
||||
return tokens.long()
|
||||
|
||||
def decode(self, tokens: torch.LongTensor):
|
||||
assert tokens.dim() == 1
|
||||
tokens = tokens.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 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) > self.model.max_positions():
|
||||
raise ValueError(
|
||||
"tokens exceeds maximum length: {} > {}".format(
|
||||
tokens.size(-1), self.model.max_positions()
|
||||
)
|
||||
)
|
||||
features, extra = self.model(
|
||||
tokens.to(device=self.device),
|
||||
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):
|
||||
features = self.extract_features(tokens.to(device=self.device))
|
||||
logits = self.model.classification_heads[head](features)
|
||||
if return_logits:
|
||||
return logits
|
||||
return F.log_softmax(logits, dim=-1)
|
||||
|
||||
def extract_features_aligned_to_words(
|
||||
self, sentence: str, return_all_hiddens: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""Extract RoBERTa features, aligned to spaCy's word-level tokenizer."""
|
||||
from fairseq.models.roberta import alignment_utils
|
||||
from spacy.tokens import Doc
|
||||
|
||||
nlp = alignment_utils.spacy_nlp()
|
||||
tokenizer = alignment_utils.spacy_tokenizer()
|
||||
|
||||
# tokenize both with GPT-2 BPE and spaCy
|
||||
bpe_toks = self.encode(sentence)
|
||||
spacy_toks = tokenizer(sentence)
|
||||
spacy_toks_ws = [t.text_with_ws for t in tokenizer(sentence)]
|
||||
alignment = alignment_utils.align_bpe_to_words(self, bpe_toks, spacy_toks_ws)
|
||||
|
||||
# extract features and align them
|
||||
features = self.extract_features(
|
||||
bpe_toks, return_all_hiddens=return_all_hiddens
|
||||
)
|
||||
features = features.squeeze(0)
|
||||
aligned_feats = alignment_utils.align_features_to_words(
|
||||
self, features, alignment
|
||||
)
|
||||
|
||||
# wrap in spaCy Doc
|
||||
doc = Doc(
|
||||
nlp.vocab,
|
||||
words=["<s>"] + [x.text for x in spacy_toks] + ["</s>"],
|
||||
spaces=[True]
|
||||
+ [x.endswith(" ") for x in spacy_toks_ws[:-1]]
|
||||
+ [True, False],
|
||||
)
|
||||
assert len(doc) == aligned_feats.size(0)
|
||||
doc.user_token_hooks["vector"] = lambda token: aligned_feats[token.i]
|
||||
return doc
|
||||
|
||||
def fill_mask(self, masked_input: str, topk: int = 5):
|
||||
masked_token = "<mask>"
|
||||
assert (
|
||||
masked_token in masked_input and masked_input.count(masked_token) == 1
|
||||
), "Please add one {0} token for the input, eg: 'He is a {0} guy'".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,
|
||||
)
|
||||
|
||||
masked_index = (tokens == self.task.mask_idx).nonzero(as_tuple=False)
|
||||
if tokens.dim() == 1:
|
||||
tokens = tokens.unsqueeze(0)
|
||||
|
||||
with utils.model_eval(self.model):
|
||||
features, extra = self.model(
|
||||
tokens.long().to(device=self.device),
|
||||
features_only=False,
|
||||
return_all_hiddens=False,
|
||||
)
|
||||
logits = features[0, masked_index, :].squeeze()
|
||||
prob = logits.softmax(dim=0)
|
||||
values, index = prob.topk(k=topk, dim=0)
|
||||
topk_predicted_token_bpe = self.task.source_dictionary.string(index)
|
||||
|
||||
topk_filled_outputs = []
|
||||
for index, predicted_token_bpe in enumerate(
|
||||
topk_predicted_token_bpe.split(" ")
|
||||
):
|
||||
predicted_token = self.bpe.decode(predicted_token_bpe)
|
||||
# Quick hack to fix https://github.com/pytorch/fairseq/issues/1306
|
||||
if predicted_token_bpe.startswith("\u2581"):
|
||||
predicted_token = " " + predicted_token
|
||||
if " {0}".format(masked_token) in masked_input:
|
||||
topk_filled_outputs.append(
|
||||
(
|
||||
masked_input.replace(
|
||||
" {0}".format(masked_token), predicted_token
|
||||
),
|
||||
values[index].item(),
|
||||
predicted_token,
|
||||
)
|
||||
)
|
||||
else:
|
||||
topk_filled_outputs.append(
|
||||
(
|
||||
masked_input.replace(masked_token, predicted_token),
|
||||
values[index].item(),
|
||||
predicted_token,
|
||||
)
|
||||
)
|
||||
return topk_filled_outputs
|
||||
|
||||
def disambiguate_pronoun(self, sentence: str) -> bool:
|
||||
"""
|
||||
Usage::
|
||||
|
||||
>>> disambiguate_pronoun('The _trophy_ would not fit in the brown suitcase because [it] was too big.')
|
||||
True
|
||||
|
||||
>>> disambiguate_pronoun('The trophy would not fit in the brown suitcase because [it] was too big.')
|
||||
'The trophy'
|
||||
"""
|
||||
assert hasattr(
|
||||
self.task, "disambiguate_pronoun"
|
||||
), "roberta.disambiguate_pronoun() requires a model trained with the WSC task."
|
||||
with utils.model_eval(self.model):
|
||||
return self.task.disambiguate_pronoun(
|
||||
self.model, sentence, use_cuda=self.device.type == "cuda"
|
||||
)
|
||||
@@ -0,0 +1,569 @@
|
||||
# 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.
|
||||
"""
|
||||
RoBERTa: A Robustly Optimized BERT Pretraining Approach.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.models import (
|
||||
FairseqEncoder,
|
||||
FairseqEncoderModel,
|
||||
register_model,
|
||||
register_model_architecture,
|
||||
)
|
||||
from fairseq.models.transformer import TransformerEncoder
|
||||
from fairseq.modules import LayerNorm
|
||||
from fairseq.modules.quant_noise import quant_noise as apply_quant_noise_
|
||||
from fairseq.modules.transformer_sentence_encoder import init_bert_params
|
||||
|
||||
from .hub_interface import RobertaHubInterface
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_model("roberta")
|
||||
class RobertaModel(FairseqEncoderModel):
|
||||
@classmethod
|
||||
def hub_models(cls):
|
||||
return {
|
||||
"roberta.base": "http://dl.fbaipublicfiles.com/fairseq/models/roberta.base.tar.gz",
|
||||
"roberta.large": "http://dl.fbaipublicfiles.com/fairseq/models/roberta.large.tar.gz",
|
||||
"roberta.large.mnli": "http://dl.fbaipublicfiles.com/fairseq/models/roberta.large.mnli.tar.gz",
|
||||
"roberta.large.wsc": "http://dl.fbaipublicfiles.com/fairseq/models/roberta.large.wsc.tar.gz",
|
||||
}
|
||||
|
||||
def __init__(self, args, encoder):
|
||||
super().__init__(encoder)
|
||||
self.args = args
|
||||
|
||||
# We follow BERT's random weight initialization
|
||||
self.apply(init_bert_params)
|
||||
|
||||
self.classification_heads = nn.ModuleDict()
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
"""Add model-specific arguments to the parser."""
|
||||
parser.add_argument(
|
||||
"--encoder-layers", type=int, metavar="L", help="num encoder layers"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--encoder-embed-dim",
|
||||
type=int,
|
||||
metavar="H",
|
||||
help="encoder embedding dimension",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--encoder-ffn-embed-dim",
|
||||
type=int,
|
||||
metavar="F",
|
||||
help="encoder embedding dimension for FFN",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--encoder-attention-heads",
|
||||
type=int,
|
||||
metavar="A",
|
||||
help="num encoder attention heads",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--activation-fn",
|
||||
choices=utils.get_available_activation_fns(),
|
||||
help="activation function to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pooler-activation-fn",
|
||||
choices=utils.get_available_activation_fns(),
|
||||
help="activation function to use for pooler layer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--encoder-normalize-before",
|
||||
action="store_true",
|
||||
help="apply layernorm before each encoder block",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layernorm-embedding",
|
||||
action="store_true",
|
||||
help="add layernorm to embedding",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dropout", type=float, metavar="D", help="dropout probability"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--attention-dropout",
|
||||
type=float,
|
||||
metavar="D",
|
||||
help="dropout probability for attention weights",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--activation-dropout",
|
||||
type=float,
|
||||
metavar="D",
|
||||
help="dropout probability after activation in FFN",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pooler-dropout",
|
||||
type=float,
|
||||
metavar="D",
|
||||
help="dropout probability in the masked_lm pooler layers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-positions", type=int, help="number of positional embeddings to learn"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--load-checkpoint-heads",
|
||||
action="store_true",
|
||||
help="(re-)register and load heads when loading checkpoints",
|
||||
)
|
||||
# args for "Reducing Transformer Depth on Demand with Structured Dropout" (Fan et al., 2019)
|
||||
parser.add_argument(
|
||||
"--encoder-layerdrop",
|
||||
type=float,
|
||||
metavar="D",
|
||||
default=0,
|
||||
help="LayerDrop probability for encoder",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--encoder-layers-to-keep",
|
||||
default=None,
|
||||
help="which layers to *keep* when pruning as a comma-separated list",
|
||||
)
|
||||
# args for Training with Quantization Noise for Extreme Model Compression ({Fan*, Stock*} et al., 2020)
|
||||
parser.add_argument(
|
||||
"--quant-noise-pq",
|
||||
type=float,
|
||||
metavar="D",
|
||||
default=0,
|
||||
help="iterative PQ quantization noise at training time",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant-noise-pq-block-size",
|
||||
type=int,
|
||||
metavar="D",
|
||||
default=8,
|
||||
help="block size of quantization noise at training time",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant-noise-scalar",
|
||||
type=float,
|
||||
metavar="D",
|
||||
default=0,
|
||||
help="scalar quantization noise and scalar quantization at training time",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--untie-weights-roberta",
|
||||
action="store_true",
|
||||
help="Untie weights between embeddings and classifiers in RoBERTa",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spectral-norm-classification-head",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Apply spectral normalization on the classification head",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_model(cls, args, task):
|
||||
"""Build a new model instance."""
|
||||
|
||||
# make sure all arguments are present
|
||||
base_architecture(args)
|
||||
|
||||
if not hasattr(args, "max_positions"):
|
||||
args.max_positions = args.tokens_per_sample
|
||||
|
||||
encoder = RobertaEncoder(args, task.source_dictionary)
|
||||
return cls(args, encoder)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
src_tokens,
|
||||
features_only=False,
|
||||
return_all_hiddens=False,
|
||||
classification_head_name=None,
|
||||
**kwargs
|
||||
):
|
||||
if classification_head_name is not None:
|
||||
features_only = True
|
||||
|
||||
x, extra = self.encoder(src_tokens, features_only, return_all_hiddens, **kwargs)
|
||||
|
||||
if classification_head_name is not None:
|
||||
x = self.classification_heads[classification_head_name](x)
|
||||
return x, extra
|
||||
|
||||
def get_normalized_probs(self, net_output, log_probs, sample=None):
|
||||
"""Get normalized probabilities (or log probs) from a net's output."""
|
||||
logits = net_output[0].float()
|
||||
if log_probs:
|
||||
return F.log_softmax(logits, dim=-1)
|
||||
else:
|
||||
return F.softmax(logits, dim=-1)
|
||||
|
||||
def register_classification_head(
|
||||
self, name, num_classes=None, inner_dim=None, **kwargs
|
||||
):
|
||||
"""Register a classification head."""
|
||||
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] = RobertaClassificationHead(
|
||||
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,
|
||||
q_noise=self.args.quant_noise_pq,
|
||||
qn_block_size=self.args.quant_noise_pq_block_size,
|
||||
do_spectral_norm=self.args.spectral_norm_classification_head,
|
||||
)
|
||||
|
||||
@property
|
||||
def supported_targets(self):
|
||||
return {"self"}
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
model_name_or_path,
|
||||
checkpoint_file="model.pt",
|
||||
data_name_or_path=".",
|
||||
bpe="gpt2",
|
||||
**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,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
logger.info(x["args"])
|
||||
return RobertaHubInterface(x["args"], x["task"], x["models"][0])
|
||||
|
||||
def upgrade_state_dict_named(self, state_dict, name):
|
||||
prefix = name + "." if name != "" else ""
|
||||
|
||||
# rename decoder -> encoder before upgrading children modules
|
||||
for k in list(state_dict.keys()):
|
||||
if k.startswith(prefix + "decoder"):
|
||||
new_k = prefix + "encoder" + k[len(prefix + "decoder") :]
|
||||
state_dict[new_k] = state_dict[k]
|
||||
del state_dict[k]
|
||||
|
||||
# rename emb_layer_norm -> layernorm_embedding
|
||||
for k in list(state_dict.keys()):
|
||||
if ".emb_layer_norm." in k:
|
||||
new_k = k.replace(".emb_layer_norm.", ".layernorm_embedding.")
|
||||
state_dict[new_k] = state_dict[k]
|
||||
del state_dict[k]
|
||||
|
||||
# upgrade children modules
|
||||
super().upgrade_state_dict_named(state_dict, name)
|
||||
|
||||
# Handle new classification heads present in the state dict.
|
||||
current_head_names = (
|
||||
[]
|
||||
if not hasattr(self, "classification_heads")
|
||||
else self.classification_heads.keys()
|
||||
)
|
||||
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]
|
||||
|
||||
# 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 RobertaLMHead(nn.Module):
|
||||
"""Head for masked language modeling."""
|
||||
|
||||
def __init__(self, embed_dim, output_dim, activation_fn, weight=None):
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(embed_dim, embed_dim)
|
||||
self.activation_fn = utils.get_activation_fn(activation_fn)
|
||||
self.layer_norm = LayerNorm(embed_dim)
|
||||
|
||||
if weight is None:
|
||||
weight = nn.Linear(embed_dim, output_dim, bias=False).weight
|
||||
self.weight = weight
|
||||
self.bias = nn.Parameter(torch.zeros(output_dim))
|
||||
|
||||
def forward(self, features, masked_tokens=None, **kwargs):
|
||||
# Only project the masked tokens while training,
|
||||
# saves both memory and computation
|
||||
if masked_tokens is not None:
|
||||
features = features[masked_tokens, :]
|
||||
|
||||
x = self.dense(features)
|
||||
x = self.activation_fn(x)
|
||||
x = self.layer_norm(x)
|
||||
# project back to size of vocabulary with bias
|
||||
x = F.linear(x, self.weight) + self.bias
|
||||
return x
|
||||
|
||||
|
||||
class RobertaClassificationHead(nn.Module):
|
||||
"""Head for sentence-level classification tasks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim,
|
||||
inner_dim,
|
||||
num_classes,
|
||||
activation_fn,
|
||||
pooler_dropout,
|
||||
q_noise=0,
|
||||
qn_block_size=8,
|
||||
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 = apply_quant_noise_(
|
||||
nn.Linear(inner_dim, num_classes), q_noise, qn_block_size
|
||||
)
|
||||
if do_spectral_norm:
|
||||
if q_noise != 0:
|
||||
raise NotImplementedError(
|
||||
"Attempting to use Spectral Normalization with Quant Noise. This is not officially supported"
|
||||
)
|
||||
self.out_proj = torch.nn.utils.spectral_norm(self.out_proj)
|
||||
|
||||
def forward(self, features, **kwargs):
|
||||
x = features[:, 0, :] # take <s> token (equiv. to [CLS])
|
||||
x = self.dropout(x)
|
||||
x = self.dense(x)
|
||||
x = self.activation_fn(x)
|
||||
x = self.dropout(x)
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class RobertaEncoder(FairseqEncoder):
|
||||
"""RoBERTa encoder."""
|
||||
|
||||
def __init__(self, args, dictionary):
|
||||
super().__init__(dictionary)
|
||||
|
||||
# set any missing default values
|
||||
base_architecture(args)
|
||||
self.args = args
|
||||
|
||||
if args.encoder_layers_to_keep:
|
||||
args.encoder_layers = len(args.encoder_layers_to_keep.split(","))
|
||||
|
||||
embed_tokens = self.build_embedding(
|
||||
len(dictionary), args.encoder_embed_dim, dictionary.pad()
|
||||
)
|
||||
|
||||
self.sentence_encoder = self.build_encoder(args, dictionary, embed_tokens)
|
||||
|
||||
self.lm_head = self.build_lm_head(
|
||||
embed_dim=args.encoder_embed_dim,
|
||||
output_dim=len(dictionary),
|
||||
activation_fn=args.activation_fn,
|
||||
weight=(
|
||||
self.sentence_encoder.embed_tokens.weight
|
||||
if not args.untie_weights_roberta
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
def build_embedding(self, vocab_size, embedding_dim, padding_idx):
|
||||
return nn.Embedding(vocab_size, embedding_dim, padding_idx)
|
||||
|
||||
def build_encoder(self, args, dictionary, embed_tokens):
|
||||
encoder = TransformerEncoder(args, dictionary, embed_tokens)
|
||||
encoder.apply(init_bert_params)
|
||||
return encoder
|
||||
|
||||
def build_lm_head(self, embed_dim, output_dim, activation_fn, weight):
|
||||
return RobertaLMHead(embed_dim, output_dim, activation_fn, weight)
|
||||
|
||||
def build_lm_head(self, embed_dim, output_dim, activation_fn, weight):
|
||||
return RobertaLMHead(embed_dim, output_dim, activation_fn, weight)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
src_tokens,
|
||||
features_only=False,
|
||||
return_all_hiddens=False,
|
||||
masked_tokens=None,
|
||||
**unused
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
src_tokens (LongTensor): input tokens of shape `(batch, src_len)`
|
||||
features_only (bool, optional): skip LM head and just return
|
||||
features. If True, the output will be of shape
|
||||
`(batch, src_len, embed_dim)`.
|
||||
return_all_hiddens (bool, optional): also return all of the
|
||||
intermediate hidden states (default: False).
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- the LM output of shape `(batch, src_len, vocab)`
|
||||
- a dictionary of additional data, where 'inner_states'
|
||||
is a list of hidden states. Note that the hidden
|
||||
states have shape `(src_len, batch, vocab)`.
|
||||
"""
|
||||
x, extra = self.extract_features(
|
||||
src_tokens, return_all_hiddens=return_all_hiddens
|
||||
)
|
||||
if not features_only:
|
||||
x = self.output_layer(x, masked_tokens=masked_tokens)
|
||||
return x, extra
|
||||
|
||||
def extract_features(self, src_tokens, return_all_hiddens=False, **kwargs):
|
||||
encoder_out = self.sentence_encoder(
|
||||
src_tokens,
|
||||
return_all_hiddens=return_all_hiddens,
|
||||
token_embeddings=kwargs.get("token_embeddings", None),
|
||||
)
|
||||
# T x B x C -> B x T x C
|
||||
features = encoder_out["encoder_out"][0].transpose(0, 1)
|
||||
inner_states = encoder_out["encoder_states"] if return_all_hiddens else None
|
||||
return features, {"inner_states": inner_states}
|
||||
|
||||
def output_layer(self, features, masked_tokens=None, **unused):
|
||||
return self.lm_head(features, masked_tokens)
|
||||
|
||||
def max_positions(self):
|
||||
"""Maximum output length supported by the encoder."""
|
||||
return self.args.max_positions
|
||||
|
||||
|
||||
@register_model_architecture("roberta", "roberta")
|
||||
def base_architecture(args):
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 12)
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 3072)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 12)
|
||||
|
||||
args.dropout = getattr(args, "dropout", 0.1)
|
||||
args.attention_dropout = getattr(args, "attention_dropout", 0.1)
|
||||
args.activation_dropout = getattr(args, "activation_dropout", 0.0)
|
||||
args.pooler_dropout = getattr(args, "pooler_dropout", 0.0)
|
||||
|
||||
args.max_source_positions = getattr(args, "max_positions", 512)
|
||||
args.no_token_positional_embeddings = getattr(
|
||||
args, "no_token_positional_embeddings", False
|
||||
)
|
||||
|
||||
# BERT has a few structural differences compared to the original Transformer
|
||||
args.encoder_learned_pos = getattr(args, "encoder_learned_pos", True)
|
||||
args.layernorm_embedding = getattr(args, "layernorm_embedding", True)
|
||||
args.no_scale_embedding = getattr(args, "no_scale_embedding", True)
|
||||
args.activation_fn = getattr(args, "activation_fn", "gelu")
|
||||
args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
|
||||
args.pooler_activation_fn = getattr(args, "pooler_activation_fn", "tanh")
|
||||
args.untie_weights_roberta = getattr(args, "untie_weights_roberta", False)
|
||||
|
||||
# Adaptive input config
|
||||
args.adaptive_input = getattr(args, "adaptive_input", False)
|
||||
|
||||
# LayerDrop config
|
||||
args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.0)
|
||||
args.encoder_layers_to_keep = getattr(args, "encoder_layers_to_keep", None)
|
||||
|
||||
# Quantization noise config
|
||||
args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)
|
||||
args.quant_noise_pq_block_size = getattr(args, "quant_noise_pq_block_size", 8)
|
||||
args.quant_noise_scalar = getattr(args, "quant_noise_scalar", 0)
|
||||
|
||||
# R4F config
|
||||
args.spectral_norm_classification_head = getattr(
|
||||
args, "spectral_norm_classification_head", False
|
||||
)
|
||||
|
||||
|
||||
@register_model_architecture("roberta", "roberta_prenorm")
|
||||
def roberta_prenorm_architecture(args):
|
||||
args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
|
||||
args.encoder_normalize_before = getattr(args, "encoder_normalize_before", True)
|
||||
base_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("roberta", "roberta_base")
|
||||
def roberta_base_architecture(args):
|
||||
base_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("roberta", "roberta_large")
|
||||
def roberta_large_architecture(args):
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 24)
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
|
||||
base_architecture(args)
|
||||
|
||||
|
||||
@register_model_architecture("roberta", "xlm")
|
||||
def xlm_architecture(args):
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 16)
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1280)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 1280 * 4)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
|
||||
base_architecture(args)
|
||||
@@ -0,0 +1,50 @@
|
||||
# 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.
|
||||
"""
|
||||
CamemBERT: a Tasty French Language Model
|
||||
"""
|
||||
|
||||
from fairseq.models import register_model
|
||||
|
||||
from .hub_interface import RobertaHubInterface
|
||||
from .model import RobertaModel
|
||||
|
||||
|
||||
@register_model("camembert")
|
||||
class CamembertModel(RobertaModel):
|
||||
@classmethod
|
||||
def hub_models(cls):
|
||||
return {
|
||||
"camembert": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base.tar.gz",
|
||||
"camembert.v0": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base.tar.gz",
|
||||
"camembert-base": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base.tar.gz",
|
||||
"camembert-large": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-large.tar.gz",
|
||||
"camembert-base-ccnet": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base-ccnet.tar.gz",
|
||||
"camembert-base-ccnet-4gb": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base-ccnet-4gb.tar.gz",
|
||||
"camembert-base-wikipedia-4gb": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base-wikipedia-4gb.tar.gz",
|
||||
"camembert-base-oscar-4gb": "http://dl.fbaipublicfiles.com/fairseq/models/camembert-base-oscar-4gb.tar.gz",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
model_name_or_path,
|
||||
checkpoint_file="model.pt",
|
||||
data_name_or_path=".",
|
||||
bpe="sentencepiece",
|
||||
**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,
|
||||
**kwargs,
|
||||
)
|
||||
return RobertaHubInterface(x["args"], x["task"], x["models"][0])
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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.
|
||||
"""
|
||||
GottBERT: a pure German Language Model
|
||||
"""
|
||||
|
||||
from fairseq.models import register_model
|
||||
|
||||
from .hub_interface import RobertaHubInterface
|
||||
from .model import RobertaModel
|
||||
|
||||
|
||||
@register_model('gottbert')
|
||||
class GottbertModel(RobertaModel):
|
||||
|
||||
@classmethod
|
||||
def hub_models(cls):
|
||||
return {
|
||||
'gottbert-base': 'https://dl.gottbert.de/fairseq/models/gottbert-base.tar.gz',
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls,
|
||||
model_name_or_path,
|
||||
checkpoint_file='model.pt',
|
||||
data_name_or_path='.',
|
||||
bpe='hf_byte_bpe',
|
||||
bpe_vocab='vocab.json',
|
||||
bpe_merges='merges.txt',
|
||||
bpe_add_prefix_space=False,
|
||||
**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,
|
||||
bpe_vocab=bpe_vocab,
|
||||
bpe_merges=bpe_merges,
|
||||
bpe_add_prefix_space=bpe_add_prefix_space,
|
||||
**kwargs,
|
||||
)
|
||||
return RobertaHubInterface(x['args'], x['task'], x['models'][0])
|
||||
@@ -0,0 +1,44 @@
|
||||
# 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.
|
||||
"""
|
||||
Unsupervised Cross-lingual Representation Learning at Scale
|
||||
"""
|
||||
|
||||
from fairseq.models import register_model
|
||||
|
||||
from .hub_interface import RobertaHubInterface
|
||||
from .model import RobertaModel
|
||||
|
||||
|
||||
@register_model("xlmr")
|
||||
class XLMRModel(RobertaModel):
|
||||
@classmethod
|
||||
def hub_models(cls):
|
||||
return {
|
||||
"xlmr.base": "http://dl.fbaipublicfiles.com/fairseq/models/xlmr.base.tar.gz",
|
||||
"xlmr.large": "http://dl.fbaipublicfiles.com/fairseq/models/xlmr.large.tar.gz",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
model_name_or_path,
|
||||
checkpoint_file="model.pt",
|
||||
data_name_or_path=".",
|
||||
bpe="sentencepiece",
|
||||
**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,
|
||||
**kwargs,
|
||||
)
|
||||
return RobertaHubInterface(x["args"], x["task"], x["models"][0])
|
||||
Reference in New Issue
Block a user