




import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


lowerCAmelCase_						=			'http://www.mocksite.com/file1.txt'
lowerCAmelCase_						=			'"text": ["foo", "foo"]'
lowerCAmelCase_						=			'6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8'
class 						_A :
  _UpperCamelCase       :      Any			=      2_0_0
  _UpperCamelCase       :      List[Any]			=      {'''Content-Length''': '''100'''}
  _UpperCamelCase       :      List[Any]			=      {}






  def  __a      (  self		:     List[str]				,						**_A		:     List[Any]				)		->   Dict:






       """simple docstring"""







       return [bytes(_A				,						'''utf-8'''				)]






def  snake_case(    *__magic_name__			,		**__magic_name__   )			->							List[str]:





     '''simple docstring'''


     return MockResponse()






@pytest.mark.parametrize('''urls_type'''			,		[str, list, dict]   )
def  snake_case(    __magic_name__			,		__magic_name__			,		__magic_name__   )			->							Union[str, Any]:





     '''simple docstring'''


     import requests

     monkeypatch.setattr(__magic_name__			,		'''request'''			,		__magic_name__   )

     lowercase :       Optional[int]       =			URL
     if issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       int       =			url
     elif issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       Tuple       =			[url]
     elif issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       Dict       =			{'''train''': url}
     lowercase :       Union[str, Any]       =			'''dummy'''
     lowercase :       Any       =			'''downloads'''
     lowercase :       str       =			tmp_path
     lowercase :       Dict       =			DownloadConfig(
         cache_dir=os.path.join(__magic_name__			,		__magic_name__   )			,		use_etag=__magic_name__			,		)
     lowercase :       str       =			DownloadManager(dataset_name=__magic_name__			,		download_config=__magic_name__   )
     lowercase :       Dict       =			dl_manager.download(__magic_name__   )
     lowercase :       Optional[Any]       =			urls
     for downloaded_paths in [downloaded_paths]:
          if isinstance(__magic_name__			,		__magic_name__   ):
               lowercase :       Dict       =			[downloaded_paths]
               lowercase :       int       =			[urls]
          elif isinstance(__magic_name__			,		__magic_name__   ):
               assert "train" in downloaded_paths.keys()
               lowercase :       Tuple       =			downloaded_paths.values()
               lowercase :       Optional[Any]       =			urls.values()
          assert downloaded_paths
          for downloaded_path, input_url in zip(__magic_name__			,		__magic_name__   ):
               assert downloaded_path == dl_manager.downloaded_paths[input_url]
               lowercase :       List[str]       =			Path(__magic_name__   )
               lowercase :       List[Any]       =			downloaded_path.parts
               assert parts[-1] == HASH
               assert parts[-2] == cache_subdir
               assert downloaded_path.exists()
               lowercase :       Tuple       =			downloaded_path.read_text()
               assert content == CONTENT
               lowercase :       Tuple       =			downloaded_path.with_suffix('''.json'''   )
               assert metadata_downloaded_path.exists()
               lowercase :       Any       =			json.loads(metadata_downloaded_path.read_text()   )
               assert metadata_content == {"url": URL, "etag": None}






@pytest.mark.parametrize('''paths_type'''			,		[str, list, dict]   )
def  snake_case(    __magic_name__			,		__magic_name__			,		__magic_name__   )			->							int:





     '''simple docstring'''


     lowercase :       Any       =			str(__magic_name__   )
     if issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       Any       =			filename
     elif issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       List[Any]       =			[filename]
     elif issubclass(__magic_name__			,		__magic_name__   ):
          lowercase :       int       =			{'''train''': filename}
     lowercase :       List[str]       =			'''dummy'''
     lowercase :       Optional[int]       =			xz_file.parent
     lowercase :       Dict       =			'''extracted'''
     lowercase :       Optional[Any]       =			DownloadConfig(
         cache_dir=__magic_name__			,		use_etag=__magic_name__			,		)
     lowercase :       str       =			DownloadManager(dataset_name=__magic_name__			,		download_config=__magic_name__   )
     lowercase :       str       =			dl_manager.extract(__magic_name__   )
     lowercase :       Dict       =			paths
     for extracted_paths in [extracted_paths]:
          if isinstance(__magic_name__			,		__magic_name__   ):
               lowercase :       Dict       =			[extracted_paths]
               lowercase :       Union[str, Any]       =			[paths]
          elif isinstance(__magic_name__			,		__magic_name__   ):
               assert "train" in extracted_paths.keys()
               lowercase :       Union[str, Any]       =			extracted_paths.values()
               lowercase :       Any       =			paths.values()
          assert extracted_paths
          for extracted_path, input_path in zip(__magic_name__			,		__magic_name__   ):
               assert extracted_path == dl_manager.extracted_paths[input_path]
               lowercase :       Tuple       =			Path(__magic_name__   )
               lowercase :       Any       =			extracted_path.parts
               assert parts[-1] == hash_url_to_filename(__magic_name__			,		etag=__magic_name__   )
               assert parts[-2] == extracted_subdir
               assert extracted_path.exists()
               lowercase :       str       =			extracted_path.read_text()
               lowercase :       Dict       =			text_file.read_text()
               assert extracted_file_content == expected_file_content






def  snake_case(    __magic_name__			,		__magic_name__   )			->							int:





     '''simple docstring'''


     assert path.endswith('''.jsonl'''   )
     for num_items, line in enumerate(__magic_name__			,		start=1   ):
          lowercase :       List[str]       =			json.loads(line.decode('''utf-8'''   )   )
          assert item.keys() == {"col_1", "col_2", "col_3"}
     assert num_items == 4






@pytest.mark.parametrize('''archive_jsonl'''			,		['''tar_jsonl_path''', '''zip_jsonl_path''']   )
def  snake_case(    __magic_name__			,		__magic_name__   )			->							List[str]:





     '''simple docstring'''


     lowercase :       Optional[Any]       =			request.getfixturevalue(__magic_name__   )
     lowercase :       List[str]       =			DownloadManager()
     for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(__magic_name__   )			,		start=1   ):
          _test_jsonl(__magic_name__			,		__magic_name__   )
     assert num_jsonl == 2






@pytest.mark.parametrize('''archive_nested_jsonl'''			,		['''tar_nested_jsonl_path''', '''zip_nested_jsonl_path''']   )
def  snake_case(    __magic_name__			,		__magic_name__   )			->							Dict:





     '''simple docstring'''


     lowercase :       Tuple       =			request.getfixturevalue(__magic_name__   )
     lowercase :       List[Any]       =			DownloadManager()
     for num_tar, (path, file) in enumerate(dl_manager.iter_archive(__magic_name__   )			,		start=1   ):
          for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(__magic_name__   )			,		start=1   ):
               _test_jsonl(__magic_name__			,		__magic_name__   )
     assert num_tar == 1
     assert num_jsonl == 2






def  snake_case(    __magic_name__   )			->							str:





     '''simple docstring'''


     lowercase :       List[str]       =			DownloadManager()
     for num_file, file in enumerate(dl_manager.iter_files(__magic_name__   )			,		start=1   ):
          assert os.path.basename(__magic_name__   ) == ("test.txt" if num_file == 1 else "train.txt")
     assert num_file == 2





import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
 from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
 from tensorflow.keras.optimizers import Adam
class 						_A (  tf.keras.optimizers.schedules.LearningRateSchedule				):






  def __init__(  self		:     str				,						_A		:     float				,						_A		:     Callable				,						_A		:     int				,						_A		:     float = 1.0				,						_A		:     str = None				,						)		->   List[str]:






       """simple docstring"""







       super().__init__()
       lowercase :       List[str]       =			initial_learning_rate
       lowercase :       List[str]       =			warmup_steps
       lowercase :       Tuple       =			power
       lowercase :       Any       =			decay_schedule_fn
       lowercase :       Union[str, Any]       =			name






  def __call__(  self		:     str				,						_A		:     Any				)		->   Optional[int]:






       """simple docstring"""







       with tf.name_scope(self.name or '''WarmUp'''				) as name:
            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
            # learning rate will be `global_step/num_warmup_steps * init_lr`.
            lowercase :       List[Any]       =			tf.cast(_A				,						tf.floataa				)
            lowercase :       Union[str, Any]       =			tf.cast(self.warmup_steps				,						tf.floataa				)
            lowercase :       List[str]       =			global_step_float / warmup_steps_float
            lowercase :       int       =			self.initial_learning_rate * tf.math.pow(_A				,						self.power				)
            return tf.cond(
                global_step_float < warmup_steps_float				,						lambda: warmup_learning_rate				,						lambda: self.decay_schedule_fn(step - self.warmup_steps				)				,						name=_A				,						)






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






       """simple docstring"""







       return {
           "initial_learning_rate": self.initial_learning_rate,
           "decay_schedule_fn": self.decay_schedule_fn,
           "warmup_steps": self.warmup_steps,
           "power": self.power,
           "name": self.name,
       }






def  snake_case(    __magic_name__			,		__magic_name__			,		__magic_name__			,		__magic_name__ = 0.0			,		__magic_name__ = 0.9			,		__magic_name__ = 0.9_9_9			,		__magic_name__ = 1e-8			,		__magic_name__ = None			,		__magic_name__ = None			,		__magic_name__ = 0.0			,		__magic_name__ = 1.0			,		__magic_name__ = None			,		)			->							int:





     '''simple docstring'''


     lowercase :       List[Any]       =			tf.keras.optimizers.schedules.PolynomialDecay(
         initial_learning_rate=__magic_name__			,		decay_steps=num_train_steps - num_warmup_steps			,		end_learning_rate=init_lr * min_lr_ratio			,		power=__magic_name__			,		)
     if num_warmup_steps:
          lowercase :       Optional[int]       =			WarmUp(
              initial_learning_rate=__magic_name__			,		decay_schedule_fn=__magic_name__			,		warmup_steps=__magic_name__			,		)
     if weight_decay_rate > 0.0:
          lowercase :       Optional[Any]       =			AdamWeightDecay(
              learning_rate=__magic_name__			,		weight_decay_rate=__magic_name__			,		beta_a=__magic_name__			,		beta_a=__magic_name__			,		epsilon=__magic_name__			,		clipnorm=__magic_name__			,		global_clipnorm=__magic_name__			,		exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']			,		include_in_weight_decay=__magic_name__			,		)
     else:
          lowercase :       str       =			tf.keras.optimizers.Adam(
              learning_rate=__magic_name__			,		beta_a=__magic_name__			,		beta_a=__magic_name__			,		epsilon=__magic_name__			,		clipnorm=__magic_name__			,		global_clipnorm=__magic_name__			,		)
     # We return the optimizer and the LR scheduler in order to better track the
     # evolution of the LR independently of the optimizer.
     return optimizer, lr_schedule
class 						_A (  _lowerCamelCase				):






  def __init__(  self		:     Optional[Any]				,						_A		:     Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001				,						_A		:     float = 0.9				,						_A		:     float = 0.999				,						_A		:     float = 1E-7				,						_A		:     bool = False				,						_A		:     float = 0.0				,						_A		:     Optional[List[str]] = None				,						_A		:     Optional[List[str]] = None				,						_A		:     str = "AdamWeightDecay"				,						**_A		:     Any				,						)		->   Optional[Any]:






       """simple docstring"""







       super().__init__(_A				,						_A				,						_A				,						_A				,						_A				,						_A				,						**_A				)
       lowercase :       Tuple       =			weight_decay_rate
       lowercase :       List[str]       =			include_in_weight_decay
       lowercase :       Optional[Any]       =			exclude_from_weight_decay






  @classmethod
  def  __a      (  cls		:     Tuple				,						_A		:     Tuple				)		->   List[str]:






       """simple docstring"""







       lowercase :       Optional[int]       =			{'''WarmUp''': WarmUp}
       return super(_A				,						cls				).from_config(_A				,						custom_objects=_A				)






  def  __a      (  self		:     Dict				,						_A		:     Tuple				,						_A		:     Dict				,						_A		:     Tuple				)		->   Tuple:






       """simple docstring"""







       super(_A				,						self				)._prepare_local(_A				,						_A				,						_A				)
       lowercase :       List[Any]       =			tf.constant(
           self.weight_decay_rate				,						name='''adam_weight_decay_rate'''				)






  def  __a      (  self		:     Tuple				,						_A		:     Optional[int]				,						_A		:     Union[str, Any]				,						_A		:     List[Any]				)		->   Any:






       """simple docstring"""







       lowercase :       str       =			self._do_use_weight_decay(var.name				)
       if do_decay:
            return var.assign_sub(
                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['''weight_decay_rate''']				,						use_locking=self._use_locking				,						)
       return tf.no_op()






  def  __a      (  self		:     Union[str, Any]				,						_A		:     Any				,						_A		:     Optional[Any]=None				,						**_A		:     Tuple				)		->   Any:






       """simple docstring"""







       lowercase      ,	lowercase :       Tuple       =			list(zip(*_A				)				)
       return super(_A				,						self				).apply_gradients(zip(_A				,						_A				)				,						name=_A				,						**_A				)






  def  __a      (  self		:     List[Any]				,						_A		:     Optional[Any]				,						_A		:     str				,						_A		:     Union[str, Any]				)		->   Union[str, Any]:






       """simple docstring"""







       if apply_state is None:
            return self._decayed_lr_t[var_dtype], {}

       lowercase :       Any       =			apply_state or {}
       lowercase :       str       =			apply_state.get((var_device, var_dtype)				)
       if coefficients is None:
            lowercase :       List[Any]       =			self._fallback_apply_state(_A				,						_A				)
            lowercase :       List[Any]       =			coefficients

       return coefficients["lr_t"], {"apply_state": apply_state}






  def  __a      (  self		:     Tuple				,						_A		:     Union[str, Any]				,						_A		:     Tuple				,						_A		:     str=None				)		->   Optional[int]:






       """simple docstring"""







       lowercase      ,	lowercase :       List[str]       =			self._get_lr(var.device				,						var.dtype.base_dtype				,						_A				)
       lowercase :       Optional[Any]       =			self._decay_weights_op(_A				,						_A				,						_A				)
       with tf.control_dependencies([decay]				):
            return super(_A				,						self				)._resource_apply_dense(_A				,						_A				,						**_A				)






  def  __a      (  self		:     Optional[int]				,						_A		:     List[Any]				,						_A		:     Dict				,						_A		:     Union[str, Any]				,						_A		:     str=None				)		->   Optional[int]:






       """simple docstring"""







       lowercase      ,	lowercase :       Optional[int]       =			self._get_lr(var.device				,						var.dtype.base_dtype				,						_A				)
       lowercase :       str       =			self._decay_weights_op(_A				,						_A				,						_A				)
       with tf.control_dependencies([decay]				):
            return super(_A				,						self				)._resource_apply_sparse(_A				,						_A				,						_A				,						**_A				)






  def  __a      (  self		:     List[Any]				)		->   str:






       """simple docstring"""







       lowercase :       Optional[Any]       =			super().get_config()
       config.update({'''weight_decay_rate''': self.weight_decay_rate}				)
       return config






  def  __a      (  self		:     str				,						_A		:     Optional[int]				)		->   Tuple:






       """simple docstring"""







       if self.weight_decay_rate == 0:
            return False

       if self._include_in_weight_decay:
            for r in self._include_in_weight_decay:
                 if re.search(_A				,						_A				) is not None:
                      return True

       if self._exclude_from_weight_decay:
            for r in self._exclude_from_weight_decay:
                 if re.search(_A				,						_A				) is not None:
                      return False
       return True


class 						_A (  _lowerCamelCase				):






  def __init__(  self		:     List[Any]				)		->   str:






       """simple docstring"""







       lowercase :       Optional[Any]       =			[]
       lowercase :       str       =			None






  @property
  def  __a      (  self		:     Any				)		->   int:






       """simple docstring"""







       if self._accum_steps is None:
            lowercase :       Optional[Any]       =			tf.Variable(
                tf.constant(0				,						dtype=tf.intaa				)				,						trainable=_A				,						synchronization=tf.VariableSynchronization.ON_READ				,						aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA				,						)

       return self._accum_steps.value()






  @property
  def  __a      (  self		:     int				)		->   List[Any]:






       """simple docstring"""







       if not self._gradients:
            raise ValueError('''The accumulator should be called first to initialize the gradients'''				)
       return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]






  def __call__(  self		:     str				,						_A		:     int				)		->   str:






       """simple docstring"""







       if not self._gradients:
            lowercase :       Optional[Any]       =			self.step  # Create the step variable.
            self._gradients.extend(
                [
                    tf.Variable(
                        tf.zeros_like(_A				)				,						trainable=_A				,						synchronization=tf.VariableSynchronization.ON_READ				,						aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA				,						)
                    if gradient is not None
                    else gradient
                    for gradient in gradients
                ]				)
       if len(_A				) != len(self._gradients				):
            raise ValueError(f"""Expected {len(self._gradients				)} gradients, but got {len(_A				)}"""				)

       for accum_gradient, gradient in zip(self._gradients				,						_A				):
            if accum_gradient is not None and gradient is not None:
                 accum_gradient.assign_add(_A				)

       self._accum_steps.assign_add(1				)






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






       """simple docstring"""







       if not self._gradients:
            return
       self._accum_steps.assign(0				)
       for gradient in self._gradients:
            if gradient is not None:
                 gradient.assign(tf.zeros_like(_A				)				)