

import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel





class  __lowerCAmelCase     :




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase=13							,   lowerCAmelCase=7							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=False							,   lowerCAmelCase=True							,   lowerCAmelCase=99							,   lowerCAmelCase=32							,   lowerCAmelCase=5							,   lowerCAmelCase=4							,   lowerCAmelCase=37							,   lowerCAmelCase="gelu"							,   lowerCAmelCase=0.1							,   lowerCAmelCase=0.1							,   lowerCAmelCase=512							,   lowerCAmelCase=16							,   lowerCAmelCase=2							,   lowerCAmelCase=0.02							,   lowerCAmelCase=3							,   lowerCAmelCase=4							,   lowerCAmelCase=None							,   )  -> Any:
        '''simple docstring'''






        _lowercase	       =parent
        _lowercase	       =batch_size
        _lowercase	       =seq_length
        _lowercase	       =is_training
        _lowercase	       =use_input_mask
        _lowercase	       =use_token_type_ids
        _lowercase	       =use_labels
        _lowercase	       =vocab_size
        _lowercase	       =hidden_size
        _lowercase	       =num_hidden_layers
        _lowercase	       =num_attention_heads
        _lowercase	       =intermediate_size
        _lowercase	       =hidden_act
        _lowercase	       =hidden_dropout_prob
        _lowercase	       =attention_probs_dropout_prob
        _lowercase	       =max_position_embeddings
        _lowercase	       =type_vocab_size
        _lowercase	       =type_sequence_label_size
        _lowercase	       =initializer_range
        _lowercase	       =num_labels
        _lowercase	       =num_choices
        _lowercase	       =scope




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






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

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

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

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

        _lowercase	       =self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        return OpenLlamaConfig(
            vocab_size=self.vocab_size							,   hidden_size=self.hidden_size							,   num_hidden_layers=self.num_hidden_layers							,   num_attention_heads=self.num_attention_heads							,   intermediate_size=self.intermediate_size							,   hidden_act=self.hidden_act							,   hidden_dropout_prob=self.hidden_dropout_prob							,   attention_probs_dropout_prob=self.attention_probs_dropout_prob							,   max_position_embeddings=self.max_position_embeddings							,   type_vocab_size=self.type_vocab_size							,   is_decoder=lowerCAmelCase							,   initializer_range=self.initializer_range							,   use_stable_embedding=lowerCAmelCase							,   )




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =OpenLlamaModel(config=lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase					)
        _lowercase	       =model(lowerCAmelCase					)
        self.parent.assertEqual(result.last_hidden_state.shape							,   (self.batch_size, self.seq_length, self.hidden_size)					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   )  -> List[Any]:
        '''simple docstring'''






        _lowercase	       =True
        _lowercase	       =OpenLlamaModel(lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(
            lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   )
        _lowercase	       =model(
            lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   )
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase					)
        self.parent.assertEqual(result.last_hidden_state.shape							,   (self.batch_size, self.seq_length, self.hidden_size)					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   )  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =OpenLlamaForCausalLM(config=lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   labels=lowerCAmelCase					)
        self.parent.assertEqual(result.logits.shape							,   (self.batch_size, self.seq_length, self.vocab_size)					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   )  -> Any:
        '''simple docstring'''






        _lowercase	       =True
        _lowercase	       =True
        _lowercase	       =OpenLlamaForCausalLM(config=lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()

        # first forward pass
        _lowercase	       =model(
            lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   use_cache=lowerCAmelCase							,   )
        _lowercase	       =outputs.past_key_values

        # create hypothetical multiple next token and extent to next_input_ids
        _lowercase	       =ids_tensor((self.batch_size, 3)							,   config.vocab_size					)
        _lowercase	       =ids_tensor((self.batch_size, 3)							,   vocab_size=2					)

        # append to next input_ids and
        _lowercase	       =torch.cat([input_ids, next_tokens]							,   dim=-1					)
        _lowercase	       =torch.cat([input_mask, next_mask]							,   dim=-1					)

        _lowercase	       =model(
            lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   output_hidden_states=lowerCAmelCase							,   )['hidden_states'][0]
        _lowercase	       =model(
            lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   past_key_values=lowerCAmelCase							,   output_hidden_states=lowerCAmelCase							,   )['hidden_states'][0]

        # select random slice
        _lowercase	       =ids_tensor((1,)							,   output_from_past.shape[-1]					).item()
        _lowercase	       =output_from_no_past[:, -3:, random_slice_idx].detach()
        _lowercase	       =output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]					)

        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(lowerCAmelCase							,   lowerCAmelCase							,   atol=1e-3					)					)




  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        _lowercase	       =self.prepare_config_and_inputs()
        (
            (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       (
            _lowercase
        )       ,       
        )	       =config_and_inputs
        _lowercase	       ={'input_ids': input_ids, 'attention_mask': input_mask}
        return config, inputs_dict





@require_torch
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			unittest.TestCase							):
  _a									=  (
      (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
  )
  _a									=  (OpenLlamaForCausalLM,) if is_torch_available() else ()
  _a									=  (
      {
          """feature-extraction""": OpenLlamaModel,
          """text-classification""": OpenLlamaForSequenceClassification,
          """text-generation""": OpenLlamaForCausalLM,
          """zero-shot""": OpenLlamaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _a									=  False
  _a									=  False




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        _lowercase	       =OpenLlamaModelTester(self					)
        _lowercase	       =ConfigTester(self							,   config_class=lowerCAmelCase							,   hidden_size=37					)




  def   A__     (				self					)  -> Any:
        '''simple docstring'''






        self.config_tester.run_common_tests()




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase					)




  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
              _lowercase	       =type
              self.model_tester.create_and_check_model(*lowerCAmelCase					)




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()
        _lowercase	       =3
        _lowercase	       =input_dict['input_ids']
        _lowercase	       =input_ids.ne(1					).to(lowerCAmelCase					)
        _lowercase	       =ids_tensor([self.model_tester.batch_size]							,   self.model_tester.type_sequence_label_size					)
        _lowercase	       =OpenLlamaForSequenceClassification(lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   labels=lowerCAmelCase					)
        self.assertEqual(result.logits.shape							,   (self.model_tester.batch_size, self.model_tester.num_labels)					)




  def   A__     (				self					)  -> List[Any]:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()
        _lowercase	       =3
        _lowercase	       ='single_label_classification'
        _lowercase	       =input_dict['input_ids']
        _lowercase	       =input_ids.ne(1					).to(lowerCAmelCase					)
        _lowercase	       =ids_tensor([self.model_tester.batch_size]							,   self.model_tester.type_sequence_label_size					)
        _lowercase	       =OpenLlamaForSequenceClassification(lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   labels=lowerCAmelCase					)
        self.assertEqual(result.logits.shape							,   (self.model_tester.batch_size, self.model_tester.num_labels)					)




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()
        _lowercase	       =3
        _lowercase	       ='multi_label_classification'
        _lowercase	       =input_dict['input_ids']
        _lowercase	       =input_ids.ne(1					).to(lowerCAmelCase					)
        _lowercase	       =ids_tensor(
            [self.model_tester.batch_size, config.num_labels]							,   self.model_tester.type_sequence_label_size					).to(torch.float					)
        _lowercase	       =OpenLlamaForSequenceClassification(lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        _lowercase	       =model(lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   labels=lowerCAmelCase					)
        self.assertEqual(result.logits.shape							,   (self.model_tester.batch_size, self.model_tester.num_labels)					)




  @unittest.skip('Open-Llama buffers include complex numbers, which breaks this test'					)
  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        pass




  @parameterized.expand([('linear',), ('dynamic',)]					)
  def   A__     (				self							,   lowerCAmelCase					)  -> List[str]:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()
        _lowercase	       =ids_tensor([1, 10]							,   config.vocab_size					)
        _lowercase	       =ids_tensor([1, int(config.max_position_embeddings * 1.5					)]							,   config.vocab_size					)

        set_seed(42					)  # Fixed seed at init time so the two models get the same random weights
        _lowercase	       =OpenLlamaModel(lowerCAmelCase					)
        original_model.to(lowerCAmelCase					)
        original_model.eval()
        _lowercase	       =original_model(lowerCAmelCase					).last_hidden_state
        _lowercase	       =original_model(lowerCAmelCase					).last_hidden_state

        set_seed(42					)  # Fixed seed at init time so the two models get the same random weights
        _lowercase	       ={'type': scaling_type, 'factor': 10.0}
        _lowercase	       =OpenLlamaModel(lowerCAmelCase					)
        scaled_model.to(lowerCAmelCase					)
        scaled_model.eval()
        _lowercase	       =scaled_model(lowerCAmelCase					).last_hidden_state
        _lowercase	       =scaled_model(lowerCAmelCase					).last_hidden_state

        # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
        # maximum sequence length, so the outputs for the short input should match.
        if scaling_type == "dynamic":
              self.assertTrue(torch.allclose(lowerCAmelCase							,   lowerCAmelCase							,   atol=1e-5					)					)
        else:
              self.assertFalse(torch.allclose(lowerCAmelCase							,   lowerCAmelCase							,   atol=1e-5					)					)

        # The output should be different for long inputs
        self.assertFalse(torch.allclose(lowerCAmelCase							,   lowerCAmelCase							,   atol=1e-5					)					)





import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
lowercase_									=   data_utils.TransfoXLTokenizer
lowercase_									=   data_utils.TransfoXLCorpus
lowercase_									=   data_utils
lowercase_									=   data_utils
def      a      (			A__	:   int      ,   A__	:   Dict      ,   A__	:   Union[str, Any]      ,   A__	:   Union[str, Any]     )    ->      List[str]:





      """simple docstring"""



      if transfo_xl_dataset_file:
            # Convert a pre-processed corpus (see original TensorFlow repo)
            with open(A__      ,   'rb'     ) as fp:
                  _lowercase	       =pickle.load(A__      ,   encoding='latin1'     )
            # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
            _lowercase	       =pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
            print(F'''Save vocabulary to {pytorch_vocab_dump_path}'''     )
            _lowercase	       =corpus.vocab.__dict__
            torch.save(A__      ,   A__     )

            _lowercase	       =corpus.__dict__
            corpus_dict_no_vocab.pop('vocab'      ,   A__     )
            _lowercase	       =pytorch_dump_folder_path + '/' + CORPUS_NAME
            print(F'''Save dataset to {pytorch_dataset_dump_path}'''     )
            torch.save(A__      ,   A__     )

      if tf_checkpoint_path:
            # Convert a pre-trained TensorFlow model
            _lowercase	       =os.path.abspath(A__     )
            _lowercase	       =os.path.abspath(A__     )

            print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'''     )
            # Initialise PyTorch model
            if transfo_xl_config_file == "":
                  _lowercase	       =TransfoXLConfig()
            else:
                  _lowercase	       =TransfoXLConfig.from_json_file(A__     )
            print(F'''Building PyTorch model from configuration: {config}'''     )
            _lowercase	       =TransfoXLLMHeadModel(A__     )

            _lowercase	       =load_tf_weights_in_transfo_xl(A__      ,   A__      ,   A__     )
            # Save pytorch-model
            _lowercase	       =os.path.join(A__      ,   A__     )
            _lowercase	       =os.path.join(A__      ,   A__     )
            print(F'''Save PyTorch model to {os.path.abspath(A__     )}'''     )
            torch.save(model.state_dict()      ,   A__     )
            print(F'''Save configuration file to {os.path.abspath(A__     )}'''     )
            with open(A__      ,   'w'      ,   encoding='utf-8'     ) as f:
                  f.write(config.to_json_string()     )


if __name__ == "__main__":
  lowercase_									=   argparse.ArgumentParser()
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default=None,
      type=str,
      required=True,
      help='Path to the folder to store the PyTorch model or dataset/vocab.',
  )
  parser.add_argument(
      '--tf_checkpoint_path',
      default='',
      type=str,
      help='An optional path to a TensorFlow checkpoint path to be converted.',
  )
  parser.add_argument(
      '--transfo_xl_config_file',
      default='',
      type=str,
      help=(
          'An optional config json file corresponding to the pre-trained BERT model. \n'
          'This specifies the model architecture.'
      ),
  )
  parser.add_argument(
      '--transfo_xl_dataset_file',
      default='',
      type=str,
      help='An optional dataset file to be converted in a vocabulary.',
  )
  lowercase_									=   parser.parse_args()
  convert_transfo_xl_checkpoint_to_pytorch(
      args.tf_checkpoint_path,
      args.transfo_xl_config_file,
      args.pytorch_dump_folder_path,
      args.transfo_xl_dataset_file,
  )


