updating readme and notebooks

README.md

@@ -1,6 +1,4 @@
-# PyTorch implementation of Google AI's BERT model with a script to load Google's pre-trained models
+# PyTorch implementation of Google AI's BERT model with Google's pre-trained models
-
-## Introduction
 
 This repository contains an op-for-op PyTorch reimplementation of [Google's TensorFlow repository for the BERT model](https://github.com/google-research/bert) that was released together with the paper [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) by Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova.
 
@@ -8,18 +6,36 @@ This implementation can load any pre-trained TensorFlow checkpoint for BERT (in
 
 The code to use, in addition, [the Multilingual and Chinese models](https://github.com/google-research/bert/blob/master/multilingual.md) will be added later this week (it's actually just the tokenization code that needs to be updated).
 
-## Installation, requirements, test
+# Documentation
 
-This code was tested on Python 3.5+. The requirements are:
+| Section | Content |
+|-|-|
+| [Installation](#installation) | How to install the package |
+| [Content](#content) | Overview of the package |
+| [Usage](#usage) | Quickstart examples |
+| [Doc](#doc) |  Detailed documentation |
+| [Examples](#examples) | Detailed examples on how to fine-tune Bert |
+| [Notebooks](#notebooks) | Introduction on the provided Jupyter Notebooks |
+| [TPU](#tup) | Notes on TPU support and pretraining scripts |
+| [Command-line interface](#Command-line-interface) | Convert a TensorFlow checkpoint in a PyTorch dump |
 
-- PyTorch (>= 0.4.1)
+# Installation
-- tqdm
 
-To install the dependencies:
+This repo was tested on Python 3.5+ and PyTorch 0.4.1
 
-````bash
+## From pip
-pip install -r ./requirements.txt
+
-````
+PyTorch pretrained bert can be installed by pip as follows:
+```bash
+pip install pytorch_pretrained_bert
+```
+
+## From source
+
+Clone the repository and run:
+```bash
+pip install [--editable] .
+```
 
 A series of tests is included in the [tests folder](https://github.com/huggingface/pytorch-pretrained-BERT/tree/master/tests) and can be run using `pytest` (install pytest if needed: `pip install pytest`).
 
@@ -28,15 +44,123 @@ You can run the tests with the command:
 python -m pytest -sv tests/
 ```
 
-## PyTorch models for BERT
+# Content
 
-We included three PyTorch models in this repository that you will find in [`modeling.py`](modeling.py):
+This package comprises the following classes that can be imported in Python and are detailed in the [Doc](#doc) section of this readme:
 
-- `BertModel` - the basic BERT Transformer model
+- Six PyTorch models (`torch.nn.Module`) for Bert with pre-trained weights:
-- `BertForSequenceClassification` - the BERT model with a sequence classification head on top
+  - `BertModel` - raw BERT Transformer model (**fully pre-trained**),
-- `BertForQuestionAnswering` - the BERT model with a token classification head on top
+  - `BertForMaskedLM` - BERT Transformer with the pre-trained masked language modeling head on top (**fully pre-trained**),
+  - `BertForNextSentencePrediction` - BERT Transformer with the pre-trained next sentence prediction classifier on top  (**fully pre-trained**),
+  - `BertForPretraining` - BERT Transformer with masked language modeling head and next sentence prediction classifier on top (**fully pre-trained**),
+  - `BertForSequenceClassification` - BERT Transformer with a sequence classification head on top (BERT Transformer is **pre-trained**, the sequence classification head **is only initialized and has to be trained**),
+  - `BertForQuestionAnswering` - BERT Transformer with a token classification head on top (BERT Transformer is **pre-trained**, the token classification head **is only initialized and has to be trained**).
 
-Here are some details on each class.
+- Three tokenizers:
+  - `BasicTokenizer` - basic tokenization (punctuation splitting, lower casing, etc.),
+  - `WordpieceTokenizer` - WordPiece tokenization,
+  - `BertTokenizer` - perform end-to-end tokenization, i.e. basic tokenization followed by WordPiece tokenization.
+
+- One optimizer:
+  - `BERTAdam` - Bert version of Adam algorithm with weight decay fix, warmup and linear decay of the learning rate.
+
+- A configuration class:
+  - `BertConfig` - Configuration class to store the configuration of a `BertModel` with utilisities to read and write from JSON configuration files.
+
+The repository further comprises:
+
+- Three examples on how to use Bert (in the [`examples` folder](./examples)):
+  - [`extract_features.py`](./examples/extract_features.py) - Show how to extract hidden states from an instance of `BertModel`,
+  - [`run_classifier.py`](./examples/run_classifier.py) - Show how to fine-tune an instance of `BertForSequenceClassification` on GLUE's MRPC task,
+  - [`run_squad.py`](./examples/run_squad.py) - Show how to fine-tune an instance of `BertForQuestionAnswering` on SQuAD v1.0 task.
+
+  These examples are detailed in the [Examples](#examples) section of this readme.
+
+- Three notebooks that were used to check that the TensorFlow and PyTorch models behave identically (in the [`notebooks` folder](./notebooks)):
+  - [`Comparing-TF-and-PT-models.ipynb`](./notebooks/Comparing-TF-and-PT-models.ipynb) - Compare the hidden states predicted by `BertModel`,
+  - [`Comparing-TF-and-PT-models-SQuAD.ipynb`](./notebooks/Comparing-TF-and-PT-models-SQuAD.ipynb) - Compare the spans predicted by  `BertForQuestionAnswering` instances,
+  - [`Comparing-TF-and-PT-models-MLM-NSP.ipynb`](./notebooks/Comparing-TF-and-PT-models-MLM-NSP.ipynb) - Compare the predictions of the `BertForPretraining` instances.
+
+  These notebooks are detailed in the [Notebooks](#notebooks) section of this readme.
+
+- A command-line interface to convert any TensorFlow checkpoint in a PyTorch dump:
+
+  This CLI is detailed in the [Command-line interface](#Command-line-interface) section of this readme.
+
+# Usage
+
+Here is a quick-start example using the `BertForMaskedLM` class with Google AI's pre-trained `Bert base uncased` model:
+
+```python
+import torch
+from pytorch_pretrained_bert import BertForMaskedLM, BertTokenizer
+
+# Load pre-trained model and tokenizer (weights and vocabulary)
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForMaskedLM.from_pretrained('bert-base-uncased')
+
+# Prepare tokenized input with a masked token
+tokenized_text = "Who was Jim Henson ? Jim Henson was a puppeteer"
+tokenized_text = tokenizer.tokenize(text)
+masked_index = 6
+tokenized_text[masked_index] = '[MASK]'
+assert tokenized_text == ['who', 'was', 'jim', 'henson', '?', 'jim', '[MASK]', 'was', 'a', 'puppet', '##eer']
+
+# Convert token to vocabulary indices
+indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
+# Assign sentence A and sentence B indices to 1st (resp 2nd) sentences
+segments_ids = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
+
+# Predict masked tokens with model
+tokens_tensor = torch.tensor([indexed_tokens])
+segments_tensors = torch.tensor([segments_ids])
+model.eval()
+predictions = model(tokens_tensor, segments_tensors)
+
+# Use model to predict
+predicted_index = torch.argmax(predictions[0, masked_index]).item()
+predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])
+assert predicted_token == 'henson'
+```
+
+# Doc
+
+Here is a detailed documentation of the classes in the package.
+
+## Loading pre-trained weigths
+
+To load Google AI's pre-trained weight, the PyTorch model classes and the tokenizer can be instantiated as
+
+```python
+model = BERT_CLASS.from_pretrain(PRE_TRAINED_MODEL_NAME_OR_PATH)
+```
+
+where
+
+- `BERT_CLASS` is either the `BertTokenizer` class (to load the vocabulary) or one of the six PyTorch model classes: `BertModel`, `BertForMaskedLM`, `BertForNextSentencePrediction`, `BertForPretraining`, `BertForSequenceClassification` or `BertForQuestionAnswering` (to load the pre-trained weights), and
+
+- `PRE_TRAINED_MODEL_NAME` is either:
+
+  - the shortcut name of a Google AI's pre-trained model selected in the list:
+
+    - `bert-base-uncased`: 12-layer, 768-hidden, 12-heads, 110M parameters
+    - `bert-large-uncased`: 24-layer, 1024-hidden, 16-heads, 340M parameters
+    - `bert-base-cased`: 12-layer, 768-hidden, 12-heads , 110M parameters
+    - `bert-base-multilingual`: 102 languages, 12-layer, 768-hidden, 12-heads, 110M parameters
+    - `bert-base-chinese`: Chinese Simplified and Traditional, 12-layer, 768-hidden, 12-heads, 110M parameters
+
+  - a path or url to a pretrained model archive containing:
+      . `bert_config.json` a configuration file for the model
+      . `pytorch_model.bin` a PyTorch dump of a pre-trained instance `BertForPreTraining` (saved with the usual `torch.save()`)
+
+If `PRE_TRAINED_MODEL_NAME` is a shortcut name, the pre-trained weights will be downloaded from AWS S3 (see the links [here](pytorch_pretrained_bert/modeling.py)) and stored in a cache folder to avoid future download (the cache folder can be found at `~/.pytorch_pretrained_bert/`).
+
+Example:
+```python
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+```
+
+## PyTorch models
 
 ### 1. `BertModel`
 
@@ -44,14 +168,14 @@ Here are some details on each class.
 
 The inputs and output are **identical to the TensorFlow model inputs and outputs**.
 
-We detail them here. This model takes as inputs:
+We detail them here. This model takes as *inputs*:
 
 - `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary (see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`), and
 - `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
 - `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences.
 - `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
 
-This model outputs a tuple composed of:
+This model *outputs* a tuple composed of:
 
 - `encoded_layers`: controled by the value of the `output_encoded_layers` argument:
 
@@ -62,7 +186,52 @@ This model outputs a tuple composed of:
 
 An example on how to use this class is given in the `extract_features.py` script which can be used to extract the hidden states of the model for a given input.
 
-### 2. `BertForSequenceClassification`
+### 2. `BertForPreTraining`
+
+`BertForPreTraining` includes the `BertModel` Transformer followed by the two pre-training heads:
+
+- the masked language modeling head, and
+- the next sentence classification head.
+
+*Inputs* comprises the inputs of the [`BertModel`](###-1.-`BertModel`) class plus two optional labels:
+
+- `masked_lm_labels`: masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss is only computed for the labels set in [0, ..., vocab_size]
+- `next_sentence_label`: next sentence classification loss: torch.LongTensor of shape [batch_size] with indices selected in [0, 1]. 0 => next sentence is the continuation, 1 => next sentence is a random sentence.
+
+*Outputs*:
+
+- if `masked_lm_labels` and `next_sentence_label` are not `None`: Outputs the total_loss which is the sum of the masked language modeling loss and the next sentence classification loss.
+- if `masked_lm_labels` or `next_sentence_label` is `None`: Outputs a tuple comprising
+  - the masked language modeling logits, and
+  - the next sentence classification logits.
+
+### 3. `BertForMaskedLM`
+
+`BertForMaskedLM` includes the `BertModel` Transformer followed by the (possibly) pre-trained  masked language modeling head.
+
+*Inputs* comprises the inputs of the [`BertModel`](###-1.-`BertModel`) class plus optional label:
+
+- `masked_lm_labels`: masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss is only computed for the labels set in [0, ..., vocab_size]
+
+*Outputs*:
+
+- if `masked_lm_labels` is not `None`: Outputs the masked language modeling loss.
+- if `masked_lm_labels` is `None`: Outputs the masked language modeling logits.
+
+### 4. `BertForNextSentencePrediction`
+
+`BertForNextSentencePrediction` includes the `BertModel` Transformer followed by the next sentence classification head.
+
+*Inputs* comprises the inputs of the [`BertModel`](###-1.-`BertModel`) class plus an optional label:
+
+- `next_sentence_label`: next sentence classification loss: torch.LongTensor of shape [batch_size] with indices selected in [0, 1]. 0 => next sentence is the continuation, 1 => next sentence is a random sentence.
+
+*Outputs*:
+
+- if `next_sentence_label` is not `None`: Outputs the next sentence classification loss.
+- if `next_sentence_label` is `None`: Outputs the next sentence classification logits.
+
+### 5. `BertForSequenceClassification`
 
 `BertForSequenceClassification` is a fine-tuning model that includes `BertModel` and a sequence-level (sequence or pair of sequences) classifier on top of the `BertModel`.
 
@@ -70,7 +239,7 @@ The sequence-level classifier is a linear layer that takes as input the last hid
 
 An example on how to use this class is given in the `run_classifier.py` script which can be used to fine-tune a single sequence (or pair of sequence) classifier using BERT, for example for the MRPC task.
 
-### 3. `BertForQuestionAnswering`
+### 6. `BertForQuestionAnswering`
 
 `BertForQuestionAnswering` is a fine-tuning model that includes `BertModel` with a token-level classifiers on top of the full sequence of last hidden states.
 
@@ -78,31 +247,54 @@ The token-level classifier takes as input the full sequence of the last hidden s
 
 An example on how to use this class is given in the `run_squad.py` script which can be used to fine-tune a token classifier using BERT, for example for the SQuAD task.
 
+## Tokenizers
 
-## Converting a TensorFlow checkpoint in a PyTorch checkpoint
+### `BertTokenizer`
 
-You can convert any TensorFlow checkpoint for BERT (in particular [the pre-trained models released by Google](https://github.com/google-research/bert#pre-trained-models)) in a PyTorch save file by using the [`convert_tf_checkpoint_to_pytorch.py`](convert_tf_checkpoint_to_pytorch.py) script.
+`BertTokenizer` perform end-to-end tokenization, i.e. basic tokenization followed by WordPiece tokenization.
 
-This script takes as input a TensorFlow checkpoint (three files starting with `bert_model.ckpt`) and the associated configuration file (`bert_config.json`), and creates a PyTorch model for this configuration, loads the weights from the TensorFlow checkpoint in the PyTorch model and saves the resulting model in a standard PyTorch save file that can be imported using `torch.load()` (see examples in `extract_features.py`, `run_classifier.py` and `run_squad.py`).
+This class has two arguments:
 
-You only need to run this conversion script **once** to get a PyTorch model. You can then disregard the TensorFlow checkpoint (the three files starting with `bert_model.ckpt`) but be sure to keep the configuration file (`bert_config.json`) and the vocabulary file (`vocab.txt`) as these are needed for the PyTorch model too.
+- `vocab_file`: path to a vocabulary file.
+- `do_lower_case`: convert text to lower-case while tokenizing. **Default = True**.
 
-To run this specific conversion script you will need to have TensorFlow and PyTorch installed (`pip install tensorflow`). The rest of the repository only requires PyTorch.
+and three methods:
 
-Here is an example of the conversion process for a pre-trained `BERT-Base Uncased` model:
+- `tokenize(text)`: convert a `str` in a list of `str` tokens by (1) performing basic tokenization and (2) WordPiece tokenization.
+- `convert_tokens_to_ids(tokens)`: convert a list of `str` tokens in a list of `int` indices in the vocabulary.
+- `convert_ids_to_tokens(tokens)`: convert a list of `int` indices in a list of `str` tokens in the vocabulary.
 
-```shell
+### `BasicTokenizer` and `WordpieceTokenizer`
-export BERT_BASE_DIR=/path/to/bert/uncased_L-12_H-768_A-12
 
-python convert_tf_checkpoint_to_pytorch.py \
+Please refer to the doc strings and code in [`tokenization.py`](./pytorch_pretrained_bert/tokenization.py) for the details of these classes. In general it is recommended to use `BertTokenizer` unless you know what you are doing.
-  --tf_checkpoint_path $BERT_BASE_DIR/bert_model.ckpt \
-  --bert_config_file $BERT_BASE_DIR/bert_config.json \
-  --pytorch_dump_path $BERT_BASE_DIR/pytorch_model.bin
-```
 
-You can download Google's pre-trained models for the conversion [here](https://github.com/google-research/bert#pre-trained-models).
+## Optimizer
 
-## Training on large batches: gradient accumulation, multi-GPU and distributed training
+### `BERTAdam`
+
+`BERTAdam` is a `torch.optimizer` adapted to be closer to the optimizer used in the TensorFlow implementation of Bert. The differences with PyTorch Adam optimizer are the following:
+
+- BERTAdam implements weight decay fix,
+- BERTAdam doesn't compensate for bias as in the regular Adam optimizer.
+
+The optimizer accepts the following arguments:
+
+- `lr` : learning rate
+- `warmup` : portion of t_total for the warmup, -1  means no warmup. Default : -1
+- `t_total` : total number of training steps for the learning
+    rate schedule, -1  means constant learning rate. Default : -1
+- `schedule` : schedule to use for the warmup (see above). Default : 'warmup_linear'
+- `b1` : Adams b1. Default : 0.9
+- `b2` : Adams b2. Default : 0.999
+- `e` : Adams epsilon. Default : 1e-6
+- `weight_decay_rate:` Weight decay. Default : 0.01
+- `max_grad_norm` : Maximum norm for the gradients (-1 means no clipping). Default : 1.0
+
+# Examples
+
+Fine-tuning the models
+
+## Training large models: introduction, tools and examples
 
 BERT-base and BERT-large are respectively 110M and 340M parameters models and it can be difficult to fine-tune them on a single GPU with the recommended batch size for good performance (in most case a batch size of 32).
 
@@ -122,26 +314,6 @@ python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=$TH
 ```
 Where `$THIS_MACHINE_INDEX` is an sequential index assigned to each of your machine (0, 1, 2...) and the machine with rank 0 has an IP address `192.168.1.1` and an open port `1234`.
 
-## TPU support and pretraining scripts
-
-TPU are not supported by the current stable release of PyTorch (0.4.1). However, the next version of PyTorch (v1.0) should support training on TPU and is expected to be released soon (see the recent [official announcement](https://cloud.google.com/blog/products/ai-machine-learning/introducing-pytorch-across-google-cloud)).
-
-We will add TPU support when this next release is published.
-
-The original TensorFlow code further comprises two scripts for pre-training BERT: [create_pretraining_data.py](https://github.com/google-research/bert/blob/master/create_pretraining_data.py) and [run_pretraining.py](https://github.com/google-research/bert/blob/master/run_pretraining.py).
-
-Since, pre-training BERT is a particularly expensive operation that basically requires one or several TPUs to be completed in a reasonable amout of time (see details [here](https://github.com/google-research/bert#pre-training-with-bert)) we have decided to wait for the inclusion of TPU support in PyTorch to convert these pre-training scripts.
-
-## Comparing the PyTorch model and the TensorFlow model predictions
-
-We also include [two Jupyter Notebooks](https://github.com/huggingface/pytorch-pretrained-BERT/tree/master/notebooks) that can be used to check that the predictions of the PyTorch model are identical to the predictions of the original TensorFlow model.
-
-- The first NoteBook ([Comparing TF and PT models.ipynb](https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/notebooks/Comparing%20TF%20and%20PT%20models.ipynb)) extracts the hidden states of a full sequence on each layers of the TensorFlow and the PyTorch models and computes the standard deviation between them. In the given example, we get a standard deviation of 1.5e-7 to 9e-7 on the various hidden state of the models.
-
-- The second NoteBook ([Comparing TF and PT models SQuAD predictions.ipynb](https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/notebooks/Comparing%20TF%20and%20PT%20models%20SQuAD%20predictions.ipynb)) compares the loss computed by the TensorFlow and the PyTorch models for identical initialization of the fine-tuning layer of the `BertForQuestionAnswering` and computes the standard deviation between them. In the given example, we get a standard deviation of 2.5e-7 between the models.
-
-Please follow the instructions given in the notebooks to run and modify them. They can also be nice example on how to use the models in a simpler way than the full fine-tuning scripts we provide.
-
 ## Fine-tuning with BERT: running the examples
 
 We showcase the same examples as [the original implementation](https://github.com/google-research/bert/): fine-tuning a sequence-level classifier on the MRPC classification corpus and a token-level classifier on the question answering dataset SQuAD.
@@ -270,3 +442,52 @@ The results were similar to the above FP32 results (actually slightly higher):
 ```bash
 {"exact_match": 84.65468306527909, "f1": 91.238669287002}
 ```
+
+# Notebooks
+
+Comparing the PyTorch model and the TensorFlow model predictions
+
+We also include [three Jupyter Notebooks](https://github.com/huggingface/pytorch-pretrained-BERT/tree/master/notebooks) that can be used to check that the predictions of the PyTorch model are identical to the predictions of the original TensorFlow model.
+
+- The first NoteBook ([Comparing TF and PT models.ipynb](https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/notebooks/Comparing%20TF%20and%20PT%20models.ipynb)) extracts the hidden states of a full sequence on each layers of the TensorFlow and the PyTorch models and computes the standard deviation between them. In the given example, we get a standard deviation of 1.5e-7 to 9e-7 on the various hidden state of the models.
+
+- The second NoteBook ([Comparing TF and PT models SQuAD predictions.ipynb](https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/notebooks/Comparing%20TF%20and%20PT%20models%20SQuAD%20predictions.ipynb)) compares the loss computed by the TensorFlow and the PyTorch models for identical initialization of the fine-tuning layer of the `BertForQuestionAnswering` and computes the standard deviation between them. In the given example, we get a standard deviation of 2.5e-7 between the models.
+
+Please follow the instructions given in the notebooks to run and modify them. They can also be nice example on how to use the models in a simpler way than the full fine-tuning scripts we provide.
+
+# Command-line interface
+
+A command-line interface is provided to convert a TensorFlow checkpoint in a PyTorch checkpoint
+
+You can convert any TensorFlow checkpoint for BERT (in particular [the pre-trained models released by Google](https://github.com/google-research/bert#pre-trained-models)) in a PyTorch save file by using the [`convert_tf_checkpoint_to_pytorch.py`](convert_tf_checkpoint_to_pytorch.py) script.
+
+This script takes as input a TensorFlow checkpoint (three files starting with `bert_model.ckpt`) and the associated configuration file (`bert_config.json`), and creates a PyTorch model for this configuration, loads the weights from the TensorFlow checkpoint in the PyTorch model and saves the resulting model in a standard PyTorch save file that can be imported using `torch.load()` (see examples in `extract_features.py`, `run_classifier.py` and `run_squad.py`).
+
+You only need to run this conversion script **once** to get a PyTorch model. You can then disregard the TensorFlow checkpoint (the three files starting with `bert_model.ckpt`) but be sure to keep the configuration file (`bert_config.json`) and the vocabulary file (`vocab.txt`) as these are needed for the PyTorch model too.
+
+To run this specific conversion script you will need to have TensorFlow and PyTorch installed (`pip install tensorflow`). The rest of the repository only requires PyTorch.
+
+Here is an example of the conversion process for a pre-trained `BERT-Base Uncased` model:
+
+```shell
+export BERT_BASE_DIR=/path/to/bert/uncased_L-12_H-768_A-12
+
+python convert_tf_checkpoint_to_pytorch.py \
+  --tf_checkpoint_path $BERT_BASE_DIR/bert_model.ckpt \
+  --bert_config_file $BERT_BASE_DIR/bert_config.json \
+  --pytorch_dump_path $BERT_BASE_DIR/pytorch_model.bin
+```
+
+You can download Google's pre-trained models for the conversion [here](https://github.com/google-research/bert#pre-trained-models).
+
+# TPU
+
+TPU support and pretraining scripts
+
+TPU are not supported by the current stable release of PyTorch (0.4.1). However, the next version of PyTorch (v1.0) should support training on TPU and is expected to be released soon (see the recent [official announcement](https://cloud.google.com/blog/products/ai-machine-learning/introducing-pytorch-across-google-cloud)).
+
+We will add TPU support when this next release is published.
+
+The original TensorFlow code further comprises two scripts for pre-training BERT: [create_pretraining_data.py](https://github.com/google-research/bert/blob/master/create_pretraining_data.py) and [run_pretraining.py](https://github.com/google-research/bert/blob/master/run_pretraining.py).
+
+Since, pre-training BERT is a particularly expensive operation that basically requires one or several TPUs to be completed in a reasonable amout of time (see details [here](https://github.com/google-research/bert#pre-training-with-bert)) we have decided to wait for the inclusion of TPU support in PyTorch to convert these pre-training scripts.

pytorch_pretrained_bert/optimization.py

@@ -42,7 +42,7 @@ SCHEDULES = {
 
 
 class BERTAdam(Optimizer):
-    """Implements BERT version of Adam algorithm with weight decay fix (and no ).
+    """Implements BERT version of Adam algorithm with weight decay fix.
     Params:
         lr: learning rate
         warmup: portion of t_total for the warmup, -1  means no warmup. Default: -1
@@ -136,7 +136,7 @@ class BERTAdam(Optimizer):
                 # the correct way of using L2 regularization/weight decay with Adam,
                 # since that will interact with the m and v parameters in strange ways.
                 #
-                # Instead we want ot decay the weights in a manner that doesn't interact
+                # Instead we want to decay the weights in a manner that doesn't interact
                 # with the m/v parameters. This is equivalent to adding the square
                 # of the weights to the loss with plain (non-momentum) SGD.
                 if group['weight_decay_rate'] > 0.0:
@@ -154,6 +154,7 @@ class BERTAdam(Optimizer):
                 state['step'] += 1
 
                 # step_size = lr_scheduled * math.sqrt(bias_correction2) / bias_correction1
+                # No bias correction
                 # bias_correction1 = 1 - beta1 ** state['step']
                 # bias_correction2 = 1 - beta2 ** state['step']