text: # coding=utf-8
  # Copyright 2020 The Google Research Authors.
  #
  # 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.
  
  """Fine-tunes an ELECTRA model on a downstream task."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import argparse
  import collections
  import json
  
  import tensorflow.compat.v1 as tf
  
  import configure_finetuning
  from finetune import preprocessing
  from finetune import task_builder
  from model import modeling
  from model import optimization
  from util import training_utils
  from util import utils
  
  
  class FinetuningModel(object):
    """Finetuning model with support for multi-task training."""
  
    def __init__(self, config: configure_finetuning.FinetuningConfig, tasks,
                 is_training, features, num_train_steps):
      # Create a shared transformer encoder
      bert_config = training_utils.get_bert_config(config)
      self.bert_config = bert_config
      if config.debug:
        bert_config.num_hidden_layers = 3
        bert_config.hidden_size = 144
        bert_config.intermediate_size = 144 * 4
        bert_config.num_attention_heads = 4
      assert config.max_seq_length <= bert_config.max_position_embeddings
      bert_model = modeling.BertModel(
          bert_config=bert_config,
          is_training=is_training,
          input_ids=features["input_ids"],
          input_mask=features["input_mask"],
          token_type_ids=features["segment_ids"],
          use_one_hot_embeddings=config.use_tpu,
          embedding_size=config.embedding_size)
      percent_done = (tf.cast(tf.train.get_or_create_global_step(), tf.float32) /
                      tf.cast(num_train_steps, tf.float32))
  
      # Add specific tasks
      self.outputs = {"task_id": features["task_id"]}
      losses = []
      for task in tasks:
        with tf.variable_scope("task_specific/" + task.name):
          task_losses, task_outputs = task.get_prediction_module(
              bert_model, features, is_training, percent_done)
          losses.append(task_losses)
          self.outputs[task.name] = task_outputs
      self.loss = tf.reduce_sum(
          tf.stack(losses, -1) *
          tf.one_hot(features["task_id"], len(config.task_names)))
  
  
  def model_fn_builder(config: configure_finetuning.FinetuningConfig, tasks,
                       num_train_steps, pretraining_config=None):
    """Returns `model_fn` closure for TPUEstimator."""
  
    def model_fn(features, labels, mode, params):
      """The `model_fn` for TPUEstimator."""
      utils.log("Building model...")
      is_training = (mode == tf.estimator.ModeKeys.TRAIN)
      model = FinetuningModel(
          config, tasks, is_training, features, num_train_steps)
  
      # Load pre-trained weights from checkpoint
      init_checkpoint = config.init_checkpoint
      if pretraining_config is not None:
        init_checkpoint = tf.train.latest_checkpoint(pretraining_config.model_dir)
        utils.log("Using checkpoint", init_checkpoint)
      tvars = tf.trainable_variables()
      scaffold_fn = None
      if init_checkpoint:
        assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
            tvars, init_checkpoint)
        if config.use_tpu:
          def tpu_scaffold():
            tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
            return tf.train.Scaffold()
          scaffold_fn = tpu_scaffold
        else:
          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
  
      # Build model for training or prediction
      if mode == tf.estimator.ModeKeys.TRAIN:
        train_op = optimization.create_optimizer(
            model.loss, config.learning_rate, num_train_steps,
            weight_decay_rate=config.weight_decay_rate,
            use_tpu=config.use_tpu,
            warmup_proportion=config.warmup_proportion,
            layerwise_lr_decay_power=config.layerwise_lr_decay,
            n_transformer_layers=model.bert_config.num_hidden_layers
        )
        output_spec = tf.estimator.tpu.TPUEstimatorSpec(
            mode=mode,
            loss=model.loss,
            train_op=train_op,
            scaffold_fn=scaffold_fn,
            training_hooks=[training_utils.ETAHook(
                {} if config.use_tpu else dict(loss=model.loss),
                num_train_steps, config.iterations_per_loop, config.use_tpu, 10)])
      else:
        assert mode == tf.estimator.ModeKeys.PREDICT
        output_spec = tf.estimator.tpu.TPUEstimatorSpec(
            mode=mode,
            predictions=utils.flatten_dict(model.outputs),
            scaffold_fn=scaffold_fn)
  
      utils.log("Building complete")
      return output_spec
  
    return model_fn
  
  
  class ModelRunner(object):
    """Fine-tunes a model on a supervised task."""
  
    def __init__(self, config: configure_finetuning.FinetuningConfig, tasks,
                 pretraining_config=None):
      self._config = config
      self._tasks = tasks
      self._preprocessor = preprocessing.Preprocessor(config, self._tasks)
  
      is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2
      tpu_cluster_resolver = None
      if config.use_tpu and config.tpu_name:
        tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
            config.tpu_name, zone=config.tpu_zone, project=config.gcp_project)
      tpu_config = tf.estimator.tpu.TPUConfig(
          iterations_per_loop=config.iterations_per_loop,
          num_shards=config.num_tpu_cores,
          per_host_input_for_training=is_per_host,
          tpu_job_name=config.tpu_job_name)
      run_config = tf.estimator.tpu.RunConfig(
          cluster=tpu_cluster_resolver,
          model_dir=config.model_dir,
          save_checkpoints_steps=config.save_checkpoints_steps,
          save_checkpoints_secs=None,
          tpu_config=tpu_config)
  
      if self._config.do_train:
        (self._train_input_fn,
         self.train_steps) = self._preprocessor.prepare_train()
      else:
        self._train_input_fn, self.train_steps = None, 0
      model_fn = model_fn_builder(
          config=config,
          tasks=self._tasks,
          num_train_steps=self.train_steps,
          pretraining_config=pretraining_config)
      self._estimator = tf.estimator.tpu.TPUEstimator(
          use_tpu=config.use_tpu,
          model_fn=model_fn,
          config=run_config,
          train_batch_size=config.train_batch_size,
          eval_batch_size=config.eval_batch_size,
          predict_batch_size=config.predict_batch_size)
  
    def train(self):
      utils.log("Training for {:} steps".format(self.train_steps))
      self._estimator.train(
          input_fn=self._train_input_fn, max_steps=self.train_steps)
  
    def evaluate(self):
      return {task.name: self.evaluate_task(task) for task in self._tasks}
  
    def evaluate_task(self, task, split="dev", return_results=True):
      """Evaluate the current model."""
      utils.log("Evaluating", task.name)
      eval_input_fn, _ = self._preprocessor.prepare_predict([task], split)
      results = self._estimator.predict(input_fn=eval_input_fn,
                                        yield_single_examples=True)
      scorer = task.get_scorer()
      for r in results:
        if r["task_id"] != len(self._tasks):  # ignore padding examples
          r = utils.nest_dict(r, self._config.task_names)
          scorer.update(r[task.name])
      if return_results:
        utils.log(task.name + ": " + scorer.results_str())
        utils.log()
        return dict(scorer.get_results())
      else:
        return scorer
  
    def write_classification_outputs(self, tasks, trial, split):
      """Write classification predictions to disk."""
      utils.log("Writing out predictions for", tasks, split)
      predict_input_fn, _ = self._preprocessor.prepare_predict(tasks, split)
      results = self._estimator.predict(input_fn=predict_input_fn,
                                        yield_single_examples=True)
      # task name -> eid -> model-logits
      logits = collections.defaultdict(dict)
      for r in results:
        if r["task_id"] != len(self._tasks):
          r = utils.nest_dict(r, self._config.task_names)
          task_name = self._config.task_names[r["task_id"]]
          logits[task_name][r[task_name]["eid"]] = (
              r[task_name]["logits"] if "logits" in r[task_name]
              else r[task_name]["predictions"])
      for task_name in logits:
        utils.log("Pickling predictions for {:} {:} examples ({:})".format(
            len(logits[task_name]), task_name, split))
        if trial <= self._config.n_writes_test:
          utils.write_pickle(logits[task_name], self._config.test_predictions(
              task_name, split, trial))
  
  
  def write_results(config: configure_finetuning.FinetuningConfig, results):
    """Write evaluation metrics to disk."""
    utils.log("Writing results to", config.results_txt)
    utils.mkdir(config.results_txt.rsplit("/", 1)[0])
    utils.write_pickle(results, config.results_pkl)
    with tf.io.gfile.GFile(config.results_txt, "w") as f:
      results_str = ""
      for trial_results in results:
        for task_name, task_results in trial_results.items():
          if task_name == "time" or task_name == "global_step":
            continue
          results_str += task_name + ": " + " - ".join(
              ["{:}: {:.2f}".format(k, v)
               for k, v in task_results.items()]) + "\n"
      f.write(results_str)
    utils.write_pickle(results, config.results_pkl)
  
  
  def run_finetuning(config: configure_finetuning.FinetuningConfig):
    """Run finetuning."""
  
    # Setup for training
    results = []
    trial = 1
    heading_info = "model={:}, trial {:}/{:}".format(
        config.model_name, trial, config.num_trials)
    heading = lambda msg: utils.heading(msg + ": " + heading_info)
    heading("Config")
    utils.log_config(config)
    generic_model_dir = config.model_dir
    tasks = task_builder.get_tasks(config)
  
    # Train and evaluate num_trials models with different random seeds
    while config.num_trials < 0 or trial <= config.num_trials:
      config.model_dir = generic_model_dir + "_" + str(trial)
      if config.do_train:
        utils.rmkdir(config.model_dir)
  
      model_runner = ModelRunner(config, tasks)
      if config.do_train:
        heading("Start training")
        model_runner.train()
        utils.log()
  
      if config.do_eval:
        heading("Run dev set evaluation")
        results.append(model_runner.evaluate())
        write_results(config, results)
        if config.write_test_outputs and trial <= config.n_writes_test:
          heading("Running on the test set and writing the predictions")
          for task in tasks:
            # Currently only writing preds for GLUE and SQuAD 2.0 is supported
            if task.name in ["cola", "mrpc", "mnli", "sst", "rte", "qnli", "qqp",
                             "sts"]:
              for split in task.get_test_splits():
                model_runner.write_classification_outputs([task], trial, split)
            elif task.name == "squad":
              scorer = model_runner.evaluate_task(task, "test", False)
              scorer.write_predictions()
              preds = utils.load_json(config.qa_preds_file("squad"))
              null_odds = utils.load_json(config.qa_na_file("squad"))
              for q, _ in preds.items():
                if null_odds[q] > config.qa_na_threshold:
                  preds[q] = ""
              utils.write_json(preds, config.test_predictions(
                  task.name, "test", trial))
            else:
              utils.log("Skipping task", task.name,
                        "- writing predictions is not supported for this task")
  
      if trial != config.num_trials and (not config.keep_all_models):
        utils.rmrf(config.model_dir)
      trial += 1
  
  
  def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", default='/home/osh7605/KoELECTRAModels/KoCharELECTRA_small_wiki_basic_prepro/models/ll',
                        help="Location of data files (model weights, etc).")
    parser.add_argument("--model-name", default='finetune',
                        help="The name of the model being fine-tuned.")
    parser.add_argument("--hparams", default="{}",
                        help="JSON dict of model hyperparameters.")
    args = parser.parse_args()
    if args.hparams.endswith(".json"):
      hparams = utils.load_json(args.hparams)
    else:
      hparams = json.loads(args.hparams)
    tf.logging.set_verbosity(tf.logging.ERROR)
    run_finetuning(configure_finetuning.FinetuningConfig(
        args.model_name, args.data_dir, **hparams))
  
  
  if __name__ == "__main__":
    main()
  