

import copy
import tempfile
import unittest

from huggingface_hub import HfFolder, delete_repo
from parameterized import parameterized
from requests.exceptions import HTTPError

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test






class _UpperCamelCase							(							unittest.TestCase      ):
			'''simple docstring'''







			@parameterized.expand([(None,), ("foo.json",)]		)
			def 							__UpperCamelCase       (	self							:       List[Any]					,					a							:       List[Any]		)							->				List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					str    =				GenerationConfig(
				    do_sample=__UpperCAmelCase					,					temperature=0.7					,					length_penalty=1.0					,					bad_words_ids=[[1, 2, 3], [4, 5]]					,					)
				with tempfile.TemporaryDirectory() as tmp_dir:
					config.save_pretrained(__UpperCAmelCase					,					config_name=__UpperCAmelCase		)
					SCREAMING_SNAKE_CASE      :					Tuple    =				GenerationConfig.from_pretrained(__UpperCAmelCase					,					config_name=__UpperCAmelCase		)

				# Checks parameters that were specified
				self.assertEqual(loaded_config.do_sample					,					__UpperCAmelCase		)
				self.assertEqual(loaded_config.temperature					,					0.7		)
				self.assertEqual(loaded_config.length_penalty					,					1.0		)
				self.assertEqual(loaded_config.bad_words_ids					,					[[1, 2, 3], [4, 5]]		)

				# Checks parameters that were not specified (defaults)
				self.assertEqual(loaded_config.top_k					,					50		)
				self.assertEqual(loaded_config.max_length					,					20		)
				self.assertEqual(loaded_config.max_time					,					__UpperCAmelCase		)
			def 							__UpperCamelCase       (	self							:       int		)							->				int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Dict    =				AutoConfig.from_pretrained("gpt2"		)
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				GenerationConfig.from_model_config(__UpperCAmelCase		)
				SCREAMING_SNAKE_CASE      :					List[Any]    =				GenerationConfig()

				# The generation config has loaded a few non-default parameters from the model config
				self.assertNotEqual(__UpperCAmelCase					,					__UpperCAmelCase		)

				# One of those parameters is eos_token_id -- check if it matches
				self.assertNotEqual(generation_config_from_model.eos_token_id					,					default_generation_config.eos_token_id		)
				self.assertEqual(generation_config_from_model.eos_token_id					,					model_config.eos_token_id		)
			def 							__UpperCamelCase       (	self							:       Optional[int]		)							->				str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				GenerationConfig()
				SCREAMING_SNAKE_CASE      :					str    =				{
				    "max_new_tokens": 1024,
				    "foo": "bar",
				}
				SCREAMING_SNAKE_CASE      :					Dict    =				copy.deepcopy(__UpperCAmelCase		)
				SCREAMING_SNAKE_CASE      :					Tuple    =				generation_config.update(**__UpperCAmelCase		)

				# update_kwargs was not modified (no side effects)
				self.assertEqual(__UpperCAmelCase					,					__UpperCAmelCase		)

				# update_kwargs was used to update the config on valid attributes
				self.assertEqual(generation_config.max_new_tokens					,					1024		)

				# `.update()` returns a dictionary of unused kwargs
				self.assertEqual(__UpperCAmelCase					,					{"foo": "bar"}		)
			def 							__UpperCamelCase       (	self							:       Optional[int]		)							->				Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				GenerationConfig()
				SCREAMING_SNAKE_CASE      :					Optional[int]    =				"bar"

				with tempfile.TemporaryDirectory("test-generation-config"		) as tmp_dir:
					generation_config.save_pretrained(__UpperCAmelCase		)

					SCREAMING_SNAKE_CASE      :					Optional[Any]    =				GenerationConfig.from_pretrained(__UpperCAmelCase		)
				# update_kwargs was used to update the config on valid attributes
				self.assertEqual(new_config.foo					,					"bar"		)

				SCREAMING_SNAKE_CASE      :					List[str]    =				GenerationConfig.from_model_config(__UpperCAmelCase		)
				assert not hasattr(__UpperCAmelCase					,					"foo"		)  # no new kwargs should be initialized if from config






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					List[Any]    =				GenerationConfig()
				self.assertEqual(default_config.temperature					,					1.0		)
				self.assertEqual(default_config.do_sample					,					__UpperCAmelCase		)
				self.assertEqual(default_config.num_beams					,					1		)

				SCREAMING_SNAKE_CASE      :					Optional[int]    =				GenerationConfig(
				    do_sample=__UpperCAmelCase					,					temperature=0.7					,					length_penalty=1.0					,					bad_words_ids=[[1, 2, 3], [4, 5]]					,					)
				self.assertEqual(config.temperature					,					0.7		)
				self.assertEqual(config.do_sample					,					__UpperCAmelCase		)
				self.assertEqual(config.num_beams					,					1		)

				with tempfile.TemporaryDirectory() as tmp_dir:
					config.save_pretrained(__UpperCAmelCase		)
					SCREAMING_SNAKE_CASE      :					Optional[Any]    =				GenerationConfig.from_pretrained(__UpperCAmelCase					,					temperature=1.0		)

				self.assertEqual(loaded_config.temperature					,					1.0		)
				self.assertEqual(loaded_config.do_sample					,					__UpperCAmelCase		)
				self.assertEqual(loaded_config.num_beams					,					1		)  # default value






@is_staging_test
class _UpperCamelCase							(							unittest.TestCase      ):
			'''simple docstring'''







			@classmethod
			def 							__UpperCamelCase       (	cls							:       Tuple		)							->				Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				TOKEN
				HfFolder.save_token(__UpperCAmelCase		)
			@classmethod
			def 							__UpperCamelCase       (	cls							:       Dict		)							->				Any:

				"""simple docstring"""

				try:
					delete_repo(token=cls._token					,					repo_id="test-generation-config"		)
				except HTTPError:
					pass

				try:
					delete_repo(token=cls._token					,					repo_id="valid_org/test-generation-config-org"		)
				except HTTPError:
					pass
			def 							__UpperCamelCase       (	self							:       str		)							->				Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Dict    =				GenerationConfig(
				    do_sample=__UpperCAmelCase					,					temperature=0.7					,					length_penalty=1.0					,					)
				config.push_to_hub("test-generation-config"					,					use_auth_token=self._token		)

				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				GenerationConfig.from_pretrained(F"{USER}/test-generation-config"		)
				for k, v in config.to_dict().items():
					if k != "transformers_version":
						self.assertEqual(__UpperCAmelCase					,					getattr(__UpperCAmelCase					,					__UpperCAmelCase		)		)

					# Reset repo
				delete_repo(token=self._token					,					repo_id="test-generation-config"		)

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
					config.save_pretrained(
					    __UpperCAmelCase					,					repo_id="test-generation-config"					,					push_to_hub=__UpperCAmelCase					,					use_auth_token=self._token		)

				SCREAMING_SNAKE_CASE      :					Any    =				GenerationConfig.from_pretrained(F"{USER}/test-generation-config"		)
				for k, v in config.to_dict().items():
					if k != "transformers_version":
						self.assertEqual(__UpperCAmelCase					,					getattr(__UpperCAmelCase					,					__UpperCAmelCase		)		)






			def 							__UpperCamelCase       (	self							:       Union[str, Any]		)							->				List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				GenerationConfig(
				    do_sample=__UpperCAmelCase					,					temperature=0.7					,					length_penalty=1.0					,					)
				config.push_to_hub("valid_org/test-generation-config-org"					,					use_auth_token=self._token		)

				SCREAMING_SNAKE_CASE      :					str    =				GenerationConfig.from_pretrained("valid_org/test-generation-config-org"		)
				for k, v in config.to_dict().items():
					if k != "transformers_version":
						self.assertEqual(__UpperCAmelCase					,					getattr(__UpperCAmelCase					,					__UpperCAmelCase		)		)

					# Reset repo
				delete_repo(token=self._token					,					repo_id="valid_org/test-generation-config-org"		)

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
					config.save_pretrained(
					    __UpperCAmelCase					,					repo_id="valid_org/test-generation-config-org"					,					push_to_hub=__UpperCAmelCase					,					use_auth_token=self._token		)

				SCREAMING_SNAKE_CASE      :					List[Any]    =				GenerationConfig.from_pretrained("valid_org/test-generation-config-org"		)
				for k, v in config.to_dict().items():
					if k != "transformers_version":
						self.assertEqual(__UpperCAmelCase					,					getattr(__UpperCAmelCase					,					__UpperCAmelCase		)		)



"""simple docstring"""







import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__lowercase    					=      datasets.utils.logging.get_logger(__name__)






@dataclass
class 						_A  (  datasets.BuilderConfig  ):







    """simple docstring"""




    UpperCAmelCase       :							int	=       1_0_0_0_0
    UpperCAmelCase       :							Optional[List[str]]	=       None
    UpperCAmelCase       :							Optional[datasets.Features]	=       None






class 						_A  (  datasets.ArrowBasedBuilder  ):







    """simple docstring"""




    UpperCAmelCase       :							str	=       ParquetConfig






    def 	__snake_case							(					self				:    Tuple):
      return datasets.DatasetInfo(features=self.config.features)






    def 	__snake_case							(					self				:    List[Any]				, __UpperCAmelCase				:    str):
      if not self.config.data_files:
        raise ValueError(f'''At least one data file must be specified, but got data_files={self.config.data_files}''')
      a						:		str      							=			dl_manager.download_and_extract(self.config.data_files)
      if isinstance(__UpperCAmelCase				, (str, list, tuple)):
        a						:		Dict      							=			data_files
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		str      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		List[Any]      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN				, gen_kwargs={"files": files})]
      a						:		Dict      							=			[]
      for split_name, files in data_files.items():
        if isinstance(__UpperCAmelCase				, __UpperCAmelCase):
          a						:		Optional[int]      							=			[files]
        # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
        a						:		Tuple      							=			[dl_manager.iter_files(__UpperCAmelCase) for file in files]
        # Infer features is they are stoed in the arrow schema
        if self.info.features is None:
          for file in itertools.chain.from_iterable(__UpperCAmelCase):
            with open(__UpperCAmelCase				, "rb") as f:
              a						:		Tuple      							=			datasets.Features.from_arrow_schema(pq.read_schema(__UpperCAmelCase))
            break
        splits.append(datasets.SplitGenerator(name=__UpperCAmelCase				, gen_kwargs={"files": files}))
      return splits






    def 	__snake_case							(					self				:    List[str]				, __UpperCAmelCase				:    pa.Table):
      if self.info.features is not None:
        # more expensive cast to support nested features with keys in a different order
        # allows str <-> int/float or str to Audio for example
        a						:		Optional[int]      							=			table_cast(__UpperCAmelCase				, self.info.features.arrow_schema)
      return pa_table






    def 	__snake_case							(					self				:    Union[str, Any]				, __UpperCAmelCase				:    int):
      a						:		Tuple      							=			self.info.features.arrow_schema if self.info.features is not None else None
      if self.info.features is not None and self.config.columns is not None:
        if sorted(field.name for field in schema) != sorted(self.config.columns):
          raise ValueError(
              f'''Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\'''')
      for file_idx, file in enumerate(itertools.chain.from_iterable(__UpperCAmelCase)):
        with open(__UpperCAmelCase				, "rb") as f:
          a						:		Tuple      							=			pq.ParquetFile(__UpperCAmelCase)
          try:
            for batch_idx, record_batch in enumerate(
                parquet_file.iter_batches(batch_size=self.config.batch_size				, columns=self.config.columns)):
              a						:		Optional[Any]      							=			pa.Table.from_batches([record_batch])
              # Uncomment for debugging (will print the Arrow table size and elements)
              # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
              # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
              yield f'''{file_idx}_{batch_idx}''', self._cast_table(__UpperCAmelCase)
          except ValueError as e:
            logger.error(f'''Failed to read file \'{file}\' with error {type(__UpperCAmelCase)}: {e}''')
            raise






