Skip to main content

Keras BERT

Version License

[中文|English]

Implementation of the BERT. Official pre-trained models could be loaded for feature extraction and prediction.

Install

pip install keras-bert

Usage

External Links

Load Official Pre-trained Models

In feature extraction demo, you should be able to get the same extraction results as the official model chinese_L-12_H-768_A-12. And in prediction demo, the missing word in the sentence could be predicted.

Run on TPU

The extraction demo shows how to convert to a model that runs on TPU.

The classification demo shows how to apply the model to simple classification tasks.

Tokenizer

The Tokenizer class is used for splitting texts and generating indices:

from keras_bert import Tokenizer

token_dict = {
    '[CLS]': 0,
    '[SEP]': 1,
    'un': 2,
    '##aff': 3,
    '##able': 4,
    '[UNK]': 5,
}
tokenizer = Tokenizer(token_dict)
print(tokenizer.tokenize('unaffable'))  # The result should be `['[CLS]', 'un', '##aff', '##able', '[SEP]']`
indices, segments = tokenizer.encode('unaffable')
print(indices)  # Should be `[0, 2, 3, 4, 1]`
print(segments)  # Should be `[0, 0, 0, 0, 0]`

print(tokenizer.tokenize(first='unaffable', second='钢'))
# The result should be `['[CLS]', 'un', '##aff', '##able', '[SEP]', '钢', '[SEP]']`
indices, segments = tokenizer.encode(first='unaffable', second='钢', max_len=10)
print(indices)  # Should be `[0, 2, 3, 4, 1, 5, 1, 0, 0, 0]`
print(segments)  # Should be `[0, 0, 0, 0, 0, 1, 1, 0, 0, 0]`

Train & Use

from tensorflow import keras
from keras_bert import get_base_dict, get_model, compile_model, gen_batch_inputs


# A toy input example
sentence_pairs = [
    [['all', 'work', 'and', 'no', 'play'], ['makes', 'jack', 'a', 'dull', 'boy']],
    [['from', 'the', 'day', 'forth'], ['my', 'arm', 'changed']],
    [['and', 'a', 'voice', 'echoed'], ['power', 'give', 'me', 'more', 'power']],
]


# Build token dictionary
token_dict = get_base_dict()  # A dict that contains some special tokens
for pairs in sentence_pairs:
    for token in pairs[0] + pairs[1]:
        if token not in token_dict:
            token_dict[token] = len(token_dict)
token_list = list(token_dict.keys())  # Used for selecting a random word


# Build & train the model
model = get_model(
    token_num=len(token_dict),
    head_num=5,
    transformer_num=12,
    embed_dim=25,
    feed_forward_dim=100,
    seq_len=20,
    pos_num=20,
    dropout_rate=0.05,
)
compile_model(model)
model.summary()

def _generator():
    while True:
        yield gen_batch_inputs(
            sentence_pairs,
            token_dict,
            token_list,
            seq_len=20,
            mask_rate=0.3,
            swap_sentence_rate=1.0,
        )

model.fit_generator(
    generator=_generator(),
    steps_per_epoch=1000,
    epochs=100,
    validation_data=_generator(),
    validation_steps=100,
    callbacks=[
        keras.callbacks.EarlyStopping(monitor='val_loss', patience=5)
    ],
)


# Use the trained model
inputs, output_layer = get_model(
    token_num=len(token_dict),
    head_num=5,
    transformer_num=12,
    embed_dim=25,
    feed_forward_dim=100,
    seq_len=20,
    pos_num=20,
    dropout_rate=0.05,
    training=False,      # The input layers and output layer will be returned if `training` is `False`
    trainable=False,     # Whether the model is trainable. The default value is the same with `training`
    output_layer_num=4,  # The number of layers whose outputs will be concatenated as a single output.
                         # Only available when `training` is `False`.
)

Use Warmup

AdamWarmup optimizer is provided for warmup and decay. The learning rate will reach lr in warmpup_steps steps, and decay to min_lr in decay_steps steps. There is a helper function calc_train_steps for calculating the two steps:

import numpy as np
from keras_bert import AdamWarmup, calc_train_steps

train_x = np.random.standard_normal((1024, 100))

total_steps, warmup_steps = calc_train_steps(
    num_example=train_x.shape[0],
    batch_size=32,
    epochs=10,
    warmup_proportion=0.1,
)

optimizer = AdamWarmup(total_steps, warmup_steps, lr=1e-3, min_lr=1e-5)

Download Pretrained Checkpoints

Several download urls has been added. You can get the downloaded and uncompressed path of a checkpoint by:

from keras_bert import get_pretrained, PretrainedList, get_checkpoint_paths

model_path = get_pretrained(PretrainedList.multi_cased_base)
paths = get_checkpoint_paths(model_path)
print(paths.config, paths.checkpoint, paths.vocab)

Extract Features

You can use helper function extract_embeddings if the features of tokens or sentences (without further tuning) are what you need. To extract the features of all tokens:

from keras_bert import extract_embeddings

model_path = 'xxx/yyy/uncased_L-12_H-768_A-12'
texts = ['all work and no play', 'makes jack a dull boy~']

embeddings = extract_embeddings(model_path, texts)

The returned result is a list with the same length as texts. Each item in the list is a numpy array truncated by the length of the input. The shapes of outputs in this example are (7, 768) and (8, 768).

When the inputs are paired-sentences, and you need the outputs of NSP and max-pooling of the last 4 layers:

from keras_bert import extract_embeddings, POOL_NSP, POOL_MAX

model_path = 'xxx/yyy/uncased_L-12_H-768_A-12'
texts = [
    ('all work and no play', 'makes jack a dull boy'),
    ('makes jack a dull boy', 'all work and no play'),
]

embeddings = extract_embeddings(model_path, texts, output_layer_num=4, poolings=[POOL_NSP, POOL_MAX])

There are no token features in the results. The outputs of NSP and max-pooling will be concatenated with the final shape (768 x 4 x 2,).

The second argument in the helper function is a generator. To extract features from file:

import codecs
from keras_bert import extract_embeddings

model_path = 'xxx/yyy/uncased_L-12_H-768_A-12'

with codecs.open('xxx.txt', 'r', 'utf8') as reader:
    texts = map(lambda x: x.strip(), reader)
    embeddings = extract_embeddings(model_path, texts)

Use tensorflow.python.keras

Add TF_KERAS=1 to environment variables to use tensorflow.python.keras.

Metadata

Release files for keras-bert 0.89.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for keras-bert 0.89.0
File Size Uploaded
keras-bert-0.89.0.tar.gz 25.6 kB Details

Release files / keras-bert-0.89.0.tar.gz

Download URL keras-bert-0.89.0.tar.gz
Size 25.6 kB
Tags Source
SHA-256 checksum
How to use checksums
d35457c91130e23fc9a97181b0eb5c7b81d6945b40edc52310f6f713ed4eaecd
BLAKE2b-256 checksum
How to use checksums
740affc65dfa4b31942ee8348e0026d2a7ee57e1769e9266c677141a3e2cac9c
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/1.13.0 pkginfo/1.5.0.1 requests/2.22.0 setuptools/41.0.1 requests-toolbelt/0.9.1 tqdm/4.32.2 CPython/3.7.4

Release history Release notifications | RSS feed

This release

0.89.0 This release

1 release file

0.88.0

1 release file

0.87.0

1 release file

0.86.0

1 release file

0.85.0

1 release file

0.84.0

1 release file

0.83.0

1 release file

0.82.0

1 release file

0.81.1

1 release file

0.81.0

1 release file

0.80.0

1 release file

0.79.0

1 release file

0.78.0

1 release file

0.77.0

1 release file

0.76.0

1 release file

0.75.0

1 release file

0.74.0

1 release file

0.73.0

1 release file

0.72.0

1 release file

0.71.0

1 release file

0.70.3

1 release file

0.70.2

1 release file

0.70.1

1 release file

0.70.0

1 release file

0.69.0

1 release file

0.68.2

1 release file

0.68.1

1 release file

0.68.0

1 release file

0.66.0

1 release file

0.65.1

1 release file

0.65.0

1 release file

0.64.0

1 release file

0.63.0

1 release file

0.62.0

1 release file

0.61.0

1 release file

0.60.0

1 release file

0.59.0

1 release file

0.58.0

1 release file

0.57.1

1 release file

0.57.0

1 release file

0.55.2

1 release file

0.55.1

1 release file

0.55.0

1 release file

0.54.0

1 release file

0.53.0

1 release file

0.52.0

1 release file

0.51.0

1 release file

0.50.0

1 release file

0.49.0

1 release file

0.47.0

1 release file

0.46.0

1 release file

0.43.0

1 release file

0.42.0

1 release file

0.41.0

1 release file

0.40.0

1 release file

0.39.0

1 release file

0.38.0

1 release file

0.37.0

1 release file

0.36.1

1 release file

0.36.0

1 release file

0.35.0

1 release file

0.34.1

1 release file

0.34.0

1 release file

0.33.1

1 release file

0.33.0

1 release file

0.32.0

1 release file

0.31.0

1 release file

0.30.0

1 release file

0.29.0

1 release file

0.28.0

1 release file

0.27.0

1 release file

0.25.0

1 release file

0.24.0

1 release file

0.23.0

1 release file

0.22.0

1 release file

0.21.0

1 release file

0.20.0

1 release file

0.19.0

1 release file

0.18.0

1 release file

0.16.0

1 release file

0.15.0

1 release file

0.14.0

1 release file

0.13.0

1 release file

0.12.0

1 release file

0.11.0

1 release file

0.10.0

1 release file

0.9.0

1 release file

0.8.0

1 release file

0.7.0

1 release file

0.6.0

1 release file

0.5.0

1 release file

0.4.0

1 release file

0.3.0

1 release file

0.2.0

1 release file

0.1.0

1 release file

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page