




'''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				:  Optional[Any] ,						_lowerCamelCase				:  Optional[Any]						)-> Dict:




     '''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				:  int ,						_lowerCamelCase				:  int ,						_lowerCamelCase				:  int						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	tmp_path / '''cache'''
     __snake_case            =	{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
          __snake_case            =	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				:  Optional[Any] ,						_lowerCamelCase				:  Tuple ,						_lowerCamelCase				:  Optional[Any]						)-> int:




     '''simple docstring'''





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




     '''simple docstring'''





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




     '''simple docstring'''





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




     '''simple docstring'''





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




     '''simple docstring'''





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







def       _UpperCamelCase			(_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Any ,						_lowerCamelCase				:  Optional[int]=("train",)						)-> Tuple:




     '''simple docstring'''





     assert isinstance(_lowerCamelCase ,						_lowerCamelCase						)
     for split in splits:
          __snake_case            =	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				:  Optional[Any] ,						_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  Union[str, Any]						)-> str:




     '''simple docstring'''





     __snake_case            =	tmp_path / '''cache'''
     __snake_case            =	{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
          __snake_case            =	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				:  Union[str, Any] ,						_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  Dict						)-> List[str]:




     '''simple docstring'''





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




     '''simple docstring'''





     if split:
          __snake_case            =	{split: jsonl_path}
     else:
          __snake_case            =	'''train'''
          __snake_case            =	{'''train''': jsonl_path, '''test''': jsonl_path}
     __snake_case            =	tmp_path / '''cache'''
     __snake_case            =	{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     __snake_case            =	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				:  List[Any]						)-> str:




     '''simple docstring'''





     return json.load(_lowerCamelCase						)







def       _UpperCamelCase			(_lowerCamelCase				:  Union[str, Any]						)-> Optional[int]:




     '''simple docstring'''





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







class      lowerCAmelCase :


       @pytest.mark.parametrize('''lines, load_json_function'''     ,	[(True, load_json_lines), (False, load_json)]	)
       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Dict:



            '''simple docstring'''




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



            '''simple docstring'''




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


       @pytest.mark.parametrize('''lines, load_json_function'''     ,	[(True, load_json_lines), (False, load_json)]	)
       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	List[str]:



            '''simple docstring'''




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



            '''simple docstring'''




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


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	Optional[Any]:



            '''simple docstring'''




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




       @pytest.mark.parametrize('''compression, extension'''     ,	[('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]	)
       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	str:



            '''simple docstring'''




            __snake_case            =	tmp_path_factory.mktemp('''data'''	) / F'''test.json.{extension}'''
            __snake_case            =	str(shared_datadir / F'''test_file.json.{extension}'''	)
            JsonDatasetWriter(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	compression=__SCREAMING_SNAKE_CASE	).write()

            with fsspec.open(__SCREAMING_SNAKE_CASE     ,	'''rb'''     ,	compression='''infer'''	) as f:
                 __snake_case            =	f.read()
            with fsspec.open(__SCREAMING_SNAKE_CASE     ,	'''rb'''     ,	compression='''infer'''	) as f:
                 __snake_case            =	f.read()
            assert exported_content == original_content














"""simple docstring"""


from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_torch_available():
		import torch

		from transformers.modeling_outputs import BaseModelOutput

		from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

__SCREAMING_SNAKE_CASE   =   logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase						)
class 						__UpperCamelCase     (     UpperCamelCase						):
					def __init__(							self							:	Union[str, Any]     , **UpperCAmelCase							:	List[Any]   )	->       int:
											super().__init__(**UpperCAmelCase   )

											if self.framework == "tf":
																	raise ValueError(f"""The {self.__class__} is only available in PyTorch."""   )

											requires_backends(self     , 'vision'   )
											self.check_model_type(UpperCAmelCase   )
					def __call__(							self							:	Optional[Any]     , UpperCAmelCase							:	Union[str, "Image.Image", List[Dict[str, Any]]]     , UpperCAmelCase							:	Union[str, List[str]] = None     , **UpperCAmelCase							:	Dict     , )	->       Any:
											if "text_queries" in kwargs:
																	lowerCAmelCase   :Dict						=			kwargs.pop('text_queries'   )

											if isinstance(UpperCAmelCase     , (str, Image.Image)   ):
																	lowerCAmelCase   :List[str]						=			{'image': image, 'candidate_labels': candidate_labels}
											else:
																	lowerCAmelCase   :Union[str, Any]						=			image
											lowerCAmelCase   :Optional[int]						=			super().__call__(UpperCAmelCase     , **UpperCAmelCase   )
											return results
					def        UpperCAmelCase__ (							self							:	List[Any]     , **UpperCAmelCase							:	Tuple   )	->       Optional[int]:
											lowerCAmelCase   :Any						=			{}
											if "threshold" in kwargs:
																	lowerCAmelCase   :int						=			kwargs['threshold']
											if "top_k" in kwargs:
																	lowerCAmelCase   :str						=			kwargs['top_k']
											return {}, {}, postprocess_params
					def        UpperCAmelCase__ (							self							:	Any     , UpperCAmelCase							:	Union[str, Any]   )	->       str:
											lowerCAmelCase   :Dict						=			load_image(inputs['image']   )
											lowerCAmelCase   :Dict						=			inputs['candidate_labels']
											if isinstance(UpperCAmelCase     , UpperCAmelCase   ):
																	lowerCAmelCase   :str						=			candidate_labels.split(','   )

											lowerCAmelCase   :str						=			torch.tensor([[image.height, image.width]]     , dtype=torch.intaa   )
											for i, candidate_label in enumerate(UpperCAmelCase   ):
																	lowerCAmelCase   :Any						=			self.tokenizer(UpperCAmelCase     , return_tensors=self.framework   )
																	lowerCAmelCase   :Optional[Any]						=			self.image_processor(UpperCAmelCase     , return_tensors=self.framework   )
																	yield {
																	    "is_last": i == len(UpperCAmelCase   ) - 1,
																	    "target_size": target_size,
																	    "candidate_label": candidate_label,
																	    **text_inputs,
																	    **image_features,
																	}
					def        UpperCAmelCase__ (							self							:	Union[str, Any]     , UpperCAmelCase							:	str   )	->       List[Any]:
											lowerCAmelCase   :Union[str, Any]						=			model_inputs.pop('target_size'   )
											lowerCAmelCase   :List[str]						=			model_inputs.pop('candidate_label'   )
											lowerCAmelCase   :Union[str, Any]						=			model_inputs.pop('is_last'   )

											lowerCAmelCase   :int						=			self.model(**UpperCAmelCase   )

											lowerCAmelCase   :Any						=			{'target_size': target_size, 'candidate_label': candidate_label, 'is_last': is_last, **outputs}
											return model_outputs
					def        UpperCAmelCase__ (							self							:	Union[str, Any]     , UpperCAmelCase							:	Optional[int]     , UpperCAmelCase							:	int=0.1     , UpperCAmelCase							:	Optional[Any]=None   )	->       str:
											lowerCAmelCase   :int						=			[]
											for model_output in model_outputs:
																	lowerCAmelCase   :int						=			model_output['candidate_label']
																	lowerCAmelCase   :List[Any]						=			BaseModelOutput(UpperCAmelCase   )
																	lowerCAmelCase   :Dict						=			self.image_processor.post_process_object_detection(
																	    outputs=UpperCAmelCase     , threshold=UpperCAmelCase     , target_sizes=model_output['target_size']   )[0]

																	for index in outputs["scores"].nonzero():
																							lowerCAmelCase   :Tuple						=			outputs['scores'][index].item()
																							lowerCAmelCase   :int						=			self._get_bounding_box(outputs['boxes'][index][0]   )

																							lowerCAmelCase   :Optional[int]						=			{'score': score, 'label': label, 'box': box}
																							results.append(UpperCAmelCase   )

											lowerCAmelCase   :Tuple						=			sorted(UpperCAmelCase     , key=lambda UpperCAmelCase   : x["score"]     , reverse=UpperCAmelCase   )
											if top_k:
																	lowerCAmelCase   :Dict						=			results[:top_k]

											return results
					def        UpperCAmelCase__ (							self							:	Optional[Any]     , UpperCAmelCase							:	"torch.Tensor"   )	->       Dict[str, int]:
											if self.framework != "pt":
																	raise ValueError('The ZeroShotObjectDetectionPipeline is only available in PyTorch.'   )
											lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Optional[int]						=			box.int().tolist()
											lowerCAmelCase   :Union[str, Any]						=			{
											    'xmin': xmin,
											    'ymin': ymin,
											    'xmax': xmax,
											    'ymax': ymax,
											}
											return bbox