BigBird

To install BigBird, clone the repository and install it in editable mode using pip.

git clone https://github.com/google-research/bigbird.git
cd bigbird
pip3 install -e .

Pretrained and fine-tuned checkpoints are available in a Google Cloud Storage bucket. You can download them using gsutil into a bigbird/ckpt directory.

Available checkpoints include:

• Pretrained BERT models: bigbr_base (base) and bigbr_large (large). These are encoder-only models with post-normalization.
• Pretrained Pegasus Encoder-Decoder: bigbp_large (large). These use pre-normalization and have separate encoder-decoder weights.
• Fine-tuned tf.SavedModel: Specifically for long document summarization, suitable for direct prediction and evaluation.

mkdir -p bigbird/ckpt
gsutil cp -r gs://bigbird-transformer/ bigbird/ckpt/

To run BigBird on Google Cloud, you can create a Compute Engine instance and a Cloud TPU instance. The following commands use europe-west4-a and the instance name bigbird as examples.

Compute Instance:

gcloud compute instances create \
  bigbird \
  --zone=europe-west4-a \
  --machine-type=n1-standard-16 \
  --boot-disk-size=50GB \
  --image-project=ml-images \
  --image-family=tf-2-3-1 \
  --maintenance-policy TERMINATE \
  --restart-on-failure \
  --scopes=cloud-platform

Cloud TPU Instance:

gcloud compute tpus create \
  bigbird \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=2.3.1

SSH into instance:

gcloud compute ssh --zone "europe-west4-a" "bigbird"

BigBird's behavior is controlled via flags and configuration parameters (defined in core/flags.py).

attention_type

Selects the attention implementation:

• original_full: Standard full attention (as in original BERT). Recommended for sequence lengths < 1024.
• simulated_sparse: Simulated sparse attention.
• block_sparse: The actual BigBird blocked sparse attention implementation.

Key Parameters

• block_size: Defines the size of the attention blocks.
• num_rand_blocks: Sets the number of random blocks.

Constraints and Requirements

• Static Shapes: The current implementation primarily supports static tensors, as it is designed for TPUs.
• Divisibility: The hidden dimension must be divisible by the number of attention heads.

To use the BigBird repository, install it directly from the GitHub repository using pip:

!pip install git+https://github.com/google-research/bigbird.git -q

To install the BigBird repository directly from the source, use the following pip command:

!pip install git+https://github.com/google-research/bigbird.git -q

To load weights from a checkpoint into an existing TransformerModel, use a tf.compat.v1.train.NewCheckpointReader. You must iterate through the model's trainable_weights and map them to the tensors in the checkpoint by name.

import tensorflow as tf
from tqdm import tqdm

ckpt_path = 'gs://bigbird-transformer/summarization/pubmed/roberta/model.ckpt-300000'
ckpt_reader = tf.compat.v1.train.NewCheckpointReader(ckpt_path)
loaded_weights = []

for v in tqdm(model.trainable_weights, position=0):
  try:
    # Attempt to get tensor by name (stripping last 2 chars from variable name)
    val = ckpt_reader.get_tensor(v.name[:-2])
  except:
    val = v.numpy()
  loaded_weights.append(val)

model.set_weights(loaded_weights)

To evaluate the quality of generated summaries, use the rouge_score library. You can iterate through a dataset, compute scores for each prediction against the ground truth, and use a BootstrapAggregator to aggregate the results.

from rouge_score import rouge_scorer
from rouge_score import scoring
from tqdm import tqdm

# Initialize scorer and aggregator
scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeLsum"], use_stemmer=True)
aggregator = scoring.BootstrapAggregator()

# Loop through dataset and aggregate scores
for ex in tqdm(dataset.take(100), position=0):
  predicted_summary = summerize(ex[0])['pred_sent'][0]
  score = scorer.score(ex[1].numpy().decode('utf-8'), predicted_summary.numpy().decode('utf-8'))
  aggregator.add_scores(score)

# Get final aggregated results
results = aggregator.aggregate()

You can run classification experiments using the scripts provided in the classifier directory. You must set your GCP project name and experiment bucket as environment variables before running the shell script.

export GCP_PROJECT_NAME=bigbird-project  # Replace by your project name
export GCP_EXP_BUCKET=gs://bigbird-transformer-training/  # Replace
sh -x bigbird/classifier/base_size.sh

To use the BigBird encoder directly in your own code (e.g., to replace a standard BERT encoder), import modeling from bigbird.core.

from bigbird.core import modeling

bigb_encoder = modeling.BertModel(...)

Alternatively, you can interact with the individual layers of the BigBird encoder stack using the encoder module:

from bigbird.core import encoder

only_layers = encoder.EncoderStack(...)