'''simple docstring'''




import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Value
from datasets.packaged_modules.json.json import Json





@pytest.fixture
def      _UpperCAmelCase    (							_UpperCamelCase    :	Union[str, Any]   )	->      Dict:
    A_     =   tmp_path / '''file.jsonl'''
    A_     =   textwrap.dedent(
        '''\
        {"col_1": -1}
        {"col_1": 1, "col_2": 2}
        {"col_1": 10, "col_2": 20}
        '''   )
    with open(_UpperCamelCase,   '''w'''   ) as f:
        f.write(_UpperCamelCase   )
    return str(_UpperCamelCase   )





@pytest.fixture
def      _UpperCAmelCase    (							_UpperCamelCase    :	Any   )	->      Optional[Any]:
    A_     =   tmp_path / '''file_utf16_encoded.jsonl'''
    A_     =   textwrap.dedent(
        '''\
        {"col_1": -1}
        {"col_1": 1, "col_2": 2}
        {"col_1": 10, "col_2": 20}
        '''   )
    with open(_UpperCamelCase,   '''w''',   encoding='''utf-16'''   ) as f:
        f.write(_UpperCamelCase   )
    return str(_UpperCamelCase   )





@pytest.fixture
def      _UpperCAmelCase    (							_UpperCamelCase    :	List[str]   )	->      Dict:
    A_     =   tmp_path / '''file_with_list_of_dicts.json'''
    A_     =   textwrap.dedent(
        '''\
        [
            {"col_1": -1},
            {"col_1": 1, "col_2": 2},
            {"col_1": 10, "col_2": 20}
        ]
        '''   )
    with open(_UpperCamelCase,   '''w'''   ) as f:
        f.write(_UpperCamelCase   )
    return str(_UpperCamelCase   )





@pytest.fixture
def      _UpperCAmelCase    (							_UpperCamelCase    :	int   )	->      int:
    A_     =   tmp_path / '''file_with_list_of_dicts_field.json'''
    A_     =   textwrap.dedent(
        '''\
        {
            "field1": 1,
            "field2": "aabb",
            "field3": [
                {"col_1": -1},
                {"col_1": 1, "col_2": 2},
                {"col_1": 10, "col_2": 20}
            ]
        }
        '''   )
    with open(_UpperCamelCase,   '''w'''   ) as f:
        f.write(_UpperCamelCase   )
    return str(_UpperCamelCase   )





@pytest.mark.parametrize(
    '''file_fixture, config_kwargs''',   [
        ('''jsonl_file''', {}),
        ('''jsonl_file_utf16_encoded''', {'''encoding''': '''utf-16'''}),
        ('''json_file_with_list_of_dicts''', {}),
        ('''json_file_with_list_of_dicts_field''', {'''field''': '''field3'''}),
    ],   )
def      _UpperCAmelCase    (							_UpperCamelCase    :	List[str],   _UpperCamelCase    :	List[Any],   _UpperCamelCase    :	List[str]   )	->      Dict:
    A_     =   Json(**_UpperCamelCase   )
    A_     =   json._generate_tables([[request.getfixturevalue(_UpperCamelCase   )]]   )
    A_     =   pa.concat_tables([table for _, table in generator]   )
    assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20]}





@pytest.mark.parametrize(
    '''file_fixture, config_kwargs''',   [
        (
            '''jsonl_file''',
            {'''features''': Features({'''col_1''': Value('''int64'''   ), '''col_2''': Value('''int64'''   ), '''missing_col''': Value('''string'''   )}   )},
        ),
        (
            '''json_file_with_list_of_dicts''',
            {'''features''': Features({'''col_1''': Value('''int64'''   ), '''col_2''': Value('''int64'''   ), '''missing_col''': Value('''string'''   )}   )},
        ),
        (
            '''json_file_with_list_of_dicts_field''',
            {
                '''field''': '''field3''',
                '''features''': Features(
                    {'''col_1''': Value('''int64'''   ), '''col_2''': Value('''int64'''   ), '''missing_col''': Value('''string'''   )}   ),
            },
        ),
    ],   )
def      _UpperCAmelCase    (							_UpperCamelCase    :	Optional[Any],   _UpperCamelCase    :	Optional[int],   _UpperCamelCase    :	Any   )	->      List[Any]:
    A_     =   Json(**_UpperCamelCase   )
    A_     =   json._generate_tables([[request.getfixturevalue(_UpperCamelCase   )]]   )
    A_     =   pa.concat_tables([table for _, table in generator]   )
    assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20], "missing_col": [None, None, None]}







'''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 	__UpperCAmelCase				:





       '''simple docstring'''



       def __init__(      self    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE=14    ,     _SCREAMING_SNAKE_CASE=7    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=False    ,     _SCREAMING_SNAKE_CASE=True    ,     _SCREAMING_SNAKE_CASE=99    ,     _SCREAMING_SNAKE_CASE=32    ,     _SCREAMING_SNAKE_CASE=4    ,     _SCREAMING_SNAKE_CASE=4    ,     _SCREAMING_SNAKE_CASE=[[["global", "local"], 2]]    ,     _SCREAMING_SNAKE_CASE=37    ,     _SCREAMING_SNAKE_CASE="gelu"    ,     _SCREAMING_SNAKE_CASE=0.1    ,     _SCREAMING_SNAKE_CASE=0.1    ,     _SCREAMING_SNAKE_CASE=512    ,     _SCREAMING_SNAKE_CASE=7    ,     _SCREAMING_SNAKE_CASE=0.02    ,     )		->     Dict:
           A_     =   parent
           A_     =   batch_size
           A_     =   seq_length
           A_     =   is_training
           A_     =   use_input_mask
           A_     =   use_token_type_ids
           A_     =   use_labels
           A_     =   vocab_size
           A_     =   hidden_size
           A_     =   num_hidden_layers
           A_     =   num_attention_heads
           A_     =   attention_types
           A_     =   intermediate_size
           A_     =   hidden_act
           A_     =   hidden_dropout_prob
           A_     =   attention_probs_dropout_prob
           A_     =   max_position_embeddings
           A_     =   window_size
           A_     =   initializer_range
           A_     =   None
           A_     =   vocab_size - 1
           A_     =   vocab_size - 1
           A_     =   vocab_size - 1

       def 					__A       (      self )		->     Optional[int]:
           A_     =   ids_tensor([self.batch_size, self.seq_length]    ,     self.vocab_size )

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

           A_     =   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=_SCREAMING_SNAKE_CASE    ,     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 					__A       (      self )		->     str:
           A_     =   self.prepare_config_and_inputs()
           A_   ,A_   ,A_     =   config_and_inputs
           A_     =   {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
           return config, inputs_dict

       def 					__A       (      self    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )		->     Union[str, Any]:
           A_     =   20
           A_     =   model_class_name(_SCREAMING_SNAKE_CASE )

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

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

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

           A_     =   model(_SCREAMING_SNAKE_CASE )

           A_     =   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    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )		->     List[Any]:
           A_     =   20
           A_     =   model_class_name(_SCREAMING_SNAKE_CASE )

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

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

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

           A_     =   model(_SCREAMING_SNAKE_CASE    ,     attention_mask=_SCREAMING_SNAKE_CASE )

           A_     =   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				(				_UpperCamelCase						,      _UpperCamelCase						,      unittest.TestCase     ):





       '''simple docstring'''



       __lowercase					:					List[str]    =				(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
       __lowercase					:					Dict    =				(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()

       def 					__A       (      self )		->     Any:
           A_     =   FlaxGPTNeoModelTester(self )

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

       def 					__A       (      self )		->     Union[str, Any]:
           for model_class_name in self.all_model_classes:
               A_   ,A_   ,A_     =   self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward_with_attn_mask(
                   _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

       @slow
       def 					__A       (      self )		->     Optional[Any]:
           A_     =   GPTaTokenizer.from_pretrained('''gpt2'''    ,     pad_token='''<|endoftext|>'''    ,     padding_side='''left''' )
           A_     =   tokenizer(['''Hello this is a long string''', '''Hey''']    ,     return_tensors='''np'''    ,     padding=_SCREAMING_SNAKE_CASE    ,     truncation=_SCREAMING_SNAKE_CASE )

           A_     =   FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M''' )
           A_     =   False
           A_     =   model.config.eos_token_id

           A_     =   jax.jit(model.generate )

           A_     =   jit_generate(
               inputs['''input_ids''']    ,     attention_mask=inputs['''attention_mask''']    ,     pad_token_id=tokenizer.pad_token_id ).sequences

           A_     =   tokenizer.batch_decode(_SCREAMING_SNAKE_CASE    ,     skip_special_tokens=_SCREAMING_SNAKE_CASE )

           A_     =   [
               '''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(_SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE )

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

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

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

                   A_     =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,     _SCREAMING_SNAKE_CASE )
                   A_     =   fx_state

                   with torch.no_grad():
                       A_     =   pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

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

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

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

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

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

                   A_     =   pt_model_class(_SCREAMING_SNAKE_CASE ).eval()
                   A_     =   model_class(_SCREAMING_SNAKE_CASE    ,     dtype=jnp.floataa )

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

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

                   with torch.no_grad():
                       A_     =   pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

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

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

                   with torch.no_grad():
                       A_     =   pt_model_loaded(**_SCREAMING_SNAKE_CASE ).to_tuple()

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

       @slow
       def 					__A       (      self )		->     List[str]:
           for model_class_name in self.all_model_classes:
               A_     =   model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M''' )
               A_     =   model(np.ones((1, 1) ) )
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE )






