bert-for-tf2 0.1.5

A TensorFlow 2.0 Keras implementation of BERT.

This repo contains a TensorFlow 2.0 Keras implementation of google-research/bert with support for loading of the original pre-trained weights, and producing activations numerically identical to the one calculated by the original model.

The implementation is build from scratch using only basic tensorflow operations, following the code in google-research/bert/modeling.py (but skipping dead code and applying some simplifications). It also utilizes kpe/params-flow to reduce common Keras boilerplate code (related to passing model and layer configuration arguments).

bert-for-tf2 should work with both TensorFlow 2.0 and TensorFlow 1.14 or newer.

LICENSE

MIT. See License File.

Install

bert-for-tf2 is on the Python Package Index (PyPI):

pip install bert-for-tf2

Usage

BERT in bert-for-tf2 is implemented as a Keras layer. You could instantiate it like this:

from bert import BertModelLayer

bert_layer = BertModelLayer(BertModelLayer.Params(
  vocab_size               = 16000,        # embedding params
  use_token_type           = True,
  use_position_embeddings  = True,
  token_type_vocab_size    = 2,

  num_layers               = 12,           # transformer encoder params
  hidden_size              = 768,
  hidden_dropout           = 0.1,
  intermediate_size        = 4*768,
  intermediate_activation  = "gelu",

  name                     = "bert"        # any other Keras layer params
))

or by using the bert_config.json from a pre-trained google model:

import os
import tensorflow as tf
from tensorflow.python import keras
from bert import BertModelLayer
from bert.loader import StockBertConfig, load_stock_weights

model_dir = ".models/uncased_L-12_H-768_A-12"

bert_config_file = os.path.join(model_dir, "bert_config.json")
bert_ckpt_file   = os.path.join(model_dir, "bert_model.ckpt")

with tf.io.gfile.GFile(bert_config_file, "r") as reader:
  stock_params = StockBertConfig.from_json_string(reader.read())
  bert_params  = stock_params.to_bert_model_layer_params()

l_bert = BertModelLayer.from_params(bert_params, name="bert")
load_stock_weights(l_bert, bert_ckpt_file)

now you can use the BERT layer in your Keras model like this:

from tensorflow.python import keras

max_seq_len = 128
l_input_ids      = keras.layers.Input(shape=(max_seq_len,), dtype='int32',
                                      name="input_ids")
l_token_type_ids = keras.layers.Input(shape=(max_seq_len,), dtype='int32',
                                      name="token_type_ids")

output = l_bert([l_input_ids, l_token_type_ids])  # [batch_size, max_seq_len, hidden_size]

N.B. see tests/test_bert_activations.py for a complete example.

Resources

• BERT - BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

• google-research/bert - the original BERT implementation

• kpe/params-flow - A Keras coding style for reducing Keras boilerplate code in custom layers by utilizing kpe/py-params