from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


SCREAMING_SNAKE_CASE			:						Optional[Any]				  =		logging.get_logger(__name__)

SCREAMING_SNAKE_CASE			:						List[str]				  =		{
    """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json""",
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}
class  A_			(							__snake_case			):
  _SCREAMING_SNAKE_CASE     					=      """blenderbot-small"""
  _SCREAMING_SNAKE_CASE     					=      ["""past_key_values"""]
  _SCREAMING_SNAKE_CASE     					=      {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}

  def __init__(		self			:   Optional[Any]  ,				__SCREAMING_SNAKE_CASE			:   Optional[int]=5_02_65  ,				__SCREAMING_SNAKE_CASE			:   Dict=5_12  ,				__SCREAMING_SNAKE_CASE			:   Any=8  ,				__SCREAMING_SNAKE_CASE			:   Dict=20_48  ,				__SCREAMING_SNAKE_CASE			:   Tuple=16  ,				__SCREAMING_SNAKE_CASE			:   List[str]=8  ,				__SCREAMING_SNAKE_CASE			:   str=20_48  ,				__SCREAMING_SNAKE_CASE			:   int=16  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]=0.0  ,				__SCREAMING_SNAKE_CASE			:   Any=0.0  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]=True  ,				__SCREAMING_SNAKE_CASE			:   Optional[int]=True  ,				__SCREAMING_SNAKE_CASE			:   List[Any]="gelu"  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]=5_12  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]=0.1  ,				__SCREAMING_SNAKE_CASE			:   List[str]=0.0  ,				__SCREAMING_SNAKE_CASE			:   Tuple=0.0  ,				__SCREAMING_SNAKE_CASE			:   Dict=0.02  ,				__SCREAMING_SNAKE_CASE			:   List[Any]=1  ,				__SCREAMING_SNAKE_CASE			:   Any=False  ,				__SCREAMING_SNAKE_CASE			:   List[str]=0  ,				__SCREAMING_SNAKE_CASE			:   Optional[Any]=1  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]=2  ,				__SCREAMING_SNAKE_CASE			:   List[Any]=2  ,				**__SCREAMING_SNAKE_CASE			:   Any  ,				):
   __a							=	vocab_size
   __a							=	max_position_embeddings
   __a							=	d_model
   __a							=	encoder_ffn_dim
   __a							=	encoder_layers
   __a							=	encoder_attention_heads
   __a							=	decoder_ffn_dim
   __a							=	decoder_layers
   __a							=	decoder_attention_heads
   __a							=	dropout
   __a							=	attention_dropout
   __a							=	activation_dropout
   __a							=	activation_function
   __a							=	init_std
   __a							=	encoder_layerdrop
   __a							=	decoder_layerdrop
   __a							=	use_cache
   __a							=	encoder_layers
   __a							=	scale_embedding  # scale factor will be sqrt(d_model) if True

   super().__init__(
       pad_token_id=__SCREAMING_SNAKE_CASE  ,				bos_token_id=__SCREAMING_SNAKE_CASE  ,				eos_token_id=__SCREAMING_SNAKE_CASE  ,				is_encoder_decoder=__SCREAMING_SNAKE_CASE  ,				decoder_start_token_id=__SCREAMING_SNAKE_CASE  ,				forced_eos_token_id=__SCREAMING_SNAKE_CASE  ,				**__SCREAMING_SNAKE_CASE  ,				)
class  A_			(							__snake_case			):

  @property
  def 		_UpperCAmelCase (		self			:   List[Any]							):
   if self.task in ["default", "seq2seq-lm"]:
    __a							=	OrderedDict(
        [
            ("input_ids", {0: "batch", 1: "encoder_sequence"}),
            ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
        ]							)

    if self.use_past:
     __a							=	{0: "batch"}
     __a							=	{0: "batch", 1: "past_decoder_sequence + sequence"}
    else:
     __a							=	{0: "batch", 1: "decoder_sequence"}
     __a							=	{0: "batch", 1: "decoder_sequence"}

    if self.use_past:
     self.fill_with_past_key_values_(__SCREAMING_SNAKE_CASE  ,				direction="inputs"							)
   elif self.task == "causal-lm":
    # TODO: figure this case out.
    __a							=	OrderedDict(
        [
            ("input_ids", {0: "batch", 1: "encoder_sequence"}),
            ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
        ]							)
    if self.use_past:
     __a							=	self.num_layers
     for i in range(__SCREAMING_SNAKE_CASE							):
      __a							=	{0: "batch", 2: "past_sequence + sequence"}
      __a							=	{0: "batch", 2: "past_sequence + sequence"}
   else:
    __a							=	OrderedDict(
        [
            ("input_ids", {0: "batch", 1: "encoder_sequence"}),
            ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
            ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
            ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
        ]							)

   return common_inputs

  @property
  def 		_UpperCAmelCase (		self			:   int							):
   if self.task in ["default", "seq2seq-lm"]:
    __a							=	super().outputs
   else:
    __a							=	super(__SCREAMING_SNAKE_CASE  ,				self							).outputs
    if self.use_past:
     __a							=	self.num_layers
     for i in range(__SCREAMING_SNAKE_CASE							):
      __a							=	{0: "batch", 2: "past_sequence + sequence"}
      __a							=	{0: "batch", 2: "past_sequence + sequence"}
   return common_outputs

  def 		_UpperCAmelCase (		self			:   Union[str, Any]  ,				__SCREAMING_SNAKE_CASE			:   PreTrainedTokenizer  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   bool = False  ,				__SCREAMING_SNAKE_CASE			:   Optional[TensorType] = None  ,				):
   __a							=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
       __SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)

   # Generate decoder inputs
   __a							=	seq_length if not self.use_past else 1
   __a							=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
       __SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)
   __a							=	{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
   __a							=	dict(**__SCREAMING_SNAKE_CASE  ,				**__SCREAMING_SNAKE_CASE							)

   if self.use_past:
    if not is_torch_available():
     raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."							)
    else:
     import torch
    __a							=	common_inputs["input_ids"].shape
    __a							=	common_inputs["decoder_input_ids"].shape[1]
    __a							=	self.num_attention_heads
    __a							=	(
        batch,
        num_encoder_attention_heads,
        encoder_seq_length,
        self._config.hidden_size // num_encoder_attention_heads,
    )
    __a							=	decoder_seq_length + 3
    __a							=	(
        batch,
        num_decoder_attention_heads,
        decoder_past_length,
        self._config.hidden_size // num_decoder_attention_heads,
    )

    __a							=	torch.cat(
        [common_inputs["decoder_attention_mask"], torch.ones(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)]  ,				dim=1							)

    __a							=	[]
    # If the number of encoder and decoder layers are present in the model configuration, both are considered
    __a							=	self.num_layers
    __a							=	min(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)
    __a							=	max(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							) - min_num_layers
    __a							=	"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

    for _ in range(__SCREAMING_SNAKE_CASE							):
     common_inputs["past_key_values"].append(
         (
             torch.zeros(__SCREAMING_SNAKE_CASE							),
             torch.zeros(__SCREAMING_SNAKE_CASE							),
             torch.zeros(__SCREAMING_SNAKE_CASE							),
             torch.zeros(__SCREAMING_SNAKE_CASE							),
         )							)
    # TODO: test this.
    __a							=	encoder_shape if remaining_side_name == "encoder" else decoder_shape
    for _ in range(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							):
     common_inputs["past_key_values"].append((torch.zeros(__SCREAMING_SNAKE_CASE							), torch.zeros(__SCREAMING_SNAKE_CASE							))							)
   return common_inputs

  def 		_UpperCAmelCase (		self			:   Union[str, Any]  ,				__SCREAMING_SNAKE_CASE			:   PreTrainedTokenizer  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   bool = False  ,				__SCREAMING_SNAKE_CASE			:   Optional[TensorType] = None  ,				):
   __a							=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
       __SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)

   if self.use_past:
    if not is_torch_available():
     raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."							)
    else:
     import torch
    __a							=	common_inputs["input_ids"].shape
    # Not using the same length for past_key_values
    __a							=	seqlen + 2
    __a							=	self.num_layers
    __a							=	self.num_attention_heads
    __a							=	(
        batch,
        num_encoder_attention_heads,
        past_key_values_length,
        self._config.hidden_size // num_encoder_attention_heads,
    )

    __a							=	common_inputs["attention_mask"].dtype
    __a							=	torch.cat(
        [common_inputs["attention_mask"], torch.ones(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				dtype=__SCREAMING_SNAKE_CASE							)]  ,				dim=1							)
    __a							=	[
        (torch.zeros(__SCREAMING_SNAKE_CASE							), torch.zeros(__SCREAMING_SNAKE_CASE							)) for _ in range(__SCREAMING_SNAKE_CASE							)
    ]
   return common_inputs

  def 		_UpperCAmelCase (		self			:   List[Any]  ,				__SCREAMING_SNAKE_CASE			:   PreTrainedTokenizer  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   bool = False  ,				__SCREAMING_SNAKE_CASE			:   Optional[TensorType] = None  ,				):
   # Copied from OnnxConfig.generate_dummy_inputs
   # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
   # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
   __a							=	compute_effective_axis_dimension(
       __SCREAMING_SNAKE_CASE  ,				fixed_dimension=OnnxConfig.default_fixed_batch  ,				num_token_to_add=0							)

   # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
   __a							=	tokenizer.num_special_tokens_to_add(__SCREAMING_SNAKE_CASE							)
   __a							=	compute_effective_axis_dimension(
       __SCREAMING_SNAKE_CASE  ,				fixed_dimension=OnnxConfig.default_fixed_sequence  ,				num_token_to_add=__SCREAMING_SNAKE_CASE							)

   # Generate dummy inputs according to compute batch and sequence
   __a							=	[" ".join([tokenizer.unk_token]							) * seq_length] * batch_size
   __a							=	dict(tokenizer(__SCREAMING_SNAKE_CASE  ,				return_tensors=__SCREAMING_SNAKE_CASE							)							)
   return common_inputs

  def 		_UpperCAmelCase (		self			:   Union[str, Any]  ,				__SCREAMING_SNAKE_CASE			:   PreTrainedTokenizer  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   int = -1  ,				__SCREAMING_SNAKE_CASE			:   bool = False  ,				__SCREAMING_SNAKE_CASE			:   Optional[TensorType] = None  ,				):
   if self.task in ["default", "seq2seq-lm"]:
    __a							=	self._generate_dummy_inputs_for_default_and_seqaseq_lm(
        __SCREAMING_SNAKE_CASE  ,				batch_size=__SCREAMING_SNAKE_CASE  ,				seq_length=__SCREAMING_SNAKE_CASE  ,				is_pair=__SCREAMING_SNAKE_CASE  ,				framework=__SCREAMING_SNAKE_CASE							)

   elif self.task == "causal-lm":
    __a							=	self._generate_dummy_inputs_for_causal_lm(
        __SCREAMING_SNAKE_CASE  ,				batch_size=__SCREAMING_SNAKE_CASE  ,				seq_length=__SCREAMING_SNAKE_CASE  ,				is_pair=__SCREAMING_SNAKE_CASE  ,				framework=__SCREAMING_SNAKE_CASE							)
   else:
    __a							=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
        __SCREAMING_SNAKE_CASE  ,				batch_size=__SCREAMING_SNAKE_CASE  ,				seq_length=__SCREAMING_SNAKE_CASE  ,				is_pair=__SCREAMING_SNAKE_CASE  ,				framework=__SCREAMING_SNAKE_CASE							)

   return common_inputs

  def 		_UpperCAmelCase (		self			:   Optional[Any]  ,				__SCREAMING_SNAKE_CASE			:   str  ,				__SCREAMING_SNAKE_CASE			:   Union[str, Any]  ,				__SCREAMING_SNAKE_CASE			:   Any  ,				__SCREAMING_SNAKE_CASE			:   List[str]							):
   if self.task in ["default", "seq2seq-lm"]:
    __a							=	super()._flatten_past_key_values_(__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)
   else:
    __a							=	super(__SCREAMING_SNAKE_CASE  ,				self							)._flatten_past_key_values_(
        __SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE  ,				__SCREAMING_SNAKE_CASE							)




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


UpperCamelCase_						:     Union[str, Any]	       =							datasets.utils.logging.get_logger(__name__)




@dataclass
class   __lowercase						(						datasets.BuilderConfig       ):
	_A       =   None
	_A       =   "utf-8"
	_A       =   None
	_A       =   None
	_A       =   True  # deprecated
	_A       =   None  # deprecated
	_A       =   10 << 20  # 10MB
	_A       =   None




class   __lowercase						(						datasets.ArrowBasedBuilder       ):
	_A       =   JsonConfig






	def       _a(self				: List[Any]    ) ->			Dict:
							if self.config.block_size is not None:
													logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"    )
													_lowercase       :	Union[str, Any]           =   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       _a(self				: Any      ,	snake_case				: Union[str, Any]    ) ->			List[str]:
							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}"""    )
							_lowercase       :	str           =   dl_manager.download_and_extract(self.config.data_files    )
							if isinstance(snake_case      ,	(str, list, tuple)    ):
													_lowercase       :	int           =   data_files
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	Union[str, Any]           =   [files]
													_lowercase       :	Union[str, Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,	gen_kwargs={"files": files}    )]
							_lowercase       :	Union[str, Any]           =   []
							for split_name, files in data_files.items():
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	List[Any]           =   [files]
													_lowercase       :	Optional[Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													splits.append(datasets.SplitGenerator(name=snake_case      ,	gen_kwargs={"files": files}    )    )
							return splits






	def       _a(self				: List[Any]      ,	snake_case				: 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    ):
																			_lowercase       :	Any           =   self.config.features.arrow_schema.field(snake_case    ).type
																			_lowercase       :	Optional[Any]           =   pa_table.append_column(snake_case      ,	pa.array([None] * len(snake_case    )      ,	type=snake_case    )    )
													# more expensive cast to support nested structures with keys in a different order
													# allows str <-> int/float or str to Audio for example
													_lowercase       :	str           =   table_cast(snake_case      ,	self.config.features.arrow_schema    )
							return pa_table






	def       _a(self				: str      ,	snake_case				: Tuple    ) ->			Tuple:
							for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case    )    ):
													# 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(snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																									_lowercase       :	Any           =   json.load(snake_case    )

																			# We keep only the field we are interested in
																			_lowercase       :	List[Any]           =   dataset[self.config.field]

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

													# If the file has one json object per line
													else:
																			with open(snake_case      ,	"rb"    ) as f:
																									_lowercase       :	Union[str, Any]           =   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
																									_lowercase       :	Dict           =   max(self.config.chunksize // 32      ,	16 << 10    )
																									_lowercase       :	int           =   (
																									    self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
																									)
																									while True:
																															_lowercase       :	Tuple           =   f.read(self.config.chunksize    )
																															if not batch:
																																					break
																															# Finish current line
																															try:
																																					batch += f.readline()
																															except (AttributeError, io.UnsupportedOperation):
																																					batch += readline(snake_case    )
																															# PyArrow only accepts utf-8 encoded bytes
																															if self.config.encoding != "utf-8":
																																					_lowercase       :	Union[str, Any]           =   batch.decode(self.config.encoding      ,	errors=snake_case    ).encode("utf-8"    )
																															try:
																																					while True:
																																											try:
																																																	_lowercase       :	Optional[Any]           =   paj.read_json(
																																																	    io.BytesIO(snake_case    )      ,	read_options=paj.ReadOptions(block_size=snake_case    )    )
																																																	break
																																											except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																	if (
																																																	    isinstance(snake_case      ,	pa.ArrowInvalid    )
																																																	    and "straddling" not in str(snake_case    )
																																																	    or block_size > len(snake_case    )
																																																	):
																																																							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(snake_case    )} 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(
																																											    snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																																																	_lowercase       :	Any           =   json.load(snake_case    )
																																					except json.JSONDecodeError:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																											raise e
																																					# If possible, parse the file as a list of json objects and exit the loop
																																					if isinstance(snake_case      ,	snake_case    ):  # list is the only sequence type supported in JSON
																																											try:
																																																	_lowercase       :	int           =   set().union(*[row.keys() for row in dataset]    )
																																																	_lowercase       :	Dict           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																																																	_lowercase       :	str           =   pa.Table.from_pydict(snake_case    )
																																											except (pa.ArrowInvalid, AttributeError) as e:
																																																	logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																																	raise ValueError(F"""Not able to read records in the JSON file at {file}."""    ) from None
																																											yield file_idx, self._cast_table(snake_case    )
																																											break
																																					else:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {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(snake_case    )
																															batch_idx += 1




