| import logging | |
| import torch | |
| import torch.nn as nn | |
| from transformers import BertModel | |
| from utils.nn_utils import gelu | |
| logger = logging.getLogger(__name__) | |
| class BertEncoder(nn.Module): | |
| """This class using pretrained `Bert` model to encode token, | |
| then fine-tuning `Bert` model | |
| """ | |
| def __init__(self, bert_model_name, trainable=False, output_size=0, activation=gelu, dropout=0.0): | |
| """This function initialize pertrained `Bert` model | |
| Arguments: | |
| bert_model_name {str} -- bert model name | |
| Keyword Arguments: | |
| output_size {float} -- output size (default: {None}) | |
| activation {nn.Module} -- activation function (default: {gelu}) | |
| dropout {float} -- dropout rate (default: {0.0}) | |
| """ | |
| super().__init__() | |
| self.bert_model = BertModel.from_pretrained(bert_model_name) | |
| logger.info("Load bert model {} successfully.".format(bert_model_name)) | |
| self.output_size = output_size | |
| if trainable: | |
| logger.info("Start fine-tuning bert model {}.".format(bert_model_name)) | |
| else: | |
| logger.info("Keep fixed bert model {}.".format(bert_model_name)) | |
| for param in self.bert_model.parameters(): | |
| param.requires_grad = trainable | |
| if self.output_size > 0: | |
| self.mlp = BertLinear(input_size=self.bert_model.config.hidden_size, | |
| output_size=self.output_size, | |
| activation=activation) | |
| else: | |
| self.output_size = self.bert_model.config.hidden_size | |
| self.mlp = lambda x: x | |
| if dropout > 0: | |
| self.dropout = nn.Dropout(p=dropout) | |
| else: | |
| self.dropout = lambda x: x | |
| def get_output_dims(self): | |
| return self.output_size | |
| def forward(self, seq_inputs, token_type_inputs=None): | |
| """forward calculates forward propagation results, get token embedding | |
| Args: | |
| seq_inputs {tensor} -- sequence inputs (tokenized) | |
| token_type_inputs (tensor, optional): token type inputs. Defaults to None. | |
| Returns: | |
| tensor: bert output for tokens | |
| """ | |
| if token_type_inputs is None: | |
| token_type_inputs = torch.zeros_like(seq_inputs) | |
| mask_inputs = (seq_inputs != 0).long() | |
| outputs = self.bert_model(input_ids=seq_inputs, attention_mask=mask_inputs, token_type_ids=token_type_inputs) | |
| last_hidden_state = outputs[0] | |
| pooled_output = outputs[1] | |
| return self.dropout(self.mlp(last_hidden_state)), self.dropout(self.mlp(pooled_output)) | |
| class BertLayerNorm(nn.Module): | |
| """This class is LayerNorm model for Bert | |
| """ | |
| def __init__(self, hidden_size, eps=1e-12): | |
| """This function sets `BertLayerNorm` parameters | |
| Arguments: | |
| hidden_size {int} -- input size | |
| Keyword Arguments: | |
| eps {float} -- epsilon (default: {1e-12}) | |
| """ | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(hidden_size)) | |
| self.bias = nn.Parameter(torch.zeros(hidden_size)) | |
| self.variance_epsilon = eps | |
| def forward(self, x): | |
| """This function propagates forwardly | |
| Arguments: | |
| x {tensor} -- input tesor | |
| Returns: | |
| tensor -- LayerNorm outputs | |
| """ | |
| u = x.mean(-1, keepdim=True) | |
| s = (x - u).pow(2).mean(-1, keepdim=True) | |
| x = (x - u) / torch.sqrt(s + self.variance_epsilon) | |
| return self.weight * x + self.bias | |
| class BertLinear(nn.Module): | |
| """This class is Linear model for Bert | |
| """ | |
| def __init__(self, input_size, output_size, activation=gelu, dropout=0.0): | |
| """This function sets `BertLinear` model parameters | |
| Arguments: | |
| input_size {int} -- input size | |
| output_size {int} -- output size | |
| Keyword Arguments: | |
| activation {function} -- activation function (default: {gelu}) | |
| dropout {float} -- dropout rate (default: {0.0}) | |
| """ | |
| super().__init__() | |
| self.input_size = input_size | |
| self.output_size = output_size | |
| self.linear = nn.Linear(input_size, output_size) | |
| self.linear.weight.data.normal_(mean=0.0, std=0.02) | |
| self.linear.bias.data.zero_() | |
| self.activation = activation | |
| self.layer_norm = BertLayerNorm(self.output_size) | |
| if dropout > 0: | |
| self.dropout = nn.Dropout(p=dropout) | |
| else: | |
| self.dropout = lambda x: x | |
| def get_input_dims(self): | |
| return self.input_size | |
| def get_output_dims(self): | |
| return self.output_size | |
| def forward(self, x): | |
| """This function propagates forwardly | |
| Arguments: | |
| x {tensor} -- input tensor | |
| Returns: | |
| tenor -- Linear outputs | |
| """ | |
| output = self.activation(self.linear(x)) | |
| return self.dropout(self.layer_norm(output)) | |