Keras BERT with pluggable tasks and encoders

What it does

Implements Google BERT and OpenAI’s Transformer LM in Keras, with support for loading their pretrained weights and fine-tuning on downstream tasks. It wraps both token-level (NER, PoS) and sentence-level (classification, next-sentence prediction) NLP tasks into a framework-independent metadata system, so the data generators and task definitions can be reused elsewhere. A Colab notebook demonstrates TPU-compatible training and inference.

The interesting bit

The training pipeline is decoupled from the encoder architecture: you can replace the Transformer with a BiLSTM or BiQRNN as long as the inputs and outputs match the expected contract. That is genuinely unusual among BERT reimplementations, which typically treat the architecture as fixed.

Key highlights

• Loads pretrained weights from both Google’s BERT and OpenAI’s Transformer models.
• Task metadata system abstracts sentence-level and token-level labels, masks, and extraction points.
• Includes TPU-compatible training and inference paths via TensorFlow backend.
• Data generators and task definitions are designed to be reusable outside Keras; the author notes they work with PyTorch too.
• TaskWeightScheduler lets you smoothly shift training emphasis from language modeling to classification.

Caveats

• Repository is archived; the author provides the code as-is and expects no further updates.
• BERT weight loading and training require TensorFlow; the Theano backend only supports OpenAI model loading, and even then not for fine-tuning.
• The tutorial notebook is explicitly described by the author as “poorly designed,” though functional.

Verdict

Worth a look if you want a Keras-native BERT with hackable task abstractions or need to study how the model’s token- and sentence-level tasks can be decoupled from the encoder. If you just need production-grade inference, modern transformers libraries have long since superseded it.