






"""simple docstring"""





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

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.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
  import torch







class 		_a       :





   '''simple docstring'''






   def __init__(   self			,				UpperCAmelCase_			,				UpperCAmelCase_=14			,				UpperCAmelCase_=7			,				UpperCAmelCase_=True			,				UpperCAmelCase_=True			,				UpperCAmelCase_=False			,				UpperCAmelCase_=True			,				UpperCAmelCase_=99			,				UpperCAmelCase_=32			,				UpperCAmelCase_=4			,				UpperCAmelCase_=4			,				UpperCAmelCase_=[[["global", "local"], 2]]			,				UpperCAmelCase_=37			,				UpperCAmelCase_="gelu"			,				UpperCAmelCase_=0.1			,				UpperCAmelCase_=0.1			,				UpperCAmelCase_=512			,				UpperCAmelCase_=7			,				UpperCAmelCase_=0.02			,				)	-> Any:




          '''simple docstring'''




          lowercase__:    List[str]        =      parent
          lowercase__:    Optional[int]        =      batch_size
          lowercase__:    List[str]        =      seq_length
          lowercase__:    Union[str, Any]        =      is_training
          lowercase__:    str        =      use_input_mask
          lowercase__:    int        =      use_token_type_ids
          lowercase__:    Optional[int]        =      use_labels
          lowercase__:    str        =      vocab_size
          lowercase__:    str        =      hidden_size
          lowercase__:    Dict        =      num_hidden_layers
          lowercase__:    List[Any]        =      num_attention_heads
          lowercase__:    Dict        =      attention_types
          lowercase__:    Tuple        =      intermediate_size
          lowercase__:    Union[str, Any]        =      hidden_act
          lowercase__:    str        =      hidden_dropout_prob
          lowercase__:    Optional[int]        =      attention_probs_dropout_prob
          lowercase__:    Optional[int]        =      max_position_embeddings
          lowercase__:    Union[str, Any]        =      window_size
          lowercase__:    Any        =      initializer_range
          lowercase__:    List[str]        =      None
          lowercase__:    Tuple        =      vocab_size - 1
          lowercase__:    List[str]        =      vocab_size - 1
          lowercase__:    Dict        =      vocab_size - 1





   def __lowercase    (   self)	-> Union[str, Any]:




          '''simple docstring'''




          lowercase__:    int        =      ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size)

          lowercase__:    List[str]        =      None
          if self.use_input_mask:
                 lowercase__:    Dict        =      random_attention_mask([self.batch_size, self.seq_length])

          lowercase__:    List[Any]        =      GPTNeoConfig(
              vocab_size=self.vocab_size			,				hidden_size=self.hidden_size			,				num_layers=self.num_hidden_layers			,				num_heads=self.num_attention_heads			,				max_position_embeddings=self.max_position_embeddings			,				use_cache=UpperCAmelCase_			,				bos_token_id=self.bos_token_id			,				eos_token_id=self.eos_token_id			,				pad_token_id=self.pad_token_id			,				window_size=self.window_size			,				attention_types=self.attention_types			,				)

          return (config, input_ids, input_mask)





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




          '''simple docstring'''




          lowercase__:    Union[str, Any]        =      self.prepare_config_and_inputs()
          lowercase__  ,     lowercase__  ,     lowercase__:    List[Any]        =      config_and_inputs
          lowercase__:    List[str]        =      {"input_ids": input_ids, "attention_mask": attention_mask}
          return config, inputs_dict





   def __lowercase    (   self			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_)	-> Optional[Any]:




          '''simple docstring'''




          lowercase__:    int        =      20
          lowercase__:    Union[str, Any]        =      model_class_name(UpperCAmelCase_)

          lowercase__:    Dict        =      model.init_cache(input_ids.shape[0]			,				UpperCAmelCase_)
          lowercase__:    Tuple        =      jnp.ones((input_ids.shape[0], max_decoder_length)			,				dtype="i4")

          lowercase__:    Optional[Any]        =      jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1)[None, :]			,				(input_ids.shape[0], input_ids.shape[-1] - 1))
          lowercase__:    int        =      model(
              input_ids[:, :-1]			,				attention_mask=UpperCAmelCase_			,				past_key_values=UpperCAmelCase_			,				position_ids=UpperCAmelCase_			,				)

          lowercase__:    Union[str, Any]        =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,				dtype="i4")
          lowercase__:    List[Any]        =      model(
              input_ids[:, -1:]			,				attention_mask=UpperCAmelCase_			,				past_key_values=outputs_cache.past_key_values			,				position_ids=UpperCAmelCase_			,				)

          lowercase__:    Dict        =      model(UpperCAmelCase_)

          lowercase__:    Optional[int]        =      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 __lowercase    (   self			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_)	-> Dict:




          '''simple docstring'''




          lowercase__:    int        =      20
          lowercase__:    List[Any]        =      model_class_name(UpperCAmelCase_)

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

          lowercase__:    List[str]        =      model.init_cache(input_ids.shape[0]			,				UpperCAmelCase_)
          lowercase__:    List[str]        =      jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1)[None, :]			,				(input_ids.shape[0], input_ids.shape[-1] - 1))

          lowercase__:    str        =      model(
              input_ids[:, :-1]			,				attention_mask=UpperCAmelCase_			,				past_key_values=UpperCAmelCase_			,				position_ids=UpperCAmelCase_			,				)
          lowercase__:    Tuple        =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,				dtype="i4")
          lowercase__:    Dict        =      model(
              input_ids[:, -1:]			,				past_key_values=outputs_cache.past_key_values			,				attention_mask=UpperCAmelCase_			,				position_ids=UpperCAmelCase_			,				)

          lowercase__:    Optional[Any]        =      model(UpperCAmelCase_			,				attention_mask=UpperCAmelCase_)

          lowercase__:    Any        =      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 		_a       (						lowercase_    ,					lowercase_    ,					unittest.TestCase   ):





   '''simple docstring'''






   UpperCamelCase__      					=    (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
   UpperCamelCase__      					=    (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()





   def __lowercase    (   self)	-> Optional[Any]:




          '''simple docstring'''




          lowercase__:    List[Any]        =      FlaxGPTNeoModelTester(self)





   def __lowercase    (   self)	-> Dict:




          '''simple docstring'''




          for model_class_name in self.all_model_classes:
                 lowercase__  ,     lowercase__  ,     lowercase__:    Dict        =      self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward(UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_)





   def __lowercase    (   self)	-> Optional[Any]:




          '''simple docstring'''




          for model_class_name in self.all_model_classes:
                 lowercase__  ,     lowercase__  ,     lowercase__:    Tuple        =      self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward_with_attn_mask(
                     UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_			,				UpperCAmelCase_)





   @slow
   def __lowercase    (   self)	-> Optional[int]:




          '''simple docstring'''




          lowercase__:    int        =      GPTaTokenizer.from_pretrained("gpt2"			,				pad_token="<|endoftext|>"			,				padding_side="left")
          lowercase__:    Union[str, Any]        =      tokenizer(["Hello this is a long string", "Hey"]			,				return_tensors="np"			,				padding=UpperCAmelCase_			,				truncation=UpperCAmelCase_)

          lowercase__:    Any        =      FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
          lowercase__:    Any        =      False
          lowercase__:    str        =      model.config.eos_token_id

          lowercase__:    Union[str, Any]        =      jax.jit(model.generate)

          lowercase__:    List[Any]        =      jit_generate(
              inputs["input_ids"]			,				attention_mask=inputs["attention_mask"]			,				pad_token_id=tokenizer.pad_token_id).sequences

          lowercase__:    Tuple        =      tokenizer.batch_decode(UpperCAmelCase_			,				skip_special_tokens=UpperCAmelCase_)

          lowercase__:    List[Any]        =      [
              "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(UpperCAmelCase_			,				UpperCAmelCase_)





   @is_pt_flax_cross_test
   def __lowercase    (   self)	-> List[Any]:




          '''simple docstring'''




          lowercase__  ,     lowercase__:    str        =      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
                        lowercase__:    Optional[Any]        =      self._prepare_for_class(UpperCAmelCase_			,				UpperCAmelCase_)
                        lowercase__:    Optional[Any]        =      {k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

                        # load corresponding PyTorch class
                        lowercase__:    Optional[Any]        =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
                        lowercase__:    List[str]        =      getattr(UpperCAmelCase_			,				UpperCAmelCase_)

                        lowercase__  ,     lowercase__:    List[Any]        =      pt_inputs["input_ids"].shape
                        lowercase__:    List[Any]        =      np.random.randint(0			,				seq_length - 1			,				size=(batch_size,))
                        for batch_idx, start_index in enumerate(UpperCAmelCase_):
                               lowercase__:    List[Any]        =      0
                               lowercase__:    Union[str, Any]        =      1
                               lowercase__:    int        =      0
                               lowercase__:    Any        =      1
                        lowercase__:    Optional[int]        =      pt_model_class(UpperCAmelCase_).eval()
                        lowercase__:    Dict        =      model_class(UpperCAmelCase_			,				dtype=jnp.floataa)

                        lowercase__:    Any        =      convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,				UpperCAmelCase_)
                        lowercase__:    Dict        =      fx_state

                        with torch.no_grad():
                               lowercase__:    Union[str, Any]        =      pt_model(**UpperCAmelCase_).to_tuple()

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

                        with tempfile.TemporaryDirectory() as tmpdirname:
                               pt_model.save_pretrained(UpperCAmelCase_)
                               lowercase__:    Any        =      model_class.from_pretrained(UpperCAmelCase_			,				from_pt=UpperCAmelCase_)

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





   @is_pt_flax_cross_test
   def __lowercase    (   self)	-> int:




          '''simple docstring'''




          lowercase__  ,     lowercase__:    List[str]        =      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
                        lowercase__:    List[Any]        =      self._prepare_for_class(UpperCAmelCase_			,				UpperCAmelCase_)
                        lowercase__:    Optional[int]        =      {k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

                        # load corresponding PyTorch class
                        lowercase__:    Union[str, Any]        =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
                        lowercase__:    Any        =      getattr(UpperCAmelCase_			,				UpperCAmelCase_)

                        lowercase__:    Union[str, Any]        =      pt_model_class(UpperCAmelCase_).eval()
                        lowercase__:    List[str]        =      model_class(UpperCAmelCase_			,				dtype=jnp.floataa)

                        lowercase__:    List[Any]        =      load_flax_weights_in_pytorch_model(UpperCAmelCase_			,				fx_model.params)
                        lowercase__  ,     lowercase__:    str        =      pt_inputs["input_ids"].shape
                        lowercase__:    Tuple        =      np.random.randint(0			,				seq_length - 1			,				size=(batch_size,))
                        for batch_idx, start_index in enumerate(UpperCAmelCase_):
                               lowercase__:    Dict        =      0
                               lowercase__:    List[str]        =      1
                               lowercase__:    Dict        =      0
                               lowercase__:    int        =      1

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

                        with torch.no_grad():
                               lowercase__:    Tuple        =      pt_model(**UpperCAmelCase_).to_tuple()

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

                        with tempfile.TemporaryDirectory() as tmpdirname:
                               fx_model.save_pretrained(UpperCAmelCase_)
                               lowercase__:    int        =      pt_model_class.from_pretrained(UpperCAmelCase_			,				from_flax=UpperCAmelCase_)

                        with torch.no_grad():
                               lowercase__:    int        =      pt_model_loaded(**UpperCAmelCase_).to_tuple()

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





   @slow
   def __lowercase    (   self)	-> Dict:




          '''simple docstring'''




          for model_class_name in self.all_model_classes:
                 lowercase__:    Union[str, Any]        =      model_class_name.from_pretrained("EleutherAI/gpt-neo-125M")
                 lowercase__:    Union[str, Any]        =      model(np.ones((1, 1)))
                 self.assertIsNotNone(UpperCAmelCase_)














"""simple docstring"""





class 		_a       :





   '''simple docstring'''






   def __init__(   self			,				UpperCAmelCase_=None)	-> List[Any]:




          '''simple docstring'''




          lowercase__:    str        =      data
          lowercase__:    Any        =      None





   def __repr__(   self)	-> Optional[Any]:




          '''simple docstring'''




          lowercase__:    Any        =      ""
          lowercase__:    Dict        =      self
          while temp:
                 string_rep += F"""<{temp.data}> ---> """
                 lowercase__:    Tuple        =      temp.next
          string_rep += "<END>"
          return string_rep







def     A(  snake_case_		):
       """simple docstring"""





       if not elements_list:
              raise Exception("The Elements List is empty"		)

       # Set first element as Head
       lowercase__:    Any        =      Node(elements_list[0]		)
       lowercase__:    Optional[int]        =      head
       # Loop through elements from position 1
       for data in elements_list[1:]:
              lowercase__:    Dict        =      Node(snake_case_		)
              lowercase__:    Dict        =      current.next
       return head


UpperCamelCase			=      [1, 3, 5, 32, 44, 12, 43]
print(F"List: {list_data}")
print("""Creating Linked List from List.""")
UpperCamelCase			=      make_linked_list(list_data)
print("""Linked List:""")







print(linked_list)






