Merged commit includes the following changes: (#7255)

258881002  by hongkuny<hongkuny@google.com>:

    Fix lint.

--
258874998  by hongkuny<hongkuny@google.com>:

    Internal

--
258872662  by hongkuny<hongkuny@google.com>:

    Fix doc

--
258871624  by hongkuny<hongkuny@google.com>:

    Internal change

PiperOrigin-RevId: 258881002

official/bert/optimization.py

@@ -135,6 +135,7 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
     return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars))
 
   def _get_lr(self, var_device, var_dtype, apply_state):
+    """Retrieves the learning rate with the given state."""
     if apply_state is None:
       return self._decayed_lr_t[var_dtype], {}
 

official/staging/shakespeare/README.md

@@ -14,16 +14,16 @@ Note that files other than shakepeare.txt can also be used to train the model to
 
 Then train the model:
 
-```
+```python
 python3 shakespeare_main.py --training_data shakespeare.txt \
     --model_dir /tmp/shakespeare
-````
+```
 
 This will place model checkpoints in `/tmp/shakespeare`, so that we can use them to make predictions.
 
 Then generate predictions:
 
-```    
+```python
 python3 shakespeare_main.py --training_data shakespeare.txt \
     --model_dir /tmp/shakespeare --notrain --predict_context=ROMEO:
 ```