





from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, 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.openai.modeling_tf_openai import (
      TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFOpenAIGPTDoubleHeadsModel,
      TFOpenAIGPTForSequenceClassification,
      TFOpenAIGPTLMHeadModel,
      TFOpenAIGPTModel,
  )



class 				__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__       							= OpenAIGPTConfig(
            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__       							= TFOpenAIGPTModel(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, input_mask]
        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__       							= TFOpenAIGPTLMHeadModel(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_      ,							*A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFOpenAIGPTDoubleHeadsModel(config=A_    )

        SCREAMING_SNAKE_CASE__       							= tf.tile(tf.expand_dims(A_      ,							1    )      ,							(1, self.num_choices, 1)    )
        SCREAMING_SNAKE_CASE__       							= tf.tile(tf.expand_dims(A_      ,							1    )      ,							(1, self.num_choices, 1)    )
        SCREAMING_SNAKE_CASE__       							= tf.tile(tf.expand_dims(A_      ,							1    )      ,							(1, self.num_choices, 1)    )

        SCREAMING_SNAKE_CASE__       							= {
            '''input_ids''': multiple_choice_inputs_ids,
            '''mc_token_ids''': mc_token_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        SCREAMING_SNAKE_CASE__       							= model(A_    )
        self.parent.assertEqual(
            result.logits.shape      ,							(self.batch_size, self.num_choices, self.seq_length, self.vocab_size)    )
        self.parent.assertEqual(result.mc_logits.shape      ,							(self.batch_size, self.num_choices)    )



    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,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
            '''labels''': sequence_labels,
        }
        SCREAMING_SNAKE_CASE__       							= TFOpenAIGPTForSequenceClassification(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__						:		List[str]							=  (
        (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
        if is_tf_available()
        else ()
    )
    lowerCamelCase__						:		Dict							=  (
        (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
    )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
    lowerCamelCase__						:		Any							=  (
        {
            """feature-extraction""": TFOpenAIGPTModel,
            """text-classification""": TFOpenAIGPTForSequenceClassification,
            """text-generation""": TFOpenAIGPTLMHeadModel,
            """zero-shot""": TFOpenAIGPTForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    lowerCamelCase__						:		List[str]							=  False
    lowerCamelCase__						:		List[Any]							=  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.
            # `OpenAIGPTConfig` 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__       							= TFOpenAIGPTModelTester(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_openai_gpt_model(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

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




    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

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





@require_tf
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= TFOpenAIGPTLMHeadModel.from_pretrained('''openai-gpt'''    )
        SCREAMING_SNAKE_CASE__       							= tf.convert_to_tensor([[4_81, 47_35, 5_44]]      ,							dtype=tf.intaa    )  # the president is
        SCREAMING_SNAKE_CASE__       							= [
            4_81,
            47_35,
            5_44,
            2_46,
            9_63,
            8_70,
            7_62,
            2_39,
            2_44,
            4_04_77,
            2_44,
            2_49,
            7_19,
            8_81,
            4_87,
            5_44,
            2_40,
            2_44,
            6_03,
            4_81,
        ]  # the president is a very good man. " \n " i\'m sure he is, " said the

        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)
