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)