Fix squad script.

PiperOrigin-RevId: 269640109

official/nlp/bert/README.md

@@ -258,8 +258,6 @@ export MODEL_DIR=gs://some_bucket/my_output_dir
 export SQUAD_VERSION=v1.1
 
 python run_squad.py \
-  --do_train=true \
-  --do_predict=true \
   --input_meta_data_path=${SQUAD_DIR}/squad_${SQUAD_VERSION}_meta_data \
   --train_data_path=${SQUAD_DIR}/squad_${SQUAD_VERSION}_train.tf_record \
   --predict_file=${SQUAD_DIR}/dev-v1.1.json \
@@ -285,8 +283,6 @@ export SQUAD_DIR=gs://some_bucket/datasets
 export SQUAD_VERSION=v1.1
 
 python run_squad.py \
-  --do_train=true \
-  --do_predict=true \
   --input_meta_data_path=${SQUAD_DIR}/squad_${SQUAD_VERSION}_meta_data \
   --train_data_path=${SQUAD_DIR}/squad_${SQUAD_VERSION}_train.tf_record \
   --predict_file=${SQUAD_DIR}/dev-v1.1.json \

official/nlp/bert/run_squad.py

@@ -41,9 +41,11 @@ from official.utils.misc import keras_utils
 from official.utils.misc import tpu_lib
 
 flags.DEFINE_enum(
-    'mode', 'train', ['train', 'predict', 'export_only'],
-    'One of {"train", "predict", "export_only"}. `train`: '
-    'trains the model and evaluates in the meantime. '
+    'mode', 'train_and_predict',
+    ['train_and_predict', 'train', 'predict', 'export_only'],
+    'One of {"train_and_predict", "train", "predict", "export_only"}. '
+    '`train_and_predict`: both train and predict to a json file. '
+    '`train`: only trains the model. '
     '`predict`: predict answers from the squad json file. '
     '`export_only`: will take the latest checkpoint inside '
     'model_dir and export a `SavedModel`.')
@@ -370,9 +372,9 @@ def main(_):
   else:
     raise ValueError('The distribution strategy type is not supported: %s' %
                      FLAGS.strategy_type)
-  if FLAGS.mode == 'train':
+  if FLAGS.mode in ('train', 'train_and_predict'):
     train_squad(strategy, input_meta_data)
-  if FLAGS.mode == 'predict':
+  if FLAGS.mode in ('predict', 'train_and_predict'):
     predict_squad(strategy, input_meta_data)