
"""simple docstring"""


import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


_lowercase							      = datasets.utils.logging.get_logger(__name__)







@dataclass
class  lowerCAmelCase_ (					datasets.BuilderConfig			):



						'''simple docstring'''


						_lowerCamelCase:			Optional[datasets.Features]					=			None
						_lowerCamelCase:			str					=			"utf-8"
						_lowerCamelCase:			Optional[str]					=			None
						_lowerCamelCase:			Optional[str]					=			None
						_lowerCamelCase:			bool					=			True  # deprecated
						_lowerCamelCase:			Optional[int]					=			None  # deprecated
						_lowerCamelCase:			int					=			10 << 20  # 10MB
						_lowerCamelCase:			Optional[bool]					=			None







class  lowerCAmelCase_ (					datasets.ArrowBasedBuilder			):



						'''simple docstring'''


						_lowerCamelCase:			Optional[int]					=			JsonConfig


						def   _SCREAMING_SNAKE_CASE					(       self   :					Optional[int]					)   ->      List[Any]:
							if self.config.block_size is not None:
								logger.warning('The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead'					)
								A        =      self.config.block_size
							if self.config.use_threads is not True:
								logger.warning(
								    'The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.'					)
							if self.config.newlines_in_values is not None:
								raise ValueError('The JSON loader parameter `newlines_in_values` is no longer supported'					)
							return datasets.DatasetInfo(features=self.config.features					)


						def   _SCREAMING_SNAKE_CASE					(       self   :					Any     ,A_   :					List[str]					)   ->      List[Any]:
							if not self.config.data_files:
								raise ValueError(F'At least one data file must be specified, but got data_files={self.config.data_files}'					)
							A        =      dl_manager.download_and_extract(self.config.data_files					)
							if isinstance(A_     ,(str, list, tuple)					):
								A        =      data_files
								if isinstance(A_     ,A_					):
									A        =      [files]
								A        =      [dl_manager.iter_files(A_					) for file in files]
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,gen_kwargs={'files': files}					)]
							A        =      []
							for split_name, files in data_files.items():
								if isinstance(A_     ,A_					):
									A        =      [files]
								A        =      [dl_manager.iter_files(A_					) for file in files]
								splits.append(datasets.SplitGenerator(name=A_     ,gen_kwargs={'files': files}					)					)
							return splits


						def   _SCREAMING_SNAKE_CASE					(       self   :					str     ,A_   :					pa.Table					)   ->      pa.Table:
							if self.config.features is not None:
								# adding missing columns
								for column_name in set(self.config.features					) - set(pa_table.column_names					):
									A        =      self.config.features.arrow_schema.field(A_					).type
									A        =      pa_table.append_column(A_     ,pa.array([None] * len(A_					)     ,type=A_					)					)
								# more expensive cast to support nested structures with keys in a different order
								# allows str <-> int/float or str to Audio for example
								A        =      table_cast(A_     ,self.config.features.arrow_schema					)
							return pa_table




						def   _SCREAMING_SNAKE_CASE					(       self   :					Any     ,A_   :					Tuple					)   ->      List[str]:
							for file_idx, file in enumerate(itertools.chain.from_iterable(A_					)					):
								# If the file is one json object and if we need to look at the list of items in one specific field
								if self.config.field is not None:
									with open(A_     ,encoding=self.config.encoding     ,errors=self.config.encoding_errors					) as f:
										A        =      json.load(A_					)

									# We keep only the field we are interested in
									A        =      dataset[self.config.field]

									# We accept two format: a list of dicts or a dict of lists
									if isinstance(A_     ,(list, tuple)					):
										A        =      set().union(*[row.keys() for row in dataset]					)
										A        =      {col: [row.get(A_					) for row in dataset] for col in keys}
									else:
										A        =      dataset
									A        =      pa.Table.from_pydict(A_					)
									yield file_idx, self._cast_table(A_					)

								# If the file has one json object per line
								else:
									with open(A_     ,'rb'					) as f:
										A        =      0
										# Use block_size equal to the chunk size divided by 32 to leverage multithreading
										# Set a default minimum value of 16kB if the chunk size is really small
										A        =      max(self.config.chunksize // 32     ,16 << 10					)
										A        =      (
										    self.config.encoding_errors if self.config.encoding_errors is not None else 'strict'
										)
										while True:
											A        =      f.read(self.config.chunksize					)
											if not batch:
												break
											# Finish current line
											try:
												batch += f.readline()
											except (AttributeError, io.UnsupportedOperation):
												batch += readline(A_					)
											# PyArrow only accepts utf-8 encoded bytes
											if self.config.encoding != "utf-8":
												A        =      batch.decode(self.config.encoding     ,errors=A_					).encode('utf-8'					)
											try:
												while True:
													try:
														A        =      paj.read_json(
														    io.BytesIO(A_					)     ,read_options=paj.ReadOptions(block_size=A_					)					)
														break
													except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
														if (
														    isinstance(A_     ,pa.ArrowInvalid					)
														    and "straddling" not in str(A_					)
														    or block_size > len(A_					)
														):
															raise
														else:
															# Increase the block size in case it was too small.
															# The block size will be reset for the next file.
															logger.debug(
															    F'Batch of {len(A_					)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.'					)
															block_size *= 2
											except pa.ArrowInvalid as e:
												try:
													with open(
													    A_     ,encoding=self.config.encoding     ,errors=self.config.encoding_errors					) as f:
														A        =      json.load(A_					)
												except json.JSONDecodeError:
													logger.error(F'Failed to read file \'{file}\' with error {type(A_					)}: {e}'					)
													raise e
												# If possible, parse the file as a list of json objects and exit the loop
												if isinstance(A_     ,A_					):  # list is the only sequence type supported in JSON
													try:
														A        =      set().union(*[row.keys() for row in dataset]					)
														A        =      {col: [row.get(A_					) for row in dataset] for col in keys}
														A        =      pa.Table.from_pydict(A_					)
													except (pa.ArrowInvalid, AttributeError) as e:
														logger.error(F'Failed to read file \'{file}\' with error {type(A_					)}: {e}'					)
														raise ValueError(F'Not able to read records in the JSON file at {file}.'					) from None
													yield file_idx, self._cast_table(A_					)
													break
												else:
													logger.error(F'Failed to read file \'{file}\' with error {type(A_					)}: {e}'					)
													raise ValueError(
													    F'Not able to read records in the JSON file at {file}. '
													    F'You should probably indicate the field of the JSON file containing your records. '
													    F'This JSON file contain the following fields: {str(list(dataset.keys()					)					)}. '
													    F'Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. '					) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
											yield (file_idx, batch_idx), self._cast_table(A_					)
											batch_idx += 1


import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel







class 						__magic_name__			:

   '''simple docstring'''




   def __init__(				self,    lowercase_,    lowercase_=13,    lowercase_=7,    lowercase_=True,    lowercase_=True,    lowercase_=False,    lowercase_=True,    lowercase_=99,    lowercase_=32,    lowercase_=5,    lowercase_=4,    lowercase_=37,    lowercase_="gelu",    lowercase_=0.1,    lowercase_=0.1,    lowercase_=512,    lowercase_=16,    lowercase_=2,    lowercase_=0.02,    lowercase_=3,    lowercase_=4,    lowercase_=None,    )							->    List[Any]:






     """simple docstring"""






     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__  =intermediate_size
     a__  =hidden_act
     a__  =hidden_dropout_prob
     a__  =attention_probs_dropout_prob
     a__  =max_position_embeddings
     a__  =type_vocab_size
     a__  =type_sequence_label_size
     a__  =initializer_range
     a__  =num_labels
     a__  =num_choices
     a__  =scope



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






     """simple docstring"""






     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__  =None
     if self.use_token_type_ids:
       a__  =ids_tensor([self.batch_size, self.seq_length],    self.type_vocab_size       )

     a__  =None
     a__  =None
     a__  =None
     if self.use_labels:
       a__  =ids_tensor([self.batch_size],    self.type_sequence_label_size       )
       a__  =ids_tensor([self.batch_size, self.seq_length],    self.num_labels       )
       a__  =ids_tensor([self.batch_size],    self.num_choices       )

     a__  =self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






     return OpenLlamaConfig(
         vocab_size=self.vocab_size,    hidden_size=self.hidden_size,    num_hidden_layers=self.num_hidden_layers,    num_attention_heads=self.num_attention_heads,    intermediate_size=self.intermediate_size,    hidden_act=self.hidden_act,    hidden_dropout_prob=self.hidden_dropout_prob,    attention_probs_dropout_prob=self.attention_probs_dropout_prob,    max_position_embeddings=self.max_position_embeddings,    type_vocab_size=self.type_vocab_size,    is_decoder=lowercase_,    initializer_range=self.initializer_range,    use_stable_embedding=lowercase_,    )



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    List[str]:






     """simple docstring"""






     a__  =OpenLlamaModel(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    attention_mask=lowercase_       )
     a__  =model(lowercase_       )
     self.parent.assertEqual(result.last_hidden_state.shape,    (self.batch_size, self.seq_length, self.hidden_size)       )



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    )							->    Any:






     """simple docstring"""






     a__  =True
     a__  =OpenLlamaModel(lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    encoder_hidden_states=lowercase_,    encoder_attention_mask=lowercase_,    )
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    encoder_hidden_states=lowercase_,    )
     a__  =model(lowercase_,    attention_mask=lowercase_       )
     self.parent.assertEqual(result.last_hidden_state.shape,    (self.batch_size, self.seq_length, self.hidden_size)       )



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    )							->    List[str]:






     """simple docstring"""






     a__  =OpenLlamaForCausalLM(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    attention_mask=lowercase_,    labels=lowercase_       )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.vocab_size)       )



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    )							->    List[Any]:






     """simple docstring"""






     a__  =True
     a__  =True
     a__  =OpenLlamaForCausalLM(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()

     # first forward pass
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    encoder_hidden_states=lowercase_,    encoder_attention_mask=lowercase_,    use_cache=lowercase_,    )
     a__  =outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     a__  =ids_tensor((self.batch_size, 3),    config.vocab_size       )
     a__  =ids_tensor((self.batch_size, 3),    vocab_size=2       )

     # append to next input_ids and
     a__  =torch.cat([input_ids, next_tokens],    dim=-1       )
     a__  =torch.cat([input_mask, next_mask],    dim=-1       )

     a__  =model(
         lowercase_,    attention_mask=lowercase_,    encoder_hidden_states=lowercase_,    encoder_attention_mask=lowercase_,    output_hidden_states=lowercase_,    )['''hidden_states'''][0]
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    encoder_hidden_states=lowercase_,    encoder_attention_mask=lowercase_,    past_key_values=lowercase_,    output_hidden_states=lowercase_,    )['''hidden_states'''][0]

     # select random slice
     a__  =ids_tensor((1,),    output_from_past.shape[-1]       ).item()
     a__  =output_from_no_past[:, -3:, random_slice_idx].detach()
     a__  =output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]       )

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(lowercase_,    lowercase_,    atol=1E-3       )       )



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






     """simple docstring"""






     a__  =self.prepare_config_and_inputs()
     (
         (
         a__
     ),      (
         a__
     ),      (
         a__
     ),      (
         a__
     ),      (
         a__
     ),      (
         a__
     ),      (
         a__
     ),      
     )  =config_and_inputs
     a__  ={'''input_ids''': input_ids, '''attention_mask''': input_mask}
     return config, inputs_dict







@require_torch
class 						__magic_name__			(							lowerCamelCase__		,       lowerCamelCase__		,       lowerCamelCase__		,       unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__ :	Optional[Any]    =   (
       (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
   )
   lowerCamelCase__ :	Optional[Any]    =   (OpenLlamaForCausalLM,) if is_torch_available() else ()
   lowerCamelCase__ :	List[str]    =   (
       {
           'feature-extraction': OpenLlamaModel,
           'text-classification': OpenLlamaForSequenceClassification,
           'text-generation': OpenLlamaForCausalLM,
           'zero-shot': OpenLlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   lowerCamelCase__ :	int    =   False
   lowerCamelCase__ :	Any    =   False



   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






     a__  =OpenLlamaModelTester(self       )
     a__  =ConfigTester(self,    config_class=lowercase_,    hidden_size=37       )



   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






     self.config_tester.run_common_tests()



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






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       a__  =type
       self.model_tester.create_and_check_model(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     a__,      a__  =self.model_tester.prepare_config_and_inputs_for_common()
     a__  =3
     a__  =input_dict['''input_ids''']
     a__  =input_ids.ne(1       ).to(lowercase_       )
     a__  =ids_tensor([self.model_tester.batch_size],    self.model_tester.type_sequence_label_size       )
     a__  =OpenLlamaForSequenceClassification(lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    attention_mask=lowercase_,    labels=lowercase_       )
     self.assertEqual(result.logits.shape,    (self.model_tester.batch_size, self.model_tester.num_labels)       )



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     a__,      a__  =self.model_tester.prepare_config_and_inputs_for_common()
     a__  =3
     a__  ='''single_label_classification'''
     a__  =input_dict['''input_ids''']
     a__  =input_ids.ne(1       ).to(lowercase_       )
     a__  =ids_tensor([self.model_tester.batch_size],    self.model_tester.type_sequence_label_size       )
     a__  =OpenLlamaForSequenceClassification(lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    attention_mask=lowercase_,    labels=lowercase_       )
     self.assertEqual(result.logits.shape,    (self.model_tester.batch_size, self.model_tester.num_labels)       )



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






     """simple docstring"""






     a__,      a__  =self.model_tester.prepare_config_and_inputs_for_common()
     a__  =3
     a__  ='''multi_label_classification'''
     a__  =input_dict['''input_ids''']
     a__  =input_ids.ne(1       ).to(lowercase_       )
     a__  =ids_tensor(
         [self.model_tester.batch_size, config.num_labels],    self.model_tester.type_sequence_label_size       ).to(torch.float       )
     a__  =OpenLlamaForSequenceClassification(lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(lowercase_,    attention_mask=lowercase_,    labels=lowercase_       )
     self.assertEqual(result.logits.shape,    (self.model_tester.batch_size, self.model_tester.num_labels)       )



   @unittest.skip('''Open-Llama buffers include complex numbers, which breaks this test'''       )
   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






     pass



   @parameterized.expand([('''linear''',), ('''dynamic''',)]       )
   def  _UpperCAmelCase (				self,    lowercase_       )							->    Optional[Any]:






     """simple docstring"""






     a__,      a__  =self.model_tester.prepare_config_and_inputs_for_common()
     a__  =ids_tensor([1, 10],    config.vocab_size       )
     a__  =ids_tensor([1, int(config.max_position_embeddings * 1.5       )],    config.vocab_size       )

     set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
     a__  =OpenLlamaModel(lowercase_       )
     original_model.to(lowercase_       )
     original_model.eval()
     a__  =original_model(lowercase_       ).last_hidden_state
     a__  =original_model(lowercase_       ).last_hidden_state

     set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
     a__  ={'''type''': scaling_type, '''factor''': 10.0}
     a__  =OpenLlamaModel(lowercase_       )
     scaled_model.to(lowercase_       )
     scaled_model.eval()
     a__  =scaled_model(lowercase_       ).last_hidden_state
     a__  =scaled_model(lowercase_       ).last_hidden_state

     # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
     # maximum sequence length, so the outputs for the short input should match.
     if scaling_type == "dynamic":
       self.assertTrue(torch.allclose(lowercase_,    lowercase_,    atol=1E-5       )       )
     else:
       self.assertFalse(torch.allclose(lowercase_,    lowercase_,    atol=1E-5       )       )

     # The output should be different for long inputs
     self.assertFalse(torch.allclose(lowercase_,    lowercase_,    atol=1E-5       )       )
