# Copyright 2023 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Runs prediction to generate submission files for SuperGLUE tasks.""" import functools import json import os import pprint from absl import app from absl import flags from absl import logging import gin import tensorflow as tf, tf_keras from official.common import distribute_utils # Imports registered experiment configs. from official.common import registry_imports # pylint: disable=unused-import from official.core import exp_factory from official.core import task_factory from official.core import train_lib from official.core import train_utils from official.modeling.hyperparams import params_dict from official.nlp.finetuning import binary_helper from official.nlp.finetuning.superglue import flags as superglue_flags # Device configs. flags.DEFINE_string('distribution_strategy', 'tpu', 'The Distribution Strategy to use for training.') flags.DEFINE_string( 'tpu', '', 'The Cloud TPU to use for training. This should be either the name ' 'used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.') flags.DEFINE_integer('num_gpus', 1, 'The number of GPUs to use at each worker.') FLAGS = flags.FLAGS EXPERIMENT_TYPE = 'bert/sentence_prediction' BEST_CHECKPOINT_EXPORT_SUBDIR = 'best_ckpt' EVAL_METRIC_MAP = { 'AX-b': 'matthews_corrcoef', 'CB': 'cls_accuracy', 'COPA': 'cls_accuracy', 'MULTIRC': 'exact_match', 'RTE': 'cls_accuracy', 'WiC': 'cls_accuracy', 'WSC': 'cls_accuracy', 'BoolQ': 'cls_accuracy', 'ReCoRD': 'cls_accuracy', 'AX-g': 'cls_accuracy', } AXG_CLASS_NAMES = ['entailment', 'not_entailment'] RTE_CLASS_NAMES = ['entailment', 'not_entailment'] CB_CLASS_NAMES = ['entailment', 'neutral', 'contradiction'] BOOLQ_CLASS_NAMES = ['True', 'False'] def _override_exp_config_by_file(exp_config, exp_config_files): """Overrides an `ExperimentConfig` object by files.""" for exp_config_file in exp_config_files: if not tf.io.gfile.exists(exp_config_file): raise ValueError('%s does not exist.' % exp_config_file) params_dict.override_params_dict( exp_config, exp_config_file, is_strict=True) return exp_config def _override_exp_config_by_flags(exp_config, input_meta_data): """Overrides an `ExperimentConfig` object by flags.""" if FLAGS.task_name in 'AX-b': override_task_cfg_fn = functools.partial( binary_helper.override_sentence_prediction_task_config, num_classes=input_meta_data['num_labels'], metric_type='matthews_corrcoef') elif FLAGS.task_name in ('CB', 'COPA', 'RTE', 'WiC', 'WSC', 'BoolQ', 'ReCoRD', 'AX-g'): override_task_cfg_fn = functools.partial( binary_helper.override_sentence_prediction_task_config, num_classes=input_meta_data['num_labels']) else: raise ValueError('Task %s not supported.' % FLAGS.task_name) binary_helper.override_trainer_cfg( exp_config.trainer, learning_rate=FLAGS.learning_rate, num_epoch=FLAGS.num_epoch, global_batch_size=FLAGS.global_batch_size, warmup_ratio=FLAGS.warmup_ratio, training_data_size=input_meta_data['train_data_size'], eval_data_size=input_meta_data['eval_data_size'], num_eval_per_epoch=FLAGS.num_eval_per_epoch, best_checkpoint_export_subdir=BEST_CHECKPOINT_EXPORT_SUBDIR, best_checkpoint_eval_metric=EVAL_METRIC_MAP[FLAGS.task_name], best_checkpoint_metric_comp='higher') override_task_cfg_fn( exp_config.task, model_config_file=FLAGS.model_config_file, init_checkpoint=FLAGS.init_checkpoint, hub_module_url=FLAGS.hub_module_url, global_batch_size=FLAGS.global_batch_size, train_input_path=FLAGS.train_input_path, validation_input_path=FLAGS.validation_input_path, seq_length=input_meta_data['max_seq_length']) return exp_config def _get_exp_config(input_meta_data, exp_config_files): """Gets an `ExperimentConfig` object.""" exp_config = exp_factory.get_exp_config(EXPERIMENT_TYPE) if exp_config_files: logging.info( 'Loading `ExperimentConfig` from file, and flags will be ignored.') exp_config = _override_exp_config_by_file(exp_config, exp_config_files) else: logging.info('Loading `ExperimentConfig` from flags.') exp_config = _override_exp_config_by_flags(exp_config, input_meta_data) exp_config.validate() exp_config.lock() pp = pprint.PrettyPrinter() logging.info('Final experiment parameters: %s', pp.pformat(exp_config.as_dict())) return exp_config def _write_submission_file(task, seq_length): """Writes submission files that can be uploaded to the leaderboard.""" tf.io.gfile.makedirs(os.path.dirname(FLAGS.test_output_path)) model = task.build_model() ckpt_file = tf.train.latest_checkpoint( os.path.join(FLAGS.model_dir, BEST_CHECKPOINT_EXPORT_SUBDIR)) logging.info('Restoring checkpoints from %s', ckpt_file) checkpoint = tf.train.Checkpoint(model=model) checkpoint.read(ckpt_file).expect_partial() write_fn = binary_helper.write_superglue_classification write_fn_map = { 'RTE': functools.partial(write_fn, class_names=RTE_CLASS_NAMES), 'AX-g': functools.partial(write_fn, class_names=AXG_CLASS_NAMES), 'CB': functools.partial(write_fn, class_names=CB_CLASS_NAMES), 'BoolQ': functools.partial(write_fn, class_names=BOOLQ_CLASS_NAMES) } logging.info('Predicting %s', FLAGS.test_input_path) write_fn_map[FLAGS.task_name]( task=task, model=model, input_file=FLAGS.test_input_path, output_file=FLAGS.test_output_path, predict_batch_size=(task.task_config.train_data.global_batch_size), seq_length=seq_length) def main(argv): if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') superglue_flags.validate_flags(FLAGS, file_exists_fn=tf.io.gfile.exists) gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params) distribution_strategy = distribute_utils.get_distribution_strategy( distribution_strategy=FLAGS.distribution_strategy, num_gpus=FLAGS.num_gpus, tpu_address=FLAGS.tpu) with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader: input_meta_data = json.loads(reader.read().decode('utf-8')) with distribution_strategy.scope(): task = None if 'train_eval' in FLAGS.mode: logging.info('Starting training and eval...') logging.info('Model dir: %s', FLAGS.model_dir) exp_config = _get_exp_config( input_meta_data=input_meta_data, exp_config_files=FLAGS.config_file) train_utils.serialize_config(exp_config, FLAGS.model_dir) task = task_factory.get_task(exp_config.task, logging_dir=FLAGS.model_dir) train_lib.run_experiment( distribution_strategy=distribution_strategy, task=task, mode='train_and_eval', params=exp_config, model_dir=FLAGS.model_dir) if 'predict' in FLAGS.mode: logging.info('Starting predict...') # When mode is `predict`, `task` will be None. if task is None: exp_config = _get_exp_config( input_meta_data=input_meta_data, exp_config_files=[os.path.join(FLAGS.model_dir, 'params.yaml')]) task = task_factory.get_task( exp_config.task, logging_dir=FLAGS.model_dir) _write_submission_file(task, input_meta_data['max_seq_length']) if __name__ == '__main__': superglue_flags.define_flags() flags.mark_flag_as_required('mode') flags.mark_flag_as_required('task_name') app.run(main)