


'''simple docstring'''

import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases






def __UpperCamelCase   (			__lowerCamelCase		: int      ,  __lowerCamelCase		: Tuple				)      ->     int:







     '''simple docstring'''






     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     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 __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Any      ,  __lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
          _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase      ,  keep_in_memory=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






@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 __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Optional[int]				)      ->     Optional[Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






@pytest.mark.parametrize(
    "features"      ,  [
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]      ,  )
def __UpperCamelCase   (			__lowerCamelCase		: str      ,  __lowerCamelCase		: Any      ,  __lowerCamelCase		: List[str]				)      ->     List[str]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_3": "float64", "col_1": "string", "col_2": "int64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     assert dataset.num_rows == 2
     assert dataset.num_columns == 3
     assert dataset.column_names == ["col_3", "col_1", "col_2"]
     for feature, expected_dtype in expected_features.items():
          assert dataset.features[feature].dtype == expected_dtype






def __UpperCamelCase   (			__lowerCamelCase		: Any      ,  __lowerCamelCase		: Optional[Any]				)      ->     int:







     '''simple docstring'''






     _a						=		{"col_2": "int64", "col_3": "float64", "col_1": "string"}
     _a						=		features.copy()
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		tmp_path / "cache"
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     assert dataset.num_rows == 2
     assert dataset.num_columns == 3
     assert dataset.column_names == ["col_2", "col_3", "col_1"]
     for feature, expected_dtype in expected_features.items():
          assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("split"      ,  [None, NamedSplit("train"				), "train", "test"]				)
def __UpperCamelCase   (			__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: Tuple				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase      ,  split=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)
     assert dataset.split == split if split else "train"






@pytest.mark.parametrize("path_type"      ,  [str, list]				)
def __UpperCamelCase   (			__lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: Optional[Any]				)      ->     Union[str, Any]:







     '''simple docstring'''






     if issubclass(__lowerCamelCase      ,  __lowerCamelCase				):
          _a						=		jsonl_path
     elif issubclass(__lowerCamelCase      ,  __lowerCamelCase				):
          _a						=		[jsonl_path]
     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






def __UpperCamelCase   (			__lowerCamelCase		: str      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Dict=("train",)				)      ->     Tuple:







     '''simple docstring'''






     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     for split in splits:
          _a						=		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 __UpperCamelCase   (			__lowerCamelCase		: int      ,  __lowerCamelCase		: int      ,  __lowerCamelCase		: str				)      ->     List[Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
          _a						=		JsonDatasetReader({"train": jsonl_path}      ,  cache_dir=__lowerCamelCase      ,  keep_in_memory=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase				)






@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 __UpperCamelCase   (			__lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: int      ,  __lowerCamelCase		: Union[str, Any]				)      ->     Dict:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader({"train": jsonl_path}      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase				)






@pytest.mark.parametrize("split"      ,  [None, NamedSplit("train"				), "train", "test"]				)
def __UpperCamelCase   (			__lowerCamelCase		: int      ,  __lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Dict				)      ->     List[str]:







     '''simple docstring'''






     if split:
          _a						=		{split: jsonl_path}
     else:
          _a						=		"train"
          _a						=		{"train": jsonl_path, "test": jsonl_path}
     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase      ,  splits=list(path.keys()				)				)
     assert all(dataset[split].split == split for split in path.keys()				)






def __UpperCamelCase   (			__lowerCamelCase		: Dict				)      ->     Tuple:







     '''simple docstring'''






     return json.load(__lowerCamelCase				)






def __UpperCamelCase   (			__lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     return [json.loads(__lowerCamelCase				) for line in buffer]






class 						__SCREAMING_SNAKE_CASE					:






     @pytest.mark.parametrize("lines, load_json_function"     ,							[(True, load_json_lines), (False, load_json)]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[Any]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase		).write()
               buffer.seek(0		)
               _a						=		load_json_function(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          assert isinstance(exported_content[0]     ,							__UpperCamelCase		)
          assert len(__UpperCamelCase		) == 10






     @pytest.mark.parametrize(
         "orient, container, keys, len_at"     ,							[
             ("records", list, {"tokens", "labels", "answers", "id"}, None),
             ("split", dict, {"columns", "data"}, "data"),
             ("index", dict, set("0123456789"		), None),
             ("columns", dict, {"tokens", "labels", "answers", "id"}, "tokens"),
             ("values", list, None, None),
             ("table", dict, {"schema", "data"}, "data"),
         ]     ,							)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    List[str]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							orient=__UpperCamelCase		).write()
               buffer.seek(0		)
               _a						=		load_json(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          if keys:
               if container is dict:
                    assert exported_content.keys() == keys
               else:
                    assert exported_content[0].keys() == keys
          else:
               assert not hasattr(__UpperCamelCase     ,							"keys"		) and not hasattr(exported_content[0]     ,							"keys"		)
          if len_at:
               assert len(exported_content[len_at]		) == 10
          else:
               assert len(__UpperCamelCase		) == 10






     @pytest.mark.parametrize("lines, load_json_function"     ,							[(True, load_json_lines), (False, load_json)]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[Any]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							num_proc=2		).write()
               buffer.seek(0		)
               _a						=		load_json_function(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          assert isinstance(exported_content[0]     ,							__UpperCamelCase		)
          assert len(__UpperCamelCase		) == 10






     @pytest.mark.parametrize(
         "orient, container, keys, len_at"     ,							[
             ("records", list, {"tokens", "labels", "answers", "id"}, None),
             ("split", dict, {"columns", "data"}, "data"),
             ("index", dict, set("0123456789"		), None),
             ("columns", dict, {"tokens", "labels", "answers", "id"}, "tokens"),
             ("values", list, None, None),
             ("table", dict, {"schema", "data"}, "data"),
         ]     ,							)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Tuple:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							orient=__UpperCamelCase     ,							num_proc=2		).write()
               buffer.seek(0		)
               _a						=		load_json(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          if keys:
               if container is dict:
                    assert exported_content.keys() == keys
               else:
                    assert exported_content[0].keys() == keys
          else:
               assert not hasattr(__UpperCamelCase     ,							"keys"		) and not hasattr(exported_content[0]     ,							"keys"		)
          if len_at:
               assert len(exported_content[len_at]		) == 10
          else:
               assert len(__UpperCamelCase		) == 10






     def  a_					(			self     ,							__UpperCamelCase		)       ->    List[Any]:
          with pytest.raises(__UpperCamelCase		):
               with io.BytesIO() as buffer:
                    JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							num_proc=0		)






     @pytest.mark.parametrize("compression, extension"     ,							[("gzip", "gz"), ("bz2", "bz2"), ("xz", "xz")]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[int]:
          _a						=		tmp_path_factory.mktemp("data"		) / f"test.json.{extension}"
          _a						=		str(shared_datadir / f"test_file.json.{extension}"		)
          JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							compression=__UpperCamelCase		).write()

          with fsspec.open(__UpperCamelCase     ,							"rb"     ,							compression="infer"		) as f:
               _a						=		f.read()
          with fsspec.open(__UpperCamelCase     ,							"rb"     ,							compression="infer"		) as f:
               _a						=		f.read()
          assert exported_content == original_content





'''simple docstring'''

import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def __UpperCamelCase   (			__lowerCamelCase		: BertModel      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: str				)      ->     List[str]:







     '''simple docstring'''






     _a						=		("dense.weight", "attention.self.query", "attention.self.key", "attention.self.value")

     _a						=		(
         ("layer.", "layer_"),
         ("word_embeddings.weight", "word_embeddings"),
         ("position_embeddings.weight", "position_embeddings"),
         ("token_type_embeddings.weight", "token_type_embeddings"),
         (".", "/"),
         ("LayerNorm/weight", "LayerNorm/gamma"),
         ("LayerNorm/bias", "LayerNorm/beta"),
         ("weight", "kernel"),
     )

     if not os.path.isdir(__lowerCamelCase				):
          os.makedirs(__lowerCamelCase				)

     _a						=		model.state_dict()

     def to_tf_var_name(__lowerCamelCase		: str				):
          for patt, repl in iter(__lowerCamelCase				):
               _a						=		name.replace(__lowerCamelCase      ,  __lowerCamelCase				)
          return F"bert/{name}"

     def create_tf_var(__lowerCamelCase		: np.ndarray      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: tf.Session				):
          _a						=		tf.dtypes.as_dtype(tensor.dtype				)
          _a						=		tf.get_variable(dtype=__lowerCamelCase      ,  shape=tensor.shape      ,  name=__lowerCamelCase      ,  initializer=tf.zeros_initializer()				)
          session.run(tf.variables_initializer([tf_var]				)				)
          session.run(__lowerCamelCase				)
          return tf_var

     tf.reset_default_graph()
     with tf.Session() as session:
          for var_name in state_dict:
               _a						=		to_tf_var_name(__lowerCamelCase				)
               _a						=		state_dict[var_name].numpy()
               if any(x in var_name for x in tensors_to_transpose				):
                    _a						=		torch_tensor.T
               _a						=		create_tf_var(tensor=__lowerCamelCase      ,  name=__lowerCamelCase      ,  session=__lowerCamelCase				)
               tf.keras.backend.set_value(__lowerCamelCase      ,  __lowerCamelCase				)
               _a						=		session.run(__lowerCamelCase				)
               print(F"Successfully created {tf_name}: {np.allclose(__lowerCamelCase      ,  __lowerCamelCase				)}"				)

          _a						=		tf.train.Saver(tf.trainable_variables()				)
          saver.save(__lowerCamelCase      ,  os.path.join(__lowerCamelCase      ,  model_name.replace("-"      ,  "_"				) + ".ckpt"				)				)






def __UpperCamelCase   (			__lowerCamelCase		: str=None				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		argparse.ArgumentParser()
     parser.add_argument("--model_name"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="model name e.g. bert-base-uncased"				)
     parser.add_argument(
         "--cache_dir"      ,  type=__lowerCamelCase      ,  default=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory containing pytorch model"				)
     parser.add_argument("--pytorch_model_path"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="/path/to/<pytorch-model-name>.bin"				)
     parser.add_argument("--tf_cache_dir"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory in which to save tensorflow model"				)
     _a						=		parser.parse_args(__lowerCamelCase				)

     _a						=		BertModel.from_pretrained(
         pretrained_model_name_or_path=args.model_name      ,  state_dict=torch.load(args.pytorch_model_path				)      ,  cache_dir=args.cache_dir      ,  )

     convert_pytorch_checkpoint_to_tf(model=__lowerCamelCase      ,  ckpt_dir=args.tf_cache_dir      ,  model_name=args.model_name				)


if __name__ == "__main__":
   main()

