






'''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


__A				:				str					=		logging.get_logger(__name__)

__A				:				Optional[Any]					=		{
    """gpt2""": """https://huggingface.co/gpt2/resolve/main/config.json""",
    """gpt2-medium""": """https://huggingface.co/gpt2-medium/resolve/main/config.json""",
    """gpt2-large""": """https://huggingface.co/gpt2-large/resolve/main/config.json""",
    """gpt2-xl""": """https://huggingface.co/gpt2-xl/resolve/main/config.json""",
    """distilgpt2""": """https://huggingface.co/distilgpt2/resolve/main/config.json""",
}
class 					lowercase							(				_lowerCamelCase   ):


    '''simple docstring'''

    lowerCAmelCase__   					=				"gpt2"
    lowerCAmelCase__   					=				["past_key_values"]
    lowerCAmelCase__   					=				{
        "hidden_size": "n_embd",
        "max_position_embeddings": "n_positions",
        "num_attention_heads": "n_head",
        "num_hidden_layers": "n_layer",
    }

    def __init__(       self  :  Union[str, Any]					,  __lowerCamelCase  :  Any=50257					,  __lowerCamelCase  :  int=1024					,  __lowerCamelCase  :  str=768					,  __lowerCamelCase  :  Union[str, Any]=12					,  __lowerCamelCase  :  Optional[int]=12					,  __lowerCamelCase  :  Optional[int]=None					,  __lowerCamelCase  :  Optional[int]="gelu_new"					,  __lowerCamelCase  :  List[Any]=0.1					,  __lowerCamelCase  :  Dict=0.1					,  __lowerCamelCase  :  Tuple=0.1					,  __lowerCamelCase  :  Optional[Any]=1E-5					,  __lowerCamelCase  :  int=0.0_2					,  __lowerCamelCase  :  List[str]="cls_index"					,  __lowerCamelCase  :  Optional[Any]=True					,  __lowerCamelCase  :  Tuple=None					,  __lowerCamelCase  :  List[Any]=True					,  __lowerCamelCase  :  Any=0.1					,  __lowerCamelCase  :  Tuple=True					,  __lowerCamelCase  :  Optional[int]=True					,  __lowerCamelCase  :  Optional[Any]=50256					,  __lowerCamelCase  :  Optional[Any]=50256					,  __lowerCamelCase  :  str=False					,  __lowerCamelCase  :  List[str]=False					,  **__lowerCamelCase  :  int					,  )				->				Any:

     '''simple docstring'''




     lowerCamelCase__	=  vocab_size
     lowerCamelCase__	=  n_positions
     lowerCamelCase__	=  n_embd
     lowerCamelCase__	=  n_layer
     lowerCamelCase__	=  n_head
     lowerCamelCase__	=  n_inner
     lowerCamelCase__	=  activation_function
     lowerCamelCase__	=  resid_pdrop
     lowerCamelCase__	=  embd_pdrop
     lowerCamelCase__	=  attn_pdrop
     lowerCamelCase__	=  layer_norm_epsilon
     lowerCamelCase__	=  initializer_range
     lowerCamelCase__	=  summary_type
     lowerCamelCase__	=  summary_use_proj
     lowerCamelCase__	=  summary_activation
     lowerCamelCase__	=  summary_first_dropout
     lowerCamelCase__	=  summary_proj_to_labels
     lowerCamelCase__	=  scale_attn_weights
     lowerCamelCase__	=  use_cache
     lowerCamelCase__	=  scale_attn_by_inverse_layer_idx
     lowerCamelCase__	=  reorder_and_upcast_attn

     lowerCamelCase__	=  bos_token_id
     lowerCamelCase__	=  eos_token_id

     super().__init__(bos_token_id=__lowerCamelCase					,  eos_token_id=__lowerCamelCase					,  **__lowerCamelCase  )
class 					lowercase							(				_lowerCamelCase   ):


    '''simple docstring'''

    def __init__(       self  :  Optional[Any]					,  __lowerCamelCase  :  PretrainedConfig					,  __lowerCamelCase  :  str = "default"					,  __lowerCamelCase  :  List[PatchingSpec] = None					,  __lowerCamelCase  :  bool = False					,  )				->				Tuple:

     '''simple docstring'''




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

    @property
    def  a__					(       self  :  Optional[Any]  )				->				Mapping[str, Mapping[int, str]]:

     '''simple docstring'''




     lowerCamelCase__	=  OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}  )
     if self.use_past:
      self.fill_with_past_key_values_(__lowerCamelCase					,  direction="inputs"  )
      lowerCamelCase__	=  {0: "batch", 1: "past_sequence + sequence"}
     else:
      lowerCamelCase__	=  {0: "batch", 1: "sequence"}

     return common_inputs

    @property
    def  a__					(       self  :  Union[str, Any]  )				->				int:

     '''simple docstring'''




     return self._config.n_layer

    @property
    def  a__					(       self  :  Optional[int]  )				->				int:

     '''simple docstring'''




     return self._config.n_head

    def  a__					(       self  :  Tuple					,  __lowerCamelCase  :  PreTrainedTokenizer					,  __lowerCamelCase  :  int = -1					,  __lowerCamelCase  :  int = -1					,  __lowerCamelCase  :  bool = False					,  __lowerCamelCase  :  Optional[TensorType] = None					,  )				->				Mapping[str, Any]:

     '''simple docstring'''




     lowerCamelCase__	=  super(__lowerCamelCase					,  self  ).generate_dummy_inputs(
         __lowerCamelCase					,  batch_size=__lowerCamelCase					,  seq_length=__lowerCamelCase					,  is_pair=__lowerCamelCase					,  framework=__lowerCamelCase  )

     # We need to order the input in the way they appears in the forward()
     lowerCamelCase__	=  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__	,						lowerCamelCase__	=  common_inputs["input_ids"].shape
       # Not using the same length for past_key_values
       lowerCamelCase__	=  seqlen + 2
       lowerCamelCase__	=  (
           batch,
           self.num_attention_heads,
           past_key_values_length,
           self._config.hidden_size // self.num_attention_heads,
       )
       lowerCamelCase__	=  [
           (torch.zeros(__lowerCamelCase  ), torch.zeros(__lowerCamelCase  )) for _ in range(self.num_layers  )
       ]

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

     return ordered_inputs




    @property
    def  a__					(       self  :  Optional[int]  )				->				int:

     '''simple docstring'''




     return 13









'''simple docstring'''





import os
import sys
from pathlib import Path

import pytest

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node

from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch






def  lowerCamelCase_  (     ):
 lowerCamelCase__	=  Dataset.from_dict({"i": range(17)})
 lowerCamelCase__	=  len(lowercase__)
 lowerCamelCase__	=  3
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(lowercase__) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






@pytest.mark.parametrize("shards_per_node"					,  [1, 2, 3])
def  lowerCamelCase_  (     lowercase__):
 def gen(lowercase__):
  for shard in shards:
   yield from ({"i": i, "shard": shard} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  shards_per_node * world_size
 lowerCamelCase__	=  {"shards": [F'''shard_{shard_idx}.txt''' for shard_idx in range(lowercase__)]}
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__					,  gen_kwargs=lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 assert full_ds.n_shards == world_size * shards_per_node
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  2
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))

 lowerCamelCase__	=  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle(seed=42)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle()

 lowerCamelCase__	=  split_dataset_by_node(full_ds.shuffle(seed=42)					,  rank=0					,  world_size=lowercase__)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(full_ds.shuffle()					,  rank=0					,  world_size=lowercase__)






@pytest.mark.parametrize("streaming"					,  [False, True])
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__):
 lowerCamelCase__	=  2
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())






@pytest.mark.parametrize(
    "nproc_per_node, num_workers"					,  [
        (2, 2),  # each node has 2 shards and each worker has 1 shards
        (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
    ]					,  )
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__					,  lowercase__):
 lowerCamelCase__	=  True
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
        --num_workers={num_workers}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())

