





'''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 				__snake_case  (     __SCREAMING_SNAKE_CASE						):


    '''simple docstring'''







    def __init__(					self      ,							A_      ,							):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= parent
        SCREAMING_SNAKE_CASE__       							= 13
        SCREAMING_SNAKE_CASE__       							= 7
        SCREAMING_SNAKE_CASE__       							= True
        SCREAMING_SNAKE_CASE__       							= True
        SCREAMING_SNAKE_CASE__       							= True
        SCREAMING_SNAKE_CASE__       							= True
        SCREAMING_SNAKE_CASE__       							= True
        SCREAMING_SNAKE_CASE__       							= 99
        SCREAMING_SNAKE_CASE__       							= 32
        SCREAMING_SNAKE_CASE__       							= 2
        SCREAMING_SNAKE_CASE__       							= 4
        SCREAMING_SNAKE_CASE__       							= 37
        SCREAMING_SNAKE_CASE__       							= '''gelu'''
        SCREAMING_SNAKE_CASE__       							= 0.1
        SCREAMING_SNAKE_CASE__       							= 0.1
        SCREAMING_SNAKE_CASE__       							= 5_12
        SCREAMING_SNAKE_CASE__       							= 16
        SCREAMING_SNAKE_CASE__       							= 2
        SCREAMING_SNAKE_CASE__       							= 0.02
        SCREAMING_SNAKE_CASE__       							= 3
        SCREAMING_SNAKE_CASE__       							= 4
        SCREAMING_SNAKE_CASE__       							= None
        SCREAMING_SNAKE_CASE__       							= self.vocab_size - 1



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size, self.seq_length]      ,							self.vocab_size    )

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

        SCREAMING_SNAKE_CASE__       							= None
        if self.use_token_type_ids:
            SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size, self.seq_length]      ,							self.type_vocab_size    )

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

        SCREAMING_SNAKE_CASE__       							= None
        SCREAMING_SNAKE_CASE__       							= None
        SCREAMING_SNAKE_CASE__       							= None
        if self.use_labels:
            SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size]      ,							self.type_sequence_label_size    )
            SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size, self.seq_length]      ,							self.num_labels    )
            SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size]      ,							self.num_choices    )

        SCREAMING_SNAKE_CASE__       							= 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      ,							)

        SCREAMING_SNAKE_CASE__       							= 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    lowercase_				(					self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							*A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFCTRLModel(config=A_    )
        SCREAMING_SNAKE_CASE__       							= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        SCREAMING_SNAKE_CASE__       							= model(A_    )

        SCREAMING_SNAKE_CASE__       							= [input_ids, None, input_mask]  # None is the input for 'past'
        SCREAMING_SNAKE_CASE__       							= model(A_    )

        SCREAMING_SNAKE_CASE__       							= model(A_    )

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



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFCTRLLMHeadModel(config=A_    )
        SCREAMING_SNAKE_CASE__       							= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        SCREAMING_SNAKE_CASE__       							= model(A_    )
        self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.seq_length, self.vocab_size)    )



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.num_labels
        SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size]      ,							self.type_sequence_label_size    )
        SCREAMING_SNAKE_CASE__       							= {
            '''input_ids''': input_ids,
            '''token_type_ids''': token_type_ids,
            '''labels''': sequence_labels,
        }
        SCREAMING_SNAKE_CASE__       							= TFCTRLForSequenceClassification(A_    )
        SCREAMING_SNAKE_CASE__       							= model(A_    )
        self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_labels)    )




    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.prepare_config_and_inputs()

        (
            (
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,		(
            SCREAMING_SNAKE_CASE__		
        )							,				
        )       							= config_and_inputs

        SCREAMING_SNAKE_CASE__       							= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict





@require_tf
class 				__snake_case  (     __SCREAMING_SNAKE_CASE ,    __SCREAMING_SNAKE_CASE ,    unittest.TestCase						):


    '''simple docstring'''







    lowerCamelCase__						:		Tuple							=  (TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
    lowerCamelCase__						:		Optional[int]							=  (TFCTRLLMHeadModel,) if is_tf_available() else ()
    lowerCamelCase__						:		Dict							=  (
        {
            """feature-extraction""": TFCTRLModel,
            """text-classification""": TFCTRLForSequenceClassification,
            """text-generation""": TFCTRLLMHeadModel,
            """zero-shot""": TFCTRLForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    lowerCamelCase__						:		Optional[int]							=  False
    lowerCamelCase__						:		int							=  False



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_    ):





        '''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    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFCTRLModelTester(self    )
        SCREAMING_SNAKE_CASE__       							= ConfigTester(self      ,							config_class=A_      ,							n_embd=37    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_model(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_lm_head(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_for_sequence_classification(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs_for_common()
        SCREAMING_SNAKE_CASE__       							= [TFCTRLLMHeadModel]
        SCREAMING_SNAKE_CASE__       							= [TFCTRLForSequenceClassification]

        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__       							= model_class(A_    )
            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:
                SCREAMING_SNAKE_CASE__       							= model.get_output_embeddings()
                assert isinstance(A_      ,							tf.keras.layers.Layer    )
                SCREAMING_SNAKE_CASE__       							= model.get_bias()
                assert isinstance(A_      ,							A_    )
                for k, v in name.items():
                    assert isinstance(A_      ,							tf.Variable    )
            elif model_class in list_other_models_with_output_ebd:
                SCREAMING_SNAKE_CASE__       							= model.get_output_embeddings()
                assert isinstance(A_      ,							tf.keras.layers.Layer    )
                SCREAMING_SNAKE_CASE__       							= model.get_bias()
                assert name is None
            else:
                SCREAMING_SNAKE_CASE__       							= model.get_output_embeddings()
                assert x is None
                SCREAMING_SNAKE_CASE__       							= model.get_bias()
                assert name is None




    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

        for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            SCREAMING_SNAKE_CASE__       							= TFCTRLModel.from_pretrained(A_    )
            self.assertIsNotNone(A_    )





@require_tf
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFCTRLLMHeadModel.from_pretrained('''ctrl'''    )
        SCREAMING_SNAKE_CASE__       							= tf.convert_to_tensor([[1_18_59, 0, 16_11, 8]]      ,							dtype=tf.intaa    )  # Legal the president is
        SCREAMING_SNAKE_CASE__       							= [
            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

        SCREAMING_SNAKE_CASE__       							= model.generate(A_      ,							do_sample=A_    )
        self.assertListEqual(output_ids[0].numpy().tolist()      ,							A_    )







import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


_A				:							Any							= get_logger(__name__)


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	str:
    if text.startswith(lowerCAmelCase_  ):
        return text[len(lowerCAmelCase_  ) :]
    return text  # or whatever


def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    return {remove_prefix(lowerCAmelCase_   ,       '''model.'''  ): v for k, v in sd.items()}


def        __snake_case       (						lowerCAmelCase_  )							->	Optional[int]:
    SCREAMING_SNAKE_CASE__       							= {}
    for k in state_dicts[0].keys():
        SCREAMING_SNAKE_CASE__       							= [sd[k] for sd in state_dicts]
        SCREAMING_SNAKE_CASE__       							= sum(lowerCAmelCase_  ) / len(lowerCAmelCase_  )
        assert isinstance(lowerCAmelCase_   ,       torch.Tensor  )
        SCREAMING_SNAKE_CASE__       							= new_t
    return new_sd


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	None:
    SCREAMING_SNAKE_CASE__       							= AutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase_  )
    if os.path.isfile(lowerCAmelCase_  ):
        SCREAMING_SNAKE_CASE__       							= [pl_ckpt_path]
    else:
        assert os.path.isdir(lowerCAmelCase_  )
        SCREAMING_SNAKE_CASE__       							= list(Path(lowerCAmelCase_  ).glob('''*.ckpt'''  )  )
        assert ckpt_files, f'''could not find any ckpt files inside the {pl_ckpt_path} directory'''

    if len(lowerCAmelCase_  ) > 1:
        logger.info(f'''averaging the weights of {ckpt_files}'''  )

    SCREAMING_SNAKE_CASE__       							= [sanitize(torch.load(lowerCAmelCase_   ,       map_location='''cpu'''  )['''state_dict''']  ) for x in ckpt_files]
    SCREAMING_SNAKE_CASE__       							= average_state_dicts(lowerCAmelCase_  )

    SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= hf_model.load_state_dict(lowerCAmelCase_   ,       strict=lowerCAmelCase_  )
    assert not missing, f'''missing keys: {missing}'''
    hf_model.save_pretrained(lowerCAmelCase_  )
    try:
        SCREAMING_SNAKE_CASE__       							= AutoTokenizer.from_pretrained(lowerCAmelCase_  )
        tok.save_pretrained(lowerCAmelCase_  )
    except Exception:
        pass
        # dont copy tokenizer if cant


if __name__ == "__main__":
  fire.Fire(convert_pl_to_hf)
