





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_					(			__a      ,       __a						)	->						List[Any]:







  """simple docstring"""





  assert isinstance(__a      ,       __a						)
  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_					(			__a      ,       __a      ,       __a						)	->						Any:







  """simple docstring"""





  lowerCamelCase__: Any         =tmp_path / "cache"
  lowerCamelCase__: Optional[int]         ={"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__: Tuple         =ParquetDatasetReader(__a      ,       cache_dir=__a      ,       keep_in_memory=__a						).read()
  _check_parquet_dataset(__a      ,       __a						)





@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_					(			__a      ,       __a      ,       __a						)	->						List[str]:







  """simple docstring"""





  lowerCamelCase__: int         =tmp_path / "cache"
  lowerCamelCase__: Tuple         ={"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCamelCase__: Union[str, Any]         =features.copy() if features else default_expected_features
  lowerCamelCase__: Optional[int]         =(
      Features({feature: Value(__a						) for feature, dtype in features.items()}						) if features is not None else None
  )
  lowerCamelCase__: int         =ParquetDatasetReader(__a      ,       features=__a      ,       cache_dir=__a						).read()
  _check_parquet_dataset(__a      ,       __a						)





@pytest.mark.parametrize("split"      ,       [None, NamedSplit("train"						), "train", "test"]						)
def 	lowerCAmelCase_					(			__a      ,       __a      ,       __a						)	->						Any:







  """simple docstring"""





  lowerCamelCase__: Any         =tmp_path / "cache"
  lowerCamelCase__: Optional[Any]         ={"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCamelCase__: Optional[Any]         =ParquetDatasetReader(__a      ,       cache_dir=__a      ,       split=__a						).read()
  _check_parquet_dataset(__a      ,       __a						)
  assert dataset.split == split if split else "train"





@pytest.mark.parametrize("path_type"      ,       [str, list]						)
def 	lowerCAmelCase_					(			__a      ,       __a      ,       __a						)	->						int:







  """simple docstring"""





  if issubclass(__a      ,       __a						):
    lowerCamelCase__: List[Any]         =parquet_path
  elif issubclass(__a      ,       __a						):
    lowerCamelCase__: str         =[parquet_path]
  lowerCamelCase__: Tuple         =tmp_path / "cache"
  lowerCamelCase__: Optional[Any]         ={"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCamelCase__: int         =ParquetDatasetReader(__a      ,       cache_dir=__a						).read()
  _check_parquet_dataset(__a      ,       __a						)





def 	lowerCAmelCase_					(			__a      ,       __a      ,       __a=("train",)						)	->						Dict:







  """simple docstring"""





  assert isinstance(__a      ,       __a						)
  for split in splits:
    lowerCamelCase__: 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_					(			__a      ,       __a      ,       __a						)	->						Any:







  """simple docstring"""





  lowerCamelCase__: List[Any]         =tmp_path / "cache"
  lowerCamelCase__: Optional[Any]         ={"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__: Tuple         =ParquetDatasetReader(
        {"train": parquet_path}      ,       cache_dir=__a      ,       keep_in_memory=__a						).read()
  _check_parquet_datasetdict(__a      ,       __a						)





@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_					(			__a      ,       __a      ,       __a						)	->						Optional[Any]:







  """simple docstring"""





  lowerCamelCase__: Tuple         =tmp_path / "cache"
  lowerCamelCase__: Optional[int]         ={"col_1": "string", "col_2": "int64", "col_3": "float64"}
  lowerCamelCase__: List[Any]         =features.copy() if features else default_expected_features
  lowerCamelCase__: int         =(
      Features({feature: Value(__a						) for feature, dtype in features.items()}						) if features is not None else None
  )
  lowerCamelCase__: Optional[Any]         =ParquetDatasetReader({"train": parquet_path}      ,       features=__a      ,       cache_dir=__a						).read()
  _check_parquet_datasetdict(__a      ,       __a						)





@pytest.mark.parametrize("split"      ,       [None, NamedSplit("train"						), "train", "test"]						)
def 	lowerCAmelCase_					(			__a      ,       __a      ,       __a						)	->						Union[str, Any]:







  """simple docstring"""





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





def 	lowerCAmelCase_					(			__a      ,       __a						)	->						int:







  """simple docstring"""





  lowerCamelCase__: List[str]         =ParquetDatasetWriter(__a      ,       tmp_path / "foo.parquet"						)
  assert writer.write() > 0
  lowerCamelCase__: List[str]         =pq.ParquetFile(tmp_path / "foo.parquet"						)
  lowerCamelCase__: List[str]         =pf.read()
  assert dataset.data.table == output_table





def 	lowerCAmelCase_					(			__a      ,       __a						)	->						List[str]:







  """simple docstring"""





  lowerCamelCase__: List[str]         =str(shared_datadir / "test_image_rgb.jpg"						)
  lowerCamelCase__: Union[str, Any]         ={"image": [image_path]}
  lowerCamelCase__: Optional[Any]         =Features({"image": Image()}						)
  lowerCamelCase__: Optional[int]         =Dataset.from_dict(__a      ,       features=__a						)
  lowerCamelCase__: Optional[int]         =ParquetDatasetWriter(__a      ,       tmp_path / "foo.parquet"						)
  assert writer.write() > 0

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

  lowerCamelCase__: Optional[Any]         =ParquetDatasetReader(str(tmp_path / "foo.parquet"						)      ,       streaming=__a						).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_					(			__a      ,       __a						)	->						Optional[Any]:







  """simple docstring"""





  assert get_writer_batch_size(__a						) == expected





import logging
import os
from dataclasses import dataclass
from typing import List, Optional, Union

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


lowerCAmelCase__    :List[str]       =				logging.getLogger(__name__)

@dataclass(frozen=UpperCAmelCase     )
class 						__a :
		_a      :    str
		_a      :    str
		_a      :    Optional[str]           =      None
		_a      :    Optional[str]           =      None
		_a      :    Optional[str]           =      None



@dataclass(frozen=UpperCAmelCase     )
class 						__a :
		_a      :    List[int]
		_a      :    Optional[List[int]]           =      None
		_a      :    Optional[List[int]]           =      None
		_a      :    Optional[Union[int, float]]           =      None
		_a      :    Optional[int]           =      None


if is_torch_available():
		import torch
		from torch.utils.data import Dataset

		class 						__a (							UpperCAmelCase     ):
				_a      :    List[InputFeatures]
				def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE = False     ,       )					->			List[Any]:




											"""simple docstring"""





											_UpperCAmelCase											=     hans_processors[task]()

											_UpperCAmelCase											=     os.path.join(
											    _SCREAMING_SNAKE_CASE     ,       'cached_{}_{}_{}_{}'.format(
											        'dev' if evaluate else 'train'     ,       tokenizer.__class__.__name__     ,       str(_SCREAMING_SNAKE_CASE		)     ,       _SCREAMING_SNAKE_CASE     ,       )     ,       )
											_UpperCAmelCase											=     processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																		# HACK(label indices are swapped in RoBERTa pretrained model)
																		_UpperCAmelCase						,			_UpperCAmelCase											=     label_list[2], label_list[1]
											_UpperCAmelCase											=     label_list

											# Make sure only the first process in distributed training processes the dataset,
											# and the others will use the cache.
											_UpperCAmelCase											=     cached_features_file + '.lock'
											with FileLock(_SCREAMING_SNAKE_CASE		):
																		if os.path.exists(_SCREAMING_SNAKE_CASE		) and not overwrite_cache:
																									logger.info(f'''Loading features from cached file {cached_features_file}'''		)
																									_UpperCAmelCase											=     torch.load(_SCREAMING_SNAKE_CASE		)
																		else:
																									logger.info(f'''Creating features from dataset file at {data_dir}'''		)

																									_UpperCAmelCase											=     (
																									    processor.get_dev_examples(_SCREAMING_SNAKE_CASE		) if evaluate else processor.get_train_examples(_SCREAMING_SNAKE_CASE		)
																									)

																									logger.info('Training examples: %s'     ,       len(_SCREAMING_SNAKE_CASE		)		)
																									_UpperCAmelCase											=     hans_convert_examples_to_features(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																									logger.info('Saving features into cached file %s'     ,       _SCREAMING_SNAKE_CASE		)
																									torch.save(self.features     ,       _SCREAMING_SNAKE_CASE		)
				def __len__(     self		)					->			List[str]:




											"""simple docstring"""





											return len(self.features		)
				def __getitem__(     self     ,       _SCREAMING_SNAKE_CASE		)					->			InputFeatures:




											"""simple docstring"""





											return self.features[i]






				def       UpperCAmelCase__							(     self		)					->			Optional[int]:




											"""simple docstring"""





											return self.label_list


if is_tf_available():
		import tensorflow as tf

		class 						__a :
				_a      :    List[InputFeatures]
				def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = 128     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE = False     ,       )					->			Tuple:




											"""simple docstring"""





											_UpperCAmelCase											=     hans_processors[task]()
											_UpperCAmelCase											=     processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																		# HACK(label indices are swapped in RoBERTa pretrained model)
																		_UpperCAmelCase						,			_UpperCAmelCase											=     label_list[2], label_list[1]
											_UpperCAmelCase											=     label_list

											_UpperCAmelCase											=     processor.get_dev_examples(_SCREAMING_SNAKE_CASE		) if evaluate else processor.get_train_examples(_SCREAMING_SNAKE_CASE		)
											_UpperCAmelCase											=     hans_convert_examples_to_features(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

											def gen():
																		for ex_index, ex in tqdm.tqdm(enumerate(self.features		)     ,       desc='convert examples to features'		):
																									if ex_index % 10000 == 0:
																																logger.info('Writing example %d of %d' % (ex_index, len(_SCREAMING_SNAKE_CASE		))		)

																									yield (
																									    {
																									        "example_id": 0,
																									        "input_ids": ex.input_ids,
																									        "attention_mask": ex.attention_mask,
																									        "token_type_ids": ex.token_type_ids,
																									    },
																									    ex.label,
																									)

											_UpperCAmelCase											=     tf.data.Dataset.from_generator(
											    _SCREAMING_SNAKE_CASE     ,       (
											        {
											            'example_id': tf.intaa,
											            'input_ids': tf.intaa,
											            'attention_mask': tf.intaa,
											            'token_type_ids': tf.intaa,
											        },
											        tf.intaa,
											    )     ,       (
											        {
											            'example_id': tf.TensorShape([]		),
											            'input_ids': tf.TensorShape([None, None]		),
											            'attention_mask': tf.TensorShape([None, None]		),
											            'token_type_ids': tf.TensorShape([None, None]		),
											        },
											        tf.TensorShape([]		),
											    )     ,       )
				def       UpperCAmelCase__							(     self		)					->			Dict:




											"""simple docstring"""





											return self.dataset
				def __len__(     self		)					->			List[Any]:




											"""simple docstring"""





											return len(self.features		)
				def __getitem__(     self     ,       _SCREAMING_SNAKE_CASE		)					->			InputFeatures:




											"""simple docstring"""





											return self.features[i]






				def       UpperCAmelCase__							(     self		)					->			List[Any]:




											"""simple docstring"""





											return self.label_list

class 						__a (							UpperCAmelCase     ):
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""simple docstring"""





									return self._create_examples(self._read_tsv(os.path.join(_SCREAMING_SNAKE_CASE     ,       'heuristics_train_set.txt'		)		)     ,       'train'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									return self._create_examples(self._read_tsv(os.path.join(_SCREAMING_SNAKE_CASE     ,       'heuristics_evaluation_set.txt'		)		)     ,       'dev'		)
		def       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									return ["contradiction", "entailment", "neutral"]






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     []
									for i, line in enumerate(_SCREAMING_SNAKE_CASE		):
																if i == 0:
																							continue
																_UpperCAmelCase											=     '%s-%s' % (set_type, line[0])
																_UpperCAmelCase											=     line[5]
																_UpperCAmelCase											=     line[6]
																_UpperCAmelCase											=     line[7][2:] if line[7].startswith('ex'		) else line[7]
																_UpperCAmelCase											=     line[0]
																examples.append(InputExample(guid=_SCREAMING_SNAKE_CASE     ,       text_a=_SCREAMING_SNAKE_CASE     ,       text_b=_SCREAMING_SNAKE_CASE     ,       label=_SCREAMING_SNAKE_CASE     ,       pairID=_SCREAMING_SNAKE_CASE		)		)
									return examples



def 			lowerCAmelCase__					(     a__:						List[InputExample]     ,    a__:						List[str]     ,    a__:						int     ,    a__:						PreTrainedTokenizer     ,    )    ->						str:


							'''simple docstring'''





							_UpperCAmelCase											=     {label: i for i, label in enumerate(a__			)}

							_UpperCAmelCase											=     []
							for ex_index, example in tqdm.tqdm(enumerate(a__			)     ,    desc='convert examples to features'			):
														if ex_index % 1_0_0_0_0 == 0:
																					logger.info('Writing example %d' % (ex_index)			)

														_UpperCAmelCase											=     tokenizer(
														    example.text_a     ,    example.text_b     ,    add_special_tokens=a__     ,    max_length=a__     ,    padding='max_length'     ,    truncation=a__     ,    return_overflowing_tokens=a__     ,    )

														_UpperCAmelCase											=     label_map[example.label] if example.label in label_map else 0

														_UpperCAmelCase											=     int(example.pairID			)

														features.append(InputFeatures(**a__     ,    label=a__     ,    pairID=a__			)			)

							for i, example in enumerate(examples[:5]			):
														logger.info('*** Example ***'			)
														logger.info(F'''guid: {example}'''			)
														logger.info(F'''features: {features[i]}'''			)

							return features


lowerCAmelCase__    :Optional[Any]       =				{
    '''hans''': 3,
}







lowerCAmelCase__    :Optional[int]       =				{
    '''hans''': HansProcessor,
}
