
"""simple docstring"""





import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


_a     :   Any	 =							get_logger(__name__)







def 			__a  (	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase = None,	):
   # Model states
   for i, state in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      Union[str, Any]      =					f"""{MODEL_NAME}.bin""" if i == 0 else f"""{MODEL_NAME}_{i}.bin"""
      UpperCAmelCase_      :      Optional[Any]      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      save(__lowerCamelCase,	__lowerCamelCase						)
      logger.info(f"""Model weights saved in {output_model_file}"""						)
   # Optimizer states
   for i, opt in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      Optional[int]      =					opt.state_dict()
      UpperCAmelCase_      :      Any      =					f"""{OPTIMIZER_NAME}.bin""" if i == 0 else f"""{OPTIMIZER_NAME}_{i}.bin"""
      UpperCAmelCase_      :      List[Any]      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      save(__lowerCamelCase,	__lowerCamelCase						)
      logger.info(f"""Optimizer state saved in {output_optimizer_file}"""						)
   # Scheduler states
   for i, scheduler in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      Optional[int]      =					scheduler.state_dict()
      UpperCAmelCase_      :      Tuple      =					f"""{SCHEDULER_NAME}.bin""" if i == 0 else f"""{SCHEDULER_NAME}_{i}.bin"""
      UpperCAmelCase_      :      Optional[Any]      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      save(__lowerCamelCase,	__lowerCamelCase						)
      logger.info(f"""Scheduler state saved in {output_scheduler_file}"""						)
   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_      :      Tuple      =					scaler.state_dict()
      UpperCAmelCase_      :      Dict      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      torch.save(__lowerCamelCase,	__lowerCamelCase						)
      logger.info(f"""Gradient scaler state saved in {output_scaler_file}"""						)
   # Random number generator states
   UpperCAmelCase_      :      List[Any]      =					{}
   UpperCAmelCase_      :      List[str]      =					f"""{RNG_STATE_NAME}_{process_index}.pkl"""
   UpperCAmelCase_      :      Any      =					random.getstate()
   UpperCAmelCase_      :      Optional[Any]      =					np.random.get_state()
   UpperCAmelCase_      :      Dict      =					torch.get_rng_state()
   if is_xpu_available():
      UpperCAmelCase_      :      Optional[int]      =					torch.xpu.get_rng_state_all()
   else:
      UpperCAmelCase_      :      int      =					torch.cuda.get_rng_state_all()
   if is_tpu_available():
      UpperCAmelCase_      :      List[str]      =					xm.get_rng_state()
   UpperCAmelCase_      :      Union[str, Any]      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
   torch.save(__lowerCamelCase,	__lowerCamelCase						)
   logger.info(f"""Random states saved in {output_states_file}"""						)
   return output_dir







def 			__a  (	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase=None,	__lowerCamelCase=None,	**__lowerCamelCase,	):
   if map_location not in [None, "cpu", "on_device"]:
      raise TypeError(
          "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"						)
   if map_location is None:
      UpperCAmelCase_      :      int      =					"cpu"
   elif map_location == "on_device":
      UpperCAmelCase_      :      List[Any]      =					PartialState().device
   # Model states
   for i, model in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      Optional[Any]      =					f"""{MODEL_NAME}.bin""" if i == 0 else f"""{MODEL_NAME}_{i}.bin"""
      UpperCAmelCase_      :      Any      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      models[i].load_state_dict(torch.load(__lowerCamelCase,	map_location=__lowerCamelCase						),	**__lowerCamelCase						)
   logger.info("All model weights loaded successfully"						)

   # Optimizer states
   for i, opt in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      Tuple      =					f"""{OPTIMIZER_NAME}.bin""" if i == 0 else f"""{OPTIMIZER_NAME}_{i}.bin"""
      UpperCAmelCase_      :      Tuple      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      UpperCAmelCase_      :      Union[str, Any]      =					torch.load(__lowerCamelCase,	map_location=__lowerCamelCase						)
      optimizers[i].load_state_dict(__lowerCamelCase						)
   logger.info("All optimizer states loaded successfully"						)

   # Scheduler states
   for i, scheduler in enumerate(__lowerCamelCase						):
      UpperCAmelCase_      :      List[str]      =					f"""{SCHEDULER_NAME}.bin""" if i == 0 else f"""{SCHEDULER_NAME}_{i}.bin"""
      UpperCAmelCase_      :      str      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      scheduler.load_state_dict(torch.load(__lowerCamelCase						)						)
   logger.info("All scheduler states loaded successfully"						)

   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_      :      str      =					os.path.join(__lowerCamelCase,	__lowerCamelCase						)
      scaler.load_state_dict(torch.load(__lowerCamelCase						)						)
      logger.info("GradScaler state loaded successfully"						)

   # Random states
   try:
      UpperCAmelCase_      :      Tuple      =					torch.load(os.path.join(__lowerCamelCase,	f"""{RNG_STATE_NAME}_{process_index}.pkl"""						)						)
      random.setstate(states["random_state"]						)
      np.random.set_state(states["numpy_random_seed"]						)
      torch.set_rng_state(states["torch_manual_seed"]						)
      if is_xpu_available():
         torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"]						)
      else:
         torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"]						)
      if is_tpu_available():
         xm.set_rng_state(states["xm_seed"]						)
      logger.info("All random states loaded successfully"						)
   except Exception:
      logger.info("Could not load random states"						)







def 			__a  (	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase = 0						):
   UpperCAmelCase_      :      Dict      =					Path(__lowerCamelCase						) / f"""custom_checkpoint_{index}.pkl"""
   logger.info(f"""Saving the state of {get_pretty_name(__lowerCamelCase						)} to {save_location}"""						)
   torch.save(obj.state_dict(),	__lowerCamelCase						)







def 			__a  (	__lowerCamelCase,	__lowerCamelCase,	__lowerCamelCase = 0						):
   UpperCAmelCase_      :      Any      =					f"""{path}/custom_checkpoint_{index}.pkl"""
   logger.info(f"""Loading the state of {get_pretty_name(__lowerCamelCase						)} from {load_location}"""						)
   obj.load_state_dict(torch.load(__lowerCamelCase,	map_location="cpu"						)						)




"""simple docstring"""





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 A_			(lowercase__   ):




  '''simple docstring'''







  def __init__(	self       ,     lowercase_       ,     ):

     """simple docstring"""
     UpperCAmelCase_      :      List[str]      =					parent
     UpperCAmelCase_      :      List[Any]      =					13
     UpperCAmelCase_      :      Dict      =					7
     UpperCAmelCase_      :      Optional[int]      =					True
     UpperCAmelCase_      :      List[Any]      =					True
     UpperCAmelCase_      :      Optional[int]      =					True
     UpperCAmelCase_      :      int      =					True
     UpperCAmelCase_      :      Any      =					True
     UpperCAmelCase_      :      Any      =					99
     UpperCAmelCase_      :      Dict      =					32
     UpperCAmelCase_      :      Dict      =					2
     UpperCAmelCase_      :      str      =					4
     UpperCAmelCase_      :      int      =					37
     UpperCAmelCase_      :      Union[str, Any]      =					"gelu"
     UpperCAmelCase_      :      List[str]      =					0.1
     UpperCAmelCase_      :      List[str]      =					0.1
     UpperCAmelCase_      :      int      =					512
     UpperCAmelCase_      :      str      =					16
     UpperCAmelCase_      :      Union[str, Any]      =					2
     UpperCAmelCase_      :      Dict      =					0.02
     UpperCAmelCase_      :      Tuple      =					3
     UpperCAmelCase_      :      Tuple      =					4
     UpperCAmelCase_      :      int      =					None
     UpperCAmelCase_      :      Dict      =					self.vocab_size - 1







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      Dict      =					ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size						)

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

     UpperCAmelCase_      :      List[Any]      =					None
     if self.use_token_type_ids:
        UpperCAmelCase_      :      Optional[int]      =					ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size						)

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

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

     UpperCAmelCase_      :      str      =					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       ,     )

     UpperCAmelCase_      :      Tuple      =					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  UpperCamelCase__      (	self       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     *lowercase_						):

     """simple docstring"""
     UpperCAmelCase_      :      int      =					TFCTRLModel(config=lowercase_						)
     UpperCAmelCase_      :      Optional[Any]      =					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
     UpperCAmelCase_      :      Union[str, Any]      =					model(lowercase_						)

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

     UpperCAmelCase_      :      str      =					model(lowercase_						)

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







  def  UpperCamelCase__      (	self       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     *lowercase_						):

     """simple docstring"""
     UpperCAmelCase_      :      Optional[int]      =					TFCTRLLMHeadModel(config=lowercase_						)
     UpperCAmelCase_      :      List[Any]      =					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
     UpperCAmelCase_      :      List[str]      =					model(lowercase_						)
     self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size)						)







  def  UpperCamelCase__      (	self       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     *lowercase_						):

     """simple docstring"""
     UpperCAmelCase_      :      List[str]      =					self.num_labels
     UpperCAmelCase_      :      str      =					ids_tensor([self.batch_size]       ,     self.type_sequence_label_size						)
     UpperCAmelCase_      :      Union[str, Any]      =					{
         "input_ids": input_ids,
         "token_type_ids": token_type_ids,
         "labels": sequence_labels,
     }
     UpperCAmelCase_      :      Union[str, Any]      =					TFCTRLForSequenceClassification(lowercase_						)
     UpperCAmelCase_      :      Tuple      =					model(lowercase_						)
     self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels)						)







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      Dict      =					self.prepare_config_and_inputs()

     (
         (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     (
         UpperCAmelCase_
     )  ,     
     )      :      Tuple      =					config_and_inputs

     UpperCAmelCase_      :      int      =					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict




@require_tf
class A_			(lowercase__       ,lowercase__       ,unittest.TestCase   ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__    :		Dict						  =	(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
  SCREAMING_SNAKE_CASE__    :		str						  =	(TFCTRLLMHeadModel,) if is_tf_available() else ()
  SCREAMING_SNAKE_CASE__    :		Any						  =	(
      {
          """feature-extraction""": TFCTRLModel,
          """text-classification""": TFCTRLForSequenceClassification,
          """text-generation""": TFCTRLLMHeadModel,
          """zero-shot""": TFCTRLForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__    :		Dict						  =	False
  SCREAMING_SNAKE_CASE__    :		List[Any]						  =	False







  def  UpperCamelCase__      (	self       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_       ,     lowercase_						):

     """simple docstring"""
     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  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      str      =					TFCTRLModelTester(self						)
     UpperCAmelCase_      :      Dict      =					ConfigTester(self       ,     config_class=lowercase_       ,     n_embd=37						)







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     self.config_tester.run_common_tests()







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      str      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_ctrl_model(*lowercase_						)







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      Union[str, Any]      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_ctrl_lm_head(*lowercase_						)







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      Optional[Any]      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_ctrl_for_sequence_classification(*lowercase_						)







  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_  ,     UpperCAmelCase_      :      List[Any]      =					self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_      :      Tuple      =					[TFCTRLLMHeadModel]
     UpperCAmelCase_      :      Dict      =					[TFCTRLForSequenceClassification]

     for model_class in self.all_model_classes:
        UpperCAmelCase_      :      Optional[int]      =					model_class(lowercase_						)
        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:
           UpperCAmelCase_      :      List[str]      =					model.get_output_embeddings()
           assert isinstance(lowercase_       ,     tf.keras.layers.Layer						)
           UpperCAmelCase_      :      List[Any]      =					model.get_bias()
           assert isinstance(lowercase_       ,     lowercase_						)
           for k, v in name.items():
              assert isinstance(lowercase_       ,     tf.Variable						)
        elif model_class in list_other_models_with_output_ebd:
           UpperCAmelCase_      :      int      =					model.get_output_embeddings()
           assert isinstance(lowercase_       ,     tf.keras.layers.Layer						)
           UpperCAmelCase_      :      Any      =					model.get_bias()
           assert name is None
        else:
           UpperCAmelCase_      :      Dict      =					model.get_output_embeddings()
           assert x is None
           UpperCAmelCase_      :      Tuple      =					model.get_bias()
           assert name is None







  @slow
  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase_      :      List[Any]      =					TFCTRLModel.from_pretrained(lowercase_						)
        self.assertIsNotNone(lowercase_						)




@require_tf
class A_			(unittest.TestCase   ):




  '''simple docstring'''







  @slow
  def  UpperCamelCase__      (	self						):

     """simple docstring"""
     UpperCAmelCase_      :      Union[str, Any]      =					TFCTRLLMHeadModel.from_pretrained("ctrl"						)
     UpperCAmelCase_      :      Union[str, Any]      =					tf.convert_to_tensor([[1_1859, 0, 1611, 8]]       ,     dtype=tf.intaa						)  # Legal the president is
     UpperCAmelCase_      :      int      =					[
         1_1859,
         0,
         1611,
         8,
         5,
         150,
         2_6449,
         2,
         19,
         348,
         469,
         3,
         2595,
         48,
         2_0740,
         24_6533,
         24_6533,
         19,
         30,
         5,
     ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

     UpperCAmelCase_      :      List[str]      =					model.generate(lowercase_       ,     do_sample=lowercase_						)
     self.assertListEqual(output_ids[0].numpy().tolist()       ,     lowercase_						)


