


import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam




class A  (   datasets.BeamBasedBuilder     ):






      def        lowerCamelCase       (					self						:     Dict )					->						Union[str, Any]:
             """simple docstring"""




             return datasets.DatasetInfo(
                 features=datasets.Features({'content': datasets.Value('string' )} )			,   supervised_keys=lowercase_			,   )






      def        lowerCamelCase       (					self						:     List[Any]			,   lowercase_						:     int			,   lowercase_						:     Optional[int] )					->						int:
             """simple docstring"""




             return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,   gen_kwargs={'examples': get_test_dummy_examples()} )]






      def        lowerCamelCase       (					self						:     Any			,   lowercase_						:     Optional[Any]			,   lowercase_						:     List[str] )					->						Optional[int]:
             """simple docstring"""




             import apache_beam as beam

             return pipeline | "Load Examples" >> beam.Create(lowercase_ )




class A  (   datasets.BeamBasedBuilder     ):






      def        lowerCamelCase       (					self						:     int )					->						str:
             """simple docstring"""




             return datasets.DatasetInfo(
                 features=datasets.Features({'a': datasets.Sequence({'b': datasets.Value('string' )} )} )			,   supervised_keys=lowercase_			,   )






      def        lowerCamelCase       (					self						:     List[Any]			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     List[str] )					->						List[str]:
             """simple docstring"""




             return [
                 datasets.SplitGenerator(name=datasets.Split.TRAIN			,   gen_kwargs={'examples': get_test_nested_examples()} )
             ]






      def        lowerCamelCase       (					self						:     Union[str, Any]			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any] )					->						Optional[int]:
             """simple docstring"""




             import apache_beam as beam

             return pipeline | "Load Examples" >> beam.Create(lowercase_ )




def 	a_      (    ):
       '''simple docstring'''





       return [(i, {"content": content}) for i, content in enumerate(['foo', 'bar', 'foobar']     )]




def 	a_      (    ):
       '''simple docstring'''





       return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['foo', 'bar', 'foobar']     )]




class A  (   UpperCamelCase_     ):






      @require_beam
      def        lowerCamelCase       (					self						:     str )					->						List[Any]:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=len(get_test_dummy_examples() )
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _lowerCamelCase       :		Tuple		=DummyBeamDataset(cache_dir=lowercase_			,   beam_runner='DirectRunner' )
                    builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   F'''{builder.name}-train.arrow''' ) ) )
                    self.assertDictEqual(builder.info.features			,   datasets.Features({'content': datasets.Value('string' )} ) )
                    _lowerCamelCase       :		Union[str, Any]		=builder.as_dataset()
                    self.assertEqual(dset['train'].num_rows			,   lowercase_ )
                    self.assertEqual(dset['train'].info.splits['train'].num_examples			,   lowercase_ )
                    self.assertDictEqual(dset['train'][0]			,   get_test_dummy_examples()[0][1] )
                    self.assertDictEqual(
                        dset['train'][expected_num_examples - 1]			,   get_test_dummy_examples()[expected_num_examples - 1][1] )
                    self.assertTrue(
                        os.path.exists(os.path.join(lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   'dataset_info.json' ) ) )
                    del dset






      @require_beam
      def        lowerCamelCase       (					self						:     List[Any] )					->						List[Any]:
             """simple docstring"""




             import apache_beam as beam

             _lowerCamelCase       :		Union[str, Any]		=beam.io.parquetio.WriteToParquet

             _lowerCamelCase       :		Dict		=len(get_test_dummy_examples() )
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _lowerCamelCase       :		Union[str, Any]		=DummyBeamDataset(cache_dir=lowercase_			,   beam_runner='DirectRunner' )
                    with patch('apache_beam.io.parquetio.WriteToParquet' ) as write_parquet_mock:
                           _lowerCamelCase       :		int		=partial(lowercase_			,   num_shards=2 )
                           builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(
                                lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   F'''{builder.name}-train-00000-of-00002.arrow''' ) ) )
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(
                                lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   F'''{builder.name}-train-00000-of-00002.arrow''' ) ) )
                    self.assertDictEqual(builder.info.features			,   datasets.Features({'content': datasets.Value('string' )} ) )
                    _lowerCamelCase       :		List[Any]		=builder.as_dataset()
                    self.assertEqual(dset['train'].num_rows			,   lowercase_ )
                    self.assertEqual(dset['train'].info.splits['train'].num_examples			,   lowercase_ )
                    # Order is not preserved when sharding, so we just check that all the elements are there
                    self.assertListEqual(sorted(dset['train']['content'] )			,   sorted(['foo', 'bar', 'foobar'] ) )
                    self.assertTrue(
                        os.path.exists(os.path.join(lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   'dataset_info.json' ) ) )
                    del dset






      @require_beam
      def        lowerCamelCase       (					self						:     Tuple )					->						Any:
             """simple docstring"""




             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _lowerCamelCase       :		str		=DummyBeamDataset(cache_dir=lowercase_ )
                    self.assertRaises(datasets.builder.MissingBeamOptions			,   builder.download_and_prepare )






      @require_beam
      def        lowerCamelCase       (					self						:     int )					->						Optional[Any]:
             """simple docstring"""




             _lowerCamelCase       :		Union[str, Any]		=len(get_test_nested_examples() )
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _lowerCamelCase       :		Dict		=NestedBeamDataset(cache_dir=lowercase_			,   beam_runner='DirectRunner' )
                    builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   F'''{builder.name}-train.arrow''' ) ) )
                    self.assertDictEqual(
                        builder.info.features			,   datasets.Features({'a': datasets.Sequence({'b': datasets.Value('string' )} )} ) )
                    _lowerCamelCase       :		Union[str, Any]		=builder.as_dataset()
                    self.assertEqual(dset['train'].num_rows			,   lowercase_ )
                    self.assertEqual(dset['train'].info.splits['train'].num_examples			,   lowercase_ )
                    self.assertDictEqual(dset['train'][0]			,   get_test_nested_examples()[0][1] )
                    self.assertDictEqual(
                        dset['train'][expected_num_examples - 1]			,   get_test_nested_examples()[expected_num_examples - 1][1] )
                    self.assertTrue(
                        os.path.exists(os.path.join(lowercase_			,   builder.name			,   'default'			,   '0.0.0'			,   'dataset_info.json' ) ) )
                    del dset






import os

import pytest
from attr import dataclass


lowerCamelCase											=       'us-east-1'  # defaults region




@dataclass
class A  :
      UpperCamelCase__							:  str
      UpperCamelCase__							:  Dict      ='arn:aws:iam::558105141721:role/sagemaker_execution_role'
      UpperCamelCase__							:  int      ={
          'task_name': 'mnli',
          'per_device_train_batch_size': 16,
          'per_device_eval_batch_size': 16,
          'do_train': True,
          'do_eval': True,
          'do_predict': True,
          'output_dir': '/opt/ml/model',
          'overwrite_output_dir': True,
          'max_steps': 500,
          'save_steps': 5500,
      }
      UpperCamelCase__							:  Optional[Any]      ={**hyperparameters, 'max_steps': 1000}






      @property
      def        lowerCamelCase       (					self						:     List[Any] )					->						str:
             """simple docstring"""




             if self.framework == "pytorch":
                    return [
                        {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
                        {"Name": "eval_accuracy", "Regex": r"eval_accuracy.*=\D*(.*?)$"},
                        {"Name": "eval_loss", "Regex": r"eval_loss.*=\D*(.*?)$"},
                    ]
             else:
                    return [
                        {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
                        {"Name": "eval_accuracy", "Regex": r"loss.*=\D*(.*?)]?$"},
                        {"Name": "eval_loss", "Regex": r"sparse_categorical_accuracy.*=\D*(.*?)]?$"},
                    ]






      @property
      def        lowerCamelCase       (					self						:     List[Any] )					->						str:
             """simple docstring"""




             return F'''{self.framework}-transfromers-test'''






      @property
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						str:
             """simple docstring"""




             return F'''./tests/sagemaker/scripts/{self.framework}'''






      @property
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						str:
             """simple docstring"""




             if self.framework == "pytorch":
                    return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04"
             else:
                    return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-training:2.4.1-transformers4.6.1-gpu-py37-cu110-ubuntu18.04"




@pytest.fixture(scope='class'     )
def 	a_      (    SCREAMING_SNAKE_CASE__					:      Tuple     ):
       '''simple docstring'''





       _lowerCamelCase       :		List[Any]		=SageMakerTestEnvironment(framework=request.cls.framework     )


