


"""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    SCREAMING_SNAKE_CASE__    (		lowercase    ):
       """simple docstring"""


       a  :    List[Any]       ="codegen"
       a  :    Union[str, Any]       ={
           "max_position_embeddings": "n_positions",
           "hidden_size": "n_embd",
           "num_attention_heads": "n_head",
           "num_hidden_layers": "n_layer",
       }

       def __init__(						self		,					snake_case__=50_400		,					snake_case__=2_048		,					snake_case__=2_048		,					snake_case__=4_096		,					snake_case__=28		,					snake_case__=16		,					snake_case__=64		,					snake_case__=None		,					snake_case__="gelu_new"		,					snake_case__=0.0		,					snake_case__=0.0		,					snake_case__=0.0		,					snake_case__=1e-5		,					snake_case__=0.02		,					snake_case__=True		,					snake_case__=50_256		,					snake_case__=50_256		,					snake_case__=False		,					**snake_case__		,					):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							vocab_size
         lowerCAmelCase					:					Optional[int]								=							n_ctx
         lowerCAmelCase					:					str								=							n_positions
         lowerCAmelCase					:					Optional[int]								=							n_embd
         lowerCAmelCase					:					Optional[int]								=							n_layer
         lowerCAmelCase					:					Optional[int]								=							n_head
         lowerCAmelCase					:					List[Any]								=							n_inner
         lowerCAmelCase					:					Optional[Any]								=							rotary_dim
         lowerCAmelCase					:					str								=							activation_function
         lowerCAmelCase					:					Tuple								=							resid_pdrop
         lowerCAmelCase					:					List[str]								=							embd_pdrop
         lowerCAmelCase					:					Tuple								=							attn_pdrop
         lowerCAmelCase					:					Dict								=							layer_norm_epsilon
         lowerCAmelCase					:					List[Any]								=							initializer_range
         lowerCAmelCase					:					Union[str, Any]								=							use_cache

         lowerCAmelCase					:					Union[str, Any]								=							bos_token_id
         lowerCAmelCase					:					Optional[Any]								=							eos_token_id

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


class    SCREAMING_SNAKE_CASE__    (		lowercase    ):
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__ = "default"		,					snake_case__ = None		,					snake_case__ = False		,					):






         """simple docstring"""





         super().__init__(snake_case__		,					task=snake_case__		,					patching_specs=snake_case__		,					use_past=snake_case__						)
         if not getattr(self._config		,					"pad_token_id"		,					snake_case__						):
           # TODO: how to do that better?
           lowerCAmelCase					:					Tuple								=							0

       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}						)
         if self.use_past:
           self.fill_with_past_key_values_(snake_case__		,					direction="inputs"						)
           lowerCAmelCase					:					Any								=							{0: "batch", 1: "past_sequence + sequence"}
         else:
           lowerCAmelCase					:					Optional[int]								=							{0: "batch", 1: "sequence"}

         return common_inputs

       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





         return self._config.n_layer

       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





         return self._config.n_head

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__ = -1		,					snake_case__ = -1		,					snake_case__ = False		,					snake_case__ = None		,					):






         """simple docstring"""





         lowerCAmelCase					:					Optional[int]								=							super(snake_case__		,					self						).generate_dummy_inputs(
             snake_case__		,					batch_size=snake_case__		,					seq_length=snake_case__		,					is_pair=snake_case__		,					framework=snake_case__						)

         # We need to order the input in the way they appears in the forward()
         lowerCAmelCase					:					List[str]								=							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

             lowerCAmelCase					:					Union[str, Any]								=							common_inputs["input_ids"].shape
             # Not using the same length for past_key_values
             lowerCAmelCase					:					Dict								=							seqlen + 2
             lowerCAmelCase					:					int								=							(
                 batch,
                 self.num_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // self.num_attention_heads,
             )
             lowerCAmelCase					:					Any								=							[
                 (torch.zeros(snake_case__						), torch.zeros(snake_case__						)) for _ in range(self.num_layers						)
             ]

         lowerCAmelCase					:					str								=							common_inputs["attention_mask"]
         if self.use_past:
           lowerCAmelCase					:					List[str]								=							ordered_inputs["attention_mask"].dtype
           lowerCAmelCase					:					Tuple								=							torch.cat(
               [ordered_inputs["attention_mask"], torch.ones(snake_case__		,					snake_case__		,					dtype=snake_case__						)]		,					dim=1						)

         return ordered_inputs







       @property
       def 	lowercase__							(						self						):






         """simple docstring"""





         return 13









"""simple docstring"""







import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
     import sqlitea

     import sqlalchemy


class 				SCREAMING_SNAKE_CASE__    (		lowercase    ):
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__		,					snake_case__ = None		,					snake_case__ = None		,					snake_case__ = False		,					**snake_case__		,					):






         """simple docstring"""





         super().__init__(features=snake_case__		,					cache_dir=snake_case__		,					keep_in_memory=snake_case__		,					**snake_case__						)
         lowerCAmelCase					:					Dict								=							Sql(
             cache_dir=snake_case__		,					features=snake_case__		,					sql=snake_case__		,					con=snake_case__		,					**snake_case__		,					)







       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							None
         lowerCAmelCase					:					Optional[int]								=							None
         lowerCAmelCase					:					str								=							None
         lowerCAmelCase					:					int								=							None

         self.builder.download_and_prepare(
             download_config=snake_case__		,					download_mode=snake_case__		,					verification_mode=snake_case__		,					base_path=snake_case__		,					)

         # Build dataset for splits
         lowerCAmelCase					:					str								=							self.builder.as_dataset(
             split="train"		,					verification_mode=snake_case__		,					in_memory=self.keep_in_memory						)
         return dataset


class 				SCREAMING_SNAKE_CASE__    :
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__ = None		,					snake_case__ = None		,					**snake_case__		,					):






         """simple docstring"""





         if num_proc is not None and num_proc <= 0:
           raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""						)

         lowerCAmelCase					:					Tuple								=							dataset
         lowerCAmelCase					:					Tuple								=							name
         lowerCAmelCase					:					List[str]								=							con
         lowerCAmelCase					:					List[Any]								=							batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
         lowerCAmelCase					:					Optional[int]								=							num_proc
         lowerCAmelCase					:					Optional[int]								=							to_sql_kwargs

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							self.to_sql_kwargs.pop("sql"		,					snake_case__						)
         lowerCAmelCase					:					List[Any]								=							self.to_sql_kwargs.pop("con"		,					snake_case__						)
         lowerCAmelCase					:					Dict								=							self.to_sql_kwargs.pop("index"		,					snake_case__						)

         lowerCAmelCase					:					Dict								=							self._write(index=snake_case__		,					**self.to_sql_kwargs						)
         return written

       def 	lowercase__							(						self		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase					:					Optional[int]								=							args
         lowerCAmelCase					:					Optional[int]								=							{**to_sql_kwargs, "if_exists": "append"} if offset > 0 else to_sql_kwargs
         lowerCAmelCase					:					List[Any]								=							query_table(
             table=self.dataset.data		,					key=slice(snake_case__		,					offset + self.batch_size						)		,					indices=self.dataset._indices		,					)
         lowerCAmelCase					:					Tuple								=							batch.to_pandas()
         lowerCAmelCase					:					str								=							df.to_sql(self.name		,					self.con		,					index=snake_case__		,					**snake_case__						)
         return num_rows or len(snake_case__						)







       def 	lowercase__							(						self		,					snake_case__		,					**snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							0

         if self.num_proc is None or self.num_proc == 1:
           for offset in logging.tqdm(
               range(0		,					len(self.dataset						)		,					self.batch_size						)		,					unit="ba"		,					disable=not logging.is_progress_bar_enabled()		,					desc="Creating SQL from Arrow format"		,					):
             written += self._batch_sql((offset, index, to_sql_kwargs)						)
         else:
           lowerCAmelCase						,  lowerCAmelCase					:					Dict								=							len(self.dataset						), self.batch_size
           with multiprocessing.Pool(self.num_proc						) as pool:
             for num_rows in logging.tqdm(
                 pool.imap(
                     self._batch_sql		,					[(offset, index, to_sql_kwargs) for offset in range(0		,					snake_case__		,					snake_case__						)]		,					)		,					total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size		,					unit="ba"		,					disable=not logging.is_progress_bar_enabled()		,					desc="Creating SQL from Arrow format"		,					):
               written += num_rows

         return written





