





"""simple docstring"""





import pyarrow.parquet as pq
import pytest

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

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Any				,		SCREAMING_SNAKE_CASE	:		Tuple     ):







      '''simple docstring'''




      assert isinstance(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )
      assert dataset.num_rows == 4
      assert dataset.num_columns == 3
      assert dataset.column_names == ["col_1", "col_2", "col_3"]
      for feature, expected_dtype in expected_features.items():
            assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""				,		[False, True]     )
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[str]				,		SCREAMING_SNAKE_CASE	:		Optional[int]				,		SCREAMING_SNAKE_CASE	:		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    		=			ParquetDatasetReader(SCREAMING_SNAKE_CASE				,		cache_dir=SCREAMING_SNAKE_CASE				,		keep_in_memory=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_dataset(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


@pytest.mark.parametrize(
    """features"""				,		[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]				,		)
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Optional[int]				,		SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		Optional[int]     ):







      '''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(SCREAMING_SNAKE_CASE     ) for feature, dtype in features.items()}     ) if features is not None else None
      )
      lowerCAmelCase    		=			ParquetDatasetReader(SCREAMING_SNAKE_CASE				,		features=SCREAMING_SNAKE_CASE				,		cache_dir=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_dataset(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


@pytest.mark.parametrize("""split"""				,		[None, NamedSplit("""train"""     ), """train""", """test"""]     )
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[Any]				,		SCREAMING_SNAKE_CASE	:		Dict				,		SCREAMING_SNAKE_CASE	:		List[str]     ):







      '''simple docstring'''




      lowerCAmelCase    		=			tmp_path / """cache"""
      lowerCAmelCase    		=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
      lowerCAmelCase    		=			ParquetDatasetReader(SCREAMING_SNAKE_CASE				,		cache_dir=SCREAMING_SNAKE_CASE				,		split=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_dataset(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )
      assert dataset.split == split if split else "train"


@pytest.mark.parametrize("""path_type"""				,		[str, list]     )
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		Union[str, Any]     ):







      '''simple docstring'''




      if issubclass(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     ):
            lowerCAmelCase    		=			parquet_path
      elif issubclass(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     ):
            lowerCAmelCase    		=			[parquet_path]
      lowerCAmelCase    		=			tmp_path / """cache"""
      lowerCAmelCase    		=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
      lowerCAmelCase    		=			ParquetDatasetReader(SCREAMING_SNAKE_CASE				,		cache_dir=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_dataset(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		Optional[int]				,		SCREAMING_SNAKE_CASE	:		List[Any]=("train",)     ):







      '''simple docstring'''




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







      '''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    		=			ParquetDatasetReader(
                {"""train""": parquet_path}				,		cache_dir=SCREAMING_SNAKE_CASE				,		keep_in_memory=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_datasetdict(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


@pytest.mark.parametrize(
    """features"""				,		[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]				,		)
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Tuple				,		SCREAMING_SNAKE_CASE	:		Optional[int]				,		SCREAMING_SNAKE_CASE	:		Any     ):







      '''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(SCREAMING_SNAKE_CASE     ) for feature, dtype in features.items()}     ) if features is not None else None
      )
      lowerCAmelCase    		=			ParquetDatasetReader({"""train""": parquet_path}				,		features=SCREAMING_SNAKE_CASE				,		cache_dir=SCREAMING_SNAKE_CASE     ).read()
      _check_parquet_datasetdict(SCREAMING_SNAKE_CASE				,		SCREAMING_SNAKE_CASE     )


@pytest.mark.parametrize("""split"""				,		[None, NamedSplit("""train"""     ), """train""", """test"""]     )
def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		Tuple				,		SCREAMING_SNAKE_CASE	:		str     ):







      '''simple docstring'''




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


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Dict				,		SCREAMING_SNAKE_CASE	:		Any     ):







      '''simple docstring'''




      lowerCAmelCase    		=			ParquetDatasetWriter(SCREAMING_SNAKE_CASE				,		tmp_path / """foo.parquet"""     )
      assert writer.write() > 0
      lowerCAmelCase    		=			pq.ParquetFile(tmp_path / """foo.parquet"""     )
      lowerCAmelCase    		=			pf.read()
      assert dataset.data.table == output_table


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Tuple				,		SCREAMING_SNAKE_CASE	:		Any     ):







      '''simple docstring'''




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

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

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






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







      '''simple docstring'''




      assert get_writer_batch_size(SCREAMING_SNAKE_CASE     ) == expected













"""simple docstring"""





import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
       import torch

       from transformers import AutoModelForImageClassification

if is_vision_available():
       from transformers import AutoImageProcessor

@require_torch
@require_vision
class 						lowercase   (   unittest.TestCase   ):





      @slow
      def        _snake_case     (			self					)						->	Any:
            lowerCAmelCase    		=			AutoImageProcessor.from_pretrained("""microsoft/dit-base-finetuned-rvlcdip"""					)
            lowerCAmelCase    		=			AutoModelForImageClassification.from_pretrained("""microsoft/dit-base-finetuned-rvlcdip"""					)
            model.to(lowercase					)

            from datasets import load_dataset

            lowerCAmelCase    		=			load_dataset("""nielsr/rvlcdip-demo"""					)

            lowerCAmelCase    		=			dataset["""train"""][0]["""image"""].convert("""RGB"""					)

            lowerCAmelCase    		=			image_processor(lowercase ,		return_tensors="""pt"""					).to(lowercase					)

            # forward pass
            with torch.no_grad():
                  lowerCAmelCase    		=			model(**lowercase					)
                  lowerCAmelCase    		=			outputs.logits

            lowerCAmelCase    		=			torch.Size((1, 16)					)
            self.assertEqual(logits.shape ,		lowercase					)

            lowerCAmelCase    		=			torch.tensor(
                [-0.4_158, -0.4_092, -0.4_347] ,		device=lowercase ,		dtype=torch.float ,		)
            self.assertTrue(torch.allclose(logits[0, :3] ,		lowercase ,		atol=1e-4					)					)






