Skip to main content

A TensorFlow 2.0 Keras implementation of BERT.

Project description

Build Status Coverage Status Version Status Python Versions Downloads

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.

NEWS

  • 25.Jul.2019 - there are now two colab notebooks under examples/ showing how to fine-tune an IMDB Movie Reviews sentiment classifier from pre-trained BERT weights using an adapter-BERT model architecture on a GPU or TPU in Google Colab.

  • 28.Jun.2019 - v.0.3.0 supports adapter-BERT (google-research/adapter-bert) for “Parameter-Efficient Transfer Learning for NLP”, i.e. fine-tuning small overlay adapter layers over BERT’s transformer encoders without changing the frozen BERT weights.

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

l_bert = 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",

  adapter_size             = None,         # see arXiv:1902.00751

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

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

import tensorflow as tf
from tensorflow import keras

from bert import BertModelLayer
from bert import params_from_pretrained_ckpt
from bert import load_stock_weights

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

bert_params = params_from_pretrained_ckpt(model_dir)
l_bert = BertModelLayer.from_params(bert_params, name="bert")

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

from tensorflow import keras

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

# using the default token_type/segment id 0
output = l_bert(l_input_ids)                              # output: [batch_size, max_seq_len, hidden_size]
model = keras.Model(inputs=l_input_ids, outputs=output)
model.build(input_shape=(None, max_seq_len))

# provide a custom token_type/segment id as a layer input
output = l_bert([l_input_ids, l_token_type_ids])          # [batch_size, max_seq_len, hidden_size]
model = keras.Model(inputs=[l_input_ids, l_token_type_ids], outputs=output)
model.build(input_shape=[(None, max_seq_len), (None, max_seq_len)])

if you choose to use adapter-BERT by setting the adapter_size parameter, you would also like to freeze all the original BERT layers by calling:

l_bert.apply_adapter_freeze()

and once the model has been build or compiled, the original pre-trained weights can be loaded in the BERT layer:

from bert import load_stock_weights

bert_ckpt_file   = os.path.join(model_dir, "bert_model.ckpt")
load_stock_weights(l_bert, bert_ckpt_file)

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

Resources

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

bert-for-tf2-0.3.11.tar.gz (30.2 kB view hashes)

Uploaded Source

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page