




from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers.models.ctrl.modeling_tf_ctrl import (
         TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFCTRLForSequenceClassification,
         TFCTRLLMHeadModel,
         TFCTRLModel,
     )




class  __lowerCAmelCase							(				lowerCAmelCase):






       def __init__( self:       Optional[Any]							,       _lowerCAmelCase:       Optional[Any]							,       ):
         lowercase     :List[str]   						=    parent
         lowercase     :Dict   						=    13
         lowercase     :int   						=    7
         lowercase     :Tuple   						=    True
         lowercase     :Dict   						=    True
         lowercase     :List[Any]   						=    True
         lowercase     :List[str]   						=    True
         lowercase     :Dict   						=    True
         lowercase     :int   						=    99
         lowercase     :int   						=    32
         lowercase     :str   						=    2
         lowercase     :List[Any]   						=    4
         lowercase     :List[Any]   						=    37
         lowercase     :List[Any]   						=    "gelu"
         lowercase     :Dict   						=    0.1
         lowercase     :List[str]   						=    0.1
         lowercase     :int   						=    5_12
         lowercase     :List[Any]   						=    16
         lowercase     :Tuple   						=    2
         lowercase     :Optional[int]   						=    0.02
         lowercase     :Union[str, Any]   						=    3
         lowercase     :int   						=    4
         lowercase     :Optional[int]   						=    None
         lowercase     :Tuple   						=    self.vocab_size - 1






       def        SCREAMING_SNAKE_CASE     ( self:       List[str]					):
         lowercase     :str   						=    ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size					)

         lowercase     :int   						=    None
         if self.use_input_mask:
           lowercase     :int   						=    random_attention_mask([self.batch_size, self.seq_length]					)

         lowercase     :List[str]   						=    None
         if self.use_token_type_ids:
           lowercase     :Dict   						=    ids_tensor([self.batch_size, self.seq_length]							,       self.type_vocab_size					)

         lowercase     :Optional[Any]   						=    None
         if self.use_mc_token_ids:
           lowercase     :Any   						=    ids_tensor([self.batch_size, self.num_choices]							,       self.seq_length					)

         lowercase     :str   						=    None
         lowercase     :List[str]   						=    None
         lowercase     :Union[str, Any]   						=    None
         if self.use_labels:
           lowercase     :int   						=    ids_tensor([self.batch_size]							,       self.type_sequence_label_size					)
           lowercase     :Tuple   						=    ids_tensor([self.batch_size, self.seq_length]							,       self.num_labels					)
           lowercase     :Union[str, Any]   						=    ids_tensor([self.batch_size]							,       self.num_choices					)

         lowercase     :Optional[Any]   						=    CTRLConfig(
             vocab_size=self.vocab_size							,       n_embd=self.hidden_size							,       n_layer=self.num_hidden_layers							,       n_head=self.num_attention_heads							,       n_positions=self.max_position_embeddings							,       pad_token_id=self.pad_token_id							,       )

         lowercase     :str   						=    ids_tensor([self.num_hidden_layers, self.num_attention_heads]							,       2					)

         return (
             config,
             input_ids,
             input_mask,
             head_mask,
             token_type_ids,
             mc_token_ids,
             sequence_labels,
             token_labels,
             choice_labels,
         )






       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       List[str]							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       str							,       *_lowerCAmelCase:       Dict					):
         lowercase     :int   						=    TFCTRLModel(config=_lowerCAmelCase					)
         lowercase     :List[Any]   						=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
         lowercase     :Tuple   						=    model(_lowerCAmelCase					)

         lowercase     :Union[str, Any]   						=    [input_ids, None, input_mask]  # None is the input for 'past'
         lowercase     :Tuple   						=    model(_lowerCAmelCase					)

         lowercase     :List[Any]   						=    model(_lowerCAmelCase					)

         self.parent.assertEqual(result.last_hidden_state.shape							,       (self.batch_size, self.seq_length, self.hidden_size)					)






       def        SCREAMING_SNAKE_CASE     ( self:       Tuple							,       _lowerCAmelCase:       Optional[int]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Optional[Any]							,       *_lowerCAmelCase:       List[Any]					):
         lowercase     :str   						=    TFCTRLLMHeadModel(config=_lowerCAmelCase					)
         lowercase     :Optional[Any]   						=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
         lowercase     :int   						=    model(_lowerCAmelCase					)
         self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.seq_length, self.vocab_size)					)






       def        SCREAMING_SNAKE_CASE     ( self:       Dict							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Any							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Any							,       _lowerCAmelCase:       Union[str, Any]							,       *_lowerCAmelCase:       Any					):
         lowercase     :Optional[int]   						=    self.num_labels
         lowercase     :List[Any]   						=    ids_tensor([self.batch_size]							,       self.type_sequence_label_size					)
         lowercase     :Tuple   						=    {
             "input_ids": input_ids,
             "token_type_ids": token_type_ids,
             "labels": sequence_labels,
         }
         lowercase     :Optional[int]   						=    TFCTRLForSequenceClassification(_lowerCAmelCase					)
         lowercase     :Optional[int]   						=    model(_lowerCAmelCase					)
         self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.num_labels)					)






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]					):
         lowercase     :List[Any]   						=    self.prepare_config_and_inputs()

         (
             (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     
         )     :str   						=    config_and_inputs

         lowercase     :Optional[int]   						=    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict







@require_tf
class  __lowerCAmelCase							(				lowerCAmelCase   ,			lowerCAmelCase   ,			unittest.TestCase):
       _a			       =    (TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
       _a			       =    (TFCTRLLMHeadModel,) if is_tf_available() else ()
       _a			       =    (
           {
               '''feature-extraction''': TFCTRLModel,
               '''text-classification''': TFCTRLForSequenceClassification,
               '''text-generation''': TFCTRLLMHeadModel,
               '''zero-shot''': TFCTRLForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       _a			       =    False
       _a			       =    False






       def        SCREAMING_SNAKE_CASE     ( self:       Dict							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       Any					):
         if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
           # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
           # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
           # config could not be created.
           return True

         return False






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]					):
         lowercase     :int   						=    TFCTRLModelTester(self					)
         lowercase     :List[str]   						=    ConfigTester(self							,       config_class=_lowerCAmelCase							,       n_embd=37					)






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[Any]					):
         self.config_tester.run_common_tests()






       def        SCREAMING_SNAKE_CASE     ( self:       Tuple					):
         lowercase     :Any   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_ctrl_model(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]					):
         lowercase     :Dict   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_ctrl_lm_head(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Any					):
         lowercase     :Dict   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_ctrl_for_sequence_classification(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[Any]					):
         lowercase						,     lowercase     :str   						=    self.model_tester.prepare_config_and_inputs_for_common()
         lowercase     :str   						=    [TFCTRLLMHeadModel]
         lowercase     :Any   						=    [TFCTRLForSequenceClassification]

         for model_class in self.all_model_classes:
           lowercase     :List[Any]   						=    model_class(_lowerCAmelCase					)
           model.build()  # may be needed for the get_bias() call below
           assert isinstance(model.get_input_embeddings()							,       tf.keras.layers.Layer					)

           if model_class in list_lm_models:
             lowercase     :Tuple   						=    model.get_output_embeddings()
             assert isinstance(_lowerCAmelCase							,       tf.keras.layers.Layer					)
             lowercase     :int   						=    model.get_bias()
             assert isinstance(_lowerCAmelCase							,       _lowerCAmelCase					)
             for k, v in name.items():
               assert isinstance(_lowerCAmelCase							,       tf.Variable					)
           elif model_class in list_other_models_with_output_ebd:
             lowercase     :Any   						=    model.get_output_embeddings()
             assert isinstance(_lowerCAmelCase							,       tf.keras.layers.Layer					)
             lowercase     :str   						=    model.get_bias()
             assert name is None
           else:
             lowercase     :Tuple   						=    model.get_output_embeddings()
             assert x is None
             lowercase     :Tuple   						=    model.get_bias()
             assert name is None






       @slow
       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]					):
         for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase     :Optional[Any]   						=    TFCTRLModel.from_pretrained(_lowerCAmelCase					)
           self.assertIsNotNone(_lowerCAmelCase					)







@require_tf
class  __lowerCAmelCase							(				unittest.TestCase):






       @slow
       def        SCREAMING_SNAKE_CASE     ( self:       int					):
         lowercase     :List[Any]   						=    TFCTRLLMHeadModel.from_pretrained("ctrl"					)
         lowercase     :str   						=    tf.convert_to_tensor([[1_18_59, 0, 16_11, 8]]							,       dtype=tf.intaa					)  # Legal the president is
         lowercase     :int   						=    [
             1_18_59,
             0,
             16_11,
             8,
             5,
             1_50,
             2_64_49,
             2,
             19,
             3_48,
             4_69,
             3,
             25_95,
             48,
             2_07_40,
             24_65_33,
             24_65_33,
             19,
             30,
             5,
         ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

         lowercase     :Optional[Any]   						=    model.generate(_lowerCAmelCase							,       do_sample=_lowerCAmelCase					)
         self.assertListEqual(output_ids[0].numpy().tolist()							,       _lowerCAmelCase					)









import argparse
import logging
import math
import time

import torch

from transformers import TransfoXLCorpus, TransfoXLLMHeadModel


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO
)
_UpperCAmelCase						:    Dict      =  logging.getLogger(__name__)



def    UpperCAmelCase__	(  ):
  lowercase     :int   						=    argparse.ArgumentParser(description="PyTorch Transformer Language Model" )
  parser.add_argument("--model_name",      type=lowerCamelCase,      default="transfo-xl-wt103",      help="pretrained model name" )
  parser.add_argument(
      "--split",      type=lowerCamelCase,      default="test",      choices=["all", "valid", "test"],      help="which split to evaluate" )
  parser.add_argument("--batch_size",      type=lowerCamelCase,      default=10,      help="batch size" )
  parser.add_argument("--tgt_len",      type=lowerCamelCase,      default=128,      help="number of tokens to predict" )
  parser.add_argument("--ext_len",      type=lowerCamelCase,      default=0,      help="length of the extended context" )
  parser.add_argument("--mem_len",      type=lowerCamelCase,      default=1600,      help="length of the retained previous heads" )
  parser.add_argument("--clamp_len",      type=lowerCamelCase,      default=1000,      help="max positional embedding index" )
  parser.add_argument("--no_cuda",      action="store_true",      help="Do not use CUDA even though CUA is available" )
  parser.add_argument("--work_dir",      type=lowerCamelCase,      required=lowerCamelCase,      help="path to the work_dir" )
  parser.add_argument("--no_log",      action="store_true",      help="do not log the eval result" )
  parser.add_argument("--same_length",      action="store_true",      help="set same length attention with masking" )
  parser.add_argument("--server_ip",      type=lowerCamelCase,      default="",      help="Can be used for distant debugging." )
  parser.add_argument("--server_port",      type=lowerCamelCase,      default="",      help="Can be used for distant debugging." )
  lowercase     :Any   						=    parser.parse_args()
  assert args.ext_len >= 0, "extended context length must be non-negative"

  if args.server_ip and args.server_port:
    # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
    import ptvsd

    print("Waiting for debugger attach" )
    ptvsd.enable_attach(address=(args.server_ip, args.server_port),      redirect_output=lowerCamelCase )
    ptvsd.wait_for_attach()

  lowercase     :Any   						=    torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu" )
  logger.info("device: {}".format(lowerCamelCase ) )

  # Load a pre-processed dataset
  # You can also build the corpus yourself using TransfoXLCorpus methods
  # The pre-processing involve computing word frequencies to prepare the Adaptive input and SoftMax
  # and tokenizing the dataset
  # The pre-processed corpus is a convertion (using the conversion script )
  lowercase     :Dict   						=    TransfoXLCorpus.from_pretrained(args.model_name )

  lowercase     :Optional[int]   						=    corpus.get_iterator("valid",      args.batch_size,      args.tgt_len,      device=lowerCamelCase,      ext_len=args.ext_len )
  lowercase     :Union[str, Any]   						=    corpus.get_iterator("test",      args.batch_size,      args.tgt_len,      device=lowerCamelCase,      ext_len=args.ext_len )

  # Load a pre-trained model
  lowercase     :List[Any]   						=    TransfoXLLMHeadModel.from_pretrained(args.model_name )
  model.to(lowerCamelCase )

  logger.info(
      "Evaluating with bsz {} tgt_len {} ext_len {} mem_len {} clamp_len {}".format(
          args.batch_size,      args.tgt_len,      args.ext_len,      args.mem_len,      args.clamp_len ) )

  model.reset_memory_length(args.mem_len )
  if args.clamp_len > 0:
    lowercase     :Tuple   						=    args.clamp_len
  if args.same_length:
    lowercase     :Tuple   						=    True

  ###############################################################################
  # Evaluation code
  ###############################################################################
  def evaluate(lowerCamelCase ):
    # Turn on evaluation mode which disables dropout.
    model.eval()
    lowercase						,     lowercase     :Optional[Any]   						=    0, 0.0
    lowercase     :Optional[Any]   						=    time.time()
    with torch.no_grad():
      lowercase     :Dict   						=    None
      for idx, (data, target, seq_len) in enumerate(lowerCamelCase ):
        lowercase     :Optional[Any]   						=    model(lowerCamelCase,      lm_labels=lowerCamelCase,      mems=lowerCamelCase )
        lowercase						,     lowercase						,     lowercase     :Optional[int]   						=    ret
        lowercase     :Optional[Any]   						=    loss.mean()
        total_loss += seq_len * loss.item()
        total_len += seq_len
      lowercase     :List[Any]   						=    time.time() - start_time
    logger.info("Time : {:.2f}s, {:.2f}ms/segment".format(lowerCamelCase,      1000 * total_time / (idx + 1) ) )
    return total_loss / total_len

  # Run on test data.
  if args.split == "all":
    lowercase     :Union[str, Any]   						=    evaluate(lowerCamelCase )
    lowercase     :List[str]   						=    evaluate(lowerCamelCase )
  elif args.split == "valid":
    lowercase     :Optional[int]   						=    evaluate(lowerCamelCase )
    lowercase     :Dict   						=    None
  elif args.split == "test":
    lowercase     :int   						=    evaluate(lowerCamelCase )
    lowercase     :Optional[Any]   						=    None

  def format_log(lowerCamelCase,      lowerCamelCase ):
    lowercase     :List[str]   						=    "| {0} loss {1:5.2f} | {0} ppl {2:9.3f} ".format(lowerCamelCase,      lowerCamelCase,      math.exp(lowerCamelCase ) )
    return log_str

  lowercase     :Optional[int]   						=    ""
  if valid_loss is not None:
    log_str += format_log(lowerCamelCase,      "valid" )
  if test_loss is not None:
    log_str += format_log(lowerCamelCase,      "test" )

  logger.info("=" * 100 )
  logger.info(lowerCamelCase )
  logger.info("=" * 100 )


if __name__ == "__main__":
     main()



