


from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


lowercase_	   =     logging.get_logger(__name__)

# General docstring
lowercase_	   =     'ResNetConfig'

# Base docstring
lowercase_	   =     'microsoft/resnet-50'
lowercase_	   =     [1, 2_0_4_8, 7, 7]

# Image classification docstring
lowercase_	   =     'microsoft/resnet-50'
lowercase_	   =     'tiger cat'

lowercase_	   =     [
    'microsoft/resnet-50',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]


class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	Optional[int]						,				a:	str						,				a:	Union[str, Any]						,				a:	Tuple = 3						,				a:	List[Any] = 1						,				a:	List[str] = "relu"    ):
       super().__init__()
       __lowerCamelCase		:							List[Any]		 =						nn.Convad(
           _SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				kernel_size=_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE						,				padding=kernel_size // 2						,				bias=_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							int		 =						nn.BatchNormad(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							List[str]		 =						ACTaFN[activation] if activation is not None else nn.Identity()




      def 					_snake_case     (    self:	int						,				a:	Optional[int]    ):
       __lowerCamelCase		:							str		 =						self.convolution(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Tuple		 =						self.normalization(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							List[str]		 =						self.activation(_SCREAMING_SNAKE_CASE    )
       return hidden_state



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	Optional[int]						,				a:	Dict    ):
       super().__init__()
       __lowerCamelCase		:							int		 =						ResNetConvLayer(
           config.num_channels						,				config.embedding_size						,				kernel_size=7						,				stride=2						,				activation=config.hidden_act    )
       __lowerCamelCase		:							Any		 =						nn.MaxPoolad(kernel_size=3						,				stride=2						,				padding=1    )
       __lowerCamelCase		:							List[Any]		 =						config.num_channels




      def 					_snake_case     (    self:	Optional[Any]						,				a:	List[str]    ):
       __lowerCamelCase		:							Optional[int]		 =						pixel_values.shape[1]
       if num_channels != self.num_channels:
        raise ValueError(
            'Make sure that the channel dimension of the pixel values match with the one set in the configuration.'    )
       __lowerCamelCase		:							Optional[int]		 =						self.embedder(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Any		 =						self.pooler(_SCREAMING_SNAKE_CASE    )
       return embedding



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	List[str]						,				a:	Any						,				a:	Dict						,				a:	List[str] = 2    ):
       super().__init__()
       __lowerCamelCase		:							Optional[int]		 =						nn.Convad(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				kernel_size=1						,				stride=_SCREAMING_SNAKE_CASE						,				bias=_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							List[Any]		 =						nn.BatchNormad(_SCREAMING_SNAKE_CASE    )




      def 					_snake_case     (    self:	str						,				a:	Optional[int]    ):
       __lowerCamelCase		:							List[str]		 =						self.convolution(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							List[str]		 =						self.normalization(_SCREAMING_SNAKE_CASE    )
       return hidden_state



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	int						,				a:	Any						,				a:	str						,				a:	List[Any] = 1						,				a:	Tuple = "relu"    ):
       super().__init__()
       __lowerCamelCase		:							Any		 =						in_channels != out_channels or stride != 1
       __lowerCamelCase		:							Tuple		 =						(
           ResNetShortCut(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE    ) if should_apply_shortcut else nn.Identity()
       )
       __lowerCamelCase		:							Tuple		 =						nn.Sequential(
           ResNetConvLayer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE    )						,				ResNetConvLayer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				activation=_SCREAMING_SNAKE_CASE    )						,				)
       __lowerCamelCase		:							List[str]		 =						ACTaFN[activation]




      def 					_snake_case     (    self:	int						,				a:	Optional[Any]    ):
       __lowerCamelCase		:							List[Any]		 =						hidden_state
       __lowerCamelCase		:							Dict		 =						self.layer(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Optional[int]		 =						self.shortcut(_SCREAMING_SNAKE_CASE    )
       hidden_state += residual
       __lowerCamelCase		:							Optional[Any]		 =						self.activation(_SCREAMING_SNAKE_CASE    )
       return hidden_state



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	Optional[Any]						,				a:	str						,				a:	Tuple						,				a:	Tuple = 1						,				a:	Optional[Any] = "relu"						,				a:	Any = 4    ):
       super().__init__()
       __lowerCamelCase		:							Any		 =						in_channels != out_channels or stride != 1
       __lowerCamelCase		:							Optional[int]		 =						out_channels // reduction
       __lowerCamelCase		:							int		 =						(
           ResNetShortCut(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE    ) if should_apply_shortcut else nn.Identity()
       )
       __lowerCamelCase		:							List[str]		 =						nn.Sequential(
           ResNetConvLayer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				kernel_size=1    )						,				ResNetConvLayer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE    )						,				ResNetConvLayer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				kernel_size=1						,				activation=_SCREAMING_SNAKE_CASE    )						,				)
       __lowerCamelCase		:							List[str]		 =						ACTaFN[activation]




      def 					_snake_case     (    self:	Any						,				a:	Union[str, Any]    ):
       __lowerCamelCase		:							Optional[int]		 =						hidden_state
       __lowerCamelCase		:							Tuple		 =						self.layer(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Any		 =						self.shortcut(_SCREAMING_SNAKE_CASE    )
       hidden_state += residual
       __lowerCamelCase		:							Optional[Any]		 =						self.activation(_SCREAMING_SNAKE_CASE    )
       return hidden_state



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	str						,				a:	str						,				a:	int						,				a:	Union[str, Any]						,				a:	Dict = 2						,				a:	Tuple = 2						,				):
       super().__init__()

       __lowerCamelCase		:							Optional[Any]		 =						ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

       __lowerCamelCase		:							Tuple		 =						nn.Sequential(
           # downsampling is done in the first layer with stride of 2
           layer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				stride=_SCREAMING_SNAKE_CASE						,				activation=config.hidden_act    )						,				*[layer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				activation=config.hidden_act    ) for _ in range(depth - 1    )]						,				)




      def 					_snake_case     (    self:	str						,				a:	int    ):
       __lowerCamelCase		:							Tuple		 =						input
       for layer in self.layers:
        __lowerCamelCase		:							str		 =						layer(_SCREAMING_SNAKE_CASE    )
       return hidden_state



class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	Any						,				a:	Union[str, Any]    ):
       super().__init__()
       __lowerCamelCase		:							Any		 =						nn.ModuleList([]    )
       # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
       self.stages.append(
           ResNetStage(
               _SCREAMING_SNAKE_CASE						,				config.embedding_size						,				config.hidden_sizes[0]						,				stride=2 if config.downsample_in_first_stage else 1						,				depth=config.depths[0]						,				)    )
       __lowerCamelCase		:							Dict		 =						zip(config.hidden_sizes						,				config.hidden_sizes[1:]    )
       for (in_channels, out_channels), depth in zip(_SCREAMING_SNAKE_CASE						,				config.depths[1:]    ):
        self.stages.append(ResNetStage(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				depth=_SCREAMING_SNAKE_CASE    )    )




      def 					_snake_case     (    self:	Optional[int]						,				a:	List[str]						,				a:	str = False						,				a:	List[str] = True    ):
       __lowerCamelCase		:							str		 =						() if output_hidden_states else None

       for stage_module in self.stages:
        if output_hidden_states:
         __lowerCamelCase		:							Optional[Any]		 =						hidden_states + (hidden_state,)

        __lowerCamelCase		:							Any		 =						stage_module(_SCREAMING_SNAKE_CASE    )

       if output_hidden_states:
        __lowerCamelCase		:							str		 =						hidden_states + (hidden_state,)

       if not return_dict:
        return tuple(v for v in [hidden_state, hidden_states] if v is not None    )

       return BaseModelOutputWithNoAttention(
           last_hidden_state=_SCREAMING_SNAKE_CASE						,				hidden_states=_SCREAMING_SNAKE_CASE						,				)



class 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      __snake_case 	=       ResNetConfig
      __snake_case 	=       """resnet"""
      __snake_case 	=       """pixel_values"""
      __snake_case 	=       True




      def 					_snake_case     (    self:	Any						,				a:	str    ):





       if isinstance(_SCREAMING_SNAKE_CASE						,				nn.Convad    ):
        nn.init.kaiming_normal_(module.weight						,				mode='fan_out'						,				nonlinearity='relu'    )
       elif isinstance(_SCREAMING_SNAKE_CASE						,				(nn.BatchNormad, nn.GroupNorm)    ):
        nn.init.constant_(module.weight						,				1    )
        nn.init.constant_(module.bias						,				0    )




      def 					_snake_case     (    self:	Any						,				a:	Tuple						,				a:	int=False    ):





       if isinstance(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE    ):
        __lowerCamelCase		:							str		 =						value



lowercase_	   =     R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

lowercase_	   =     R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'


@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top.""" ,				__UpperCamelCase ,				)
class 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	List[str]						,				a:	Union[str, Any]    ):
       super().__init__(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Any		 =						config
       __lowerCamelCase		:							Dict		 =						ResNetEmbeddings(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Optional[Any]		 =						ResNetEncoder(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Optional[Any]		 =						nn.AdaptiveAvgPoolad((1, 1)    )
       # Initialize weights and apply final processing
       self.post_init()




      @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE    )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC						,				output_type=_SCREAMING_SNAKE_CASE						,				config_class=_CONFIG_FOR_DOC						,				modality='vision'						,				expected_output=_EXPECTED_OUTPUT_SHAPE						,				)
      def 					_snake_case     (    self:	Optional[int]						,				a:	Optional[Any]						,				a:	str = None						,				a:	str = None    ):
       __lowerCamelCase		:							List[str]		 =						(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       __lowerCamelCase		:							List[Any]		 =						return_dict if return_dict is not None else self.config.use_return_dict

       __lowerCamelCase		:							int		 =						self.embedder(_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Optional[int]		 =						self.encoder(
           _SCREAMING_SNAKE_CASE						,				output_hidden_states=_SCREAMING_SNAKE_CASE						,				return_dict=_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Union[str, Any]		 =						encoder_outputs[0]

       __lowerCamelCase		:							Optional[Any]		 =						self.pooler(_SCREAMING_SNAKE_CASE    )

       if not return_dict:
        return (last_hidden_state, pooled_output) + encoder_outputs[1:]

       return BaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=_SCREAMING_SNAKE_CASE						,				pooler_output=_SCREAMING_SNAKE_CASE						,				hidden_states=encoder_outputs.hidden_states						,				)



@add_start_docstrings(
    """\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    """ ,				__UpperCamelCase ,				)
class 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	List[str]						,				a:	Any    ):
       super().__init__(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Union[str, Any]		 =						config.num_labels
       __lowerCamelCase		:							List[str]		 =						ResNetModel(_SCREAMING_SNAKE_CASE    )
       # classification head
       __lowerCamelCase		:							Optional[int]		 =						nn.Sequential(
           nn.Flatten()						,				nn.Linear(config.hidden_sizes[-1]						,				config.num_labels    ) if config.num_labels > 0 else nn.Identity()						,				)
       # initialize weights and apply final processing
       self.post_init()




      @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE    )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT						,				output_type=_SCREAMING_SNAKE_CASE						,				config_class=_CONFIG_FOR_DOC						,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT						,				)
      def 					_snake_case     (    self:	Optional[int]						,				a:	List[str] = None						,				a:	Union[str, Any] = None						,				a:	Any = None						,				a:	int = None						,				):
       __lowerCamelCase		:							Tuple		 =						return_dict if return_dict is not None else self.config.use_return_dict

       __lowerCamelCase		:							int		 =						self.resnet(_SCREAMING_SNAKE_CASE						,				output_hidden_states=_SCREAMING_SNAKE_CASE						,				return_dict=_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Tuple		 =						outputs.pooler_output if return_dict else outputs[1]

       __lowerCamelCase		:							List[str]		 =						self.classifier(_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Any		 =						None

       if labels is not None:
        if self.config.problem_type is None:
         if self.num_labels == 1:
          __lowerCamelCase		:							Optional[Any]		 =						'''regression'''
         elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
          __lowerCamelCase		:							int		 =						'''single_label_classification'''
         else:
          __lowerCamelCase		:							Optional[Any]		 =						'''multi_label_classification'''
        if self.config.problem_type == "regression":
         __lowerCamelCase		:							Any		 =						MSELoss()
         if self.num_labels == 1:
          __lowerCamelCase		:							Any		 =						loss_fct(logits.squeeze()						,				labels.squeeze()    )
         else:
          __lowerCamelCase		:							Optional[Any]		 =						loss_fct(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE    )
        elif self.config.problem_type == "single_label_classification":
         __lowerCamelCase		:							Optional[int]		 =						CrossEntropyLoss()
         __lowerCamelCase		:							List[str]		 =						loss_fct(logits.view(-1						,				self.num_labels    )						,				labels.view(-1    )    )
        elif self.config.problem_type == "multi_label_classification":
         __lowerCamelCase		:							List[str]		 =						BCEWithLogitsLoss()
         __lowerCamelCase		:							List[Any]		 =						loss_fct(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE    )

       if not return_dict:
        __lowerCamelCase		:							Tuple		 =						(logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output

       return ImageClassifierOutputWithNoAttention(loss=_SCREAMING_SNAKE_CASE						,				logits=_SCREAMING_SNAKE_CASE						,				hidden_states=outputs.hidden_states    )



@add_start_docstrings(
    """\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    """ ,				__UpperCamelCase ,				)
class 	A_ (	__UpperCamelCase ,				__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	Union[str, Any]						,				a:	Dict    ):
       super().__init__(_SCREAMING_SNAKE_CASE    )
       super()._init_backbone(_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Optional[Any]		 =						[config.embedding_size] + config.hidden_sizes
       __lowerCamelCase		:							Optional[int]		 =						ResNetEmbeddings(_SCREAMING_SNAKE_CASE    )
       __lowerCamelCase		:							Union[str, Any]		 =						ResNetEncoder(_SCREAMING_SNAKE_CASE    )

       # initialize weights and apply final processing
       self.post_init()




      @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE    )
      @replace_return_docstrings(output_type=_SCREAMING_SNAKE_CASE						,				config_class=_CONFIG_FOR_DOC    )
      def 					_snake_case     (    self:	Optional[Any]						,				a:	List[Any]						,				a:	Union[str, Any] = None						,				a:	int = None    ):
       __lowerCamelCase		:							Optional[Any]		 =						return_dict if return_dict is not None else self.config.use_return_dict
       __lowerCamelCase		:							int		 =						(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )

       __lowerCamelCase		:							Any		 =						self.embedder(_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							Any		 =						self.encoder(_SCREAMING_SNAKE_CASE						,				output_hidden_states=_SCREAMING_SNAKE_CASE						,				return_dict=_SCREAMING_SNAKE_CASE    )

       __lowerCamelCase		:							str		 =						outputs.hidden_states

       __lowerCamelCase		:							Dict		 =						()
       for idx, stage in enumerate(self.stage_names    ):
        if stage in self.out_features:
         feature_maps += (hidden_states[idx],)

       if not return_dict:
        __lowerCamelCase		:							str		 =						(feature_maps,)
        if output_hidden_states:
         output += (outputs.hidden_states,)
        return output

       return BackboneOutput(
           feature_maps=_SCREAMING_SNAKE_CASE						,				hidden_states=outputs.hidden_states if output_hidden_states else None						,				attentions=_SCREAMING_SNAKE_CASE						,				)





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

