class RoBERTaEncoder(RoBERTaEncoderBase): """A PyTorch RoBERTa implementation""" class Config(RoBERTaEncoderBase.Config): num_encoder_layers: int = 12 num_attention_heads: int = 12 model_path: str = "manifold://pytext_training/tree/static/models/roberta_base_torch.pt" def __init__(self, config: Config, output_encoded_layers: bool, **kwarg) -> None: super().__init__(config, output_encoded_layers=output_encoded_layers) # assert config.pretrained_encoder.load_path, "Load path cannot be empty." self.encoder = SentenceEncoder(transformer=Transformer( embedding_dim=config.embedding_dim, layers=[ TransformerLayer( embedding_dim=config.embedding_dim, attention=MultiheadSelfAttention( config.embedding_dim, config.num_attention_heads), ) for _ in range(config.num_encoder_layers) ], )) roberta_state = torch.load( config.model_path, map_location=lambda s, l: default_restore_location(s, "cpu"), ) self.encoder.load_roberta_state_dict(roberta_state["model"]) self.representation_dim = self.encoder.transformer.token_embedding.weight.size( -1)
class RoBERTaEncoder(RoBERTaEncoderBase): """A PyTorch RoBERTa implementation""" class Config(RoBERTaEncoderBase.Config): embedding_dim: int = 768 vocab_size: int = 50265 num_encoder_layers: int = 12 num_attention_heads: int = 12 model_path: str = ( "manifold://pytext_training/tree/static/models/roberta_base_torch.pt" ) # Loading the state dict of the model depends on whether the model was # previously finetuned in PyText or not. If it was finetuned then we # dont need to translate the state dict and can just load it` # directly. is_finetuned: bool = False def __init__(self, config: Config, output_encoded_layers: bool, **kwarg) -> None: super().__init__(config, output_encoded_layers=output_encoded_layers) # map to the real model_path config.model_path = (resources.roberta.RESOURCE_MAP[config.model_path] if config.model_path in resources.roberta.RESOURCE_MAP else config.model_path) # assert config.pretrained_encoder.load_path, "Load path cannot be empty." self.encoder = SentenceEncoder(transformer=Transformer( vocab_size=config.vocab_size, embedding_dim=config.embedding_dim, layers=[ TransformerLayer( embedding_dim=config.embedding_dim, attention=MultiheadSelfAttention( config.embedding_dim, config.num_attention_heads), ) for _ in range(config.num_encoder_layers) ], )) self.apply(init_params) if config.model_path: with PathManager.open(config.model_path, "rb") as f: roberta_state = torch.load(f, map_location=lambda s, l: default_restore_location(s, "cpu")) # In case the model has previously been loaded in PyText and finetuned, # then we dont need to do the special state dict translation. Load # it directly if not config.is_finetuned: self.encoder.load_roberta_state_dict(roberta_state["model"]) else: self.load_state_dict(roberta_state) self.representation_dim = self._embedding().weight.size(-1) log_class_usage(__class__) def _embedding(self): # used to tie weights in MaskedLM model return self.encoder.transformer.token_embedding