


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


lowercase_	   =     logging.get_logger(__name__)

lowercase_	   =     {
    "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_ (	_UpperCAmelCase   ):




      '''simple docstring'''


      __snake_case 	=       """blenderbot-small"""
      __snake_case 	=       ["""past_key_values"""]
      __snake_case 	=       {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




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

       super().__init__(
           pad_token_id=lowercase_						,				bos_token_id=lowercase_						,				eos_token_id=lowercase_						,				is_encoder_decoder=lowercase_						,				decoder_start_token_id=lowercase_						,				forced_eos_token_id=lowercase_						,				**lowercase_						,				)



class 	A_ (	_UpperCAmelCase   ):




      '''simple docstring'''


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

        if self.use_past:
         __lowerCamelCase		:							Any		 =						{0: """batch"""}
         __lowerCamelCase		:							List[Any]		 =						{0: """batch""", 1: """past_decoder_sequence + sequence"""}
        else:
         __lowerCamelCase		:							List[str]		 =						{0: """batch""", 1: """decoder_sequence"""}
         __lowerCamelCase		:							Optional[int]		 =						{0: """batch""", 1: """decoder_sequence"""}

        if self.use_past:
         self.fill_with_past_key_values_(lowercase_						,				direction='inputs'    )
       elif self.task == "causal-lm":
        # TODO: figure this case out.
        __lowerCamelCase		:							List[str]		 =						OrderedDict(
            [
                ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
            ]    )
        if self.use_past:
         __lowerCamelCase		:							Optional[Any]		 =						self.num_layers
         for i in range(lowercase_    ):
          __lowerCamelCase		:							List[Any]		 =						{0: """batch""", 2: """past_sequence + sequence"""}
          __lowerCamelCase		:							Any		 =						{0: """batch""", 2: """past_sequence + sequence"""}
       else:
        __lowerCamelCase		:							str		 =						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 					_snake_case     (    self:	str    ):
       if self.task in ["default", "seq2seq-lm"]:
        __lowerCamelCase		:							List[str]		 =						super().outputs
       else:
        __lowerCamelCase		:							str		 =						super(lowercase_						,				self    ).outputs
        if self.use_past:
         __lowerCamelCase		:							int		 =						self.num_layers
         for i in range(lowercase_    ):
          __lowerCamelCase		:							str		 =						{0: """batch""", 2: """past_sequence + sequence"""}
          __lowerCamelCase		:							str		 =						{0: """batch""", 2: """past_sequence + sequence"""}
       return common_outputs




      def 					_snake_case     (    self:	int						,				a:	PreTrainedTokenizer						,				a:	int = -1						,				a:	int = -1						,				a:	bool = False						,				a:	Optional[TensorType] = None						,				):
       __lowerCamelCase		:							List[Any]		 =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           lowercase_						,				lowercase_						,				lowercase_						,				lowercase_						,				lowercase_    )

       # Generate decoder inputs
       __lowerCamelCase		:							Dict		 =						seq_length if not self.use_past else 1
       __lowerCamelCase		:							Union[str, Any]		 =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           lowercase_						,				lowercase_						,				lowercase_						,				lowercase_						,				lowercase_    )
       __lowerCamelCase		:							Any		 =						{F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
       __lowerCamelCase		:							Optional[int]		 =						dict(**lowercase_						,				**lowercase_    )

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

        __lowerCamelCase		:							Optional[int]		 =						torch.cat(
            [common_inputs['decoder_attention_mask'], torch.ones(lowercase_						,				lowercase_    )]						,				dim=1    )

        __lowerCamelCase		:							Union[str, Any]		 =						[]
        # If the number of encoder and decoder layers are present in the model configuration, both are considered
        __lowerCamelCase		:							Dict		 =						self.num_layers
        __lowerCamelCase		:							Any		 =						min(lowercase_						,				lowercase_    )
        __lowerCamelCase		:							Optional[Any]		 =						max(lowercase_						,				lowercase_    ) - min_num_layers
        __lowerCamelCase		:							int		 =						"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

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




      def 					_snake_case     (    self:	Any						,				a:	PreTrainedTokenizer						,				a:	int = -1						,				a:	int = -1						,				a:	bool = False						,				a:	Optional[TensorType] = None						,				):
       __lowerCamelCase		:							Dict		 =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           lowercase_						,				lowercase_						,				lowercase_						,				lowercase_						,				lowercase_    )

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

        __lowerCamelCase		:							Optional[Any]		 =						common_inputs["""attention_mask"""].dtype
        __lowerCamelCase		:							List[Any]		 =						torch.cat(
            [common_inputs['attention_mask'], torch.ones(lowercase_						,				lowercase_						,				dtype=lowercase_    )]						,				dim=1    )
        __lowerCamelCase		:							List[Any]		 =						[
            (torch.zeros(lowercase_    ), torch.zeros(lowercase_    )) for _ in range(lowercase_    )
        ]
       return common_inputs




      def 					_snake_case     (    self:	List[str]						,				a:	PreTrainedTokenizer						,				a:	int = -1						,				a:	int = -1						,				a:	bool = False						,				a:	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
       __lowerCamelCase		:							Tuple		 =						compute_effective_axis_dimension(
           lowercase_						,				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
       __lowerCamelCase		:							Dict		 =						tokenizer.num_special_tokens_to_add(lowercase_    )
       __lowerCamelCase		:							Optional[Any]		 =						compute_effective_axis_dimension(
           lowercase_						,				fixed_dimension=OnnxConfig.default_fixed_sequence						,				num_token_to_add=lowercase_    )

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




      def 					_snake_case     (    self:	int						,				a:	PreTrainedTokenizer						,				a:	int = -1						,				a:	int = -1						,				a:	bool = False						,				a:	Optional[TensorType] = None						,				):
       if self.task in ["default", "seq2seq-lm"]:
        __lowerCamelCase		:							Tuple		 =						self._generate_dummy_inputs_for_default_and_seqaseq_lm(
            lowercase_						,				batch_size=lowercase_						,				seq_length=lowercase_						,				is_pair=lowercase_						,				framework=lowercase_    )

       elif self.task == "causal-lm":
        __lowerCamelCase		:							List[Any]		 =						self._generate_dummy_inputs_for_causal_lm(
            lowercase_						,				batch_size=lowercase_						,				seq_length=lowercase_						,				is_pair=lowercase_						,				framework=lowercase_    )
       else:
        __lowerCamelCase		:							Union[str, Any]		 =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase_						,				batch_size=lowercase_						,				seq_length=lowercase_						,				is_pair=lowercase_						,				framework=lowercase_    )

       return common_inputs




      def 					_snake_case     (    self:	int						,				a:	Union[str, Any]						,				a:	List[Any]						,				a:	Union[str, Any]						,				a:	str    ):
       if self.task in ["default", "seq2seq-lm"]:
        __lowerCamelCase		:							Tuple		 =						super()._flatten_past_key_values_(lowercase_						,				lowercase_						,				lowercase_						,				lowercase_    )
       else:
        __lowerCamelCase		:							List[Any]		 =						super(lowercase_						,				self    )._flatten_past_key_values_(
            lowercase_						,				lowercase_						,				lowercase_						,				lowercase_    )





import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							List[str]		 =						np.inf

 def set_batch_size(SCREAMING_SNAKE_CASE__						) -> None:
  nonlocal batch_size
  if isinstance(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							Union[str, Any]		 =						min(SCREAMING_SNAKE_CASE__      ,	config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS						)
  elif isinstance(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							Optional[Any]		 =						min(SCREAMING_SNAKE_CASE__      ,	config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS						)
  elif isinstance(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						) and feature.dtype == "binary":
   __lowerCamelCase		:							List[str]		 =						min(SCREAMING_SNAKE_CASE__      ,	config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS						)

 _visit(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)

 return None if batch_size is np.inf else batch_size


class 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	Tuple						,				a:	NestedDataStructureLike[PathLike]						,				a:	Optional[NamedSplit] = None						,				a:	Optional[Features] = None						,				a:	str = None						,				a:	bool = False						,				a:	bool = False						,				a:	Optional[int] = None						,				**a:	Optional[Any]						,				):
       super().__init__(
           a						,				split=a						,				features=a						,				cache_dir=a						,				keep_in_memory=a						,				streaming=a						,				num_proc=a						,				**a						,				)
       __lowerCamelCase		:							List[Any]		 =						path_or_paths if isinstance(a						,				a    ) else {self.split: path_or_paths}
       __lowerCamelCase		:							Optional[Any]		 =						_PACKAGED_DATASETS_MODULES['parquet'][1]
       __lowerCamelCase		:							List[str]		 =						Parquet(
           cache_dir=a						,				data_files=a						,				features=a						,				hash=a						,				**a						,				)




      def 					_snake_case     (    self:	List[str]    ):
       # Build iterable dataset
       if self.streaming:
        __lowerCamelCase		:							str		 =						self.builder.as_streaming_dataset(split=self.split    )
       # Build regular (map-style) dataset
       else:
        __lowerCamelCase		:							str		 =						None
        __lowerCamelCase		:							Optional[Any]		 =						None
        __lowerCamelCase		:							Union[str, Any]		 =						None
        __lowerCamelCase		:							int		 =						None

        self.builder.download_and_prepare(
            download_config=a						,				download_mode=a						,				verification_mode=a						,				base_path=a						,				num_proc=self.num_proc						,				)
        __lowerCamelCase		:							Tuple		 =						self.builder.as_dataset(
            split=self.split						,				verification_mode=a						,				in_memory=self.keep_in_memory    )
       return dataset



class 	A_ :




      '''simple docstring'''


      def __init__(    self:	Optional[int]						,				a:	Dataset						,				a:	Union[PathLike, BinaryIO]						,				a:	Optional[int] = None						,				**a:	List[Any]						,				):
       __lowerCamelCase		:							Optional[int]		 =						dataset
       __lowerCamelCase		:							List[Any]		 =						path_or_buf
       __lowerCamelCase		:							List[str]		 =						batch_size or get_writer_batch_size(dataset.features    )
       __lowerCamelCase		:							List[Any]		 =						parquet_writer_kwargs




      def 					_snake_case     (    self:	Union[str, Any]    ):
       __lowerCamelCase		:							Optional[int]		 =						self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

       if isinstance(self.path_or_buf						,				(str, bytes, os.PathLike)    ):
        with open(self.path_or_buf						,				'wb+'    ) as buffer:
         __lowerCamelCase		:							Optional[int]		 =						self._write(file_obj=a						,				batch_size=a						,				**self.parquet_writer_kwargs    )
       else:
        __lowerCamelCase		:							Any		 =						self._write(file_obj=self.path_or_buf						,				batch_size=a						,				**self.parquet_writer_kwargs    )
       return written




      def 					_snake_case     (    self:	Optional[int]						,				a:	BinaryIO						,				a:	int						,				**a:	str    ):
       __lowerCamelCase		:							Dict		 =						0
       __lowerCamelCase		:							Union[str, Any]		 =						parquet_writer_kwargs.pop('path_or_buf'						,				a    )
       __lowerCamelCase		:							str		 =						self.dataset.features.arrow_schema

       __lowerCamelCase		:							Any		 =						pq.ParquetWriter(a						,				schema=a						,				**a    )

       for offset in logging.tqdm(
           range(0						,				len(self.dataset    )						,				a    )						,				unit='ba'						,				disable=not logging.is_progress_bar_enabled()						,				desc='Creating parquet from Arrow format'						,				):
        __lowerCamelCase		:							Any		 =						query_table(
            table=self.dataset._data						,				key=slice(a						,				offset + batch_size    )						,				indices=self.dataset._indices if self.dataset._indices is not None else None						,				)
        writer.write_table(a    )
        written += batch.nbytes
       writer.close()
       return written

