






"""simple docstring"""


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 __a   :



  def __init__(self  :		int				,			lowercase_  :		Tuple				,			lowercase_  :		Union[str, Any]=14				,			lowercase_  :		Optional[int]=7				,			lowercase_  :		Any=True				,			lowercase_  :		Optional[Any]=True				,			lowercase_  :		Any=False				,			lowercase_  :		Dict=True				,			lowercase_  :		int=99				,			lowercase_  :		int=32				,			lowercase_  :		Optional[Any]=5				,			lowercase_  :		int=4				,			lowercase_  :		Tuple=37				,			lowercase_  :		Optional[Any]="gelu"				,			lowercase_  :		List[Any]=0.1				,			lowercase_  :		str=0.1				,			lowercase_  :		str=512				,			lowercase_  :		Dict=0.02				,			)		->							str:
   __A	:       str       =	parent
   __A	:       str       =	batch_size
   __A	:       Dict       =	seq_length
   __A	:       Tuple       =	is_training
   __A	:       int       =	use_input_mask
   __A	:       Optional[Any]       =	use_token_type_ids
   __A	:       Union[str, Any]       =	use_labels
   __A	:       List[str]       =	vocab_size
   __A	:       Union[str, Any]       =	hidden_size
   __A	:       Optional[Any]       =	num_hidden_layers
   __A	:       Optional[int]       =	num_attention_heads
   __A	:       List[str]       =	intermediate_size
   __A	:       List[str]       =	hidden_act
   __A	:       int       =	hidden_dropout_prob
   __A	:       Dict       =	attention_probs_dropout_prob
   __A	:       Tuple       =	max_position_embeddings
   __A	:       Optional[int]       =	initializer_range
   __A	:       int       =	None
   __A	:       Any       =	vocab_size - 1
   __A	:       List[Any]       =	vocab_size - 1
   __A	:       Any       =	vocab_size - 1



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

   __A	:       Tuple       =	None
   if self.use_input_mask:
    __A	:       int       =	random_attention_mask([self.batch_size, self.seq_length]   )

   __A	:       List[str]       =	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=_lowerCamelCase				,			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  :		Any   )		->							List[Any]:
   __A	:       Dict       =	self.prepare_config_and_inputs()
   __A	:       Dict       =	config_and_inputs
   __A	:       Tuple       =	{"""input_ids""": input_ids, """attention_mask""": attention_mask}
   return config, inputs_dict



  def 						A_				(self  :		str   )		->							Optional[int]:
   __A	:       Optional[int]       =	self.prepare_config_and_inputs()

   __A	:       Optional[int]       =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
   __A	:       List[str]       =	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  :		List[Any]				,			lowercase_  :		Optional[int]				,			lowercase_  :		Any				,			lowercase_  :		Any				,			lowercase_  :		Dict   )		->							Optional[Any]:
   __A	:       List[str]       =	20
   __A	:       Any       =	model_class_name(_lowerCamelCase   )

   __A	:       int       =	model.init_cache(input_ids.shape[0]				,			_lowerCamelCase   )
   __A	:       Optional[Any]       =	jnp.ones((input_ids.shape[0], max_decoder_length)				,			dtype='i4'   )

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

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

   __A	:       List[str]       =	model(_lowerCamelCase   )

   __A	:       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 						A_				(self  :		int				,			lowercase_  :		Any				,			lowercase_  :		Any				,			lowercase_  :		Union[str, Any]				,			lowercase_  :		Any   )		->							int:
   __A	:       str       =	20
   __A	:       Tuple       =	model_class_name(_lowerCamelCase   )

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

   __A	:       Tuple       =	model.init_cache(input_ids.shape[0]				,			_lowerCamelCase   )
   __A	:       List[Any]       =	jnp.broadcast_to(
       jnp.arange(input_ids.shape[-1] - 1   )[None, :]				,			(input_ids.shape[0], input_ids.shape[-1] - 1)   )

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

   __A	:       Optional[Any]       =	model(_lowerCamelCase				,			attention_mask=_lowerCamelCase   )

   __A	:       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}"   )






@require_flax
class __a   (   lowerCAmelCase__ ,lowerCAmelCase__ ,unittest.TestCase				):
  __lowercase    :    int						=					(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
  __lowercase    :    List[str]						=					(FlaxGPTaLMHeadModel,) if is_flax_available() else ()



  def 						A_				(self  :		List[Any]   )		->							Optional[int]:
   __A	:       int       =	FlaxGPTaModelTester(self   )



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

   for model_class_name in self.all_model_classes:
    __A	:       Optional[Any]       =	self.model_tester.prepare_config_and_inputs()
    self.model_tester.check_use_cache_forward(_lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase   )



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

   for model_class_name in self.all_model_classes:
    __A	:       Tuple       =	self.model_tester.prepare_config_and_inputs()
    self.model_tester.check_use_cache_forward_with_attn_mask(
        _lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase				,			_lowerCamelCase   )



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

   __A	:       List[str]       =	FlaxGPTaLMHeadModel.from_pretrained('gpt2'   )
   __A	:       Optional[int]       =	False
   __A	:       List[Any]       =	model.config.eos_token_id

   __A	:       Optional[int]       =	jax.jit(model.generate   )

   __A	:       Any       =	jit_generate(inputs['input_ids']				,			attention_mask=inputs['attention_mask']   ).sequences

   __A	:       int       =	tokenizer.batch_decode(_lowerCamelCase				,			skip_special_tokens=_lowerCamelCase   )

   __A	:       str       =	[
       """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(_lowerCamelCase				,			_lowerCamelCase   )



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

     # load corresponding PyTorch class
     __A	:       Dict       =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
     __A	:       List[str]       =	getattr(_lowerCamelCase				,			_lowerCamelCase   )

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

     __A	:       str       =	convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,			_lowerCamelCase   )
     __A	:       List[str]       =	fx_state

     with torch.no_grad():
      __A	:       str       =	pt_model(**_lowerCamelCase   ).to_tuple()

     __A	:       Tuple       =	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   )
      __A	:       Optional[int]       =	model_class.from_pretrained(_lowerCamelCase				,			from_pt=_lowerCamelCase   )

     __A	:       Tuple       =	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 						A_				(self  :		Tuple   )		->							Any:
   __A	:       int       =	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	:       str       =	self._prepare_for_class(_lowerCamelCase				,			_lowerCamelCase   )
     __A	:       Dict       =	{k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

     # load corresponding PyTorch class
     __A	:       Optional[Any]       =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
     __A	:       Dict       =	getattr(_lowerCamelCase				,			_lowerCamelCase   )

     __A	:       str       =	pt_model_class(_lowerCamelCase   ).eval()
     __A	:       Dict       =	model_class(_lowerCamelCase				,			dtype=jnp.floataa   )

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

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

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

     __A	:       Union[str, Any]       =	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   )
      __A	:       Optional[int]       =	pt_model_class.from_pretrained(_lowerCamelCase				,			from_flax=_lowerCamelCase   )

     with torch.no_grad():
      __A	:       int       =	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   )



  @slow
  def 						A_				(self  :		List[Any]   )		->							Union[str, Any]:

   for model_class_name in self.all_model_classes:
    __A	:       Optional[int]       =	model_class_name.from_pretrained('gpt2'				,			from_pt=_lowerCamelCase   )
    __A	:       List[Any]       =	model(np.ones((1, 1)   )   )
    self.assertIsNotNone(_lowerCamelCase   )













"""simple docstring"""


import csv
import json
import os

import datasets


# TODO: Add BibTeX citation
# Find for instance the citation on arxiv or on the dataset repo/website
_UpperCamelCase   :						List[Any]     						= """\
@InProceedings{huggingface:dataset,
title = {A great new dataset},
author={huggingface, Inc.
},
year={2020}
}
"""

# TODO: Add description of the dataset here
# You can copy an official description
_UpperCamelCase   :						Any     						= """\
This new dataset is designed to solve this great NLP task and is crafted with a lot of care.
"""

# TODO: Add a link to an official homepage for the dataset here
_UpperCamelCase   :						Dict     						= """"""

# TODO: Add the licence for the dataset here if you can find it
_UpperCamelCase   :						Tuple     						= """"""

# TODO: Add link to the official dataset URLs here
# The HuggingFace Datasets library doesn't host the datasets but only points to the original files.
# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)
_UpperCamelCase   :						str     						= {
    """first_domain""": """https://huggingface.co/great-new-dataset-first_domain.zip""",
    """second_domain""": """https://huggingface.co/great-new-dataset-second_domain.zip""",
}






class __a   (   datasets.GeneratorBasedBuilder				):
  __lowercase    :    int						=					datasets.Version('1.1.0'				)

  # This is an example of a dataset with multiple configurations.
  # If you don't want/need to define several sub-sets in your dataset,
  # just remove the BUILDER_CONFIG_CLASS and the BUILDER_CONFIGS attributes.

  # If you need to make complex sub-parts in the datasets with configurable options
  # You can create your own builder configuration class to store attribute, inheriting from datasets.BuilderConfig
  # BUILDER_CONFIG_CLASS = MyBuilderConfig

  # You will be able to load one or the other configurations in the following list with
  # data = datasets.load_dataset('my_dataset', 'first_domain')
  # data = datasets.load_dataset('my_dataset', 'second_domain')
  __lowercase    :    Optional[Any]						=					[
      datasets.BuilderConfig(name='first_domain' ,version=snake_case_ ,description='This part of my dataset covers a first domain'				),
      datasets.BuilderConfig(name='second_domain' ,version=snake_case_ ,description='This part of my dataset covers a second domain'				),
  ]

  __lowercase    :    Any						=					'first_domain'  # It's not mandatory to have a default configuration. Just use one if it make sense.



  def 						A_				(self  :		Optional[Any]   )		->							List[Any]:
   # TODO: This method specifies the datasets.DatasetInfo object which contains informations and typings for the dataset
   if self.config.name == "first_domain":  # This is the name of the configuration selected in BUILDER_CONFIGS above
    __A	:       Dict       =	datasets.Features(
        {
            'sentence': datasets.Value('string'   ),
            'option1': datasets.Value('string'   ),
            'answer': datasets.Value('string'   )
            # These are the features of your dataset like images, labels ...
        }   )
   else:  # This is an example to show how to have different features for "first_domain" and "second_domain"
    __A	:       Optional[int]       =	datasets.Features(
        {
            'sentence': datasets.Value('string'   ),
            'option2': datasets.Value('string'   ),
            'second_domain_answer': datasets.Value('string'   )
            # These are the features of your dataset like images, labels ...
        }   )
   return datasets.DatasetInfo(
       # This is the description that will appear on the datasets page.
       description=_DESCRIPTION				,			features=lowercase_				,			homepage=_HOMEPAGE				,			license=_LICENSE				,			citation=_CITATION				,			)



  def 						A_				(self  :		Dict				,			lowercase_  :		Union[str, Any]   )		->							Union[str, Any]:
   # TODO: This method is tasked with downloading/extracting the data and defining the splits depending on the configuration
   # If several configurations are possible (listed in BUILDER_CONFIGS), the configuration selected by the user is in self.config.name

   # dl_manager is a datasets.download.DownloadManager that can be used to download and extract URLS
   # It can accept any type or nested list/dict and will give back the same structure with the url replaced with path to local files.
   # By default the archives will be extracted and a path to a cached folder where they are extracted is returned instead of the archive
   __A	:       Optional[int]       =	_URLS[self.config.name]
   __A	:       Any       =	dl_manager.download_and_extract(lowercase_   )
   return [
       datasets.SplitGenerator(
           name=datasets.Split.TRAIN				,			gen_kwargs={
               'filepath': os.path.join(lowercase_				,			'train.jsonl'   ),
               'split': 'train',
           }				,			),
       datasets.SplitGenerator(
           name=datasets.Split.VALIDATION				,			gen_kwargs={
               'filepath': os.path.join(lowercase_				,			'dev.jsonl'   ),
               'split': 'dev',
           }				,			),
       datasets.SplitGenerator(
           name=datasets.Split.TEST				,			gen_kwargs={
               'filepath': os.path.join(lowercase_				,			'test.jsonl'   ),
               'split': 'test'
           }				,			),
   ]



  def 						A_				(self  :		Tuple				,			lowercase_  :		Dict				,			lowercase_  :		Optional[Any]   )		->							Tuple:
   # TODO: This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
   # The `key` is for legacy reasons (tfds) and is not important in itself, but must be unique for each example.
   with open(lowercase_				,			encoding='utf-8'   ) as f:
    for key, row in enumerate(lowercase_   ):
     __A	:       int       =	json.loads(lowercase_   )
     if self.config.name == "first_domain":
      # Yields examples as (key, example) tuples
      yield key, {
          "sentence": data["sentence"],
          "option1": data["option1"],
          "answer": "" if split == "test" else data["answer"],
      }
     else:
      yield key, {
          "sentence": data["sentence"],
          "option2": data["option2"],
          "second_domain_answer": "" if split == "test" else data["second_domain_answer"],
      }





