# ------------------------------------------------------------------------------
# Adapted from https://github.com/lonePatient/BERT-NER-Pytorch
# Original licence: Copyright (c) 2020 Weitang Liu, under the MIT License.
# ------------------------------------------------------------------------------
import math
import torch
import torch.nn as nn
from mmocr.models.builder import build_activation_layer
class BertModel(nn.Module):
"""Implement Bert model for named entity recognition task.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch
Args:
num_hidden_layers (int): The number of hidden layers.
initializer_range (float):
vocab_size (int): Number of words supported.
hidden_size (int): Hidden size.
max_position_embeddings (int): Max positionsembedding size.
type_vocab_size (int): The size of type_vocab.
layer_norm_eps (float): eps.
hidden_dropout_prob (float): The dropout probability of hidden layer.
output_attentions (bool): Whether use the attentions in output
output_hidden_states (bool): Whether use the hidden_states in output.
num_attention_heads (int): The number of attention heads.
attention_probs_dropout_prob (float): The dropout probability
for the attention probabilities normalized from
the attention scores.
intermediate_size (int): The size of intermediate layer.
hidden_act_cfg (str): hidden layer activation
"""
def __init__(self,
num_hidden_layers=12,
initializer_range=0.02,
vocab_size=21128,
hidden_size=768,
max_position_embeddings=128,
type_vocab_size=2,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1,
output_attentions=False,
output_hidden_states=False,
num_attention_heads=12,
attention_probs_dropout_prob=0.1,
intermediate_size=3072,
hidden_act_cfg=dict(type='GeluNew')):
super().__init__()
self.embeddings = BertEmbeddings(
vocab_size=vocab_size,
hidden_size=hidden_size,
max_position_embeddings=max_position_embeddings,
type_vocab_size=type_vocab_size,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob)
self.encoder = BertEncoder(
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
num_hidden_layers=num_hidden_layers,
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
attention_probs_dropout_prob=attention_probs_dropout_prob,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob,
intermediate_size=intermediate_size,
hidden_act_cfg=hidden_act_cfg)
self.pooler = BertPooler(hidden_size=hidden_size)
self.num_hidden_layers = num_hidden_layers
self.initializer_range = initializer_range
self.init_weights()
def _resize_token_embeddings(self, new_num_tokens):
old_embeddings = self.embeddings.word_embeddings
new_embeddings = self._get_resized_embeddings(old_embeddings,
new_num_tokens)
self.embeddings.word_embeddings = new_embeddings
return self.embeddings.word_embeddings
def forward(self,
input_ids,
attention_masks=None,
token_type_ids=None,
position_ids=None,
head_mask=None):
if attention_masks is None:
attention_masks = torch.ones_like(input_ids)
if token_type_ids is None:
token_type_ids = torch.zeros_like(input_ids)
attention_masks = attention_masks[:, None, None]
attention_masks = attention_masks.to(
dtype=next(self.parameters()).dtype)
attention_masks = (1.0 - attention_masks) * -10000.0
if head_mask is not None:
if head_mask.dim() == 1:
head_mask = head_mask[None, None, :, None, None]
elif head_mask.dim() == 2:
head_mask = head_mask[None, :, None, None]
head_mask = head_mask.to(dtype=next(self.parameters()).dtype)
else:
head_mask = [None] * self.num_hidden_layers
embedding_output = self.embeddings(
input_ids,
position_ids=position_ids,
token_type_ids=token_type_ids)
sequence_output, *encoder_outputs = self.encoder(
embedding_output, attention_masks, head_mask=head_mask)
# sequence_output = encoder_outputs[0]
pooled_output = self.pooler(sequence_output)
# add hidden_states and attentions if they are here
# sequence_output, pooled_output, (hidden_states), (attentions)
outputs = (
sequence_output,
pooled_output,
) + tuple(encoder_outputs)
return outputs
def _init_weights(self, module):
"""Initialize the weights."""
if isinstance(module, (nn.Linear, nn.Embedding)):
# Slightly different from the TF version which
# uses truncated_normal for initialization
# cf https://github.com/pytorch/pytorch/pull/5617
module.weight.data.normal_(mean=0.0, std=self.initializer_range)
elif isinstance(module, torch.nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
def init_weights(self):
"""Initialize and prunes weights if needed."""
# Initialize weights
self.apply(self._init_weights)
class BertEmbeddings(nn.Module):
"""Construct the embeddings from word, position and token_type embeddings.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
Args:
vocab_size (int): Number of words supported.
hidden_size (int): Hidden size.
max_position_embeddings (int): Max positions embedding size.
type_vocab_size (int): The size of type_vocab.
layer_norm_eps (float): eps.
hidden_dropout_prob (float): The dropout probability of hidden layer.
"""
def __init__(self,
vocab_size=21128,
hidden_size=768,
max_position_embeddings=128,
type_vocab_size=2,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1):
super().__init__()
self.word_embeddings = nn.Embedding(
vocab_size, hidden_size, padding_idx=0)
self.position_embeddings = nn.Embedding(max_position_embeddings,
hidden_size)
self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size)
# self.LayerNorm is not snake-cased to stick with
# TensorFlow model variable name and be able to load
# any TensorFlow checkpoint file
self.LayerNorm = torch.nn.LayerNorm(hidden_size, eps=layer_norm_eps)
self.dropout = nn.Dropout(hidden_dropout_prob)
def forward(self, input_ids, token_type_ids=None, position_ids=None):
seq_length = input_ids.size(1)
if position_ids is None:
position_ids = torch.arange(
seq_length, dtype=torch.long, device=input_ids.device)
position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
if token_type_ids is None:
token_type_ids = torch.zeros_like(input_ids)
words_emb = self.word_embeddings(input_ids)
position_emb = self.position_embeddings(position_ids)
token_type_emb = self.token_type_embeddings(token_type_ids)
embeddings = words_emb + position_emb + token_type_emb
embeddings = self.LayerNorm(embeddings)
embeddings = self.dropout(embeddings)
return embeddings
class BertEncoder(nn.Module):
"""The code is adapted from https://github.com/lonePatient/BERT-NER-
Pytorch."""
def __init__(self,
output_attentions=False,
output_hidden_states=False,
num_hidden_layers=12,
hidden_size=768,
num_attention_heads=12,
attention_probs_dropout_prob=0.1,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1,
intermediate_size=3072,
hidden_act_cfg=dict(type='GeluNew')):
super().__init__()
self.output_attentions = output_attentions
self.output_hidden_states = output_hidden_states
self.layer = nn.ModuleList([
BertLayer(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
output_attentions=output_attentions,
attention_probs_dropout_prob=attention_probs_dropout_prob,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob,
intermediate_size=intermediate_size,
hidden_act_cfg=hidden_act_cfg)
for _ in range(num_hidden_layers)
])
def forward(self, hidden_states, attention_mask=None, head_mask=None):
all_hidden_states = ()
all_attentions = ()
for i, layer_module in enumerate(self.layer):
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states, )
layer_outputs = layer_module(hidden_states, attention_mask,
head_mask[i])
hidden_states = layer_outputs[0]
if self.output_attentions:
all_attentions = all_attentions + (layer_outputs[1], )
# Add last layer
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states, )
outputs = (hidden_states, )
if self.output_hidden_states:
outputs = outputs + (all_hidden_states, )
if self.output_attentions:
outputs = outputs + (all_attentions, )
# last-layer hidden state, (all hidden states), (all attentions)
return outputs
class BertPooler(nn.Module):
def __init__(self, hidden_size=768):
super().__init__()
self.dense = nn.Linear(hidden_size, hidden_size)
self.activation = nn.Tanh()
def forward(self, hidden_states):
# We "pool" the model by simply taking the hidden state corresponding
# to the first token.
first_token_tensor = hidden_states[:, 0]
pooled_output = self.dense(first_token_tensor)
pooled_output = self.activation(pooled_output)
return pooled_output
class BertLayer(nn.Module):
"""Bert layer.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
hidden_size=768,
num_attention_heads=12,
output_attentions=False,
attention_probs_dropout_prob=0.1,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1,
intermediate_size=3072,
hidden_act_cfg=dict(type='GeluNew')):
super().__init__()
self.attention = BertAttention(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
output_attentions=output_attentions,
attention_probs_dropout_prob=attention_probs_dropout_prob,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob)
self.intermediate = BertIntermediate(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act_cfg=hidden_act_cfg)
self.output = BertOutput(
intermediate_size=intermediate_size,
hidden_size=hidden_size,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob)
def forward(self, hidden_states, attention_mask=None, head_mask=None):
attention_outputs = self.attention(hidden_states, attention_mask,
head_mask)
attention_output = attention_outputs[0]
intermediate_output = self.intermediate(attention_output)
layer_output = self.output(intermediate_output, attention_output)
outputs = (layer_output, ) + attention_outputs[
1:] # add attentions if we output them
return outputs
class BertSelfAttention(nn.Module):
"""Bert self attention module.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
hidden_size=768,
num_attention_heads=12,
output_attentions=False,
attention_probs_dropout_prob=0.1):
super().__init__()
if hidden_size % num_attention_heads != 0:
raise ValueError('The hidden size (%d) is not a multiple of'
'the number of attention heads (%d)' %
(hidden_size, num_attention_heads))
self.output_attentions = output_attentions
self.num_attention_heads = num_attention_heads
self.att_head_size = int(hidden_size / num_attention_heads)
self.all_head_size = self.num_attention_heads * self.att_head_size
self.query = nn.Linear(hidden_size, self.all_head_size)
self.key = nn.Linear(hidden_size, self.all_head_size)
self.value = nn.Linear(hidden_size, self.all_head_size)
self.dropout = nn.Dropout(attention_probs_dropout_prob)
def transpose_for_scores(self, x):
new_x_shape = x.size()[:-1] + (self.num_attention_heads,
self.att_head_size)
x = x.view(*new_x_shape)
return x.permute(0, 2, 1, 3)
def forward(self, hidden_states, attention_mask=None, head_mask=None):
mixed_query_layer = self.query(hidden_states)
mixed_key_layer = self.key(hidden_states)
mixed_value_layer = self.value(hidden_states)
query_layer = self.transpose_for_scores(mixed_query_layer)
key_layer = self.transpose_for_scores(mixed_key_layer)
value_layer = self.transpose_for_scores(mixed_value_layer)
# Take the dot product between "query" and
# "key" to get the raw attention scores.
attention_scores = torch.matmul(query_layer,
key_layer.transpose(-1, -2))
attention_scores = attention_scores / math.sqrt(self.att_head_size)
if attention_mask is not None:
# Apply the attention mask is precomputed for
# all layers in BertModel forward() function.
attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities.
attention_probs = nn.Softmax(dim=-1)(attention_scores)
# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs = self.dropout(attention_probs)
# Mask heads if we want to.
if head_mask is not None:
attention_probs = attention_probs * head_mask
context_layer = torch.matmul(attention_probs, value_layer)
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
new_context_layer_shape = context_layer.size()[:-2] + (
self.all_head_size, )
context_layer = context_layer.view(*new_context_layer_shape)
outputs = (context_layer,
attention_probs) if self.output_attentions else (
context_layer, )
return outputs
class BertSelfOutput(nn.Module):
"""Bert self output.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
hidden_size=768,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1):
super().__init__()
self.dense = nn.Linear(hidden_size, hidden_size)
self.LayerNorm = torch.nn.LayerNorm(hidden_size, eps=layer_norm_eps)
self.dropout = nn.Dropout(hidden_dropout_prob)
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = self.dropout(hidden_states)
hidden_states = self.LayerNorm(hidden_states + input_tensor)
return hidden_states
class BertAttention(nn.Module):
"""Bert Attention module implementation.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
hidden_size=768,
num_attention_heads=12,
output_attentions=False,
attention_probs_dropout_prob=0.1,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1):
super().__init__()
self.self = BertSelfAttention(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
output_attentions=output_attentions,
attention_probs_dropout_prob=attention_probs_dropout_prob)
self.output = BertSelfOutput(
hidden_size=hidden_size,
layer_norm_eps=layer_norm_eps,
hidden_dropout_prob=hidden_dropout_prob)
def forward(self, input_tensor, attention_mask=None, head_mask=None):
self_outputs = self.self(input_tensor, attention_mask, head_mask)
attention_output = self.output(self_outputs[0], input_tensor)
outputs = (attention_output,
) + self_outputs[1:] # add attentions if we output them
return outputs
class BertIntermediate(nn.Module):
"""Bert BertIntermediate module implementation.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
hidden_size=768,
intermediate_size=3072,
hidden_act_cfg=dict(type='GeluNew')):
super().__init__()
self.dense = nn.Linear(hidden_size, intermediate_size)
self.intermediate_act_fn = build_activation_layer(hidden_act_cfg)
def forward(self, hidden_states):
hidden_states = self.dense(hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states)
return hidden_states
class BertOutput(nn.Module):
"""Bert output module.
The code is adapted from https://github.com/lonePatient/BERT-NER-Pytorch.
"""
def __init__(self,
intermediate_size=3072,
hidden_size=768,
layer_norm_eps=1e-12,
hidden_dropout_prob=0.1):
super().__init__()
self.dense = nn.Linear(intermediate_size, hidden_size)
self.LayerNorm = torch.nn.LayerNorm(hidden_size, eps=layer_norm_eps)
self.dropout = nn.Dropout(hidden_dropout_prob)
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = self.dropout(hidden_states)
hidden_states = self.LayerNorm(hidden_states + input_tensor)
return hidden_states