





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      import jax
      import jax.numpy as jnp

      from transformers.modeling_flax_pytorch_utils import (
          convert_pytorch_state_dict_to_flax,
          load_flax_weights_in_pytorch_model,
      )
      from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
      import torch

class 		UpperCamelCase__ :



     '''simple docstring'''

     def __init__(							self						:    Tuple						,lowerCamelCase__						:    Union[str, Any]						,lowerCamelCase__						:    Optional[int]=14						,lowerCamelCase__						:    str=7						,lowerCamelCase__						:    List[Any]=True						,lowerCamelCase__						:    List[Any]=True						,lowerCamelCase__						:    str=False						,lowerCamelCase__						:    List[str]=True						,lowerCamelCase__						:    Any=99						,lowerCamelCase__						:    Tuple=32						,lowerCamelCase__						:    Optional[Any]=4						,lowerCamelCase__						:    Tuple=4						,lowerCamelCase__						:    str=4						,lowerCamelCase__						:    List[Any]=37						,lowerCamelCase__						:    Optional[int]="gelu"						,lowerCamelCase__						:    Optional[int]=0.1						,lowerCamelCase__						:    Optional[int]=0.1						,lowerCamelCase__						:    List[Any]=512						,lowerCamelCase__						:    List[str]=0.02						,)     ->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						parent
           SCREAMING_SNAKE_CASE		=						batch_size
           SCREAMING_SNAKE_CASE		=						seq_length
           SCREAMING_SNAKE_CASE		=						is_training
           SCREAMING_SNAKE_CASE		=						use_input_mask
           SCREAMING_SNAKE_CASE		=						use_token_type_ids
           SCREAMING_SNAKE_CASE		=						use_labels
           SCREAMING_SNAKE_CASE		=						vocab_size
           SCREAMING_SNAKE_CASE		=						hidden_size
           SCREAMING_SNAKE_CASE		=						rotary_dim
           SCREAMING_SNAKE_CASE		=						num_hidden_layers
           SCREAMING_SNAKE_CASE		=						num_attention_heads
           SCREAMING_SNAKE_CASE		=						intermediate_size
           SCREAMING_SNAKE_CASE		=						hidden_act
           SCREAMING_SNAKE_CASE		=						hidden_dropout_prob
           SCREAMING_SNAKE_CASE		=						attention_probs_dropout_prob
           SCREAMING_SNAKE_CASE		=						max_position_embeddings
           SCREAMING_SNAKE_CASE		=						initializer_range
           SCREAMING_SNAKE_CASE		=						None
           SCREAMING_SNAKE_CASE		=						vocab_size - 1
           SCREAMING_SNAKE_CASE		=						vocab_size - 1
           SCREAMING_SNAKE_CASE		=						vocab_size - 1






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		List[str]:




           '''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		=						GPTJConfig(
               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						,use_cache=lowerCamelCase__						,bos_token_id=self.bos_token_id						,eos_token_id=self.eos_token_id						,pad_token_id=self.pad_token_id						,rotary_dim=self.rotary_dim						,)

           return (config, input_ids, input_mask)






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.prepare_config_and_inputs()
           SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						config_and_inputs
           SCREAMING_SNAKE_CASE		=						{"""input_ids""": input_ids, """attention_mask""": attention_mask}
           return config, inputs_dict






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Tuple						,lowerCamelCase__						:    Dict						,lowerCamelCase__						:    List[str]						,lowerCamelCase__						:    List[Any]						,lowerCamelCase__						:    Optional[Any]       )     ->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						20
           SCREAMING_SNAKE_CASE		=						model_class_name(lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						model.init_cache(input_ids.shape[0]						,lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						jnp.ones((input_ids.shape[0], max_decoder_length)						,dtype="""i4"""       )

           SCREAMING_SNAKE_CASE		=						jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,(input_ids.shape[0], input_ids.shape[-1] - 1)       )
           SCREAMING_SNAKE_CASE		=						model(
               input_ids[:, :-1]						,attention_mask=lowerCamelCase__						,past_key_values=lowerCamelCase__						,position_ids=lowerCamelCase__						,)

           SCREAMING_SNAKE_CASE		=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,dtype="""i4"""       )
           SCREAMING_SNAKE_CASE		=						model(
               input_ids[:, -1:]						,attention_mask=lowerCamelCase__						,past_key_values=outputs_cache.past_key_values						,position_ids=lowerCamelCase__						,)

           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
           self.parent.assertTrue(diff < 1e-3						,msg=F"""Max diff is {diff}"""       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]						,lowerCamelCase__						:    Optional[Any]						,lowerCamelCase__						:    Any						,lowerCamelCase__						:    List[str]						,lowerCamelCase__						:    Optional[int]       )     ->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						20
           SCREAMING_SNAKE_CASE		=						model_class_name(lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						jnp.concatenate(
               [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])       )]						,axis=-1						,)

           SCREAMING_SNAKE_CASE		=						model.init_cache(input_ids.shape[0]						,lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,(input_ids.shape[0], input_ids.shape[-1] - 1)       )

           SCREAMING_SNAKE_CASE		=						model(
               input_ids[:, :-1]						,attention_mask=lowerCamelCase__						,past_key_values=lowerCamelCase__						,position_ids=lowerCamelCase__						,)
           SCREAMING_SNAKE_CASE		=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,dtype="""i4"""       )
           SCREAMING_SNAKE_CASE		=						model(
               input_ids[:, -1:]						,past_key_values=outputs_cache.past_key_values						,attention_mask=lowerCamelCase__						,position_ids=lowerCamelCase__						,)

           SCREAMING_SNAKE_CASE		=						model(lowerCamelCase__						,attention_mask=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
           self.parent.assertTrue(diff < 1e-3						,msg=F"""Max diff is {diff}"""       )

@require_flax
class 		UpperCamelCase__ (      lowerCAmelCase_							,      lowerCAmelCase_							,      unittest.TestCase   ):



     '''simple docstring'''

     __snake_case    :	List[Any]			=			(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
     __snake_case    :	List[str]			=			(FlaxGPTJForCausalLM,) if is_flax_available() else ()






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						FlaxGPTJModelTester(self       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		Any:




           '''simple docstring'''


           for model_class_name in self.all_model_classes:
                 SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Optional[int]:




           '''simple docstring'''


           for model_class_name in self.all_model_classes:
                 SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward_with_attn_mask(
                     lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__       )






     @tooslow
     def 				SCREAMING_SNAKE_CASE__  (							self						:    str       )     ->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						GPTaTokenizer.from_pretrained("""gpt2"""						,pad_token="""<|endoftext|>"""						,padding_side="""left"""       )
           SCREAMING_SNAKE_CASE		=						tokenizer(["""Hello this is a long string""", """Hey"""]						,return_tensors="""np"""						,padding=lowerCamelCase__						,truncation=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						FlaxGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""       )
           SCREAMING_SNAKE_CASE		=						False
           SCREAMING_SNAKE_CASE		=						model.config.eos_token_id

           SCREAMING_SNAKE_CASE		=						jax.jit(model.generate       )

           SCREAMING_SNAKE_CASE		=						jit_generate(
               inputs["""input_ids"""]						,attention_mask=inputs["""attention_mask"""]						,pad_token_id=tokenizer.pad_token_id       ).sequences

           SCREAMING_SNAKE_CASE		=						tokenizer.batch_decode(lowerCamelCase__						,skip_special_tokens=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						[
               """Hello this is a long string of text.\n\nI'm trying to get the text of the""",
               """Hey, I'm a little late to the party. I'm going to""",
           ]

           self.assertListEqual(lowerCamelCase__						,lowerCamelCase__       )






     @is_pt_flax_cross_test
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Dict       )     ->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__       ):
                       # prepare inputs
                       SCREAMING_SNAKE_CASE		=						self._prepare_for_class(lowerCamelCase__						,lowerCamelCase__       )
                       SCREAMING_SNAKE_CASE		=						{k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

                       # load corresponding PyTorch class
                       SCREAMING_SNAKE_CASE		=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                       SCREAMING_SNAKE_CASE		=						getattr(lowerCamelCase__						,lowerCamelCase__       )

                       SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						pt_inputs["""input_ids"""].shape
                       SCREAMING_SNAKE_CASE		=						np.random.randint(0						,seq_length - 1						,size=(batch_size,)       )
                       for batch_idx, start_index in enumerate(lowerCamelCase__       ):
                             SCREAMING_SNAKE_CASE		=						0
                             SCREAMING_SNAKE_CASE		=						1
                             SCREAMING_SNAKE_CASE		=						0
                             SCREAMING_SNAKE_CASE		=						1
                       SCREAMING_SNAKE_CASE		=						pt_model_class(lowerCamelCase__       ).eval()
                       SCREAMING_SNAKE_CASE		=						model_class(lowerCamelCase__						,dtype=jnp.floataa       )

                       SCREAMING_SNAKE_CASE		=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,lowerCamelCase__       )
                       SCREAMING_SNAKE_CASE		=						fx_state

                       with torch.no_grad():
                             SCREAMING_SNAKE_CASE		=						pt_model(**lowerCamelCase__       ).to_tuple()

                       SCREAMING_SNAKE_CASE		=						fx_model(**lowerCamelCase__       ).to_tuple()
                       self.assertEqual(len(lowerCamelCase__       )						,len(lowerCamelCase__       )						,"""Output lengths differ between Flax and PyTorch"""       )
                       for fx_output, pt_output in zip(lowerCamelCase__						,lowerCamelCase__       ):
                             self.assert_almost_equals(fx_output[:, -1]						,pt_output[:, -1].numpy()						,4e-2       )

                       with tempfile.TemporaryDirectory() as tmpdirname:
                             pt_model.save_pretrained(lowerCamelCase__       )
                             SCREAMING_SNAKE_CASE		=						model_class.from_pretrained(lowerCamelCase__						,from_pt=lowerCamelCase__       )

                       SCREAMING_SNAKE_CASE		=						fx_model_loaded(**lowerCamelCase__       ).to_tuple()
                       self.assertEqual(
                           len(lowerCamelCase__       )						,len(lowerCamelCase__       )						,"""Output lengths differ between Flax and PyTorch"""       )
                       for fx_output_loaded, pt_output in zip(lowerCamelCase__						,lowerCamelCase__       ):
                             self.assert_almost_equals(fx_output_loaded[:, -1]						,pt_output[:, -1].numpy()						,4e-2       )






     @is_pt_flax_cross_test
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Union[str, Any]       )     ->		str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__       ):
                       # prepare inputs
                       SCREAMING_SNAKE_CASE		=						self._prepare_for_class(lowerCamelCase__						,lowerCamelCase__       )
                       SCREAMING_SNAKE_CASE		=						{k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

                       # load corresponding PyTorch class
                       SCREAMING_SNAKE_CASE		=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                       SCREAMING_SNAKE_CASE		=						getattr(lowerCamelCase__						,lowerCamelCase__       )

                       SCREAMING_SNAKE_CASE		=						pt_model_class(lowerCamelCase__       ).eval()
                       SCREAMING_SNAKE_CASE		=						model_class(lowerCamelCase__						,dtype=jnp.floataa       )

                       SCREAMING_SNAKE_CASE		=						load_flax_weights_in_pytorch_model(lowerCamelCase__						,fx_model.params       )
                       SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						pt_inputs["""input_ids"""].shape
                       SCREAMING_SNAKE_CASE		=						np.random.randint(0						,seq_length - 1						,size=(batch_size,)       )
                       for batch_idx, start_index in enumerate(lowerCamelCase__       ):
                             SCREAMING_SNAKE_CASE		=						0
                             SCREAMING_SNAKE_CASE		=						1
                             SCREAMING_SNAKE_CASE		=						0
                             SCREAMING_SNAKE_CASE		=						1

                       # make sure weights are tied in PyTorch
                       pt_model.tie_weights()

                       with torch.no_grad():
                             SCREAMING_SNAKE_CASE		=						pt_model(**lowerCamelCase__       ).to_tuple()

                       SCREAMING_SNAKE_CASE		=						fx_model(**lowerCamelCase__       ).to_tuple()
                       self.assertEqual(len(lowerCamelCase__       )						,len(lowerCamelCase__       )						,"""Output lengths differ between Flax and PyTorch"""       )
                       for fx_output, pt_output in zip(lowerCamelCase__						,lowerCamelCase__       ):
                             self.assert_almost_equals(fx_output[:, -1]						,pt_output[:, -1].numpy()						,4e-2       )

                       with tempfile.TemporaryDirectory() as tmpdirname:
                             fx_model.save_pretrained(lowerCamelCase__       )
                             SCREAMING_SNAKE_CASE		=						pt_model_class.from_pretrained(lowerCamelCase__						,from_flax=lowerCamelCase__       )

                       with torch.no_grad():
                             SCREAMING_SNAKE_CASE		=						pt_model_loaded(**lowerCamelCase__       ).to_tuple()

                       self.assertEqual(
                           len(lowerCamelCase__       )						,len(lowerCamelCase__       )						,"""Output lengths differ between Flax and PyTorch"""       )
                       for fx_output, pt_output in zip(lowerCamelCase__						,lowerCamelCase__       ):
                             self.assert_almost_equals(fx_output[:, -1]						,pt_output[:, -1].numpy()						,4e-2       )






     @tooslow
     def 				SCREAMING_SNAKE_CASE__  (							self						:    str       )     ->		str:




           '''simple docstring'''


           for model_class_name in self.all_model_classes:
                 SCREAMING_SNAKE_CASE		=						model_class_name.from_pretrained("""EleutherAI/gpt-j-6B"""       )
                 SCREAMING_SNAKE_CASE		=						model(np.ones((1, 1)       )       )
                 self.assertIsNotNone(lowerCamelCase__       )









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE_						   =       {
    """configuration_llama""": ["""LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LlamaConfig"""],
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						   =       ["""LlamaTokenizer"""]

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						   =       ["""LlamaTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						   =       [
          """LlamaForCausalLM""",
          """LlamaModel""",
          """LlamaPreTrainedModel""",
          """LlamaForSequenceClassification""",
      ]


if TYPE_CHECKING:
      from .configuration_llama import LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, LlamaConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_llama import LlamaTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_llama_fast import LlamaTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel


else:
      import sys

      SCREAMING_SNAKE_CASE_						   =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


