def __init__(
self,
model_name: str,
*,
max_length: int = None,
sub_module: str = None,
train_parameters: bool = True,
last_layer_only: bool = True,
override_weights_file: Optional[str] = None,
override_weights_strip_prefix: Optional[str] = None,
masked_language_modeling: bool = True,
) -> None:
TokenEmbedder.__init__(self) # Call the base class constructor
tokenizer = PretrainedTransformerTokenizer(model_name)
self.masked_language_modeling = masked_language_modeling
if self.masked_language_modeling:
self.config = AutoConfig.from_pretrained(model_name,
output_hidden_states=True)
# We only need access to the HF tokenizer if we are masked language modeling
self.tokenizer = tokenizer.tokenizer
# The only differences when masked language modeling are:
# 1) `output_hidden_states` must be True to get access to token embeddings.
# 2) We need to use `AutoModelForMaskedLM` to get the correct model
self.transformer_model = AutoModelForMaskedLM.from_pretrained(
model_name, config=self.config)
# Eveything after the if statement (including the else) is copied directly from:
# https://github.com/allenai/allennlp/blob/master/allennlp/modules/token_embedders/pretrained_transformer_embedder.py
else:
from allennlp.common import cached_transformers
self.transformer_model = cached_transformers.get(
model_name, True, override_weights_file,
override_weights_strip_prefix)
self.config = self.transformer_model.config
if sub_module:
assert hasattr(self.transformer_model, sub_module)
self.transformer_model = getattr(self.transformer_model,
sub_module)
self._max_length = max_length
# I'm not sure if this works for all models; open an issue on github if you find a case
# where it doesn't work.
self.output_dim = self.config.hidden_size
self._scalar_mix: Optional[ScalarMix] = None
if not last_layer_only:
self._scalar_mix = ScalarMix(self.config.num_hidden_layers)
self.config.output_hidden_states = True
self._num_added_start_tokens = len(
tokenizer.single_sequence_start_tokens)
self._num_added_end_tokens = len(tokenizer.single_sequence_end_tokens)
self._num_added_tokens = self._num_added_start_tokens + self._num_added_end_tokens
if not train_parameters:
for param in self.transformer_model.parameters():
param.requires_grad = False
def __init__(
self,
model_name: str,
*,
max_length: int = None,
sub_module: str = None,
train_parameters: bool = True,
last_layer_only: bool = True,
override_weights_file: Optional[str] = None,
override_weights_strip_prefix: Optional[str] = None,
gradient_checkpointing: Optional[bool] = None,
tokenizer_kwargs: Optional[Dict[str, Any]] = None,
transformer_kwargs: Optional[Dict[str, Any]] = None,
masked_language_modeling: bool = True,
load_directory: Optional[str] = None
) -> None:
TokenEmbedder.__init__(self) # Call the base class constructor
tokenizer = PretrainedTransformerTokenizer(model_name, tokenizer_kwargs=tokenizer_kwargs)
self.masked_language_modeling = masked_language_modeling
if self.masked_language_modeling:
self.config = AutoConfig.from_pretrained(model_name, output_hidden_states=True)
# We only need access to the HF tokenizer if we are masked language modeling
self.tokenizer = tokenizer.tokenizer
# The only differences when masked language modeling are:
# 1) `output_hidden_states` must be True to get access to token embeddings.
# 2) We need to use `AutoModelForMaskedLM` to get the correct model
self.transformer_model = AutoModelForMaskedLM.from_pretrained(
# self.transformer_model = RobertaForAugment.from_pretrained()
model_name, config=self.config, **(transformer_kwargs or {})
)
if load_directory is not None:
print("Loading Model from:", load_directory)
state = torch.load(load_directory)
model_dict = self.transformer_model.state_dict()
# ckpt__dict = state['state_dict']
state = {k: v for k, v in state.items() if k in model_dict}
model_dict.update(state)
self.transformer_model.load_state_dict(model_dict, strict=False)
print("Loading Model from:", load_directory, "...Finished.")
# Eveything after the if statement (including the else) is copied directly from:
# https://github.com/allenai/allennlp/blob/master/allennlp/modules/token_embedders/pretrained_transformer_embedder.py
else:
from allennlp.common import cached_transformers
self.transformer_model = cached_transformers.get(
model_name, True, override_weights_file, override_weights_strip_prefix
)
self.config = self.transformer_model.config
if gradient_checkpointing is not None:
self.transformer_model.config.update({"gradient_checkpointing": gradient_checkpointing})
if sub_module:
assert hasattr(self.transformer_model, sub_module)
self.transformer_model = getattr(self.transformer_model, sub_module)
# print("max_length", max_length)
self._max_length = max_length
# I'm not sure if this works for all models; open an issue on github if you find a case
# where it doesn't work.
self.output_dim = self.config.hidden_size
self._scalar_mix: Optional[ScalarMix] = None
if not last_layer_only:
self._scalar_mix = ScalarMix(self.config.num_hidden_layers)
self.config.output_hidden_states = True
self._num_added_start_tokens = len(tokenizer.single_sequence_start_tokens)
self._num_added_end_tokens = len(tokenizer.single_sequence_end_tokens)
self._num_added_tokens = self._num_added_start_tokens + self._num_added_end_tokens
self.encoder = BertEncoder(self.config)
self.layer = torch.nn.ModuleList([BertLayer(self.config)
for _ in range(self.config.num_hidden_layers)])
self.embeddings = BertEmbeddings(self.config)
self.output_hidden_states = self.config.output_hidden_states
if not train_parameters:
for param in self.transformer_model.parameters():
param.requires_grad = False