
'''simple docstring'''



import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FlaubertForMultipleChoice,
         FlaubertForQuestionAnswering,
         FlaubertForQuestionAnsweringSimple,
         FlaubertForSequenceClassification,
         FlaubertForTokenClassification,
         FlaubertModel,
         FlaubertWithLMHeadModel,
     )
     from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class    _a				(_lowerCamelCase):





       """simple docstring"""




       def __init__(							self       ,							A__       ,							A__=13       ,							A__=7       ,							A__=True       ,							A__=True       ,							A__=True       ,							A__=True       ,							A__=True       ,							A__=False       ,							A__=False       ,							A__=False       ,							A__=2       ,							A__=99       ,							A__=0       ,							A__=32       ,							A__=5       ,							A__=4       ,							A__=0.1       ,							A__=0.1       ,							A__=5_12       ,							A__=12       ,							A__=2       ,							A__=0.02       ,							A__=3       ,							A__=4       ,							A__="last"       ,							A__=None       ,							A__=None       ,							)    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				parent
          _SCREAMING_SNAKE_CASE      =				batch_size
          _SCREAMING_SNAKE_CASE      =				seq_length
          _SCREAMING_SNAKE_CASE      =				is_training
          _SCREAMING_SNAKE_CASE      =				use_input_lengths
          _SCREAMING_SNAKE_CASE      =				use_token_type_ids
          _SCREAMING_SNAKE_CASE      =				use_labels
          _SCREAMING_SNAKE_CASE      =				gelu_activation
          _SCREAMING_SNAKE_CASE      =				sinusoidal_embeddings
          _SCREAMING_SNAKE_CASE      =				causal
          _SCREAMING_SNAKE_CASE      =				asm
          _SCREAMING_SNAKE_CASE      =				n_langs
          _SCREAMING_SNAKE_CASE      =				vocab_size
          _SCREAMING_SNAKE_CASE      =				n_special
          _SCREAMING_SNAKE_CASE      =				hidden_size
          _SCREAMING_SNAKE_CASE      =				num_hidden_layers
          _SCREAMING_SNAKE_CASE      =				num_attention_heads
          _SCREAMING_SNAKE_CASE      =				hidden_dropout_prob
          _SCREAMING_SNAKE_CASE      =				attention_probs_dropout_prob
          _SCREAMING_SNAKE_CASE      =				max_position_embeddings
          _SCREAMING_SNAKE_CASE      =				type_vocab_size
          _SCREAMING_SNAKE_CASE      =				type_sequence_label_size
          _SCREAMING_SNAKE_CASE      =				initializer_range
          _SCREAMING_SNAKE_CASE      =				num_labels
          _SCREAMING_SNAKE_CASE      =				num_choices
          _SCREAMING_SNAKE_CASE      =				summary_type
          _SCREAMING_SNAKE_CASE      =				use_proj
          _SCREAMING_SNAKE_CASE      =				scope






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size      )
          _SCREAMING_SNAKE_CASE      =				random_attention_mask([self.batch_size, self.seq_length]      )

          _SCREAMING_SNAKE_CASE      =				None
          if self.use_input_lengths:
             _SCREAMING_SNAKE_CASE      =				(
                 ids_tensor([self.batch_size]       ,							vocab_size=2      ) + self.seq_length - 2
             )  # small variation of seq_length

          _SCREAMING_SNAKE_CASE      =				None
          if self.use_token_type_ids:
             _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size, self.seq_length]       ,							self.n_langs      )

          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				None
          if self.use_labels:
             _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size]       ,							self.type_sequence_label_size      )
             _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size, self.seq_length]       ,							self.num_labels      )
             _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size]       ,							2      ).float()
             _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size]       ,							self.num_choices      )

          _SCREAMING_SNAKE_CASE      =				self.get_config()

          return (
              config,
              input_ids,
              token_type_ids,
              input_lengths,
              sequence_labels,
              token_labels,
              is_impossible_labels,
              choice_labels,
              input_mask,
          )






       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          return FlaubertConfig(
              vocab_size=self.vocab_size       ,							n_special=self.n_special       ,							emb_dim=self.hidden_size       ,							n_layers=self.num_hidden_layers       ,							n_heads=self.num_attention_heads       ,							dropout=self.hidden_dropout_prob       ,							attention_dropout=self.attention_probs_dropout_prob       ,							gelu_activation=self.gelu_activation       ,							sinusoidal_embeddings=self.sinusoidal_embeddings       ,							asm=self.asm       ,							causal=self.causal       ,							n_langs=self.n_langs       ,							max_position_embeddings=self.max_position_embeddings       ,							initializer_range=self.initializer_range       ,							summary_type=self.summary_type       ,							use_proj=self.use_proj       ,							)






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				FlaubertModel(config=A__      )
          model.to(A__      )
          model.eval()
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							lengths=A__       ,							langs=A__      )
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							langs=A__      )
          _SCREAMING_SNAKE_CASE      =				model(A__      )
          self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				FlaubertWithLMHeadModel(A__      )
          model.to(A__      )
          model.eval()

          _SCREAMING_SNAKE_CASE      =				model(A__       ,							token_type_ids=A__       ,							labels=A__      )
          self.parent.assertEqual(result.loss.shape       ,							()      )
          self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				FlaubertForQuestionAnsweringSimple(A__      )
          model.to(A__      )
          model.eval()

          _SCREAMING_SNAKE_CASE      =				model(A__      )

          _SCREAMING_SNAKE_CASE      =				model(A__       ,							start_positions=A__       ,							end_positions=A__      )
          self.parent.assertEqual(result.start_logits.shape       ,							(self.batch_size, self.seq_length)      )
          self.parent.assertEqual(result.end_logits.shape       ,							(self.batch_size, self.seq_length)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				FlaubertForQuestionAnswering(A__      )
          model.to(A__      )
          model.eval()

          _SCREAMING_SNAKE_CASE      =				model(A__      )

          _SCREAMING_SNAKE_CASE      =				model(
              A__       ,							start_positions=A__       ,							end_positions=A__       ,							cls_index=A__       ,							is_impossible=A__       ,							p_mask=A__       ,							)

          _SCREAMING_SNAKE_CASE      =				model(
              A__       ,							start_positions=A__       ,							end_positions=A__       ,							cls_index=A__       ,							is_impossible=A__       ,							)

          ((_SCREAMING_SNAKE_CASE) ,							)      =				result_with_labels.to_tuple()

          _SCREAMING_SNAKE_CASE      =				model(A__       ,							start_positions=A__       ,							end_positions=A__      )

          ((_SCREAMING_SNAKE_CASE) ,							)      =				result_with_labels.to_tuple()

          self.parent.assertEqual(result_with_labels.loss.shape       ,							()      )
          self.parent.assertEqual(result.start_top_log_probs.shape       ,							(self.batch_size, model.config.start_n_top)      )
          self.parent.assertEqual(result.start_top_index.shape       ,							(self.batch_size, model.config.start_n_top)      )
          self.parent.assertEqual(
              result.end_top_log_probs.shape       ,							(self.batch_size, model.config.start_n_top * model.config.end_n_top)      )
          self.parent.assertEqual(
              result.end_top_index.shape       ,							(self.batch_size, model.config.start_n_top * model.config.end_n_top)      )
          self.parent.assertEqual(result.cls_logits.shape       ,							(self.batch_size,)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				FlaubertForSequenceClassification(A__      )
          model.to(A__      )
          model.eval()

          _SCREAMING_SNAKE_CASE      =				model(A__      )
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							labels=A__      )

          self.parent.assertEqual(result.loss.shape       ,							()      )
          self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.type_sequence_label_size)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				self.num_labels
          _SCREAMING_SNAKE_CASE      =				FlaubertForTokenClassification(A__      )
          model.to(A__      )
          model.eval()

          _SCREAMING_SNAKE_CASE      =				model(A__       ,							attention_mask=A__       ,							labels=A__      )
          self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__       ,							)    ->			int:
          _SCREAMING_SNAKE_CASE      =				self.num_choices
          _SCREAMING_SNAKE_CASE      =				FlaubertForMultipleChoice(config=A__      )
          model.to(A__      )
          model.eval()
          _SCREAMING_SNAKE_CASE      =				input_ids.unsqueeze(1      ).expand(-1       ,							self.num_choices       ,							-1      ).contiguous()
          _SCREAMING_SNAKE_CASE      =				token_type_ids.unsqueeze(1      ).expand(-1       ,							self.num_choices       ,							-1      ).contiguous()
          _SCREAMING_SNAKE_CASE      =				input_mask.unsqueeze(1      ).expand(-1       ,							self.num_choices       ,							-1      ).contiguous()
          _SCREAMING_SNAKE_CASE      =				model(
              A__       ,							attention_mask=A__       ,							token_type_ids=A__       ,							labels=A__       ,							)
          self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_choices)      )






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.prepare_config_and_inputs()
          (
              (
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							(
              _SCREAMING_SNAKE_CASE
          ) ,							
          )      =				config_and_inputs
          _SCREAMING_SNAKE_CASE      =				{
              """input_ids""": input_ids,
              """token_type_ids""": token_type_ids,
              """lengths""": input_lengths,
              """attention_mask""": input_mask,
          }
          return config, inputs_dict






@require_torch
class    _a				(_lowerCamelCase  ,					_lowerCamelCase  ,					unittest.TestCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  (
           (
               FlaubertModel,
               FlaubertWithLMHeadModel,
               FlaubertForQuestionAnswering,
               FlaubertForQuestionAnsweringSimple,
               FlaubertForSequenceClassification,
               FlaubertForTokenClassification,
               FlaubertForMultipleChoice,
           )
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE      =  (
           {
               'feature-extraction': FlaubertModel,
               'fill-mask': FlaubertWithLMHeadModel,
               'question-answering': FlaubertForQuestionAnsweringSimple,
               'text-classification': FlaubertForSequenceClassification,
               'token-classification': FlaubertForTokenClassification,
               'zero-shot': FlaubertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__       ,							A__       ,							A__      )    ->			Optional[Any]:
          if (
              pipeline_test_casse_name == "QAPipelineTests"
              and tokenizer_name is not None
              and not tokenizer_name.endswith("""Fast"""      )
          ):
             # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
             # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
             # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
             return True

          return False






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__=False      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				super()._prepare_for_class(A__       ,							A__       ,							return_labels=A__      )

          if return_labels:
             if model_class.__name__ == "FlaubertForQuestionAnswering":
                _SCREAMING_SNAKE_CASE      =				torch.zeros(
                    self.model_tester.batch_size       ,							dtype=torch.long       ,							device=A__      )
                _SCREAMING_SNAKE_CASE      =				torch.zeros(
                    self.model_tester.batch_size       ,							dtype=torch.long       ,							device=A__      )

          return inputs_dict






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				FlaubertModelTester(self      )
          _SCREAMING_SNAKE_CASE      =				ConfigTester(self       ,							config_class=A__       ,							emb_dim=37      )






       def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
          self.config_tester.run_common_tests()






       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_model(*A__      )






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_lm_head(*A__      )






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_simple_qa(*A__      )






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_qa(*A__      )






       def        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_sequence_classif(*A__      )






       def        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_token_classif(*A__      )






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_multiple_choice(*A__      )






       @slow
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _SCREAMING_SNAKE_CASE      =				FlaubertModel.from_pretrained(A__      )
             self.assertIsNotNone(A__      )






       @slow
       @require_torch_gpu
       def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()
          for model_class in self.all_model_classes:
             # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
             if model_class == FlaubertForMultipleChoice:
                return

             _SCREAMING_SNAKE_CASE      =				True
             _SCREAMING_SNAKE_CASE      =				model_class(config=A__      )

             _SCREAMING_SNAKE_CASE      =				self._prepare_for_class(A__       ,							A__      )
             _SCREAMING_SNAKE_CASE      =				torch.jit.trace(
                 A__       ,							(inputs_dict["""input_ids"""].to("""cpu"""      ), inputs_dict["""attention_mask"""].to("""cpu"""      ))      )

             with tempfile.TemporaryDirectory() as tmp:
                torch.jit.save(A__       ,							os.path.join(A__       ,							"""traced_model.pt"""      )      )
                _SCREAMING_SNAKE_CASE      =				torch.jit.load(os.path.join(A__       ,							"""traced_model.pt"""      )       ,							map_location=A__      )
                loaded(inputs_dict["""input_ids"""].to(A__      )       ,							inputs_dict["""attention_mask"""].to(A__      )      )






@require_torch
class    _a				(unittest.TestCase):





       """simple docstring"""




       @slow
       def        UpperCamelCase				(							self      )    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased"""      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]      )
          with torch.no_grad():
             _SCREAMING_SNAKE_CASE      =				model(A__      )[0]
          _SCREAMING_SNAKE_CASE      =				torch.Size((1, 11, 7_68)      )
          self.assertEqual(output.shape       ,							A__      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(
              [[[-2.6251, -1.4298, -0.0227], [-2.8510, -1.6387, 0.2258], [-2.8114, -1.1832, -0.3066]]]      )

          self.assertTrue(torch.allclose(output[:, :3, :3]       ,							A__       ,							atol=1E-4      )      )








'''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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[int]:




   """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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     List[Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				tmp_path / """cache"""
   _SCREAMING_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():
      _SCREAMING_SNAKE_CASE      =				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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[int]:




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




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				tmp_path / """cache"""
   _SCREAMING_SNAKE_CASE      =				{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   _SCREAMING_SNAKE_CASE      =				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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[int]:




   """simple docstring"""
   if issubclass(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						):
      _SCREAMING_SNAKE_CASE      =				parquet_path
   elif issubclass(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						):
      _SCREAMING_SNAKE_CASE      =				[parquet_path]
   _SCREAMING_SNAKE_CASE      =				tmp_path / """cache"""
   _SCREAMING_SNAKE_CASE      =				{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   _SCREAMING_SNAKE_CASE      =				ParquetDatasetReader(SCREAMING_SNAKE_CASE_    ,      cache_dir=SCREAMING_SNAKE_CASE_						).read()
   _check_parquet_dataset(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_=("train",)						)     ->     List[str]:




   """simple docstring"""
   assert isinstance(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)
   for split in splits:
      _SCREAMING_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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     List[Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				tmp_path / """cache"""
   _SCREAMING_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():
      _SCREAMING_SNAKE_CASE      =				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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[Any]:




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




   """simple docstring"""
   if split:
      _SCREAMING_SNAKE_CASE      =				{split: parquet_path}
   else:
      _SCREAMING_SNAKE_CASE      =				"""train"""
      _SCREAMING_SNAKE_CASE      =				{"""train""": parquet_path, """test""": parquet_path}
   _SCREAMING_SNAKE_CASE      =				tmp_path / """cache"""
   _SCREAMING_SNAKE_CASE      =				{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   _SCREAMING_SNAKE_CASE      =				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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Dict:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				ParquetDatasetWriter(SCREAMING_SNAKE_CASE_    ,      tmp_path / """foo.parquet"""						)
   assert writer.write() > 0
   _SCREAMING_SNAKE_CASE      =				pq.ParquetFile(tmp_path / """foo.parquet"""						)
   _SCREAMING_SNAKE_CASE      =				pf.read()
   assert dataset.data.table == output_table
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[Any]:




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

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

   _SCREAMING_SNAKE_CASE      =				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    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     str:




   """simple docstring"""
   assert get_writer_batch_size(SCREAMING_SNAKE_CASE_						) == expected






