

from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_SCREAMING_SNAKE_CASE  :						Tuple         =    logging.get_logger(__name__)
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     if isinstance(_A   ,	np.ndarray       ):
          return list(tensor.shape       )

     SCREAMING_SNAKE_CASE__					=					tf.shape(_A       )

     if tensor.shape == tf.TensorShape(_A       ):
          return dynamic

     SCREAMING_SNAKE_CASE__					=					tensor.shape.as_list()

     return [dynamic[i] if s is None else s for i, s in enumerate(_A       )]
def        UpperCAmelCase_     (		_A   ,	_A = None   ,	_A = None       ):





     '''simple docstring'''




     return tf.nn.softmax(logits=logits + 1e-9   ,	axis=_A   ,	name=_A       )
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A=1e-5   ,	_A=-1       ):





     '''simple docstring'''




     if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(_A   ,	_A       ):
          raise NotImplementedError('''Only 1D weight and bias tensors are supported for now, with only a single axis.'''       )

     # Get mean and variance on the axis to be normalized
     SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					tf.nn.moments(_A   ,	axes=[axis]   ,	keepdims=_A       )

     if axis != -1:
          # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
          # on every dimension except axis
          SCREAMING_SNAKE_CASE__					=					[1] * inputs.shape.rank
          SCREAMING_SNAKE_CASE__					=					shape_list(_A       )[axis]
          SCREAMING_SNAKE_CASE__					=					tf.reshape(_A   ,	_A       )
          SCREAMING_SNAKE_CASE__					=					tf.reshape(_A   ,	_A       )

     # Compute layer normalization using the batch_normalization
     # function.
     SCREAMING_SNAKE_CASE__					=					tf.nn.batch_normalization(
         _A   ,	_A   ,	_A   ,	offset=_A   ,	scale=_A   ,	variance_epsilon=_A   ,	)
     return outputs
def        UpperCAmelCase_     (		_A   ,	_A=0   ,	_A=-1       ):





     '''simple docstring'''




     if end_dim < 0:
          end_dim += input.shape.rank
     if start_dim < 0:
          start_dim += input.shape.rank

     if start_dim == end_dim:
          return input

     SCREAMING_SNAKE_CASE__					=					tf.shape(_A       )
     SCREAMING_SNAKE_CASE__					=					tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]       )
     SCREAMING_SNAKE_CASE__					=					tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]   ,	axis=0       )
     return tf.reshape(_A   ,	_A       )
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     if not isinstance(_A   ,	tf.Tensor       ):
          SCREAMING_SNAKE_CASE__					=					tf.convert_to_tensor(_A       )  # Catches stray NumPy inputs
     if encoder_attention_mask.shape.rank == 3:
          SCREAMING_SNAKE_CASE__					=					encoder_attention_mask[:, None, :, :]
     if encoder_attention_mask.shape.rank == 2:
          SCREAMING_SNAKE_CASE__					=					encoder_attention_mask[:, None, None, :]
     # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
     # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
     # /transformer/transformer_layers.py#L270
     # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
     # encoder_extended_attention_mask.transpose(-1, -2))
     SCREAMING_SNAKE_CASE__					=					(
         tf.cast(1   ,	encoder_attention_mask.dtype       ) - encoder_extended_attention_mask
     ) * encoder_extended_attention_mask.dtype.min

     return encoder_extended_attention_mask
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A = "input_ids"       ):





     '''simple docstring'''




     tf.debugging.assert_less(
         _A   ,	tf.cast(_A   ,	dtype=tensor.dtype       )   ,	message=(
             F'''The maximum value of {tensor_name} ({tf.math.reduce_max(_A       )}) must be smaller than the embedding '''
             F'''layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'''
         )   ,	)
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					6_45_12
     # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
     # because in that case even chunking the array would not make the saving
     # possible.
     SCREAMING_SNAKE_CASE__					=					[x for x in data if len(_A       ) > HDF5_OBJECT_HEADER_LIMIT]

     # Expecting this to never be true.
     if bad_attributes:
          raise RuntimeError(
              '''The following attributes cannot be saved to HDF5 file because '''
              F'''they are larger than {HDF5_OBJECT_HEADER_LIMIT} '''
              F'''bytes: {bad_attributes}'''       )

     SCREAMING_SNAKE_CASE__					=					np.asarray(_A       )

     SCREAMING_SNAKE_CASE__					=					1
     SCREAMING_SNAKE_CASE__					=					np.array_split(_A   ,	_A       )

     # This will never loop forever thanks to the test above.
     while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data       ):
          num_chunks += 1
          SCREAMING_SNAKE_CASE__					=					np.array_split(_A   ,	_A       )

     if num_chunks > 1:
          for chunk_id, chunk_data in enumerate(_A       ):
               SCREAMING_SNAKE_CASE__					=					chunk_data
     else:
          SCREAMING_SNAKE_CASE__					=					data
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     if name in group.attrs:
          SCREAMING_SNAKE_CASE__					=					[n.decode('''utf8'''       ) if hasattr(_A   ,	'''decode'''       ) else n for n in group.attrs[name]]
     else:
          SCREAMING_SNAKE_CASE__					=					[]
          SCREAMING_SNAKE_CASE__					=					0
          while "%s%d" % (name, chunk_id) in group.attrs:
               data.extend(
                   [n.decode('''utf8'''       ) if hasattr(_A   ,	'''decode'''       ) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]       )
               chunk_id += 1
     return data







def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     def _expand_single_ad_tensor(_A       ):
          if isinstance(_A   ,	tf.Tensor       ) and t.shape.rank == 1:
               return tf.expand_dims(_A   ,	axis=-1       )
          return t

     return tf.nest.map_structure(_expand_single_ad_tensor   ,	_A       )





from dataclasses import dataclass, field
from typing import Optional






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Model name or path of model to be trained."} )
  a     =      field(
      default="./"		,  metadata={"help": "Save dir where model repo is cloned and models updates are saved to."} )
  a     =      field(
      default="codeparrot/codeparrot-clean-train"		,  metadata={"help": "Name or path of training dataset."} )
  a     =      field(
      default="codeparrot/codeparrot-clean-valid"		,  metadata={"help": "Name or path of validation dataset."} )
  a     =      field(default=2		,  metadata={"help": "Batch size for training."} )
  a     =      field(default=2		,  metadata={"help": "Batch size for evaluation."} )
  a     =      field(default=0.1		,  metadata={"help": "Value of weight decay."} )
  a     =      field(
      default=1_00_00		,  metadata={"help": "Size of buffer used to shuffle streaming dataset."} )
  a     =      field(default=2e-4		,  metadata={"help": "Learning rate fo training."} )
  a     =      field(default="cosine"		,  metadata={"help": "Learning rate."} )
  a     =      field(
      default=7_50		,  metadata={"help": "Number of warmup steps in the learning rate schedule."} )
  a     =      field(
      default=16		,  metadata={"help": "Number of gradient accumulation steps."} )
  a     =      field(
      default=A__		,  metadata={"help": "Use gradient checkpointing to reduce memory footprint."} )
  a     =      field(default=5_00_00		,  metadata={"help": "Maximum number of training steps."} )
  a     =      field(
      default=-1		,  metadata={"help": "Maximum number of evaluation steps. If -1 the full dataset is evaluated."} )
  a     =      field(default=10_24		,  metadata={"help": "Sequence lengths used for training."} )
  a     =      field(default=1		,  metadata={"help": "Training seed."} )
  a     =      field(
      default=10_24		,  metadata={"help": "Interval to save checkpoints. Measured as number of forward passes not training steps."}		,  )
  a     =      field(
      default=A__		,  metadata={"help": "States path if the training should continue from a checkpoint folder."} )
  a     =      field(default=A__		,  metadata={"help": "If True the data is pretokenized."} )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Model name or path of model to be evaluated."} )
  a     =      field(
      default="codeparrot/codeparrot-clean-valid"		,  metadata={"help": "Name or path of validation dataset."} )
  a     =      field(default=2		,  metadata={"help": "Batch size used for evaluation."} )
  a     =      field(
      default=-1		,  metadata={"help": "Maximum number of evaluation steps. If -1 the full dataset is evaluated."} )
  a     =      field(default=10_24		,  metadata={"help": "Length of sequences to be evaluated."} )
  a     =      field(default=1		,  metadata={"help": "Random seed used for evaluation."} )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Model name or path of model to be evaluated."} )
  a     =      field(default=A__		,  metadata={"help": "Number of workers used for code evaluation."} )
  a     =      field(
      default=A__		,  metadata={"help": "The number of human-eval tasks to run. If not included all tasks are evaluated."}		,  )
  a     =      field(
      default=A__		,  metadata={"help": "Sample from the language model's output distribution."} )
  a     =      field(default=0.2		,  metadata={"help": "Sampling temperature used for generation."} )
  a     =      field(default=2_56		,  metadata={"help": "Maximum number of newly generated tokens."} )
  a     =      field(default=0		,  metadata={"help": "Top-k parameter used for generation."} )
  a     =      field(default=0.9_5		,  metadata={"help": "Top-p parameter used for nucleus sampling."} )
  a     =      field(default=10		,  metadata={"help": "Number of generations to run in parallel."} )
  a     =      field(
      default=2_00		,  metadata={"help": "Number of completions to generate for each sample."} )
  a     =      field(default=1		,  metadata={"help": "Random seed used for evaluation."} )
  a     =      field(
      default="eval_results.json"		,  metadata={"help": "Random seed used for evaluation."} )
  a     =      field(
      default="0"		,  metadata={"help": "Allow `code_eval` to execute Python code on machine"} )
  a     =      field(
      default=-1		,  metadata={
          "help": (
              "Determine which device to run the `text-generation` Pipeline on. -1 is CPU and any zero or positive"
              " number corresponds to which GPU device id to run on."
          )
      }		,  )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default=A__		,  metadata={
          "help": "The number of CPU cores to use for parallel preprocessing. Default uses the maximum available."
      }		,  )
  a     =      field(
      default="transformersbook/codeparrot"		,  metadata={"help": "Folder or name of dataset to process."} )
  a     =      field(
      default="codeparrot-clean"		,  metadata={"help": "Folder to save processed processed dataset."} )
  a     =      field(
      default=10_00_00		,  metadata={"help": "Number of files to save per JSON output file."} )
  a     =      field(default="content"		,  metadata={"help": "Column containing text data to process."} )
  a     =      field(
      default=10_00		,  metadata={"help": "Maximum line length in file, otherwise file is filtered."} )
  a     =      field(
      default=1_00		,  metadata={"help": "Maximum mean line length in file, otherwise file is filtered."} )
  a     =      field(
      default=0.2_5		,  metadata={"help": "Maximum fraction of non-alphanumeric characters, otherwise file is filtered."} )
  a     =      field(
      default=1.5		,  metadata={"help": "Minimum character token ratio for the file, otherwise file is filtered."} )
  a     =      field(
      default=0.7		,  metadata={"help": "Probability for filtering config, test and uncommon files."} )
  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Name or path to the tokenizer."}		,  )
  a     =      field(
      default=A__		,  metadata={"help": "If True, near-duplicate samples are removed."} )
  a     =      field(
      default=0.8_5		,  metadata={"help": "Jaccard threshold for near-duplicate samples."} )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="gpt2"		,  metadata={"help": "Base tokenizer to build new tokenizer from."} )
  a     =      field(
      default="transformersbook/codeparrot-train"		,  metadata={"help": "Dataset to train tokenizer on."} )
  a     =      field(default="content"		,  metadata={"help": "Column containing text data to process."} )
  a     =      field(default=20_00_00		,  metadata={"help": "Number of examples to train tokenizer on."} )
  a     =      field(
      default=3_27_68		,  metadata={"help": "Number of examples to train the tokenizer on."} )
  a     =      field(default="codeparrot"		,  metadata={"help": "Name of new tokenizer."} )
  a     =      field(default=A__		,  metadata={"help": "Push saved tokenizer to the hub."} )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Name or path to the tokenizer."} )
  a     =      field(
      default="codeparrot/codeparrot-clean-train"		,  metadata={"help": "Name or path to the dataset to pretokenize."} )
  a     =      field(
      default="tokenized-codeparrot-train"		,  metadata={"help": "Repo name of the pretokenized data."} )
  a     =      field(default=A__		,  metadata={"help": "Number of workers used for code evaluation."} )






@dataclass
class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      field(
      default="gpt2-large"		,  metadata={"help": "Configuration to use for model initialization."} )
  a     =      field(
      default="codeparrot/codeparrot"		,  metadata={"help": "Tokenizer attached to model."} )
  a     =      field(default="codeparrot"		,  metadata={"help": "Name of the created model."} )
  a     =      field(default=A__		,  metadata={"help": "Push saved tokenizer to the hub."} )


