


import json
import os
import pickle
import subprocess
from hashlib import mda
from pathlib import Path
from tempfile import gettempdir
from textwrap import dedent
from types import CodeType, FunctionType
from unittest import TestCase
from unittest.mock import patch

import pytest
from multiprocess import Pool
from packaging import version

import datasets
from datasets.fingerprint import Hasher, fingerprint_transform
from datasets.table import InMemoryTable

from .utils import (
    require_regex,
    require_spacy,
    require_spacy_model,
    require_tiktoken,
    require_torch,
    require_transformers,
)







class    lowerCamelCase__     :





     def __init__(self	,  UpperCAmelCase )  ->     List[str]:
      _lowercase							   =foo





     def __call__(self )  ->     List[str]:
      return self.foo







class    lowerCamelCase__     ( datasets.Dataset):





     @fingerprint_transform(inplace=UpperCAmelCase )
     def   __A      (self	,  UpperCAmelCase	,  *UpperCAmelCase	,  **UpperCAmelCase )  ->     Dict:
      return DatasetChild(self.data	,  fingerprint=UpperCAmelCase )





     @fingerprint_transform(inplace=UpperCAmelCase )
     def   __A      (self	,  UpperCAmelCase	,  *UpperCAmelCase	,  **UpperCAmelCase )  ->     Dict:
      return DatasetChild(self.data	,  fingerprint=UpperCAmelCase )







class    lowerCamelCase__     :





     def __init__(self	,  UpperCAmelCase )  ->     List[Any]:
      _lowercase							   =callable





     def __call__(self	,  *UpperCAmelCase	,  **UpperCAmelCase )  ->     Optional[int]:
      if self.callable is not None:
       return self.callable(*UpperCAmelCase	,  **UpperCAmelCase )





     def __getstate__(self )  ->     str:
      raise pickle.PicklingError()







class    lowerCamelCase__     ( lowerCAmelCase):





     @require_transformers
     @pytest.mark.integration
     def   __A      (self )  ->     Any:
      from transformers import AutoTokenizer

      def encode(UpperCAmelCase ):
       return tokenizer(UpperCAmelCase )

      # TODO: add hash consistency tests across sessions
      _lowercase							   =AutoTokenizer.from_pretrained('''bert-base-uncased''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase : tokenizer(UpperCAmelCase ) ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =AutoTokenizer.from_pretrained('''bert-base-cased''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase : tokenizer(UpperCAmelCase ) ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =AutoTokenizer.from_pretrained('''bert-base-uncased''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(lambda UpperCAmelCase : tokenizer(UpperCAmelCase ) ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     @require_transformers
     @pytest.mark.integration
     def   __A      (self )  ->     Any:
      from transformers import AutoTokenizer

      _lowercase							   =AutoTokenizer.from_pretrained('''gpt2''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      tokenizer('''Hello world !''' )  # call once to change the tokenizer's cache
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )





     @require_regex
     def   __A      (self )  ->     List[Any]:
      import regex

      _lowercase							   =regex.Regex('''foo''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =regex.Regex('''bar''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =regex.Regex('''foo''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )







class    lowerCamelCase__     ( lowerCAmelCase):





     def   __A      (self )  ->     List[Any]:
      def func():
       return foo

      _lowercase							   =[0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =[1]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =[0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Union[str, Any]:
      def func():
       pass

      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()

      def func():
       pass

      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     List[str]:
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([0] ) ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([1] ) ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([0] ) ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([0] ).__call__ ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([1] ).__call__ ) ).hexdigest()
      _lowercase							   =mda(datasets.utils.py_utils.dumps(Foo([0] ).__call__ ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Optional[int]:
      _lowercase							   =(
          '''co_argcount''',
          '''co_kwonlyargcount''',
          '''co_nlocals''',
          '''co_stacksize''',
          '''co_flags''',
          '''co_code''',
          '''co_consts''',
          '''co_names''',
          '''co_varnames''',
          '''co_filename''',
          '''co_name''',
          '''co_firstlineno''',
          '''co_lnotab''',
          '''co_freevars''',
          '''co_cellvars''',
      )

      def _create_code_pyaa(*UpperCAmelCase ):
       if hasattr(args[-3]	,  '''encode''' ):
        _lowercase							   =list(UpperCAmelCase )
        _lowercase							   =args[-3].encode()  # co_lnotab
        _lowercase							   =args[-1_0].encode()  # co_code
       if hasattr(UpperCAmelCase	,  '''co_posonlyargcount''' ):
        if len(UpperCAmelCase ) == 1_6:
         return CodeType(*UpperCAmelCase )
        elif len(UpperCAmelCase ) == 1_5:
         return CodeType(args[0]	,  0	,  *args[1:] )
        return CodeType(args[0]	,  0	,  0	,  *args[1:] )
       elif hasattr(UpperCAmelCase	,  '''co_kwonlyargcount''' ):
        if len(UpperCAmelCase ) == 1_6:
         return CodeType(args[0]	,  *args[2:] )
        elif len(UpperCAmelCase ) == 1_5:
         return CodeType(*UpperCAmelCase )
        return CodeType(args[0]	,  0	,  *args[1:] )
       if len(UpperCAmelCase ) == 1_6:
        return CodeType(args[0]	,  *args[3:] )
       elif len(UpperCAmelCase ) == 1_5:
        return CodeType(args[0]	,  *args[2:] )
       return CodeType(*UpperCAmelCase )

      def create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ):
       def func():
        return returned_obj

       _lowercase							   =func.__code__
       # Use _create_code from dill in order to make it work for different python versions
       if datasets.config.PY_VERSION < version.parse('''3.8''' ):
        _lowercase							   =_create_code_pyaa(
            *[getattr(UpperCAmelCase	,  UpperCAmelCase ) if k != '''co_filename''' else co_filename for k in code_args_pyaa] )
       else:
        _lowercase							   =code.replace(co_filename=UpperCAmelCase )
       return FunctionType(UpperCAmelCase	,  func.__globals__	,  func.__name__	,  func.__defaults__	,  func.__closure__ )

      _lowercase   ,					_lowercase							   ='''<ipython-input-2-e0383a102aae>''', [0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      _lowercase   ,					_lowercase							   ='''<ipython-input-2-e0383a102aae>''', [1]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      _lowercase   ,					_lowercase							   ='''<ipython-input-5-713f6613acf3>''', [0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )

      _lowercase   ,					_lowercase							   =os.path.join(gettempdir()	,  '''ipykernel_12345'''	,  '''321456789.py''' ), [0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      _lowercase   ,					_lowercase							   =os.path.join(gettempdir()	,  '''ipykernel_12345'''	,  '''321456789.py''' ), [1]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      _lowercase   ,					_lowercase							   =os.path.join(gettempdir()	,  '''ipykernel_12345'''	,  '''654123987.py''' ), [0]
      _lowercase							   =mda(datasets.utils.py_utils.dumps(create_ipython_func(UpperCAmelCase	,  UpperCAmelCase ) ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     List[str]:
      _lowercase   ,					_lowercase							   =[0], [1]

      def func():
       return foo, bar

      _lowercase							   ='''__main__'''

      def globalvars_mocka_side_effect(UpperCAmelCase	,  *UpperCAmelCase	,  **UpperCAmelCase ):
       return {"foo": foo, "bar": bar}

      def globalvars_mocka_side_effect(UpperCAmelCase	,  *UpperCAmelCase	,  **UpperCAmelCase ):
       return {"bar": bar, "foo": foo}

      with patch('''dill.detect.globalvars'''	,  side_effect=UpperCAmelCase ) as globalvars_mocka:
       _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
       self.assertGreater(globalvars_mocka.call_count	,  0 )
      with patch('''dill.detect.globalvars'''	,  side_effect=UpperCAmelCase ) as globalvars_mocka:
       _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
       self.assertGreater(globalvars_mocka.call_count	,  0 )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )







class    lowerCamelCase__     ( lowerCAmelCase):





     def   __A      (self )  ->     Dict:
      _lowercase							   =Hasher.hash('''hello''' )
      _lowercase							   =Hasher.hash('''hello''' )
      _lowercase							   =Hasher.hash('''there''' )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =Hasher.hash(Foo('''hello''' ) )
      _lowercase							   =Hasher.hash(Foo('''hello''' ) )
      _lowercase							   =Hasher.hash(Foo('''there''' ) )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   =Hasher()
      for x in ["hello", Foo('''hello''' )]:
       hasher.update(UpperCAmelCase )
      _lowercase							   =hasher.hexdigest()
      _lowercase							   =Hasher()
      for x in ["hello", Foo('''hello''' )]:
       hasher.update(UpperCAmelCase )
      _lowercase							   =hasher.hexdigest()
      _lowercase							   =Hasher()
      for x in ["there", Foo('''there''' )]:
       hasher.update(UpperCAmelCase )
      _lowercase							   =hasher.hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Any:
      with self.assertRaises(pickle.PicklingError ):
       Hasher.hash(UnpicklableCallable(Foo('''hello''' ) ) )





     def   __A      (self )  ->     List[Any]:
      _lowercase							   ='''abc'''
      _lowercase							   =[string, string]  # two strings have the same ids
      _lowercase							   =[string, string]
      _lowercase							   =json.loads(f"[\"{string}\", \"{string}\"]" )  # two strings have different ids
      self.assertIs(obja[0]	,  UpperCAmelCase )
      self.assertIs(obja[0]	,  obja[1] )
      self.assertIs(obja[0]	,  UpperCAmelCase )
      self.assertIs(obja[0]	,  obja[1] )
      self.assertIsNot(obja[0]	,  UpperCAmelCase )
      self.assertIsNot(obja[0]	,  obja[1] )
      _lowercase							   =Hasher.hash(UpperCAmelCase )
      _lowercase							   =Hasher.hash(UpperCAmelCase )
      _lowercase							   =Hasher.hash(UpperCAmelCase )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )





     @require_tiktoken
     def   __A      (self )  ->     Union[str, Any]:
      import tiktoken

      _lowercase							   =tiktoken.get_encoding('''gpt2''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =tiktoken.get_encoding('''r50k_base''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =tiktoken.get_encoding('''gpt2''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     @require_torch
     def   __A      (self )  ->     Optional[int]:
      import torch

      _lowercase							   =torch.tensor([1.0] )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =torch.tensor([2.0] )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =torch.tensor([1.0] )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )





     @require_spacy
     @require_spacy_model('''en_core_web_sm''' )
     @require_spacy_model('''fr_core_news_sm''' )
     @pytest.mark.integration
     def   __A      (self )  ->     str:
      import spacy

      _lowercase							   =spacy.load('''en_core_web_sm''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =spacy.load('''fr_core_news_sm''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      _lowercase							   =spacy.load('''en_core_web_sm''' )
      _lowercase							   =mda(datasets.utils.py_utils.dumps(UpperCAmelCase ) ).hexdigest()
      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )
      self.assertNotEqual(UpperCAmelCase	,  UpperCAmelCase )



@pytest.mark.integration
def        UpperCAmelCase_      (							__snake_case		)  ->    int:







 """simple docstring"""


 _lowercase							   =tmp_path / '''dir1'''
 _lowercase							   =tmp_path / '''dir2'''
 dira.mkdir()
 dira.mkdir()
 _lowercase							   ='''script.py'''
 _lowercase							   =dedent(
     '''
    from datasets.fingerprint import Hasher
    def foo():
        pass
    print(Hasher.hash(foo))
    '''		)
 _lowercase							   =dira / script_filename
 _lowercase							   =dira / script_filename
 with script_patha.open('''w'''		) as f:
  f.write(__snake_case		)
 with script_patha.open('''w'''		) as f:
  f.write(__snake_case		)
 _lowercase							   =subprocess.check_output(['''python''', str(__snake_case		)]		)
 _lowercase							   =subprocess.check_output(['''python''', str(__snake_case		)]		)
 assert fingerprinta == fingerprinta
def        UpperCAmelCase_      (							)  ->    Optional[Any]:







 """simple docstring"""


 _lowercase							   ={'''a''': [0, 1, 2]}
 _lowercase							   =DatasetChild(InMemoryTable.from_pydict(__snake_case		)		)
 _lowercase							   =dataset.funca()._fingerprint
 assert expected_fingerprint == dataset.funca()._fingerprint
 assert expected_fingerprint != dataset.funca()._fingerprint

 with Pool(2		) as p:
  assert expected_fingerprint == p.apply_async(dataset.funca		).get()._fingerprint
  assert expected_fingerprint != p.apply_async(dataset.funca		).get()._fingerprint
def        UpperCAmelCase_      (							)  ->    List[Any]:







 """simple docstring"""


 _lowercase							   ={'''a''': [0, 1, 2]}







 class    lowerCamelCase__     ( datasets.Dataset):





      @fingerprint_transform(inplace=UpperCAmelCase )
      def   __A      (self	,  UpperCAmelCase )  ->     Optional[int]:
       return DummyDatasetChild(self.data	,  fingerprint=UpperCAmelCase )



 _lowercase							   =DummyDatasetChild(InMemoryTable.from_pydict(__snake_case		)		).func()







 class    lowerCamelCase__     ( datasets.Dataset):





      @fingerprint_transform(inplace=UpperCAmelCase	,  version='''1.0.0''' )
      def   __A      (self	,  UpperCAmelCase )  ->     List[Any]:
       return DummyDatasetChild(self.data	,  fingerprint=UpperCAmelCase )



 _lowercase							   =DummyDatasetChild(InMemoryTable.from_pydict(__snake_case		)		).func()







 class    lowerCamelCase__     ( datasets.Dataset):





      @fingerprint_transform(inplace=UpperCAmelCase	,  version='''2.0.0''' )
      def   __A      (self	,  UpperCAmelCase )  ->     List[str]:
       return DummyDatasetChild(self.data	,  fingerprint=UpperCAmelCase )



 _lowercase							   =DummyDatasetChild(InMemoryTable.from_pydict(__snake_case		)		).func()

 assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}		) == 3



def        UpperCAmelCase_      (							)  ->    Any:







 """simple docstring"""


 _lowercase							   =Hasher()
 hasher.update(lambda __snake_case		: x		)





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, 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, floats_tensor, 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.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
  import torch







class    lowerCamelCase__     :





     def __init__(self	,  UpperCAmelCase	,  UpperCAmelCase=1_4	,  UpperCAmelCase=7	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=False	,  UpperCAmelCase=True	,  UpperCAmelCase=9_9	,  UpperCAmelCase=3_2	,  UpperCAmelCase=5	,  UpperCAmelCase=4	,  UpperCAmelCase=3_7	,  UpperCAmelCase="gelu"	,  UpperCAmelCase=0.1	,  UpperCAmelCase=0.1	,  UpperCAmelCase=5_1_2	,  UpperCAmelCase=0.02	,  )  ->     Any:
      _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							   =initializer_range
      _lowercase							   =None
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1





     def   __A      (self )  ->     int:
      _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							   =GPTaConfig(
          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=UpperCAmelCase	,  bos_token_id=self.bos_token_id	,  eos_token_id=self.eos_token_id	,  pad_token_id=self.pad_token_id	,  )

      return (config, input_ids, input_mask)





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =self.prepare_config_and_inputs()
      _lowercase   ,					_lowercase   ,					_lowercase							   =config_and_inputs
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
      return config, inputs_dict





     def   __A      (self )  ->     Tuple:
      _lowercase   ,					_lowercase   ,					_lowercase							   =self.prepare_config_and_inputs()

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

      return (
          config,
          input_ids,
          attention_mask,
          encoder_hidden_states,
          encoder_attention_mask,
      )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     Dict:
      _lowercase							   =2_0
      _lowercase							   =model_class_name(UpperCAmelCase )

      _lowercase							   =model.init_cache(input_ids.shape[0]	,  UpperCAmelCase )
      _lowercase							   =jnp.ones((input_ids.shape[0], max_decoder_length)	,  dtype='''i4''' )

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

      _lowercase							   =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,  dtype='''i4''' )
      _lowercase							   =model(
          input_ids[:, -1:]	,  attention_mask=UpperCAmelCase	,  past_key_values=outputs_cache.past_key_values	,  position_ids=UpperCAmelCase	,  )

      _lowercase							   =model(UpperCAmelCase )

      _lowercase							   =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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     List[str]:
      _lowercase							   =2_0
      _lowercase							   =model_class_name(UpperCAmelCase )

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

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

      _lowercase							   =model(
          input_ids[:, :-1]	,  attention_mask=UpperCAmelCase	,  past_key_values=UpperCAmelCase	,  position_ids=UpperCAmelCase	,  )
      _lowercase							   =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,  dtype='''i4''' )
      _lowercase							   =model(
          input_ids[:, -1:]	,  past_key_values=outputs_cache.past_key_values	,  attention_mask=UpperCAmelCase	,  position_ids=UpperCAmelCase	,  )

      _lowercase							   =model(UpperCAmelCase	,  attention_mask=UpperCAmelCase )

      _lowercase							   =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    lowerCamelCase__     ( lowerCAmelCase     ,							lowerCAmelCase     ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
     SCREAMING_SNAKE_CASE__       =    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()





     def   __A      (self )  ->     str:
      _lowercase							   =FlaxGPTaModelTester(self )





     def   __A      (self )  ->     Optional[int]:
      for model_class_name in self.all_model_classes:
       _lowercase   ,					_lowercase   ,					_lowercase							   =self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Dict:
      for model_class_name in self.all_model_classes:
       _lowercase   ,					_lowercase   ,					_lowercase							   =self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward_with_attn_mask(
           UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )





     @slow
     def   __A      (self )  ->     Dict:
      _lowercase							   =GPTaTokenizer.from_pretrained('''gpt2'''	,  pad_token='''</s>'''	,  padding_side='''left''' )
      _lowercase							   =tokenizer(['''Hello this is a long string''', '''Hey''']	,  return_tensors='''np'''	,  padding=UpperCAmelCase	,  truncation=UpperCAmelCase )

      _lowercase							   =FlaxGPTaLMHeadModel.from_pretrained('''gpt2''' )
      _lowercase							   =False
      _lowercase							   =model.config.eos_token_id

      _lowercase							   =jax.jit(model.generate )

      _lowercase							   =jit_generate(inputs['''input_ids''']	,  attention_mask=inputs['''attention_mask'''] ).sequences

      _lowercase							   =tokenizer.batch_decode(UpperCAmelCase	,  skip_special_tokens=UpperCAmelCase )

      _lowercase							   =[
          '''Hello this is a long string of words. I\'m going to start with the first one.\n''',
          '''Hey, I\'m not sure if I\'m going to be able to do''',
      ]

      self.assertListEqual(UpperCAmelCase	,  UpperCAmelCase )





     @is_pt_flax_cross_test
     def   __A      (self )  ->     Optional[Any]:
      _lowercase   ,					_lowercase							   =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							   =self._prepare_for_class(UpperCAmelCase	,  UpperCAmelCase )
        _lowercase							   ={k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

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

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

        _lowercase							   =convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,  UpperCAmelCase )
        _lowercase							   =fx_state

        with torch.no_grad():
         _lowercase							   =pt_model(**UpperCAmelCase ).to_tuple()

        _lowercase							   =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							   =model_class.from_pretrained(UpperCAmelCase	,  from_pt=UpperCAmelCase )

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

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

        _lowercase							   =pt_model_class(UpperCAmelCase ).eval()
        _lowercase							   =model_class(UpperCAmelCase	,  dtype=jnp.floataa )

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

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

        with torch.no_grad():
         _lowercase							   =pt_model(**UpperCAmelCase ).to_tuple()

        _lowercase							   =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							   =pt_model_class.from_pretrained(UpperCAmelCase	,  from_flax=UpperCAmelCase )

        with torch.no_grad():
         _lowercase							   =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   __A      (self )  ->     Union[str, Any]:
      for model_class_name in self.all_model_classes:
       _lowercase							   =model_class_name.from_pretrained('''gpt2'''	,  from_pt=UpperCAmelCase )
       _lowercase							   =model(np.ones((1, 1) ) )
       self.assertIsNotNone(UpperCAmelCase )

