Source code for texar.torch.modules.pretrained.roberta

# Copyright 2019 The Texar Authors. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
Utils of RoBERTa Modules.

import os
from abc import ABC
from typing import Any, Dict

import torch

from texar.torch.modules.pretrained.pretrained_base import PretrainedMixin

__all__ = [


[docs]class PretrainedRoBERTaMixin(PretrainedMixin, ABC): r"""A mixin class to support loading pre-trained checkpoints for modules that implement the RoBERTa model. The RoBERTa model was proposed in (`Liu et al`. 2019) `RoBERTa: A Robustly Optimized BERT Pretraining Approach`_. As a variant of the standard BERT model, RoBERTa trains for more iterations on more data with a larger batch size as well as other tweaks in pre-training. Differing from the standard BERT, the RoBERTa model does not use segmentation embedding. Available model names include: * ``roberta-base``: RoBERTa using the BERT-base architecture, 125M parameters. * ``roberta-large``: RoBERTa using the BERT-large architecture, 355M parameters. We provide the following RoBERTa classes: * :class:`~texar.torch.modules.RoBERTaEncoder` for text encoding. * :class:`~texar.torch.modules.RoBERTaClassifier` for text classification and sequence tagging. .. _`RoBERTa: A Robustly Optimized BERT Pretraining Approach`: """ _MODEL_NAME = "RoBERTa" _MODEL2URL = { 'roberta-base': _ROBERTA_PATH + "roberta.base.tar.gz", 'roberta-large': _ROBERTA_PATH + "roberta.large.tar.gz", } @classmethod def _transform_config(cls, pretrained_model_name: str, cache_dir: str) -> Dict[str, Any]: info = list(os.walk(cache_dir)) root, _, files = info[0] config_path = None for file in files: if file.endswith(''): config_path = os.path.join(root, file) args = torch.load(config_path, map_location="cpu")['args'] hidden_dim = args.encoder_embed_dim vocab_size = 50265 position_size = args.max_positions + 2 embedding_dropout = args.dropout num_blocks = args.encoder_layers num_heads = args.encoder_attention_heads dropout_rate = args.attention_dropout residual_dropout = args.dropout intermediate_size = args.encoder_ffn_embed_dim hidden_act = args.activation_fn if config_path is None: raise ValueError(f"Cannot find the config file in {cache_dir}") configs = { 'hidden_size': hidden_dim, 'embed': { 'name': 'word_embeddings', 'dim': hidden_dim }, 'vocab_size': vocab_size, 'position_embed': { 'name': 'position_embeddings', 'dim': hidden_dim }, 'position_size': position_size, 'encoder': { 'name': 'encoder', 'embedding_dropout': embedding_dropout, 'num_blocks': num_blocks, 'multihead_attention': { 'use_bias': True, 'num_units': hidden_dim, 'num_heads': num_heads, 'output_dim': hidden_dim, 'dropout_rate': dropout_rate, 'name': 'self' }, 'residual_dropout': residual_dropout, 'dim': hidden_dim, 'eps': 1e-12, 'use_bert_config': True, 'poswise_feedforward': { "layers": [{ 'type': 'Linear', 'kwargs': { 'in_features': hidden_dim, 'out_features': intermediate_size, 'bias': True, } }, { 'type': 'Bert' + hidden_act.upper() }, { 'type': 'Linear', 'kwargs': { 'in_features': intermediate_size, 'out_features': hidden_dim, 'bias': True, } }], }, } } return configs def _init_from_checkpoint(self, pretrained_model_name: str, cache_dir: str, **kwargs): global_tensor_map = { 'decoder.sentence_encoder.embed_tokens.weight': 'word_embedder._embedding', 'decoder.sentence_encoder.embed_positions.weight': 'position_embedder._embedding', 'decoder.sentence_encoder.emb_layer_norm.weight': 'encoder.input_normalizer.weight', 'decoder.sentence_encoder.emb_layer_norm.bias': 'encoder.input_normalizer.bias', } attention_tensor_map = { 'final_layer_norm.weight': 'encoder.output_layer_norm.{}.weight', 'final_layer_norm.bias': 'encoder.output_layer_norm.{}.bias', 'fc1.weight': 'encoder.poswise_networks.{}._layers.0.weight', 'fc1.bias': 'encoder.poswise_networks.{}._layers.0.bias', 'fc2.weight': 'encoder.poswise_networks.{}._layers.2.weight', 'fc2.bias': 'encoder.poswise_networks.{}._layers.2.bias', 'self_attn_layer_norm.weight': 'encoder.poswise_layer_norm.{}.weight', 'self_attn_layer_norm.bias': 'encoder.poswise_layer_norm.{}.bias', 'self_attn.out_proj.weight': 'encoder.self_attns.{}.O_dense.weight', 'self_attn.out_proj.bias': 'encoder.self_attns.{}.O_dense.bias', 'self_attn.in_proj_weight': [ 'encoder.self_attns.{}.Q_dense.weight', 'encoder.self_attns.{}.K_dense.weight', 'encoder.self_attns.{}.V_dense.weight', ], 'self_attn.in_proj_bias': [ 'encoder.self_attns.{}.Q_dense.bias', 'encoder.self_attns.{}.K_dense.bias', 'encoder.self_attns.{}.V_dense.bias' ], } checkpoint_path = os.path.abspath(os.path.join(cache_dir, '')) device = next(self.parameters()).device params = torch.load(checkpoint_path, map_location=device)['model'] for name, tensor in params.items(): if name in global_tensor_map: v_name = global_tensor_map[name] pointer = self._name_to_variable(v_name) assert pointer.shape == tensor.shape = elif name.startswith('decoder.sentence_encoder.layers.'): name = name.lstrip('decoder.sentence_encoder.layers.') layer_num, layer_name = name.split('.', 1) if layer_name in attention_tensor_map: v_names = attention_tensor_map[layer_name] if isinstance(v_names, str): pointer = self._name_to_variable( v_names.format(layer_num)) assert pointer.shape == tensor.shape = else: # Q, K, V in self-attention tensors = torch.chunk(tensor, chunks=3, dim=0) for i in range(3): pointer = self._name_to_variable( v_names[i].format(layer_num)) assert pointer.shape == tensors[i].shape = tensors[i].data.type(pointer.dtype) else: raise NameError(f"Layer name '{layer_name}' not found")