



import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse('''0.8.3'''):
	raise Exception('''requires gluonnlp == 0.8.3''')

if version.parse(mx.__version__) != version.parse('''1.5.0'''):
	raise Exception('''requires mxnet == 1.5.0''')

logging.set_verbosity_info()
lowerCAmelCase	:		List[str]				=   logging.get_logger(__name__)

lowerCAmelCase	:		Any				=   '''The Nymphenburg Palace is a beautiful palace in Munich!'''



def   __lowerCAmelCase					(							lowerCamelCase   :						Tuple				,		lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  {
	    "attention_cell": "multi_head",
	    "num_layers": 4,
	    "units": 10_24,
	    "hidden_size": 7_68,
	    "max_length": 5_12,
	    "num_heads": 8,
	    "scaled": True,
	    "dropout": 0.1,
	    "use_residual": True,
	    "embed_size": 10_24,
	    "embed_dropout": 0.1,
	    "word_embed": None,
	    "layer_norm_eps": 1e-5,
	    "token_type_vocab_size": 2,
	}

	__lowerCAmelCase												=  bort_4_8_768_1024_hparams

	# Let's construct the original Bort model here
	# Taken from official BERT implementation, see:
	# https://github.com/alexa/bort/blob/master/bort/bort.py
	__lowerCAmelCase												=  BERTEncoder(
	    attention_cell=predefined_args["attention_cell"]				,		num_layers=predefined_args["num_layers"]				,		units=predefined_args["units"]				,		hidden_size=predefined_args["hidden_size"]				,		max_length=predefined_args["max_length"]				,		num_heads=predefined_args["num_heads"]				,		scaled=predefined_args["scaled"]				,		dropout=predefined_args["dropout"]				,		output_attention=lowercase__				,		output_all_encodings=lowercase__				,		use_residual=predefined_args["use_residual"]				,		activation=predefined_args.get("activation"				,		"gelu"					)				,		layer_norm_eps=predefined_args.get("layer_norm_eps"				,		lowercase__					)				,		)

	# Vocab information needs to be fetched first
	# It's the same as RoBERTa, so RobertaTokenizer can be used later
	__lowerCAmelCase												=  "openwebtext_ccnews_stories_books_cased"

	# Specify download folder to Gluonnlp's vocab
	__lowerCAmelCase												=  os.path.join(get_home_dir()				,		"models"					)
	__lowerCAmelCase												=  _load_vocab(lowercase__				,		lowercase__				,		lowercase__				,		cls=lowercase__					)

	__lowerCAmelCase												=  nlp.model.BERTModel(
	    lowercase__				,		len(lowercase__					)				,		units=predefined_args["units"]				,		embed_size=predefined_args["embed_size"]				,		embed_dropout=predefined_args["embed_dropout"]				,		word_embed=predefined_args["word_embed"]				,		use_pooler=lowercase__				,		use_token_type_embed=lowercase__				,		token_type_vocab_size=predefined_args["token_type_vocab_size"]				,		use_classifier=lowercase__				,		use_decoder=lowercase__				,		)

	original_bort.load_parameters(lowercase__				,		cast_dtype=lowercase__				,		ignore_extra=lowercase__					)
	__lowerCAmelCase												=  original_bort._collect_params_with_prefix()

	# Build our config 🤗
	__lowerCAmelCase												=  {
	    "architectures": ["BertForMaskedLM"],
	    "attention_probs_dropout_prob": predefined_args["dropout"],
	    "hidden_act": "gelu",
	    "hidden_dropout_prob": predefined_args["dropout"],
	    "hidden_size": predefined_args["embed_size"],
	    "initializer_range": 0.0_2,
	    "intermediate_size": predefined_args["hidden_size"],
	    "layer_norm_eps": predefined_args["layer_norm_eps"],
	    "max_position_embeddings": predefined_args["max_length"],
	    "model_type": "bort",
	    "num_attention_heads": predefined_args["num_heads"],
	    "num_hidden_layers": predefined_args["num_layers"],
	    "pad_token_id": 1,  # 2 = BERT, 1 = RoBERTa
	    "type_vocab_size": 1,  # 2 = BERT, 1 = RoBERTa
	    "vocab_size": len(lowercase__					),
	}

	__lowerCAmelCase												=  BertConfig.from_dict(lowercase__					)
	__lowerCAmelCase												=  BertForMaskedLM(lowercase__					)
	hf_bort_model.eval()

	# Parameter mapping table (Gluonnlp to Transformers)
	# * denotes layer index
	#
	# | Gluon Parameter                                                | Transformers Parameter
	# | -------------------------------------------------------------- | ----------------------
	# | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
	# | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
	# | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
	# | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
	# | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
	# | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
	# | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
	# | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
	# | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
	# | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
	# | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
	# | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
	# | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
	# | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

	# Helper function to convert MXNET Arrays to PyTorch
	def to_torch(lowerCamelCase   :						Any					) -> nn.Parameter:
		return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()					)					)

	# Check param shapes and map new HF param back
	def check_and_map_params(lowerCamelCase   :						Tuple				,		lowerCamelCase   :						Optional[int]					):
		__lowerCAmelCase												=  hf_param.shape

		__lowerCAmelCase												=  to_torch(params[gluon_param]					)
		__lowerCAmelCase												=  gluon_param.shape

		assert (
		    shape_hf == shape_gluon
		), f'''The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'''

		return gluon_param

	__lowerCAmelCase												=  check_and_map_params(
	    hf_bort_model.bert.embeddings.word_embeddings.weight				,		"word_embed.0.weight"					)
	__lowerCAmelCase												=  check_and_map_params(
	    hf_bort_model.bert.embeddings.position_embeddings.weight				,		"encoder.position_weight"					)
	__lowerCAmelCase												=  check_and_map_params(
	    hf_bort_model.bert.embeddings.LayerNorm.bias				,		"encoder.layer_norm.beta"					)
	__lowerCAmelCase												=  check_and_map_params(
	    hf_bort_model.bert.embeddings.LayerNorm.weight				,		"encoder.layer_norm.gamma"					)

	# Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
	__lowerCAmelCase												=  torch.zeros_like(
	    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data					)

	for i in range(hf_bort_config.num_hidden_layers					):
		__lowerCAmelCase												=  hf_bort_model.bert.encoder.layer[i]

		# self attention
		__lowerCAmelCase												=  layer.attention.self

		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.key.bias.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_key.bias'''					)

		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.key.weight.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_key.weight'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.query.bias.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_query.bias'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.query.weight.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_query.weight'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.value.bias.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_value.bias'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_attn.value.weight.data				,		f'''encoder.transformer_cells.{i}.attention_cell.proj_value.weight'''					)

		# self attention output
		__lowerCAmelCase												=  layer.attention.output

		__lowerCAmelCase												=  check_and_map_params(
		    self_output.dense.bias				,		f'''encoder.transformer_cells.{i}.proj.bias'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_output.dense.weight				,		f'''encoder.transformer_cells.{i}.proj.weight'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_output.LayerNorm.bias				,		f'''encoder.transformer_cells.{i}.layer_norm.beta'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    self_output.LayerNorm.weight				,		f'''encoder.transformer_cells.{i}.layer_norm.gamma'''					)

		# intermediate
		__lowerCAmelCase												=  layer.intermediate

		__lowerCAmelCase												=  check_and_map_params(
		    intermediate.dense.bias				,		f'''encoder.transformer_cells.{i}.ffn.ffn_1.bias'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    intermediate.dense.weight				,		f'''encoder.transformer_cells.{i}.ffn.ffn_1.weight'''					)

		# output
		__lowerCAmelCase												=  layer.output

		__lowerCAmelCase												=  check_and_map_params(
		    bert_output.dense.bias				,		f'''encoder.transformer_cells.{i}.ffn.ffn_2.bias'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    bert_output.dense.weight				,		f'''encoder.transformer_cells.{i}.ffn.ffn_2.weight'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    bert_output.LayerNorm.bias				,		f'''encoder.transformer_cells.{i}.ffn.layer_norm.beta'''					)
		__lowerCAmelCase												=  check_and_map_params(
		    bert_output.LayerNorm.weight				,		f'''encoder.transformer_cells.{i}.ffn.layer_norm.gamma'''					)

	# Save space and energy 🎄
	hf_bort_model.half()

	# Compare output of both models
	__lowerCAmelCase												=  RobertaTokenizer.from_pretrained("roberta-base"					)

	__lowerCAmelCase												=  tokenizer.encode_plus(lowercase__					)["input_ids"]

	# Get gluon output
	__lowerCAmelCase												=  mx.nd.array([input_ids]					)
	__lowerCAmelCase												=  original_bort(inputs=lowercase__				,		token_types=[]					)

	# Get Transformer output (save and reload model again)
	hf_bort_model.save_pretrained(lowercase__					)
	__lowerCAmelCase												=  BertModel.from_pretrained(lowercase__					)
	hf_bort_model.eval()

	__lowerCAmelCase												=  tokenizer.encode_plus(lowercase__				,		return_tensors="pt"					)
	__lowerCAmelCase												=  hf_bort_model(**lowercase__					)[0]

	__lowerCAmelCase												=  output_gluon[0].asnumpy()
	__lowerCAmelCase												=  output_hf[0].detach().numpy()

	__lowerCAmelCase												=  np.max(np.abs(hf_layer - gluon_layer					)					).item()
	__lowerCAmelCase												=  np.allclose(lowercase__				,		lowercase__				,		atol=1e-3					)

	if success:
		print("✔️ Both model do output the same tensors"					)
	else:
		print("❌ Both model do **NOT** output the same tensors"					)
		print("Absolute difference is:"				,		lowercase__					)


if __name__ == "__main__":
	lowerCAmelCase	:		List[Any]				=   argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--bort_checkpoint_path''', default=None, type=str, required=True, help='''Path the official Bort params file.'''
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
	)
	lowerCAmelCase	:		List[str]				=   parser.parse_args()
	convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)




'''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   __lowerCAmelCase					(							lowerCamelCase   :						str				,		lowerCamelCase   :						Union[str, Any]					):






	'''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   __lowerCAmelCase					(							lowerCamelCase   :						Tuple				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"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												=  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   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_3": "float64", "col_1": "string", "col_2": "int64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						Any				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	__lowerCAmelCase												=  {"col_2": "int64", "col_3": "float64", "col_1": "string"}
	__lowerCAmelCase												=  features.copy()
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Optional[Any]				,		lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						Any				,		lowerCamelCase   :						List[str]				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




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



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						str=("train",)					):






	'''simple docstring'''




	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	for split in splits:
		__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						Dict				,		lowerCamelCase   :						int				,		lowerCamelCase   :						List[str]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"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												=  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   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Any				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int					):






	'''simple docstring'''




	if split:
		__lowerCAmelCase												=  {split: jsonl_path}
	else:
		__lowerCAmelCase												=  "train"
		__lowerCAmelCase												=  {"train": jsonl_path, "test": jsonl_path}
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						Optional[Any]					):






	'''simple docstring'''




	return json.load(lowerCamelCase					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Tuple					):






	'''simple docstring'''




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






class 					UpperCAmelCase__	:



					@pytest.mark.parametrize("lines, load_json_function"			,				[(True, load_json_lines), (False, load_json)]     )
					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			str:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Union[str, Any]:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Tuple:
						__lowerCAmelCase												=  tmp_path_factory.mktemp("data"     ) / F'''test.json.{extension}'''
						__lowerCAmelCase												=  str(shared_datadir / F'''test_file.json.{extension}'''     )
						JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				compression=UpperCamelCase     ).write()

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