

'''simple docstring'''







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__     (   UpperCAmelCase__					)      ->      Optional[int]:
       A_						       =   np.inf

       def set_batch_size(UpperCAmelCase__					) -> None:
              nonlocal batch_size
              if isinstance(UpperCAmelCase__,							UpperCAmelCase__					):
                     A_						       =   min(UpperCAmelCase__,							config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS					)
              elif isinstance(UpperCAmelCase__,							UpperCAmelCase__					):
                     A_						       =   min(UpperCAmelCase__,							config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS					)
              elif isinstance(UpperCAmelCase__,							UpperCAmelCase__					) and feature.dtype == "binary":
                     A_						       =   min(UpperCAmelCase__,							config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS					)

       _visit(UpperCAmelCase__,							UpperCAmelCase__					)

       return None if batch_size is np.inf else batch_size


class        A__   (   _snake_case      ):



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					UpperCamelCase__ = False		,					UpperCamelCase__ = False		,					UpperCamelCase__ = None		,					**UpperCamelCase__		,					)    ->   List[str]:





           '''simple docstring'''
           super().__init__(
               UpperCamelCase__		,					split=UpperCamelCase__		,					features=UpperCamelCase__		,					cache_dir=UpperCamelCase__		,					keep_in_memory=UpperCamelCase__		,					streaming=UpperCamelCase__		,					num_proc=UpperCamelCase__		,					**UpperCamelCase__		,					)
           A_						       =   path_or_paths if isinstance(UpperCamelCase__		,					UpperCamelCase__							) else {self.split: path_or_paths}
           A_						       =   _PACKAGED_DATASETS_MODULES["""parquet"""][1]
           A_						       =   Parquet(
               cache_dir=UpperCamelCase__		,					data_files=UpperCamelCase__		,					features=UpperCamelCase__		,					hash=UpperCamelCase__		,					**UpperCamelCase__		,					)







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





           '''simple docstring'''
           # Build iterable dataset
           if self.streaming:
                  A_						       =   self.builder.as_streaming_dataset(split=self.split							)
           # Build regular (map-style) dataset
           else:
                  A_						       =   None
                  A_						       =   None
                  A_						       =   None
                  A_						       =   None

                  self.builder.download_and_prepare(
                      download_config=UpperCamelCase__		,					download_mode=UpperCamelCase__		,					verification_mode=UpperCamelCase__		,					base_path=UpperCamelCase__		,					num_proc=self.num_proc		,					)
                  A_						       =   self.builder.as_dataset(
                      split=self.split		,					verification_mode=UpperCamelCase__		,					in_memory=self.keep_in_memory							)
           return dataset


class        A__   :



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__ = None		,					**UpperCamelCase__		,					)    ->   Dict:





           '''simple docstring'''
           A_						       =   dataset
           A_						       =   path_or_buf
           A_						       =   batch_size or get_writer_batch_size(dataset.features							)
           A_						       =   parquet_writer_kwargs



    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           A_						       =   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:
                         A_						       =   self._write(file_obj=UpperCamelCase__		,					batch_size=UpperCamelCase__		,					**self.parquet_writer_kwargs							)
           else:
                  A_						       =   self._write(file_obj=self.path_or_buf		,					batch_size=UpperCamelCase__		,					**self.parquet_writer_kwargs							)
           return written







    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   int:





           '''simple docstring'''
           A_						       =   0
           A_						       =   parquet_writer_kwargs.pop("""path_or_buf"""		,					UpperCamelCase__							)
           A_						       =   self.dataset.features.arrow_schema

           A_						       =   pq.ParquetWriter(UpperCamelCase__		,					schema=UpperCamelCase__		,					**UpperCamelCase__							)

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



'''simple docstring'''







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

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


__lowerCamelCase									=     logging.get_logger(__name__)


__lowerCamelCase									=     {
    '''Salesforce/codegen-350M-nl''': '''https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json''',
    '''Salesforce/codegen-350M-multi''': '''https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json''',
    '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json''',
    '''Salesforce/codegen-2B-nl''': '''https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-2B-multi''': '''https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-2B-mono''': '''https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-6B-nl''': '''https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-6B-multi''': '''https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-6B-mono''': '''https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-16B-nl''': '''https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-16B-multi''': '''https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-16B-mono''': '''https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json''',
}


class        A__   (   _snake_case      ):
    lowercase						       =		"codegen"
    lowercase						       =		{
        "max_position_embeddings": "n_positions",
        "hidden_size": "n_embd",
        "num_attention_heads": "n_head",
        "num_hidden_layers": "n_layer",
    }



    def __init__(       self		,					UpperCamelCase__=50400		,					UpperCamelCase__=2048		,					UpperCamelCase__=2048		,					UpperCamelCase__=4096		,					UpperCamelCase__=28		,					UpperCamelCase__=16		,					UpperCamelCase__=64		,					UpperCamelCase__=None		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=1e-5		,					UpperCamelCase__=0.02		,					UpperCamelCase__=True		,					UpperCamelCase__=50256		,					UpperCamelCase__=50256		,					UpperCamelCase__=False		,					**UpperCamelCase__		,					)    ->   Tuple:





           '''simple docstring'''
           A_						       =   vocab_size
           A_						       =   n_ctx
           A_						       =   n_positions
           A_						       =   n_embd
           A_						       =   n_layer
           A_						       =   n_head
           A_						       =   n_inner
           A_						       =   rotary_dim
           A_						       =   activation_function
           A_						       =   resid_pdrop
           A_						       =   embd_pdrop
           A_						       =   attn_pdrop
           A_						       =   layer_norm_epsilon
           A_						       =   initializer_range
           A_						       =   use_cache

           A_						       =   bos_token_id
           A_						       =   eos_token_id

           super().__init__(
               bos_token_id=UpperCamelCase__		,					eos_token_id=UpperCamelCase__		,					tie_word_embeddings=UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__ = "default"		,					UpperCamelCase__ = None		,					UpperCamelCase__ = False		,					)    ->   Tuple:





           '''simple docstring'''
           super().__init__(UpperCamelCase__		,					task=UpperCamelCase__		,					patching_specs=UpperCamelCase__		,					use_past=UpperCamelCase__							)
           if not getattr(self._config		,					"""pad_token_id"""		,					UpperCamelCase__							):
                  # TODO: how to do that better?
                  A_						       =   0



    @property
    def        snake_case_      (       self							)    ->   Mapping[str, Mapping[int, str]]:





           '''simple docstring'''
           A_						       =   OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}							)
           if self.use_past:
                  self.fill_with_past_key_values_(UpperCamelCase__		,					direction="""inputs"""							)
                  A_						       =   {0: """batch""", 1: """past_sequence + sequence"""}
           else:
                  A_						       =   {0: """batch""", 1: """sequence"""}

           return common_inputs



    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return self._config.n_layer



    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return self._config.n_head



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__ = -1		,					UpperCamelCase__ = -1		,					UpperCamelCase__ = False		,					UpperCamelCase__ = None		,					)    ->   Mapping[str, Any]:





           '''simple docstring'''
           A_						       =   super(UpperCamelCase__		,					self							).generate_dummy_inputs(
               UpperCamelCase__		,					batch_size=UpperCamelCase__		,					seq_length=UpperCamelCase__		,					is_pair=UpperCamelCase__		,					framework=UpperCamelCase__							)

           # We need to order the input in the way they appears in the forward()
           A_						       =   OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}							)

           # Need to add the past_keys
           if self.use_past:
                  if not is_torch_available():
                         raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
                  else:
                         import torch

                         A_       ,      A_						       =   common_inputs["""input_ids"""].shape
                         # Not using the same length for past_key_values
                         A_						       =   seqlen + 2
                         A_						       =   (
                             batch,
                             self.num_attention_heads,
                             past_key_values_length,
                             self._config.hidden_size // self.num_attention_heads,
                         )
                         A_						       =   [
                             (torch.zeros(UpperCamelCase__							), torch.zeros(UpperCamelCase__							)) for _ in range(self.num_layers							)
                         ]

           A_						       =   common_inputs["""attention_mask"""]
           if self.use_past:
                  A_						       =   ordered_inputs["""attention_mask"""].dtype
                  A_						       =   torch.cat(
                      [ordered_inputs["""attention_mask"""], torch.ones(UpperCamelCase__		,					UpperCamelCase__		,					dtype=UpperCamelCase__							)]		,					dim=1							)

           return ordered_inputs







    @property
    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           return 13
