text: # coding=utf-8
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  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.
  """
  Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
  GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
  using a masked language modeling (MLM) loss.
  """
  
  import logging
  import math
  import os
  import re
  from dataclasses import dataclass, field
  from typing import Optional
  
  import torch
  import torch.nn.functional as F
  from transformers import (
      CONFIG_MAPPING,
      MODEL_WITH_LM_HEAD_MAPPING,
      AutoConfig,
      AutoModelWithLMHead,
      AutoTokenizer,
      DataCollatorForLanguageModeling,
      HfArgumentParser,
      LineByLineTextDataset,
      PreTrainedTokenizer,
      TextDataset,
      Trainer,
      TrainingArguments,
      set_seed,
  )
  
  from models.affect_lm import AffectGPT2LMHeadModel
  from scripts.evaluation.webtext_dataset import WebTextPretokenizedDataset
  
  logger = logging.getLogger(__name__)
  
  MODEL_CONFIG_CLASSES = list(MODEL_WITH_LM_HEAD_MAPPING.keys())
  MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
  
  # Controllable model settings
  AFFECT_LABEL_IDX = 0
  AFFECT_NUM_LABELS = 2
  
  TARGET_CTRL_CODE = '<|nontoxic|>'
  CTRL_CODES = ['<|nontoxic|>', '<|toxic|>']
  assert TARGET_CTRL_CODE in CTRL_CODES
  
  
  @dataclass
  class ModelArguments:
      """
      Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
      """
  
      model_name_or_path: Optional[str] = field(
          default=None,
          metadata={
              "help": "The model checkpoint for weights initialization. Leave None if you want to train a model from scratch."
          },
      )
      model_type: Optional[str] = field(
          default=None,
          metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
      )
      controllable_model: Optional[str] = field(
          default=None,
          metadata={"help": "Use a controllable model for evaluation"},
      )
      config_name: Optional[str] = field(
          default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
      )
      tokenizer_name: Optional[str] = field(
          default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
      )
      cache_dir: Optional[str] = field(
          default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
      )
  
  
  @dataclass
  class DataTrainingArguments:
      """
      Arguments pertaining to what data we are going to input our model for training and eval.
      """
  
      train_data_file: Optional[str] = field(
          default=None, metadata={"help": "The input training data file (a text file)."}
      )
      eval_data_file: Optional[str] = field(
          default=None,
          metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
      )
      line_by_line: bool = field(
          default=False,
          metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
      )
      webtext: bool = field(
          default=False,
          metadata={"help": "Use webtext eval set"},
      )
  
      mlm: bool = field(
          default=False, metadata={"help": "Train with masked-language modeling loss instead of language modeling."}
      )
      mlm_probability: float = field(
          default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
      )
  
      block_size: int = field(
          default=-1,
          metadata={
              "help": "Optional input sequence length after tokenization."
                      "The training dataset will be truncated in block of this size for training."
                      "Default to the model max input length for single sentence inputs (take into account special tokens)."
          },
      )
      overwrite_cache: bool = field(
          default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
      )
  
  
  def get_dataset(args: DataTrainingArguments, tokenizer: PreTrainedTokenizer, inline_meta: str = None, local_rank=-1):
      file_path = args.eval_data_file
      if args.webtext:
          return WebTextPretokenizedDataset(
              tokenizer=tokenizer, file_path=file_path, block_size=args.block_size, inline_meta=inline_meta,
              local_rank=local_rank
          )
      elif args.line_by_line:
          return LineByLineTextDataset(
              tokenizer=tokenizer, file_path=file_path, block_size=args.block_size, local_rank=local_rank
          )
      else:
          return TextDataset(
              tokenizer=tokenizer, file_path=file_path, block_size=args.block_size, local_rank=local_rank,
          )
  
  
  def main():
      # See all possible arguments in src/transformers/training_args.py
      # or by passing the --help flag to this script.
      # We now keep distinct sets of args, for a cleaner separation of concerns.
  
      parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
      model_args, data_args, training_args = parser.parse_args_into_dataclasses()
  
      if training_args.do_train:
          raise ValueError(
              "Training not supported with this script."
          )
  
      if data_args.eval_data_file is None:
          raise ValueError(
              "Cannot do evaluation without an evaluation data file."
          )
  
      # Setup logging
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
      )
      logger.warning(
          "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
          training_args.local_rank,
          training_args.device,
          training_args.n_gpu,
          bool(training_args.local_rank != -1),
          training_args.fp16,
      )
      logger.info("Training/evaluation parameters %s", training_args)
  
      # Set seed
      set_seed(training_args.seed)
  
      # Load pretrained model and tokenizer
      #
      # Distributed training:
      # The .from_pretrained methods guarantee that only one local process can concurrently
      # download model & vocab.
  
      if model_args.config_name:
          config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
      elif model_args.model_name_or_path:
          config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
      else:
          config = CONFIG_MAPPING[model_args.model_type]()
          logger.warning("You are instantiating a new config instance from scratch.")
  
      if model_args.tokenizer_name:
          tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, cache_dir=model_args.cache_dir)
      elif model_args.model_name_or_path:
          tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
      else:
          raise ValueError(
              "You are instantiating a new tokenizer from scratch. This is not supported, but you can do it from another script, save it,"
              "and load it from here, using --tokenizer_name"
          )
  
      inline_meta = None
      if model_args.controllable_model == 'ctrl-gpt2':
          print("Loaded added tokens:", tokenizer.added_tokens_encoder)
          inline_meta = TARGET_CTRL_CODE
          assert inline_meta in tokenizer.get_vocab().keys()
          print("Using inline meta:", inline_meta)
  
      if model_args.model_name_or_path:
          controllable_model = model_args.controllable_model
          if controllable_model and controllable_model.startswith('affect-gpt2'):
              logger.info("Using AffectLM for evaluation")
              model = AffectGPT2LMHeadModel.from_pretrained(
                  model_args.model_name_or_path,
                  from_tf=bool(".ckpt" in model_args.model_name_or_path),
                  config=config,
                  cache_dir=model_args.cache_dir,
              )
  
              # Evaluate with non-toxic affect
              beta = int(re.search(r'beta-([0-9]+)', controllable_model).group(1))
              logging.info(f"Setting AffectLM beta to {beta}")
              model.affect.beta = beta
              affect_label = F.one_hot(torch.LongTensor([AFFECT_LABEL_IDX]),
                                       num_classes=AFFECT_NUM_LABELS).float().to(training_args.device)
              model.set_affect_labels(affect_label)
          else:
              model = AutoModelWithLMHead.from_pretrained(
                  model_args.model_name_or_path,
                  from_tf=bool(".ckpt" in model_args.model_name_or_path),
                  config=config,
                  cache_dir=model_args.cache_dir,
              )
      else:
          logger.info("Training new model from scratch")
          model = AutoModelWithLMHead.from_config(config)
  
      model.resize_token_embeddings(len(tokenizer))
  
      if config.model_type in ["bert", "roberta", "distilbert", "camembert"] and not data_args.mlm:
          raise ValueError(
              "BERT and RoBERTa-like models do not have LM heads but masked LM heads. They must be run using the --mlm "
              "flag (masked language modeling)."
          )
  
      if data_args.block_size <= 0:
          data_args.block_size = tokenizer.max_len
          # Our input block size will be the max possible for the model
      else:
          data_args.block_size = min(data_args.block_size, tokenizer.max_len)
  
      # Get datasets
      eval_dataset = get_dataset(
          data_args, tokenizer=tokenizer, local_rank=training_args.local_rank, inline_meta=inline_meta
      )
      data_collator = DataCollatorForLanguageModeling(
          tokenizer=tokenizer, mlm=data_args.mlm, mlm_probability=data_args.mlm_probability
      )
  
      # Initialize our Trainer
      trainer = Trainer(
          model=model,
          args=training_args,
          data_collator=data_collator,
          eval_dataset=eval_dataset,
          prediction_loss_only=True,
      )
  
      # Evaluation
      results = {}
      if training_args.local_rank in [-1, 0]:
          logger.info("*** Evaluate ***")
  
          eval_output = trainer.evaluate()
  
          perplexity = math.exp(eval_output["loss"])
          result = {"perplexity": perplexity}
  
          output_eval_file = os.path.join(training_args.output_dir, "eval_results_lm.txt")
          with open(output_eval_file, "w") as writer:
              logger.info("***** Eval results *****")
              for key in sorted(result.keys()):
                  logger.info("  %s = %s", key, str(result[key]))
                  writer.write("%s = %s\n" % (key, str(result[key])))
  
          results.update(result)
  
      return results
  
  
  if __name__ == "__main__":
      main()
  