SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity
Files
12e013dbac308f96b723f82095130d1a4416c155
HuggingFace_transformer/bert_model.py
thomwolf 12e013dbac added wordpiece - updated readme
2018-10-30 23:09:09 +01:00

475 lines
20 KiB
Python
Raw Blame History

"""
A PyTorch implementation of Google's BERT Model.
From "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding"
By Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova
Link: http://arxiv.org/abs/1810.04805
Adapted from HuggingFace's OpenAI PyTorch code and its adaptation by AllenNLP.
"""
# pylint: disable=invalid-name,arguments-differ
from typing import NamedTuple, List
import copy
import io
import json
import math
import pathlib
import re
import tarfile
import numpy as np
import torch
from torch.nn import Parameter
# pylint: disable=line-too-long
_PARAMETER_NAMES = ["model/we:0",
"model/h0/attn/c_attn/w:0", "model/h0/attn/c_attn/b:0", "model/h0/attn/c_proj/w:0",
"model/h0/attn/c_proj/b:0", "model/h0/ln_1/g:0", "model/h0/ln_1/b:0", "model/h0/mlp/c_fc/w:0",
"model/h0/mlp/c_fc/b:0", "model/h0/mlp/c_proj/w:0", "model/h0/mlp/c_proj/b:0", "model/h0/ln_2/g:0",
"model/h0/ln_2/b:0", "model/h1/attn/c_attn/w:0", "model/h1/attn/c_attn/b:0", "model/h1/attn/c_proj/w:0",
"model/h1/attn/c_proj/b:0", "model/h1/ln_1/g:0", "model/h1/ln_1/b:0", "model/h1/mlp/c_fc/w:0",
"model/h1/mlp/c_fc/b:0", "model/h1/mlp/c_proj/w:0", "model/h1/mlp/c_proj/b:0", "model/h1/ln_2/g:0",
"model/h1/ln_2/b:0", "model/h2/attn/c_attn/w:0", "model/h2/attn/c_attn/b:0", "model/h2/attn/c_proj/w:0",
"model/h2/attn/c_proj/b:0", "model/h2/ln_1/g:0", "model/h2/ln_1/b:0", "model/h2/mlp/c_fc/w:0",
"model/h2/mlp/c_fc/b:0", "model/h2/mlp/c_proj/w:0", "model/h2/mlp/c_proj/b:0", "model/h2/ln_2/g:0",
"model/h2/ln_2/b:0", "model/h3/attn/c_attn/w:0", "model/h3/attn/c_attn/b:0", "model/h3/attn/c_proj/w:0",
"model/h3/attn/c_proj/b:0", "model/h3/ln_1/g:0", "model/h3/ln_1/b:0", "model/h3/mlp/c_fc/w:0",
"model/h3/mlp/c_fc/b:0", "model/h3/mlp/c_proj/w:0", "model/h3/mlp/c_proj/b:0", "model/h3/ln_2/g:0",
"model/h3/ln_2/b:0", "model/h4/attn/c_attn/w:0", "model/h4/attn/c_attn/b:0", "model/h4/attn/c_proj/w:0",
"model/h4/attn/c_proj/b:0", "model/h4/ln_1/g:0", "model/h4/ln_1/b:0", "model/h4/mlp/c_fc/w:0",
"model/h4/mlp/c_fc/b:0", "model/h4/mlp/c_proj/w:0", "model/h4/mlp/c_proj/b:0", "model/h4/ln_2/g:0",
"model/h4/ln_2/b:0", "model/h5/attn/c_attn/w:0", "model/h5/attn/c_attn/b:0", "model/h5/attn/c_proj/w:0",
"model/h5/attn/c_proj/b:0", "model/h5/ln_1/g:0", "model/h5/ln_1/b:0", "model/h5/mlp/c_fc/w:0",
"model/h5/mlp/c_fc/b:0", "model/h5/mlp/c_proj/w:0", "model/h5/mlp/c_proj/b:0", "model/h5/ln_2/g:0",
"model/h5/ln_2/b:0", "model/h6/attn/c_attn/w:0", "model/h6/attn/c_attn/b:0", "model/h6/attn/c_proj/w:0",
"model/h6/attn/c_proj/b:0", "model/h6/ln_1/g:0", "model/h6/ln_1/b:0", "model/h6/mlp/c_fc/w:0",
"model/h6/mlp/c_fc/b:0", "model/h6/mlp/c_proj/w:0", "model/h6/mlp/c_proj/b:0", "model/h6/ln_2/g:0",
"model/h6/ln_2/b:0", "model/h7/attn/c_attn/w:0", "model/h7/attn/c_attn/b:0", "model/h7/attn/c_proj/w:0",
"model/h7/attn/c_proj/b:0", "model/h7/ln_1/g:0", "model/h7/ln_1/b:0", "model/h7/mlp/c_fc/w:0",
"model/h7/mlp/c_fc/b:0", "model/h7/mlp/c_proj/w:0", "model/h7/mlp/c_proj/b:0", "model/h7/ln_2/g:0",
"model/h7/ln_2/b:0", "model/h8/attn/c_attn/w:0", "model/h8/attn/c_attn/b:0", "model/h8/attn/c_proj/w:0",
"model/h8/attn/c_proj/b:0", "model/h8/ln_1/g:0", "model/h8/ln_1/b:0", "model/h8/mlp/c_fc/w:0",
"model/h8/mlp/c_fc/b:0", "model/h8/mlp/c_proj/w:0", "model/h8/mlp/c_proj/b:0", "model/h8/ln_2/g:0",
"model/h8/ln_2/b:0", "model/h9/attn/c_attn/w:0", "model/h9/attn/c_attn/b:0", "model/h9/attn/c_proj/w:0",
"model/h9/attn/c_proj/b:0", "model/h9/ln_1/g:0", "model/h9/ln_1/b:0", "model/h9/mlp/c_fc/w:0",
"model/h9/mlp/c_fc/b:0", "model/h9/mlp/c_proj/w:0", "model/h9/mlp/c_proj/b:0", "model/h9/ln_2/g:0",
"model/h9/ln_2/b:0", "model/h10/attn/c_attn/w:0", "model/h10/attn/c_attn/b:0", "model/h10/attn/c_proj/w:0",
"model/h10/attn/c_proj/b:0", "model/h10/ln_1/g:0", "model/h10/ln_1/b:0", "model/h10/mlp/c_fc/w:0",
"model/h10/mlp/c_fc/b:0", "model/h10/mlp/c_proj/w:0", "model/h10/mlp/c_proj/b:0", "model/h10/ln_2/g:0",
"model/h10/ln_2/b:0", "model/h11/attn/c_attn/w:0", "model/h11/attn/c_attn/b:0", "model/h11/attn/c_proj/w:0",
"model/h11/attn/c_proj/b:0", "model/h11/ln_1/g:0", "model/h11/ln_1/b:0", "model/h11/mlp/c_fc/w:0",
"model/h11/mlp/c_fc/b:0", "model/h11/mlp/c_proj/w:0", "model/h11/mlp/c_proj/b:0", "model/h11/ln_2/g:0",
"model/h11/ln_2/b:0", "model/clf/w:0", "model/clf/b:0"]
# pylint: enable=line-too-long
def gelu(x: torch.Tensor) -> torch.Tensor:
return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
class BERTConfig(NamedTuple):
"""
BERT's hyper-parameters
"""
embedding_dim: int = 768
num_heads: int = 12
dropout: float = 0.1
class LayerNorm(torch.nn.Module):
"""
A layernorm module in the Tensorflow style (with the epsilon inside the square root).
"""
def __init__(self, n_state, e=1e-5):
super().__init__()
self.g = torch.nn.Parameter(torch.ones(n_state))
self.b = torch.nn.Parameter(torch.zeros(n_state))
self.e = e
def forward(self, x):
u = x.mean(-1, keepdim=True)
s = (x - u).pow(2).mean(-1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.e)
return self.g * x + self.b
class Conv1D(torch.nn.Module):
"""
A batched linear layer using torch.addmm
"""
def __init__(self, nf: int, rf: int, nx: int) -> None:
super().__init__()
self.rf = rf
self.nf = nf
w = torch.empty(nx, nf)
torch.nn.init.normal_(w, std=0.02)
self.w = Parameter(w)
self.b = Parameter(torch.zeros(nf))
def forward(self, x: torch.Tensor) -> torch.Tensor:
size_out = x.size()[:-1] + (self.nf,)
x = torch.addmm(self.b, x.view(-1, x.size(-1)), self.w)
x = x.view(*size_out)
return x
class Attention(torch.nn.Module):
"""
A self-attention layer comprising a sequence of:
- a linear layer: instance of the `Conv1D` class,
- spliting the inputs in key, value, query tensors (x.split),
- reshaping key, value, query tensors according to the number of head (self.split_heads)
- appying self attention (self._attn)
- merging back the heads results (self.merge_heads)
- a linear layer: instance of the `Conv1D` class,
- a dropout layer: instance of `torch.nn.Dropout` class.
See above for the details of Conv1D.
"""
def __init__(self,
nx: int,
n_ctx: int,
config: BERTConfig,
scale: bool = False) -> None:
super().__init__()
n_state = nx # in Attention: n_state=768 (nx=n_embd)
# [switch nx => n_state from Block to Attention to keep identical to TF implem]
assert n_state % config.num_heads == 0
self.register_buffer('b', torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
self.n_head = config.num_heads
self.split_size = n_state
self.scale = scale
self.c_attn = Conv1D(n_state * 3, 1, nx)
self.c_proj = Conv1D(n_state, 1, nx)
self.attn_dropout = torch.nn.Dropout(config.dropout)
self.resid_dropout = torch.nn.Dropout(config.dropout)
def _attn(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
w = torch.matmul(q, k)
if self.scale:
w = w / math.sqrt(v.size(-1))
w = w * self.b + -1e9 * (1 - self.b) # TF implem method: mask_attn_weights
w = torch.nn.Softmax(dim=-1)(w)
w = self.attn_dropout(w)
return torch.matmul(w, v)
def merge_heads(self, x: torch.Tensor):
# pylint: disable=no-self-use
x = x.permute(0, 2, 1, 3).contiguous()
new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
return x.view(*new_x_shape) # in Tensorflow implem: fct merge_states
def split_heads(self, x: torch.Tensor, k: bool = False):
new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
x = x.view(*new_x_shape) # in Tensorflow implem: fct split_states
if k:
return x.permute(0, 2, 3, 1)
else:
return x.permute(0, 2, 1, 3)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.c_attn(x)
query, key, value = x.split(self.split_size, dim=2)
query = self.split_heads(query)
key = self.split_heads(key, k=True)
value = self.split_heads(value)
a = self._attn(query, key, value)
a = self.merge_heads(a)
a = self.c_proj(a)
a = self.resid_dropout(a)
return a
class MLP(torch.nn.Module):
"""
A multi-layer perceptron layer comprising a sequence of:
- a linear layer: instance of the `Conv1D` class,
- an activation function: the `gelu` function,
- another linear layer: instance of the `Conv1D` class,
- a dropout layer: instance of `torch.nn.Dropout` class.
See above for the details of Conv1D and the gelu function.
"""
def __init__(self, n_state: int, config: BERTConfig) -> None: # in MLP: n_state=3072 (4 * n_embd)
super().__init__()
self.c_fc = Conv1D(n_state, 1, config.embedding_dim)
self.c_proj = Conv1D(config.embedding_dim, 1, n_state)
self.act = gelu
self.dropout = torch.nn.Dropout(config.dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.act(self.c_fc(x))
h2 = self.c_proj(h)
return self.dropout(h2)
class Block(torch.nn.Module):
"""
A Transformer Block comprising a sequence of:
- a self-attention layer: instance of the `Attention` class,
- a Layer Normalization layer: instance of the `LayerNorm` class,
- a Multi-layer perceptron layer: instance of the `MLP` class,
- another Layer Normalization layer: instance of the `LayerNorm` class.
See above for the details of these classes.
"""
def __init__(self,
n_ctx: int,
config: BERTConfig,
scale: bool = False) -> None:
super().__init__()
nx = config.embedding_dim
self.attn = Attention(nx, n_ctx, config, scale)
self.ln_1 = LayerNorm(nx)
self.mlp = MLP(4 * nx, config)
self.ln_2 = LayerNorm(nx)
def forward(self, x: torch.Tensor) -> torch.Tensor:
a = self.attn(x)
n = self.ln_1(x + a)
m = self.mlp(n)
h = self.ln_2(n + m)
return h
class BERT(torch.nn.Module):
"""
Google's BERT Model.
Default parameters are the ones for Google's pretrained model.
Parameters
----------
vocab_size: ``int`` (optional, default: 40478)
The size of the vocabulary (number of byte pair embeddings)
excluding the n_special embeddings (if any), and the positional embeddings.
n_ctx: ``int`` (optional, default: 512)
The number of positional encodings to use for evaluation.
embedding_dim: ``int`` (optional, default: 768)
The dimension of the output embeddings.
num_heads: ``int`` (optional, default: 12)
How many "heads" the attention has.
num_layers: ``int`` (optional, default: 12)
How many layers of "blocks" the transformer has.
dropout_probability: ``float`` (optional, default: 0.1)
Dropout for all layers.
model_path: ``str`` (optional, default: ``None``)
A tar.gz file containing serialized model weights. If supplied,
the weights will be loaded from that file.
requires_grad: ``bool`` (optional, default: ``False``)
If true, the transformer will be fine-tuneable.
n_special: ``int`` (optional, default: ``-1``)
The number of special tokens added to the byte pair vocabulary
"""
def __init__(self,
vocab_size: int = 40478,
n_ctx: int = 512,
embedding_dim: int = 768,
num_heads: int = 12,
num_layers: int = 12,
dropout_probability: float = 0.1,
model_path: str = None,
requires_grad: bool = False,
n_special: int = -1) -> None:
super().__init__()
config = BERTConfig(
embedding_dim,
num_heads,
dropout_probability,
)
# the embedding size is vocab_size + n_special embeddings + n_ctx
embedding_size = vocab_size + max(n_special, 0) + n_ctx
self.vocab_size = embedding_size
self.n_ctx = n_ctx
self.n_special = n_special
self.num_output_layers = 1 + num_layers
self.embed = torch.nn.Embedding(embedding_size, embedding_dim)
self.drop = torch.nn.Dropout(dropout_probability)
block = Block(n_ctx, config, scale=True)
self.h = torch.nn.ModuleList([copy.deepcopy(block) for _ in range(num_layers)])
self.decoder = torch.nn.Linear(embedding_dim, embedding_size, bias=False)
self.decoder.weight = self.embed.weight # Tied weights
# To reproduce the noise_shape parameter of TF implementation
torch.nn.init.normal_(self.embed.weight, std=0.02)
for parameter in self.parameters():
parameter.requires_grad = requires_grad
if model_path:
self.load_weights(model_path, n_special=n_special, n_ctx=n_ctx)
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
#x = x.view(-1, x.size(2), x.size(3))
# x is (batch_size, sequence_length) tensor of byte-pair ids
# e is (batch_size, sequence_length, 2, embedding_dim) tensor of embeddings
e = self.embed(x)
# h is (batch_size, sequence_length, embedding_dim)
h = e.sum(dim=2)
all_layers = [h]
for block in self.h:
h = block(h)
all_layers.append(h)
# result is list of (batch_size, sequence_length, embedding_dim)
return all_layers
def load_weights(self,
bert_model_path: str,
n_ctx: int = -1,
n_special: int = -1,
n_transfer: int = 12,
n_embd: int = 768,
names: List[str] = _PARAMETER_NAMES) -> None:
# pylint: disable=dangerous-default-value
with tarfile.open(bert_model_path) as tmp:
num_params_files = len([member for member in tmp.getmembers() if member.name.endswith('.npy')])
shapesfile = tmp.extractfile('model/params_shapes.json')
if shapesfile:
shapes = json.loads(shapesfile.read())
else:
raise Exception("unable to find model/params_shapes.json in the archive")
# numpy can't read from a tarfile directly, so we need a workaround
# https://github.com/numpy/numpy/issues/7989#issuecomment-341656702
init_params: List[np.ndarray] = []
for n in range(num_params_files):
array_file = io.BytesIO()
array_file.write(tmp.extractfile(f'model/params_{n}.npy').read())
array_file.seek(0)
# each np.load is a (11653478,) numpy array
init_params.append(np.load(array_file))
# init_params is a list of 10 arrays of size (11653578,)
# shapes are [[512, 768], [40478, 768], [1, 768, 2304], [2304], ... # 146 elts
# products are [512 * 768, 40478 * 768, ...]
# offsets is [512 * 768, 512 * 768 + 40478 * 768, ...]
offsets = np.cumsum([np.prod(shape) for shape in shapes])
# split into the 146 subarrays corresponding to shapes
init_params = np.split(np.concatenate(init_params, 0), offsets)[:-1]
# reshape
init_params = [param.reshape(shape) for param, shape in zip(init_params, shapes)]
# truncate if necessary
if n_ctx > 0:
# positional embeddings?
# init_params[0] is (512, 768) = (max_chars, embedding_dim)
init_params[0] = init_params[0][:n_ctx]
# combine init_params[1] and init_params[0]
if n_special > 0:
# init_params[1] is (40478, 768)
# special is (n_special, 768)
# init_params[0] is (512, 768)
# result is (40990 + n_special, 768)
init_params[0] = np.concatenate(
[init_params[1],
(np.random.randn(n_special, n_embd) * 0.02).astype(np.float32),
init_params[0]],
0
)
else:
# result is (40990, 768)
init_params[0] = np.concatenate([init_params[1], init_params[0]], 0)
del init_params[1]
# number of dimensions to transfer, 12 per layer, plus one extra
if n_transfer == -1:
n_transfer = 0
else:
n_transfer = 1 + n_transfer * 12
# squeeze?
init_params = [arr.squeeze() for arr in init_params]
# embedding.weight is (vocab_size, embedding_dim)
# make sure init_params[0] has the same shape
try:
assert self.embed.weight.shape == init_params[0].shape
except AssertionError as e:
e.args += (self.embed.weight.shape, init_params[0].shape)
raise
# and then assign it
self.embed.weight.data = torch.from_numpy(init_params[0])
self.decoder.weight = self.embed.weight
# for each (name, array) pair to transfer over
for name, ip in zip(names[1:n_transfer], init_params[1:n_transfer]):
# "model/h0/attn/c_attn/w:0"
name = name[6:] # "h0/attn/c_attn/w:0"
assert name[-2:] == ":0"
name = name[:-2] # "h0/attn/c_attn/w"
name_parts = name.split('/') # ['h0', 'attn', 'c_attn', 'w']
pointer = self
for m_name in name_parts:
if re.fullmatch(r'[A-Za-z]+\d+', m_name):
l = re.split(r'(\d+)', m_name) # ['h', '0', '']
else:
l = [m_name] # ['attn']
pointer = getattr(pointer, l[0])
if len(l) >= 2:
num = int(l[1])
pointer = pointer[num]
try:
assert pointer.shape == ip.shape
except AssertionError as e:
e.args += (pointer.shape, ip.shape)
raise
pointer.data = torch.from_numpy(ip) # pylint: disable=attribute-defined-outside-init
def dump_weights(self,
output_dir: str,
num_pieces: int = 10) -> None:
output_path = pathlib.Path(output_dir) / 'model'
output_path.mkdir(exist_ok=True, parents=True) # pylint: disable=no-member
named_parameters = list(self.named_parameters())
# embedding weights get special treatment
_, array = named_parameters[0]
num_bpe = self.vocab_size - self.n_ctx
byte_pair_embeddings = array[:num_bpe]
positional_embeddings = array[num_bpe:]
arrays = [positional_embeddings.numpy().ravel(), byte_pair_embeddings.numpy().ravel()]
shapes = [positional_embeddings.shape, byte_pair_embeddings.shape]
names = ["model/we:0"]
for param_name, tensor in named_parameters[1:]:
param_name = f'h{param_name}' # 'h0.attn.c_attn.w'
parts = param_name.split(".") # ['h0', 'attn', 'c_attn', 'w']
name = "model/" + '/'.join(parts) + ':0' # 'model/h0/attn/c_attn/w:0'
array = tensor.numpy().ravel()
arrays.append(array)
shapes.append(list(tensor.shape))
names.append(name)
# write out the arrays
big_array = np.concatenate(arrays)
total_size = len(big_array)
batch_size = math.ceil(total_size / num_pieces)
for i in range(num_pieces):
filename = output_path / f"params_{i}.npy"
start = i * batch_size
end = start + batch_size
subarray = big_array[start:end]
np.save(filename, subarray)
# write out the shapes
with open(output_path / 'params_shapes.json', 'w') as shapes_file:
json.dump(shapes, shapes_file)
Reference in New Issue View Git Blame Copy Permalink