


"""simple docstring"""







import json
import os
import shutil
import warnings
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import List

from ..utils import logging
from . import BaseTransformersCLICommand


try:
     from cookiecutter.main import cookiecutter

     lowerCAmelCase__				    =       True
except ImportError:
     lowerCAmelCase__				    =       False

lowerCAmelCase__				    =       logging.get_logger(__name__)  # pylint: disable=invalid-name
def  a__      (    SCREAMING_SNAKE_CASE   :   Namespace    ):


  '''simple docstring'''







  return AddNewModelCommand(args.testing    ,				args.testing_file    ,				path=args.path    )





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


       @staticmethod
       def 	lowercase__							(						snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							parser.add_parser("add-new-model"						)
         add_new_model_parser.add_argument("--testing"		,					action="store_true"		,					help="If in testing mode."						)
         add_new_model_parser.add_argument("--testing_file"		,					type=_a		,					help="Configuration file on which to run."						)
         add_new_model_parser.add_argument(
             "--path"		,					type=_a		,					help="Path to cookiecutter. Should only be used for testing purposes."						)
         add_new_model_parser.set_defaults(func=_a						)

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






         """simple docstring"""





         lowerCAmelCase					:					Dict								=							testing
         lowerCAmelCase					:					Tuple								=							testing_file
         lowerCAmelCase					:					int								=							path







       def 	lowercase__							(						self						):






         """simple docstring"""





         warnings.warn(
             "The command `transformers-cli add-new-model` is deprecated and will be removed in v5 of Transformers. "
             "It is not actively maintained anymore, so might give a result that won't pass all tests and quality "
             "checks, you should use `transformers-cli add-new-model-like` instead."						)
         if not _has_cookiecutter:
           raise ImportError(
               "Model creation dependencies are required to use the `add_new_model` command. Install them by running "
               "the following at the root of your `transformers` clone:\n\n\t$ pip install -e .[modelcreation]\n"						)
         # Ensure that there is no other `cookiecutter-template-xxx` directory in the current working directory
         lowerCAmelCase					:					Any								=							[directory for directory in os.listdir() if "cookiecutter-template-" == directory[:22]]
         if len(_a						) > 0:
           raise ValueError(
               "Several directories starting with `cookiecutter-template-` in current working directory. "
               "Please clean your directory by removing all folders starting with `cookiecutter-template-` or "
               "change your working directory."						)

         lowerCAmelCase					:					Tuple								=							(
             Path(_a						).parent.parent.parent.parent if self._path is None else Path(self._path						).parent.parent
         )
         lowerCAmelCase					:					List[Any]								=							path_to_transformer_root / "templates" / "adding_a_new_model"

         # Execute cookiecutter
         if not self._testing:
           cookiecutter(str(_a						)						)
         else:
           with open(self._testing_file		,					"r"						) as configuration_file:
             lowerCAmelCase					:					int								=							json.load(_a						)

           cookiecutter(
               str(path_to_cookiecutter if self._path is None else self._path						)		,					no_input=_a		,					extra_context=_a		,					)

         lowerCAmelCase					:					Optional[Any]								=							[directory for directory in os.listdir() if "cookiecutter-template-" in directory[:22]][0]

         # Retrieve configuration
         with open(directory + "/configuration.json"		,					"r"						) as configuration_file:
           lowerCAmelCase					:					Any								=							json.load(_a						)

         lowerCAmelCase					:					Tuple								=							configuration["lowercase_modelname"]
         lowerCAmelCase					:					Optional[Any]								=							configuration["generate_tensorflow_pytorch_and_flax"]
         os.remove(f"""{directory}/configuration.json"""						)

         lowerCAmelCase					:					List[Any]								=							"PyTorch" in generate_tensorflow_pytorch_and_flax
         lowerCAmelCase					:					int								=							"TensorFlow" in generate_tensorflow_pytorch_and_flax
         lowerCAmelCase					:					List[Any]								=							"Flax" in generate_tensorflow_pytorch_and_flax

         lowerCAmelCase					:					Tuple								=							f"""{path_to_transformer_root}/src/transformers/models/{lowercase_model_name}"""
         os.makedirs(_a		,					exist_ok=_a						)
         os.makedirs(f"""{path_to_transformer_root}/tests/models/{lowercase_model_name}"""		,					exist_ok=_a						)

         # Tests require submodules as they have parent imports
         with open(f"""{path_to_transformer_root}/tests/models/{lowercase_model_name}/__init__.py"""		,					"w"						):
           pass

         shutil.move(
             f"""{directory}/__init__.py"""		,					f"""{model_dir}/__init__.py"""		,					)
         shutil.move(
             f"""{directory}/configuration_{lowercase_model_name}.py"""		,					f"""{model_dir}/configuration_{lowercase_model_name}.py"""		,					)

         def remove_copy_lines(snake_case__						):
           with open(_a		,					"r"						) as f:
             lowerCAmelCase					:					int								=							f.readlines()
           with open(_a		,					"w"						) as f:
             for line in lines:
               if "# Copied from transformers." not in line:
                 f.write(_a						)

         if output_pytorch:
           if not self._testing:
             remove_copy_lines(f"""{directory}/modeling_{lowercase_model_name}.py"""						)

           shutil.move(
               f"""{directory}/modeling_{lowercase_model_name}.py"""		,					f"""{model_dir}/modeling_{lowercase_model_name}.py"""		,					)

           shutil.move(
               f"""{directory}/test_modeling_{lowercase_model_name}.py"""		,					f"""{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_{lowercase_model_name}.py"""		,					)
         else:
           os.remove(f"""{directory}/modeling_{lowercase_model_name}.py"""						)
           os.remove(f"""{directory}/test_modeling_{lowercase_model_name}.py"""						)

         if output_tensorflow:
           if not self._testing:
             remove_copy_lines(f"""{directory}/modeling_tf_{lowercase_model_name}.py"""						)

           shutil.move(
               f"""{directory}/modeling_tf_{lowercase_model_name}.py"""		,					f"""{model_dir}/modeling_tf_{lowercase_model_name}.py"""		,					)

           shutil.move(
               f"""{directory}/test_modeling_tf_{lowercase_model_name}.py"""		,					f"""{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_tf_{lowercase_model_name}.py"""		,					)
         else:
           os.remove(f"""{directory}/modeling_tf_{lowercase_model_name}.py"""						)
           os.remove(f"""{directory}/test_modeling_tf_{lowercase_model_name}.py"""						)

         if output_flax:
           if not self._testing:
             remove_copy_lines(f"""{directory}/modeling_flax_{lowercase_model_name}.py"""						)

           shutil.move(
               f"""{directory}/modeling_flax_{lowercase_model_name}.py"""		,					f"""{model_dir}/modeling_flax_{lowercase_model_name}.py"""		,					)

           shutil.move(
               f"""{directory}/test_modeling_flax_{lowercase_model_name}.py"""		,					f"""{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_flax_{lowercase_model_name}.py"""		,					)
         else:
           os.remove(f"""{directory}/modeling_flax_{lowercase_model_name}.py"""						)
           os.remove(f"""{directory}/test_modeling_flax_{lowercase_model_name}.py"""						)

         shutil.move(
             f"""{directory}/{lowercase_model_name}.md"""		,					f"""{path_to_transformer_root}/docs/source/en/model_doc/{lowercase_model_name}.md"""		,					)

         shutil.move(
             f"""{directory}/tokenization_{lowercase_model_name}.py"""		,					f"""{model_dir}/tokenization_{lowercase_model_name}.py"""		,					)

         shutil.move(
             f"""{directory}/tokenization_fast_{lowercase_model_name}.py"""		,					f"""{model_dir}/tokenization_{lowercase_model_name}_fast.py"""		,					)

         from os import fdopen, remove
         from shutil import copymode, move
         from tempfile import mkstemp

         def replace(snake_case__		,					snake_case__		,					snake_case__						):
           # Create temp file
           lowerCAmelCase					:					List[str]								=							mkstemp()
           lowerCAmelCase					:					Tuple								=							False
           with fdopen(_a		,					"w"						) as new_file:
             with open(_a						) as old_file:
               for line in old_file:
                 new_file.write(_a						)
                 if line_to_copy_below in line:
                   lowerCAmelCase					:					str								=							True
                   for line_to_copy in lines_to_copy:
                     new_file.write(_a						)

           if not line_found:
             raise ValueError(f"""Line {line_to_copy_below} was not found in file."""						)

           # Copy the file permissions from the old file to the new file
           copymode(_a		,					_a						)
           # Remove original file
           remove(_a						)
           # Move new file
           move(_a		,					_a						)

         def skip_units(snake_case__						):
           return (
               ("generating PyTorch" in line and not output_pytorch)
               or ("generating TensorFlow" in line and not output_tensorflow)
               or ("generating Flax" in line and not output_flax)
           )

         def replace_in_files(snake_case__						):
           with open(_a						) as datafile:
             lowerCAmelCase					:					Tuple								=							[]
             lowerCAmelCase					:					str								=							False
             lowerCAmelCase					:					Tuple								=							False
             for line in datafile:
               if "# To replace in: " in line and "##" not in line:
                 lowerCAmelCase					:					int								=							line.split("\""						)[1]
                 lowerCAmelCase					:					Optional[Any]								=							skip_units(_a						)
               elif "# Below: " in line and "##" not in line:
                 lowerCAmelCase					:					Any								=							line.split("\""						)[1]
                 lowerCAmelCase					:					str								=							skip_units(_a						)
               elif "# End." in line and "##" not in line:
                 if not skip_file and not skip_snippet:
                   replace(_a		,					_a		,					_a						)

                 lowerCAmelCase					:					Optional[Any]								=							[]
               elif "# Replace with" in line and "##" not in line:
                 lowerCAmelCase					:					Dict								=							[]
               elif "##" not in line:
                 lines_to_copy.append(_a						)

           remove(_a						)

         replace_in_files(f"""{directory}/to_replace_{lowercase_model_name}.py"""						)
         os.rmdir(_a						)









"""simple docstring"""







import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
def  a__      (    SCREAMING_SNAKE_CASE   :   str    ,				SCREAMING_SNAKE_CASE   :   Any    ):


  '''simple docstring'''







  assert isinstance(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
  assert dataset.num_rows == 4
  assert dataset.num_columns == 3
  assert dataset.column_names == ["col_1", "col_2", "col_3"]
  for feature, expected_dtype in expected_features.items():
    assert dataset.features[feature].dtype == expected_dtype
@pytest.mark.parametrize("keep_in_memory"    ,				[False, True]    )
def  a__      (    SCREAMING_SNAKE_CASE   :   Any    ,				SCREAMING_SNAKE_CASE   :   List[str]    ,				SCREAMING_SNAKE_CASE   :   int    ):


  '''simple docstring'''







  lowerCAmelCase					:					List[Any]								=							tmp_path / "cache"
  lowerCAmelCase					:					Optional[int]								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    lowerCAmelCase					:					str								=							ParquetDatasetReader(SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ,				keep_in_memory=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_dataset(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
@pytest.mark.parametrize(
    "features"    ,				[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]    ,				)
def  a__      (    SCREAMING_SNAKE_CASE   :   Tuple    ,				SCREAMING_SNAKE_CASE   :   Optional[Any]    ,				SCREAMING_SNAKE_CASE   :   Tuple    ):


  '''simple docstring'''







  lowerCAmelCase					:					Optional[int]								=							tmp_path / "cache"
  lowerCAmelCase					:					Optional[Any]								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCAmelCase					:					Tuple								=							features.copy() if features else default_expected_features
  lowerCAmelCase					:					List[str]								=							(
      Features({feature: Value(SCREAMING_SNAKE_CASE    ) for feature, dtype in features.items()}    ) if features is not None else None
  )
  lowerCAmelCase					:					List[str]								=							ParquetDatasetReader(SCREAMING_SNAKE_CASE    ,				features=SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_dataset(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
@pytest.mark.parametrize("split"    ,				[None, NamedSplit("train"    ), "train", "test"]    )
def  a__      (    SCREAMING_SNAKE_CASE   :   Union[str, Any]    ,				SCREAMING_SNAKE_CASE   :   Union[str, Any]    ,				SCREAMING_SNAKE_CASE   :   Optional[int]    ):


  '''simple docstring'''







  lowerCAmelCase					:					List[Any]								=							tmp_path / "cache"
  lowerCAmelCase					:					int								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCAmelCase					:					Tuple								=							ParquetDatasetReader(SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ,				split=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_dataset(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
  assert dataset.split == split if split else "train"
@pytest.mark.parametrize("path_type"    ,				[str, list]    )
def  a__      (    SCREAMING_SNAKE_CASE   :   List[Any]    ,				SCREAMING_SNAKE_CASE   :   Dict    ,				SCREAMING_SNAKE_CASE   :   List[str]    ):


  '''simple docstring'''







  if issubclass(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    ):
    lowerCAmelCase					:					Optional[Any]								=							parquet_path
  elif issubclass(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    ):
    lowerCAmelCase					:					Dict								=							[parquet_path]
  lowerCAmelCase					:					Optional[Any]								=							tmp_path / "cache"
  lowerCAmelCase					:					Any								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCAmelCase					:					Optional[Any]								=							ParquetDatasetReader(SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_dataset(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
def  a__      (    SCREAMING_SNAKE_CASE   :   Dict    ,				SCREAMING_SNAKE_CASE   :   List[str]    ,				SCREAMING_SNAKE_CASE   :   Tuple=("train",)    ):


  '''simple docstring'''







  assert isinstance(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
  for split in splits:
    lowerCAmelCase					:					List[Any]								=							dataset_dict[split]
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]
    for feature, expected_dtype in expected_features.items():
      assert dataset.features[feature].dtype == expected_dtype
@pytest.mark.parametrize("keep_in_memory"    ,				[False, True]    )
def  a__      (    SCREAMING_SNAKE_CASE   :   Optional[int]    ,				SCREAMING_SNAKE_CASE   :   Tuple    ,				SCREAMING_SNAKE_CASE   :   Tuple    ):


  '''simple docstring'''







  lowerCAmelCase					:					Tuple								=							tmp_path / "cache"
  lowerCAmelCase					:					List[Any]								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    lowerCAmelCase					:					Optional[Any]								=							ParquetDatasetReader(
        {"train": parquet_path}    ,				cache_dir=SCREAMING_SNAKE_CASE    ,				keep_in_memory=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_datasetdict(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
@pytest.mark.parametrize(
    "features"    ,				[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]    ,				)
def  a__      (    SCREAMING_SNAKE_CASE   :   Optional[int]    ,				SCREAMING_SNAKE_CASE   :   Dict    ,				SCREAMING_SNAKE_CASE   :   Optional[int]    ):


  '''simple docstring'''







  lowerCAmelCase					:					Any								=							tmp_path / "cache"
  lowerCAmelCase					:					Union[str, Any]								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCAmelCase					:					List[Any]								=							features.copy() if features else default_expected_features
  lowerCAmelCase					:					List[Any]								=							(
      Features({feature: Value(SCREAMING_SNAKE_CASE    ) for feature, dtype in features.items()}    ) if features is not None else None
  )
  lowerCAmelCase					:					Optional[int]								=							ParquetDatasetReader({"train": parquet_path}    ,				features=SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_datasetdict(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
@pytest.mark.parametrize("split"    ,				[None, NamedSplit("train"    ), "train", "test"]    )
def  a__      (    SCREAMING_SNAKE_CASE   :   str    ,				SCREAMING_SNAKE_CASE   :   Dict    ,				SCREAMING_SNAKE_CASE   :   List[str]    ):


  '''simple docstring'''







  if split:
    lowerCAmelCase					:					List[str]								=							{split: parquet_path}
  else:
    lowerCAmelCase					:					List[str]								=							"train"
    lowerCAmelCase					:					str								=							{"train": parquet_path, "test": parquet_path}
  lowerCAmelCase					:					Optional[int]								=							tmp_path / "cache"
  lowerCAmelCase					:					List[Any]								=							{"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCAmelCase					:					List[str]								=							ParquetDatasetReader(SCREAMING_SNAKE_CASE    ,				cache_dir=SCREAMING_SNAKE_CASE    ).read()
  _check_parquet_datasetdict(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    ,				splits=list(path.keys()    )    )
  assert all(dataset[split].split == split for split in path.keys()    )
def  a__      (    SCREAMING_SNAKE_CASE   :   Optional[Any]    ,				SCREAMING_SNAKE_CASE   :   Any    ):


  '''simple docstring'''







  lowerCAmelCase					:					Optional[int]								=							ParquetDatasetWriter(SCREAMING_SNAKE_CASE    ,				tmp_path / "foo.parquet"    )
  assert writer.write() > 0
  lowerCAmelCase					:					List[str]								=							pq.ParquetFile(tmp_path / "foo.parquet"    )
  lowerCAmelCase					:					Union[str, Any]								=							pf.read()
  assert dataset.data.table == output_table
def  a__      (    SCREAMING_SNAKE_CASE   :   List[Any]    ,				SCREAMING_SNAKE_CASE   :   List[str]    ):


  '''simple docstring'''







  lowerCAmelCase					:					str								=							str(shared_datadir / "test_image_rgb.jpg"    )
  lowerCAmelCase					:					List[str]								=							{"image": [image_path]}
  lowerCAmelCase					:					str								=							Features({"image": Image()}    )
  lowerCAmelCase					:					Optional[Any]								=							Dataset.from_dict(SCREAMING_SNAKE_CASE    ,				features=SCREAMING_SNAKE_CASE    )
  lowerCAmelCase					:					List[Any]								=							ParquetDatasetWriter(SCREAMING_SNAKE_CASE    ,				tmp_path / "foo.parquet"    )
  assert writer.write() > 0

  lowerCAmelCase					:					Any								=							Dataset.from_parquet(str(tmp_path / "foo.parquet"    )    )
  assert dataset.features == reloaded_dataset.features

  lowerCAmelCase					:					Optional[int]								=							ParquetDatasetReader(str(tmp_path / "foo.parquet"    )    ,				streaming=SCREAMING_SNAKE_CASE    ).read()
  assert dataset.features == reloaded_iterable_dataset.features





@pytest.mark.parametrize(
    "feature, expected"    ,				[
        (Features({"foo": Value("int32"    )}    ), None),
        (Features({"image": Image(), "foo": Value("int32"    )}    ), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({"nested": Sequence(Audio()    )}    ), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ]    ,				)
def  a__      (    SCREAMING_SNAKE_CASE   :   str    ,				SCREAMING_SNAKE_CASE   :   List[str]    ):


  '''simple docstring'''







  assert get_writer_batch_size(SCREAMING_SNAKE_CASE    ) == expected





