When Google published “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding” in October 2018, the paper landed together with a reference implementation — the google-research/bert repository. That repo is the artifact this tour reads. It is written for TensorFlow 1.x in graph mode, built on the Estimator API, and it carries the unmistakable texture of late-2018 research code: flag parsing at the top of every script, TPUEstimator wiring, and comments that explain not just what the code does but why the authors chose that design.
What strikes most readers first is the size. There is no package directory, no src/ tree, no build system. The whole thing is roughly a dozen files at the repository root, and the ones that matter are eight Python scripts: modeling.py, tokenization.py, run_classifier.py, run_squad.py, optimization.py, extract_features.py, create_pretraining_data.py, and run_pretraining.py. That is the complete surface area of the most influential language model of its era — the codebase that other teams ported to PyTorch, that spawned the fine-tuning ecosystem, and that turned “pre-train once, fine-tune everywhere” from an idea into an industry default.
Reading it top-down is rewarding precisely because the graph-mode style forces every tensor transformation to be explicit. Nothing hides behind an nn.Module with implicit state; you can trace a token from a raw string all the way to a probability without leaving a handful of files. This post walks that path at two zoom levels, mapped by the two diagrams below.
The overview graph: four task runners on the right pull from the tokenizer and the shared encoder; optimization is a small, self-contained module that every trainer reuses.
Reading the overview from left to right:
- BertModel, L107 is the heart of the repo. Its constructor wires the embedding lookup, the attention mask, the encoder stack, and a tanh pooler into one computation graph, then exposes accessors like
get_sequence_output()andget_pooled_output()for downstream heads. - BertConfig, L31 is the hyperparameter record that separates BERT-Base (12 layers, 768 hidden units, 12 heads) from BERT-Large (24 layers, 1024 hidden, 16 heads), including dropout rates and the GELU activation choice.
- FullTokenizer, L161 is the single entry point for raw text. It chains a rule-based cleaner with a subword splitter and finally maps tokens to vocabulary indices.
- WordpieceTokenizer, L300 implements the greedy longest-match search that turns words missing from the vocabulary into
##-prefixed fragments — the detail that let a 30k-token vocabulary cover English nearly losslessly. - model_fn_builder, L619 returns the
model_fnclosure thatTPUEstimatorcalls; it is the junction where the encoder, the task head, the loss, and the training op meet. - create_optimizer, L25 assembles the learning-rate schedule — linear warmup followed by polynomial decay — and clips gradients to a global norm of 1.0.
- AdamWeightDecayOptimizer, L87 is the Adam variant with decoupled weight decay, deliberately skipping LayerNorm kernels and biases — a recipe later papers copied almost verbatim.
- write_predictions, L741 closes the question-answering loop, joining raw start/end logits with token features to emit n-best answer strings in original text.
One File, One Encoder: modeling.py
Everything in modeling.py exists to answer one question: given input_ids, input_mask, and token_type_ids, what tensors come out? The constructor of BertModel reads like a checklist of the architecture. First, embedding_lookup, L380 maps integer WordPiece ids to dense vectors, with a one-hot matmul path kept for TPU efficiency. Then embedding_postprocessor, L428 adds learned position embeddings and segment (token-type) embeddings on top, applies LayerNorm and dropout — this is where BERT differs from the original Transformer by using learned positions rather than sinusoids.
Attention itself is split into two functions. attention_layer, L558 does one round of multi-head scaled dot-product attention: it reshapes the 768-dim hidden states into per-head slices, computes Q·Kᵀ/√d, adds the mask as an additive bias before the softmax, and projects the concatenated heads back to the model width. transformer_model, L754 is the loop that runs num_hidden_layers rounds of that attention plus a position-wise feed-forward network activated by gelu, L264. The attention mask starts life as a simple 2D padding indicator and is expanded into a 3D bias by create_attention_mask_from_input_mask, L524.
Finally, the pooler. Lines 224 through 232 squeeze out the hidden state of the first token (the [CLS] slot), push it through a dense tanh layer, and stash the result; get_pooled_output, L234 hands it to task heads, while get_sequence_output() returns the full per-token sequence for tagging and extraction.
WordPiece, With the Rules in Plain Sight: tokenization.py
The tokenizer file is the most portable code in the repo — pure Python, no TensorFlow at all. load_vocab, L121 reads vocab.txt, one token per line, into an ordered dictionary. FullTokenizer, L161 then composes two stages. Stage one is BasicTokenizer, L185: optional lower-casing and accent stripping, splitting CJK characters into singletons, and breaking punctuation into separate tokens. Stage two is WordpieceTokenizer, L300: for each word, take the longest prefix that exists in the vocabulary, mark it, then repeat on the remainder with a ## continuation marker, giving up after 100 characters and emitting the [UNK] token instead.
The output mapping happens in convert_tokens_to_ids, L144, a trivial dictionary lookup that nonetheless defines the exact integer contract every downstream consumer depends on. Because the logic is deterministic and dependency-free, this file was lifted almost unchanged into later reimplementations; its modern, fast-Rust descendant lives on in tokenizers, which still ships BERT’s original vocab format.
Two Task Scripts, One Pattern
run_classifier.py and run_squad.py share an identical skeleton: read examples, convert to fixed-length features, build a model on top of BertModel, wrap it in an Estimator, train, predict. Learning one teaches you the other, and both teach you the TF1 idiom.
In the classifier script, GLUE datasets enter through DataProcessor subclasses (MrpcProcessor, L296, XnliProcessor, L207) that yield InputExample records. convert_single_example, L377 tokenizes a sentence pair and calls _truncate_seq_pair to fit max_seq_length by trimming the longer sentence from the right; file_based_convert_examples_to_features, L479 serializes the InputFeatures, L161 records into TFRecord protos that file_based_input_fn_builder, L509 parses back inside the Estimator input function. The head is deliberately boring: create_model, L574 takes the pooled output, applies dropout, and does a single matrix multiply into a softmax.
The SQuAD script adds span logic. read_squad_examples, L227 aligns answer character offsets to whitespace tokens; convert_examples_to_features, L309 slides a doc_stride window across long passages, marking which window holds the true answer via _check_is_max_context; the model at create_model, L550 emits two logits vectors for start and end positions. Prediction time in write_predictions, L741 reassembles WordPiece output into readable text with get_final_text, L927. A third runner, extract_features.py (entry at main, L343), skips heads entirely and dumps contextual embeddings from any requested layer — still a popular way to get fixed BERT features for downstream models.
From Install to Your First Model
This is TensorFlow 1-era code, and honesty matters here: the paths below describe how the repo was meant to be run, and what it costs. The requirements.txt asks for tensorflow >= 1.11.0 (CPU build; the GPU build was a separate tensorflow-gpu package). On a modern machine the pragmatic route is a dedicated Python 3.7-or-earlier environment with the last TF1 release:
conda create -n bert-tf1 python=3.7 -y
conda activate bert-tf1
pip install tensorflow==1.15
git clone https://github.com/google-research/bert.git
python run_classifier.py --help # sanity-check the flag surface
Then download the uncased_L-12_H-768_A-12 BERT-Base checkpoint, grab the GLUE MRPC data, and run the README’s fine-tuning example:
export BERT_BASE_DIR=/path/to/uncased_L-12_H-768_A-12
export GLUE_DIR=/path/to/glue
python run_classifier.py \
--task_name=MRPC \
--do_train=true \
--do_eval=true \
--data_dir=$GLUE_DIR/MRPC \
--vocab_file=$BERT_BASE_DIR/vocab.txt \
--bert_config_file=$BERT_BASE_DIR/bert_config.json \
--init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
--max_seq_length=128 \
--train_batch_size=32 \
--learning_rate=2e-5 \
--num_train_epochs=3.0 \
--output_dir=/tmp/mrpc_output/
MRPC is tiny — about 3,600 examples — so this converges in roughly 343 steps and a few minutes on most GPUs, landing near 84–88% dev accuracy. The README is candid about hardware: everything runs on CPU, GPU, or Cloud TPU; BERT-Base fine-tuning fits comfortably in 12 GB of GPU memory, while BERT-Large wants 12–16 GB and will still out-of-memory on SQuAD even at batch size 1 on a 12 GB card. Pretraining from scratch is another world entirely — the paper’s BERT-Base replication took about an hour on a 16-cloud-TPU pod, and run_pretraining.py (entry at main, L406) assumes that scale, with create_pretraining_data.py (main, L436) masking 15% of tokens up front.
Honest Limits
Treat this repository as an archive, not a dependency. It is frozen TF1 code: tf.layers, Estimator, and tf.contrib calls that modern TensorFlow removed, plus Python 2/3 straddling helpers like convert_to_unicode. Nothing here has been maintained for years, and Google has long since moved on. The architecture itself was also superseded within two years — RoBERTa, ELECTRA, and DeBERTa each kept the skeleton and replaced the training recipe or the head design.
For real work today, the ecosystem lives elsewhere. Hugging Face transformers reimplements BERT in PyTorch, loads these exact checkpoints, and has become the de facto interface for the entire encoder family. The WordPiece algorithm survived as the tokenizers library in fast Rust. If you need the original frozen graph out of the TF1 world, exporting through ONNX remains the cleanest exit. What this repo still offers that no replacement can, is legibility: eight files where every design decision of the foundational era — learned positions, the [CLS] pooler, decoupled weight decay, warmup schedules — sits in plain sight with a comment above it.
The detail graph: from FullTokenizer's two stages through the encoder stack, into the classifier and SQuAD heads, with the pretraining scripts sharing the same optimizer.
If you want to go deeper, start at Enjoyed this post? Never miss out on future posts by following us BertModel and follow the tensors; every symbol above links to the exact line in the upstream repository, and the repository’s README remains the best companion for checkpoint formats and hardware expectations.