def get_output_from_encoder(encoder, input_ids, segment_ids,
input_mask) -> EncoderOutput:
"""Pass inputs to encoder, return encoder output.
Args:
encoder: bare model outputting raw hidden-states without any specific head.
input_ids: token indices (see huggingface.co/transformers/glossary.html#input-ids).
segment_ids: token type ids (see huggingface.co/transformers/glossary.html#token-type-ids).
input_mask: attention mask (see huggingface.co/transformers/glossary.html#attention-mask).
Raises:
RuntimeError if encoder output contains less than 2 elements.
Returns:
EncoderOutput containing pooled and unpooled model outputs as well as any other outputs.
"""
model_arch = ModelArchitectures.from_encoder(encoder)
if model_arch in [
ModelArchitectures.BERT,
ModelArchitectures.ROBERTA,
ModelArchitectures.ALBERT,
ModelArchitectures.XLM_ROBERTA,
]:
pooled, unpooled, other = get_output_from_standard_transformer_models(
encoder=encoder,
input_ids=input_ids,
segment_ids=segment_ids,
input_mask=input_mask,
)
elif model_arch == ModelArchitectures.ELECTRA:
pooled, unpooled, other = get_output_from_electra(
encoder=encoder,
input_ids=input_ids,
segment_ids=segment_ids,
input_mask=input_mask,
)
elif model_arch in [
ModelArchitectures.BART,
ModelArchitectures.MBART,
]:
pooled, unpooled, other = get_output_from_bart_models(
encoder=encoder,
input_ids=input_ids,
input_mask=input_mask,
)
elif model_arch == ModelArchitectures.DISTILBERT:
pooled, unpooled, other = get_output_from_distilbert(
encoder=encoder,
input_ids=input_ids,
input_mask=input_mask,
)
else:
raise KeyError(model_arch)
# Extend later with attention, hidden_acts, etc
if other:
return EncoderOutput(pooled=pooled, unpooled=unpooled, other=other)
else:
return EncoderOutput(pooled=pooled, unpooled=unpooled)
def load_encoder_from_transformers_weights(encoder: nn.Module,
weights_dict: dict,
return_remainder=False):
"""Find encoder weights in weights dict, load them into encoder, return any remaining weights.
TODO: clarify how we know the encoder weights will be prefixed by transformer name.
Args:
encoder (PreTrainedModel): Transformer w/o heads (embedding layer + self-attention layer).
weights_dict (Dict): model weights.
return_remainder (bool): If True, return any leftover weights.
Returns:
Dict containing any leftover weights.
"""
remainder_weights_dict = {}
load_weights_dict = {}
model_arch = ModelArchitectures.from_encoder(encoder=encoder)
encoder_prefix = MODEL_PREFIX[model_arch] + "."
# Encoder
for k, v in weights_dict.items():
if k.startswith(encoder_prefix):
load_weights_dict[strings.remove_prefix(k, encoder_prefix)] = v
else:
remainder_weights_dict[k] = v
encoder.load_state_dict(load_weights_dict)
if return_remainder:
return remainder_weights_dict
def get_model_arch_from_jiant_model(
jiant_model: nn.Module) -> ModelArchitectures:
return ModelArchitectures.from_encoder(encoder=jiant_model.encoder)