






'''simple docstring'''
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			):




    @parameterized.expand([(None,), ("foo.json",)]			)
    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     GenerationConfig(
           do_sample=__lowerCamelCase  ,     temperature=0.7  ,     length_penalty=1.0  ,     bad_words_ids=[[1, 2, 3], [4, 5]]  ,     )
       with tempfile.TemporaryDirectory() as tmp_dir:
          config.save_pretrained(__lowerCamelCase  ,     config_name=__lowerCamelCase			)
          lowercase_	:    Optional[int]           =     GenerationConfig.from_pretrained(__lowerCamelCase  ,     config_name=__lowerCamelCase			)

       # Checks parameters that were specified
       self.assertEqual(loaded_config.do_sample  ,     __lowerCamelCase			)
       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  ,     5_0			)
       self.assertEqual(loaded_config.max_length  ,     2_0			)
       self.assertEqual(loaded_config.max_time  ,     __lowerCamelCase			)




    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     AutoConfig.from_pretrained("gpt2"			)
       lowercase_	:    List[str]           =     GenerationConfig.from_model_config(__lowerCamelCase			)
       lowercase_	:    Optional[int]           =     GenerationConfig()

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

       # 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 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     GenerationConfig()
       lowercase_	:    Optional[int]           =     {
           '''max_new_tokens''': 1_0_2_4,
           '''foo''': '''bar''',
       }
       lowercase_	:    Dict           =     copy.deepcopy(__lowerCamelCase			)
       lowercase_	:    str           =     generation_config.update(**__lowerCamelCase			)

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

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

       # `.update()` returns a dictionary of unused kwargs
       self.assertEqual(__lowerCamelCase  ,     {"foo": "bar"}			)




    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     GenerationConfig()
       lowercase_	:    List[str]           =     '''bar'''

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

          lowercase_	:    Dict           =     GenerationConfig.from_pretrained(__lowerCamelCase			)
       # update_kwargs was used to update the config on valid attributes
       self.assertEqual(new_config.foo  ,     "bar"			)

       lowercase_	:    Union[str, Any]           =     GenerationConfig.from_model_config(__lowerCamelCase			)
       assert not hasattr(__lowerCamelCase  ,     "foo"			)  # no new kwargs should be initialized if from config





    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     GenerationConfig()
       self.assertEqual(default_config.temperature  ,     1.0			)
       self.assertEqual(default_config.do_sample  ,     __lowerCamelCase			)
       self.assertEqual(default_config.num_beams  ,     1			)

       lowercase_	:    List[Any]           =     GenerationConfig(
           do_sample=__lowerCamelCase  ,     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  ,     __lowerCamelCase			)
       self.assertEqual(config.num_beams  ,     1			)

       with tempfile.TemporaryDirectory() as tmp_dir:
          config.save_pretrained(__lowerCamelCase			)
          lowercase_	:    Dict           =     GenerationConfig.from_pretrained(__lowerCamelCase  ,     temperature=1.0			)

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


@is_staging_test
class        _UpperCAmelCase						(		unittest.TestCase			):




    @classmethod
    def 				lowerCAmelCase__				(			cls  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     TOKEN
       HfFolder.save_token(__lowerCamelCase			)




    @classmethod
    def 				lowerCAmelCase__				(			cls  :       Dict			):






       '''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 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     GenerationConfig(
           do_sample=__lowerCamelCase  ,     temperature=0.7  ,     length_penalty=1.0  ,     )
       config.push_to_hub("test-generation-config"  ,     use_auth_token=self._token			)

       lowercase_	:    Tuple           =     GenerationConfig.from_pretrained(f"""{USER}/test-generation-config"""			)
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(__lowerCamelCase  ,     getattr(__lowerCamelCase  ,     __lowerCamelCase			)			)

        # 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(
              __lowerCamelCase  ,     repo_id="test-generation-config"  ,     push_to_hub=__lowerCamelCase  ,     use_auth_token=self._token			)

       lowercase_	:    int           =     GenerationConfig.from_pretrained(f"""{USER}/test-generation-config"""			)
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(__lowerCamelCase  ,     getattr(__lowerCamelCase  ,     __lowerCamelCase			)			)





    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    str           =     GenerationConfig(
           do_sample=__lowerCamelCase  ,     temperature=0.7  ,     length_penalty=1.0  ,     )
       config.push_to_hub("valid_org/test-generation-config-org"  ,     use_auth_token=self._token			)

       lowercase_	:    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(__lowerCamelCase  ,     getattr(__lowerCamelCase  ,     __lowerCamelCase			)			)

        # 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(
              __lowerCamelCase  ,     repo_id="valid_org/test-generation-config-org"  ,     push_to_hub=__lowerCamelCase  ,     use_auth_token=self._token			)

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












'''simple docstring'''
import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


class        _UpperCAmelCase						(		snake_case			):




    def __init__(			self  :       Tuple  ,     a  :       NestedDataStructureLike[PathLike]  ,     a  :       Optional[NamedSplit] = None  ,     a  :       Optional[Features] = None  ,     a  :       str = None  ,     a  :       bool = False  ,     a  :       bool = False  ,     a  :       Optional[str] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       super().__init__(
           a  ,     split=a  ,     features=a  ,     cache_dir=a  ,     keep_in_memory=a  ,     streaming=a  ,     num_proc=a  ,     **a  ,     )
       lowercase_	:    str           =     field
       lowercase_	:    Optional[Any]           =     path_or_paths if isinstance(a  ,     a			) else {self.split: path_or_paths}
       lowercase_	:    Any           =     Json(
           cache_dir=a  ,     data_files=a  ,     features=a  ,     field=a  ,     **a  ,     )





    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       if self.streaming:
          lowercase_	:    Any           =     self.builder.as_streaming_dataset(split=self.split			)
       # Build regular (map-style) dataset
       else:
          lowercase_	:    Dict           =     None
          lowercase_	:    Optional[int]           =     None
          lowercase_	:    str           =     None
          lowercase_	:    str           =     None

          self.builder.download_and_prepare(
              download_config=a  ,     download_mode=a  ,     verification_mode=a  ,     base_path=a  ,     num_proc=self.num_proc  ,     )
          lowercase_	:    int           =     self.builder.as_dataset(
              split=self.split  ,     verification_mode=a  ,     in_memory=self.keep_in_memory			)
       return dataset


class        _UpperCAmelCase						:




    def __init__(			self  :       str  ,     a  :       Dataset  ,     a  :       Union[PathLike, BinaryIO]  ,     a  :       Optional[int] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       if num_proc is not None and num_proc <= 0:
          raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""			)

       lowercase_	:    Dict           =     dataset
       lowercase_	:    Optional[int]           =     path_or_buf
       lowercase_	:    List[Any]           =     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       lowercase_	:    Optional[Any]           =     num_proc
       lowercase_	:    List[Any]           =     "utf-8"
       lowercase_	:    List[str]           =     to_json_kwargs




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    str           =     self.to_json_kwargs.pop("path_or_buf"  ,     a			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("orient"  ,     "records"			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("lines"  ,     True if orient == "records" else False			)
       lowercase_	:    List[str]           =     self.to_json_kwargs.pop("index"  ,     False if orient in ["split", "table"] else True			)
       lowercase_	:    int           =     self.to_json_kwargs.pop("compression"  ,     a			)

       if compression not in [None, "infer", "gzip", "bz2", "xz"]:
          raise NotImplementedError(f"""`datasets` currently does not support {compression} compression"""			)

       if isinstance(self.path_or_buf  ,     (str, bytes, os.PathLike)			):
          with fsspec.open(self.path_or_buf  ,     "wb"  ,     compression=a			) as buffer:
             lowercase_	:    Dict           =     self._write(file_obj=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       else:
          if compression:
             raise NotImplementedError(
                 f"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
                 " was passed. Please provide a local path instead."			)
          lowercase_	:    Dict           =     self._write(
              file_obj=self.path_or_buf  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       return written




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       List[str]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    List[Any]           =     args

       lowercase_	:    Optional[int]           =     query_table(
           table=self.dataset.data  ,     key=slice(a  ,     offset + self.batch_size			)  ,     indices=self.dataset._indices  ,     )
       lowercase_	:    Dict           =     batch.to_pandas().to_json(
           path_or_buf=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **a			)
       if not json_str.endswith("\n"			):
          json_str += "\n"
       return json_str.encode(self.encoding			)





    def 				lowerCAmelCase__				(			self  :       int  ,     a  :       BinaryIO  ,     a  :       int  ,     a  :       str  ,     a  :       Union[str, Any]  ,     **a  :       str  ,     ):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     0

       if self.num_proc is None or self.num_proc == 1:
          for offset in logging.tqdm(
              range(0  ,     len(self.dataset			)  ,     self.batch_size			)  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
             lowercase_	:    Dict           =     self._batch_json((offset, orient, lines, index, to_json_kwargs)			)
             written += file_obj.write(a			)
       else:
          lowercase_    ,		lowercase_	:    Any           =     len(self.dataset			), self.batch_size
          with multiprocessing.Pool(self.num_proc			) as pool:
             for json_str in logging.tqdm(
                 pool.imap(
                     self._batch_json  ,     [(offset, orient, lines, index, to_json_kwargs) for offset in range(0  ,     a  ,     a			)]  ,     )  ,     total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
                written += file_obj.write(a			)

       return written




