





import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase							:       str					 =	logging.get_logger(__name__)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Tuple       ,			UpperCamelCase__   :   List[Any]					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				int           =   []
						for i in range(encoder_config.num_hidden_layers					):
												# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.norm1.weight""", F"""encoder.encoder.layer.{i}.layernorm_before.weight""")					)
												rename_keys.append((F"""encoder.deit.blocks.{i}.norm1.bias""", F"""encoder.encoder.layer.{i}.layernorm_before.bias""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.attn.proj.weight""", F"""encoder.encoder.layer.{i}.attention.output.dense.weight""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.attn.proj.bias""", F"""encoder.encoder.layer.{i}.attention.output.dense.bias""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.norm2.weight""", F"""encoder.encoder.layer.{i}.layernorm_after.weight""")					)
												rename_keys.append((F"""encoder.deit.blocks.{i}.norm2.bias""", F"""encoder.encoder.layer.{i}.layernorm_after.bias""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.mlp.fc1.weight""", F"""encoder.encoder.layer.{i}.intermediate.dense.weight""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.mlp.fc1.bias""", F"""encoder.encoder.layer.{i}.intermediate.dense.bias""")					)
												rename_keys.append(
												    (F"""encoder.deit.blocks.{i}.mlp.fc2.weight""", F"""encoder.encoder.layer.{i}.output.dense.weight""")					)
												rename_keys.append((F"""encoder.deit.blocks.{i}.mlp.fc2.bias""", F"""encoder.encoder.layer.{i}.output.dense.bias""")					)

						# cls token, position embeddings and patch embeddings of encoder
						rename_keys.extend(
						    [
						        ('''encoder.deit.cls_token''', '''encoder.embeddings.cls_token'''),
						        ('''encoder.deit.pos_embed''', '''encoder.embeddings.position_embeddings'''),
						        ('''encoder.deit.patch_embed.proj.weight''', '''encoder.embeddings.patch_embeddings.projection.weight'''),
						        ('''encoder.deit.patch_embed.proj.bias''', '''encoder.embeddings.patch_embeddings.projection.bias'''),
						        ('''encoder.deit.norm.weight''', '''encoder.layernorm.weight'''),
						        ('''encoder.deit.norm.bias''', '''encoder.layernorm.bias'''),
						    ]					)

						return rename_keys






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   List[Any]					)      ->    Dict:

						"""simple docstring"""





						for i in range(encoder_config.num_hidden_layers					):
												# queries, keys and values (only weights, no biases)
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   state_dict.pop(F"""encoder.deit.blocks.{i}.attn.qkv.weight"""					)

												__SCREAMING_SNAKE_CASE:				List[str]           =   in_proj_weight[
												    : encoder_config.hidden_size, :
												]
												__SCREAMING_SNAKE_CASE:				str           =   in_proj_weight[
												    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
												]
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   in_proj_weight[
												    -encoder_config.hidden_size :, :
												]






def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Optional[Any]					)      ->    List[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   dct.pop(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				int           =   val






def 		lowerCAmelCase  (	UpperCamelCase__   :   str					)      ->    Union[str, Any]:

						"""simple docstring"""





						if "handwritten" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   '''https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg'''  # industry
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
						elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				str           =   '''https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'''
						__SCREAMING_SNAKE_CASE:				List[str]           =   Image.open(requests.get(UpperCamelCase__       ,			stream=UpperCamelCase__					).raw					).convert('''RGB'''					)
						return im






@torch.no_grad()
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Union[str, Any]					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   ViTConfig(image_size=384       ,			qkv_bias=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   TrOCRConfig()

						# size of the architecture
						if "base" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   768
						elif "large" in checkpoint_url:
												# use ViT-large encoder
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   1_024
												__SCREAMING_SNAKE_CASE:				Any           =   4_096
												__SCREAMING_SNAKE_CASE:				str           =   24
												__SCREAMING_SNAKE_CASE:				List[Any]           =   16
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   1_024
						else:
												raise ValueError('''Should either find \'base\' or \'large\' in checkpoint URL'''					)

						# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
						if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   False
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   '''relu'''
												__SCREAMING_SNAKE_CASE:				Dict           =   1_024
												__SCREAMING_SNAKE_CASE:				str           =   True
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   False
												__SCREAMING_SNAKE_CASE:				int           =   False

						# load HuggingFace model
						__SCREAMING_SNAKE_CASE:				int           =   ViTModel(UpperCamelCase__       ,			add_pooling_layer=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				int           =   TrOCRForCausalLM(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   VisionEncoderDecoderModel(encoder=UpperCamelCase__       ,			decoder=UpperCamelCase__					)
						model.eval()

						# load state_dict of original model, rename some keys
						__SCREAMING_SNAKE_CASE:				List[str]           =   torch.hub.load_state_dict_from_url(UpperCamelCase__       ,			map_location='''cpu'''       ,			check_hash=UpperCamelCase__					)['''model''']

						__SCREAMING_SNAKE_CASE:				Optional[int]           =   create_rename_keys(UpperCamelCase__       ,			UpperCamelCase__					)
						for src, dest in rename_keys:
												rename_key(UpperCamelCase__       ,			UpperCamelCase__       ,			UpperCamelCase__					)
						read_in_q_k_v(UpperCamelCase__       ,			UpperCamelCase__					)

						# remove parameters we don't need
						del state_dict["encoder.deit.head.weight"]
						del state_dict["encoder.deit.head.bias"]
						del state_dict["decoder.version"]

						# add prefix to decoder keys
						for key, val in state_dict.copy().items():
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   state_dict.pop(UpperCamelCase__					)
												if key.startswith('''decoder'''					) and "output_projection" not in key:
																		__SCREAMING_SNAKE_CASE:				Optional[Any]           =   val
												else:
																		__SCREAMING_SNAKE_CASE:				Dict           =   val

    # load state dict
						model.load_state_dict(UpperCamelCase__					)

						# Check outputs on an image
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   ViTImageProcessor(size=encoder_config.image_size					)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   RobertaTokenizer.from_pretrained('''roberta-large'''					)
						__SCREAMING_SNAKE_CASE:				List[Any]           =   TrOCRProcessor(UpperCamelCase__       ,			UpperCamelCase__					)

						__SCREAMING_SNAKE_CASE:				int           =   processor(images=prepare_img(UpperCamelCase__					)       ,			return_tensors='''pt'''					).pixel_values

						# verify logits
						__SCREAMING_SNAKE_CASE:				int           =   torch.tensor([[model.config.decoder.decoder_start_token_id]]					)
						__SCREAMING_SNAKE_CASE:				List[str]           =   model(pixel_values=UpperCamelCase__       ,			decoder_input_ids=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Any           =   outputs.logits

						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.Size([1, 1, 50_265]					)
						if "trocr-base-handwritten" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				Tuple           =   torch.tensor(
												    [-1.45_02, -4.66_83, -0.53_47, -2.92_91, 9.14_35, -3.05_71, 8.97_64, 1.75_60, 8.73_58, -1.53_11]					)
						elif "trocr-large-handwritten" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				List[str]           =   torch.tensor(
												    [-2.64_37, -1.31_29, -2.25_96, -5.34_55, 6.35_39, 1.76_04, 5.49_91, 1.47_02, 5.61_13, 2.01_70]					)
						elif "trocr-base-printed" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				Any           =   torch.tensor(
												    [-5.68_16, -5.83_88, 1.13_98, -6.90_34, 6.85_05, -2.43_93, 1.22_84, -1.02_32, -1.96_61, -3.92_10]					)
						elif "trocr-large-printed" in checkpoint_url:
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   torch.tensor(
												    [-6.01_62, -7.09_59, 4.41_55, -5.10_63, 7.04_68, -3.16_31, 2.64_66, -0.30_81, -0.81_06, -1.75_35]					)

						if "stage1" not in checkpoint_url:
												assert logits.shape == expected_shape, "Shape of logits not as expected"
												assert torch.allclose(logits[0, 0, :10]       ,			UpperCamelCase__       ,			atol=1E-3					), "First elements of logits not as expected"

						Path(UpperCamelCase__					).mkdir(exist_ok=UpperCamelCase__					)
						print(F"""Saving model to {pytorch_dump_folder_path}"""					)
						model.save_pretrained(UpperCamelCase__					)
						print(F"""Saving processor to {pytorch_dump_folder_path}"""					)
						processor.save_pretrained(UpperCamelCase__					)


if __name__ == "__main__":
							lowerCAmelCase							:       Union[str, Any]					 =	argparse.ArgumentParser()

							parser.add_argument(
							    """--checkpoint_url""",
							    default="""https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt""",
							    type=str,
							    help="""URL to the original PyTorch checkpoint (.pth file).""",
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
							)
							lowerCAmelCase							:       List[Any]					 =	parser.parse_args()
							convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










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  (	UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   Optional[Any]					)      ->    str:

						"""simple docstring"""





						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						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  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   List[Any]					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Dict           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Dict           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__       ,			keep_in_memory=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)






@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  (	UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Optional[int]					)      ->    int:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[str]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Tuple           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Any           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)






@pytest.mark.parametrize(
    '''features'''       ,			[
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ]       ,			)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   Optional[int]					)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[int]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Dict           =   {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
						__SCREAMING_SNAKE_CASE:				int           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   JsonDatasetReader(UpperCamelCase__       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						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  (	UpperCamelCase__   :   str       ,			UpperCamelCase__   :   Optional[int]					)      ->    str:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				str           =   {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
						__SCREAMING_SNAKE_CASE:				Dict           =   features.copy()
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						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  (	UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   str       ,			UpperCamelCase__   :   int					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__       ,			split=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)
						assert dataset.split == split if split else "train"






@pytest.mark.parametrize('''path_type'''       ,			[str, list]					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   str					)      ->    Any:

						"""simple docstring"""





						if issubclass(UpperCamelCase__       ,			UpperCamelCase__					):
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   jsonl_path
						elif issubclass(UpperCamelCase__       ,			UpperCamelCase__					):
												__SCREAMING_SNAKE_CASE:				List[str]           =   [jsonl_path]
						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Union[str, Any]=("train",)					)      ->    List[str]:

						"""simple docstring"""





						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						for split in splits:
												__SCREAMING_SNAKE_CASE:				Tuple           =   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  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   Tuple					)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				List[str]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader({'''train''': jsonl_path}       ,			cache_dir=UpperCamelCase__       ,			keep_in_memory=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__					)






@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  (	UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   Tuple       ,			UpperCamelCase__   :   List[str]					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				str           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader({'''train''': jsonl_path}       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__					)






@pytest.mark.parametrize('''split'''       ,			[None, NamedSplit('''train'''					), '''train''', '''test''']					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   int					)      ->    Optional[Any]:

						"""simple docstring"""





						if split:
												__SCREAMING_SNAKE_CASE:				str           =   {split: jsonl_path}
						else:
												__SCREAMING_SNAKE_CASE:				Dict           =   '''train'''
												__SCREAMING_SNAKE_CASE:				Tuple           =   {'''train''': jsonl_path, '''test''': jsonl_path}
						__SCREAMING_SNAKE_CASE:				Tuple           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				str           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				int           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__       ,			splits=list(path.keys()					)					)
						assert all(dataset[split].split == split for split in path.keys()					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Optional[int]:

						"""simple docstring"""





						return json.load(UpperCamelCase__					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Any					)      ->    Any:

						"""simple docstring"""





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





class 				a			:






				@pytest.mark.parametrize('''lines, load_json_function'''			, [(True, load_json_lines), (False, load_json)]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				str           =   load_json_function(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										assert isinstance(exported_content[0]			, _lowerCAmelCase					)
										assert len(_lowerCAmelCase					) == 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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






				@pytest.mark.parametrize('''lines, load_json_function'''			, [(True, load_json_lines), (False, load_json)]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, num_proc=2					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   load_json_function(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										assert isinstance(exported_content[0]			, _lowerCAmelCase					)
										assert len(_lowerCAmelCase					) == 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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase			, num_proc=2					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Dict           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






				def 			snake_case_    (  self			, _lowerCAmelCase					):

										"""simple docstring"""




										with pytest.raises(_lowerCAmelCase					):
																with io.BytesIO() as buffer:
																						JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=0					)






				@pytest.mark.parametrize('''compression, extension'''			, [('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Any           =   tmp_path_factory.mktemp('''data'''					) / f"""test.json.{extension}"""
										__SCREAMING_SNAKE_CASE:				Any           =   str(shared_datadir / f"""test_file.json.{extension}"""					)
										JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, compression=_lowerCAmelCase					).write()

										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   f.read()
										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Any           =   f.read()
										assert exported_content == original_content



