"""simple docstring"""





import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):
      # Construct model
      if openai_config_file == "":
            UpperCAmelCase_      			=							OpenAIGPTConfig()
      else:
            UpperCAmelCase_      			=							OpenAIGPTConfig.from_json_file(lowerCAmelCase__				)
      UpperCAmelCase_      			=							OpenAIGPTModel(lowerCAmelCase__				)

      # Load weights from numpy
      load_tf_weights_in_openai_gpt(lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				)

      # Save pytorch-model
      UpperCAmelCase_      			=							pytorch_dump_folder_path + "/" + WEIGHTS_NAME
      UpperCAmelCase_      			=							pytorch_dump_folder_path + "/" + CONFIG_NAME
      print(f"""Save PyTorch model to {pytorch_weights_dump_path}"""				)
      torch.save(model.state_dict()     ,      lowerCAmelCase__				)
      print(f"""Save configuration file to {pytorch_config_dump_path}"""				)
      with open(lowerCAmelCase__     ,      "w"     ,      encoding="utf-8"				) as f:
            f.write(config.to_json_string()				)


if __name__ == "__main__":
  lowerCamelCase       	=						argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--openai_checkpoint_folder_path""",
      default=None,
      type=str,
      required=True,
      help="""Path to the TensorFlow checkpoint path.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--openai_config_file""",
      default="""""",
      type=str,
      help=(
          """An optional config json file corresponding to the pre-trained OpenAI model. \n"""
          """This specifies the model architecture."""
      ),
  )
  lowerCamelCase       	=						parser.parse_args()
  convert_openai_checkpoint_to_pytorch(
      args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
  )









"""simple docstring"""





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCamelCase       	=						datasets.utils.logging.get_logger(__name__)
@dataclass
class 				lowercase__       (   datasets.BuilderConfig				):







       '''simple docstring'''


       UpperCamelCase						=					None
       UpperCamelCase						=					"utf-8"
       UpperCamelCase						=					None
       UpperCamelCase						=					None
       UpperCamelCase						=					True  # deprecated
       UpperCamelCase						=					None  # deprecated
       UpperCamelCase						=					10 << 20  # 10MB
       UpperCamelCase						=					None



class 				lowercase__       (   datasets.ArrowBasedBuilder				):







       '''simple docstring'''


       UpperCamelCase						=					JsonConfig



       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             if self.config.block_size is not None:
                   logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"   )
                   UpperCAmelCase_      			=							self.config.block_size
             if self.config.use_threads is not True:
                   logger.warning(
                       "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."   )
             if self.config.newlines_in_values is not None:
                   raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"   )
             return datasets.DatasetInfo(features=self.config.features   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  Optional[int]   )			->						Optional[int]:


             '''simple docstring'''




             if not self.config.data_files:
                   raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""   )
             UpperCAmelCase_      			=							dl_manager.download_and_extract(self.config.data_files   )
             if isinstance(_UpperCAmelCase		,    (str, list, tuple)   ):
                   UpperCAmelCase_      			=							data_files
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,    gen_kwargs={"files": files}   )]
             UpperCAmelCase_      			=							[]
             for split_name, files in data_files.items():
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   splits.append(datasets.SplitGenerator(name=_UpperCAmelCase		,    gen_kwargs={"files": files}   )   )
             return splits



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  pa.Table   )			->						pa.Table:


             '''simple docstring'''




             if self.config.features is not None:
                   # adding missing columns
                   for column_name in set(self.config.features   ) - set(pa_table.column_names   ):
                         UpperCAmelCase_      			=							self.config.features.arrow_schema.field(_UpperCAmelCase   ).type
                         UpperCAmelCase_      			=							pa_table.append_column(_UpperCAmelCase		,    pa.array([None] * len(_UpperCAmelCase   )		,    type=_UpperCAmelCase   )   )
                   # more expensive cast to support nested structures with keys in a different order
                   # allows str <-> int/float or str to Audio for example
                   UpperCAmelCase_      			=							table_cast(_UpperCAmelCase		,    self.config.features.arrow_schema   )
             return pa_table







       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  List[Any]   )			->						str:


             '''simple docstring'''




             for file_idx, file in enumerate(itertools.chain.from_iterable(_UpperCAmelCase   )   ):
                   # If the file is one json object and if we need to look at the list of items in one specific field
                   if self.config.field is not None:
                         with open(_UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                               UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )

                         # We keep only the field we are interested in
                         UpperCAmelCase_      			=							dataset[self.config.field]

                         # We accept two format: a list of dicts or a dict of lists
                         if isinstance(_UpperCAmelCase		,    (list, tuple)   ):
                               UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                               UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                         else:
                               UpperCAmelCase_      			=							dataset
                         UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                         yield file_idx, self._cast_table(_UpperCAmelCase   )

                   # If the file has one json object per line
                   else:
                         with open(_UpperCAmelCase		,    "rb"   ) as f:
                               UpperCAmelCase_      			=							0
                               # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                               # Set a default minimum value of 16kB if the chunk size is really small
                               UpperCAmelCase_      			=							max(self.config.chunksize // 32		,    16 << 10   )
                               UpperCAmelCase_      			=							(
                                   self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
                               )
                               while True:
                                     UpperCAmelCase_      			=							f.read(self.config.chunksize   )
                                     if not batch:
                                           break
                                     # Finish current line
                                     try:
                                           batch += f.readline()
                                     except (AttributeError, io.UnsupportedOperation):
                                           batch += readline(_UpperCAmelCase   )
                                     # PyArrow only accepts utf-8 encoded bytes
                                     if self.config.encoding != "utf-8":
                                           UpperCAmelCase_      			=							batch.decode(self.config.encoding		,    errors=_UpperCAmelCase   ).encode("utf-8"   )
                                     try:
                                           while True:
                                                 try:
                                                       UpperCAmelCase_      			=							paj.read_json(
                                                           io.BytesIO(_UpperCAmelCase   )		,    read_options=paj.ReadOptions(block_size=_UpperCAmelCase   )   )
                                                       break
                                                 except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                                       if (
                                                           isinstance(_UpperCAmelCase		,    pa.ArrowInvalid   )
                                                           and "straddling" not in str(_UpperCAmelCase   )
                                                           or block_size > len(_UpperCAmelCase   )
                                                       ):
                                                             raise
                                                       else:
                                                             # Increase the block size in case it was too small.
                                                             # The block size will be reset for the next file.
                                                             logger.debug(
                                                                 F"""Batch of {len(_UpperCAmelCase   )} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."""   )
                                                             block_size *= 2
                                     except pa.ArrowInvalid as e:
                                           try:
                                                 with open(
                                                     _UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                                                       UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )
                                           except json.JSONDecodeError:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise e
                                           # If possible, parse the file as a list of json objects and exit the loop
                                           if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):  # list is the only sequence type supported in JSON
                                                 try:
                                                       UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                                                       UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                                                       UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                                                 except (pa.ArrowInvalid, AttributeError) as e:
                                                       logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                       raise ValueError(F"""Not able to read records in the JSON file at {file}."""   ) from None
                                                 yield file_idx, self._cast_table(_UpperCAmelCase   )
                                                 break
                                           else:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise ValueError(
                                                     F"""Not able to read records in the JSON file at {file}. """
                                                     F"""You should probably indicate the field of the JSON file containing your records. """
                                                     F"""This JSON file contain the following fields: {str(list(dataset.keys()   )   )}. """
                                                     F"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """   ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                     yield (file_idx, batch_idx), self._cast_table(_UpperCAmelCase   )
                                     batch_idx += 1




