

import contextlib
import os
import sqlitea

import pytest

from datasets import Dataset, Features, Value
from datasets.io.sql import SqlDatasetReader, SqlDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases, require_sqlalchemy
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     List[str]   ,      SCREAMING_SNAKE_CASE__      :     Union[str, Any]				)							->				Dict:





				'''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
@require_sqlalchemy
@pytest.mark.parametrize('keep_in_memory'   ,      [False, True]				)
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Union[str, Any]   ,      SCREAMING_SNAKE_CASE__      :     str   ,      SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     str				)							->				Dict:





				'''simple docstring'''






				A__   				=       tmp_path / 'cache'
				A__   				=       {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								A__   				=       SqlDatasetReader(
								    'dataset'   ,      'sqlite:///' + sqlite_path   ,      cache_dir=SCREAMING_SNAKE_CASE__   ,      keep_in_memory=SCREAMING_SNAKE_CASE__				).read()
				_check_sql_dataset(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)
@require_sqlalchemy
@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 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Any				)							->				Union[str, Any]:





				'''simple docstring'''






				A__   				=       tmp_path / 'cache'
				A__   				=       {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
				A__   				=       features.copy() if features else default_expected_features
				A__   				=       (
				    Features({feature: Value(SCREAMING_SNAKE_CASE__				) for feature, dtype in features.items()}				) if features is not None else None
				)
				A__   				=       SqlDatasetReader('dataset'   ,      'sqlite:///' + sqlite_path   ,      features=SCREAMING_SNAKE_CASE__   ,      cache_dir=SCREAMING_SNAKE_CASE__				).read()
				_check_sql_dataset(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     List[Any]				)							->				Optional[int]:





				'''simple docstring'''






				with contextlib.closing(sqlitea.connect(SCREAMING_SNAKE_CASE__				)				) as con:
								A__   				=       con.cursor()
								cur.execute('SELECT * FROM dataset'				)
								for row in cur:
												yield row
@require_sqlalchemy
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     List[str]   ,      SCREAMING_SNAKE_CASE__      :     List[Any]   ,      SCREAMING_SNAKE_CASE__      :     Union[str, Any]				)							->				Dict:





				'''simple docstring'''






				A__   				=       tmp_path / 'cache'
				A__   				=       os.path.join(SCREAMING_SNAKE_CASE__   ,      'tmp.sql'				)
				A__   				=       SqlDatasetReader('dataset'   ,      'sqlite:///' + sqlite_path   ,      cache_dir=SCREAMING_SNAKE_CASE__				).read()
				SqlDatasetWriter(SCREAMING_SNAKE_CASE__   ,      'dataset'   ,      'sqlite:///' + output_sqlite_path   ,      num_proc=1				).write()

				A__   				=       iter_sql_file(SCREAMING_SNAKE_CASE__				)
				A__   				=       iter_sql_file(SCREAMING_SNAKE_CASE__				)

				for rowa, rowa in zip(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				):
								assert rowa == rowa
@require_sqlalchemy
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     str   ,      SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Dict				)							->				List[str]:





				'''simple docstring'''






				A__   				=       tmp_path / 'cache'
				A__   				=       os.path.join(SCREAMING_SNAKE_CASE__   ,      'tmp.sql'				)
				A__   				=       SqlDatasetReader('dataset'   ,      'sqlite:///' + sqlite_path   ,      cache_dir=SCREAMING_SNAKE_CASE__				).read()
				SqlDatasetWriter(SCREAMING_SNAKE_CASE__   ,      'dataset'   ,      'sqlite:///' + output_sqlite_path   ,      num_proc=2				).write()

				A__   				=       iter_sql_file(SCREAMING_SNAKE_CASE__				)
				A__   				=       iter_sql_file(SCREAMING_SNAKE_CASE__				)

				for rowa, rowa in zip(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				):
								assert rowa == rowa

@require_sqlalchemy
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     List[str]   ,      SCREAMING_SNAKE_CASE__      :     Optional[int]				)							->				Tuple:





				'''simple docstring'''






				A__   				=       tmp_path / 'cache'
				A__   				=       os.path.join(SCREAMING_SNAKE_CASE__   ,      'tmp.sql'				)
				A__   				=       SqlDatasetReader('dataset'   ,      'sqlite:///' + sqlite_path   ,      cache_dir=SCREAMING_SNAKE_CASE__				).read()
				with pytest.raises(SCREAMING_SNAKE_CASE__				):
								SqlDatasetWriter(SCREAMING_SNAKE_CASE__   ,      'dataset'   ,      'sqlite:///' + output_sqlite_path   ,      num_proc=0				).write()




import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
     from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments


@require_torch
class 			A						(							unittest.TestCase			):







    """simple docstring"""





    def 						snake_case__					(				self			:	List[Any],lowercase_			:	str		)->   List[Any]:





        '''simple docstring'''







        for model_result in results.values():
            for batch_size, sequence_length in zip(model_result['bs'],model_result['ss']		):
                A__   				=       model_result['result'][batch_size][sequence_length]
                self.assertIsNotNone(lowercase_		)


    def 						snake_case__					(				self			:	Dict		)->   List[str]:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


    def 						snake_case__					(				self			:	Dict		)->   List[str]:





        '''simple docstring'''







        A__   				=       'sgugger/tiny-distilbert-classification'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,only_pretrain_model=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


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





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,torchscript=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


    @unittest.skipIf(torch_device == 'cpu','Cant do half precision'		)
    def 						snake_case__					(				self			:	Any		)->   Dict:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,fpaa=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


    def 						snake_case__					(				self			:	Any		)->   Optional[Any]:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       AutoConfig.from_pretrained(lowercase_		)
        # set architectures equal to `None`
        A__   				=       None
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_,configs=[config]		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


    def 						snake_case__					(				self			:	Union[str, Any]		)->   int:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_train_result		)
        self.check_results_dict_not_empty(results.memory_train_result		)


    @unittest.skipIf(torch_device == 'cpu','Can\'t do half precision'		)
    def 						snake_case__					(				self			:	List[Any]		)->   Dict:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],fpaa=lowercase_,multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_train_result		)
        self.check_results_dict_not_empty(results.memory_train_result		)


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





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       AutoConfig.from_pretrained(lowercase_		)
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_,configs=[config]		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


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





        '''simple docstring'''







        A__   				=       'sshleifer/tinier_bart'
        A__   				=       AutoConfig.from_pretrained(lowercase_		)
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_,configs=[config]		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_inference_result		)
        self.check_results_dict_not_empty(results.memory_inference_result		)


    def 						snake_case__					(				self			:	List[str]		)->   List[str]:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        A__   				=       AutoConfig.from_pretrained(lowercase_		)
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_,configs=[config]		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_train_result		)
        self.check_results_dict_not_empty(results.memory_train_result		)


    def 						snake_case__					(				self			:	int		)->   Union[str, Any]:





        '''simple docstring'''







        A__   				=       'sshleifer/tinier_bart'
        A__   				=       AutoConfig.from_pretrained(lowercase_		)
        A__   				=       PyTorchBenchmarkArguments(
            models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],multi_process=lowercase_,)
        A__   				=       PyTorchBenchmark(lowercase_,configs=[config]		)
        A__   				=       benchmark.run()
        self.check_results_dict_not_empty(results.time_train_result		)
        self.check_results_dict_not_empty(results.memory_train_result		)


    def 						snake_case__					(				self			:	Optional[Any]		)->   Tuple:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'
        with tempfile.TemporaryDirectory() as tmp_dir:
            A__   				=       PyTorchBenchmarkArguments(
                models=[MODEL_ID],training=lowercase_,inference=lowercase_,save_to_csv=lowercase_,sequence_lengths=[8],batch_sizes=[1],inference_time_csv_file=os.path.join(lowercase_,'inf_time.csv'		),train_memory_csv_file=os.path.join(lowercase_,'train_mem.csv'		),inference_memory_csv_file=os.path.join(lowercase_,'inf_mem.csv'		),train_time_csv_file=os.path.join(lowercase_,'train_time.csv'		),env_info_csv_file=os.path.join(lowercase_,'env.csv'		),multi_process=lowercase_,)
            A__   				=       PyTorchBenchmark(lowercase_		)
            benchmark.run()
            self.assertTrue(Path(os.path.join(lowercase_,'inf_time.csv'		)		).exists()		)
            self.assertTrue(Path(os.path.join(lowercase_,'train_time.csv'		)		).exists()		)
            self.assertTrue(Path(os.path.join(lowercase_,'inf_mem.csv'		)		).exists()		)
            self.assertTrue(Path(os.path.join(lowercase_,'train_mem.csv'		)		).exists()		)
            self.assertTrue(Path(os.path.join(lowercase_,'env.csv'		)		).exists()		)


    def 						snake_case__					(				self			:	Tuple		)->   str:





        '''simple docstring'''







        A__   				=       'sshleifer/tiny-gpt2'

        def _check_summary_is_not_empty(lowercase_			:	Optional[Any]		):
            self.assertTrue(hasattr(lowercase_,'sequential'		)		)
            self.assertTrue(hasattr(lowercase_,'cumulative'		)		)
            self.assertTrue(hasattr(lowercase_,'current'		)		)
            self.assertTrue(hasattr(lowercase_,'total'		)		)

        with tempfile.TemporaryDirectory() as tmp_dir:
            A__   				=       PyTorchBenchmarkArguments(
                models=[MODEL_ID],training=lowercase_,inference=lowercase_,sequence_lengths=[8],batch_sizes=[1],log_filename=os.path.join(lowercase_,'log.txt'		),log_print=lowercase_,trace_memory_line_by_line=lowercase_,multi_process=lowercase_,)
            A__   				=       PyTorchBenchmark(lowercase_		)
            A__   				=       benchmark.run()
            _check_summary_is_not_empty(result.inference_summary		)
            _check_summary_is_not_empty(result.train_summary		)
            self.assertTrue(Path(os.path.join(lowercase_,'log.txt'		)		).exists()		)

