





"""simple docstring"""





import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
       import torch
       from torch import nn

       from transformers import (
           Adafactor,
           AdamW,
           get_constant_schedule,
           get_constant_schedule_with_warmup,
           get_cosine_schedule_with_warmup,
           get_cosine_with_hard_restarts_schedule_with_warmup,
           get_inverse_sqrt_schedule,
           get_linear_schedule_with_warmup,
           get_polynomial_decay_schedule_with_warmup,
       )


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str				,		SCREAMING_SNAKE_CASE	:		Dict=10     ):







      '''simple docstring'''




      lowerCAmelCase    		=			[]
      for _ in range(SCREAMING_SNAKE_CASE     ):
            lrs.append(scheduler.get_lr()[0]     )
            scheduler.step()
      return lrs


def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		List[Any]				,		SCREAMING_SNAKE_CASE	:		Optional[int]=10     ):







      '''simple docstring'''




      lowerCAmelCase    		=			[]
      for step in range(SCREAMING_SNAKE_CASE     ):
            lrs.append(scheduler.get_lr()[0]     )
            scheduler.step()
            if step == num_steps // 2:
                  with tempfile.TemporaryDirectory() as tmpdirname:
                        lowerCAmelCase    		=			os.path.join(SCREAMING_SNAKE_CASE				,		"""schedule.bin"""     )
                        torch.save(scheduler.state_dict()				,		SCREAMING_SNAKE_CASE     )

                        lowerCAmelCase    		=			torch.load(SCREAMING_SNAKE_CASE     )
                        scheduler.load_state_dict(SCREAMING_SNAKE_CASE     )
      return lrs

@require_torch
class 						lowercase   (   unittest.TestCase   ):





      def        _snake_case     (			self ,		lowercase ,		lowercase ,		lowercase					)						->	Tuple:
            self.assertEqual(len(__lowerCamelCase					) ,		len(__lowerCamelCase					)					)
            for a, b in zip(__lowerCamelCase ,		__lowerCamelCase					):
                  self.assertAlmostEqual(__lowerCamelCase ,		__lowerCamelCase ,		delta=__lowerCamelCase					)





      def        _snake_case     (			self					)						->	str:
            lowerCAmelCase    		=			torch.tensor([0.1, -0.2, -0.1] ,		requires_grad=__lowerCamelCase					)
            lowerCAmelCase    		=			torch.tensor([0.4, 0.2, -0.5]					)
            lowerCAmelCase    		=			nn.MSELoss()
            # No warmup, constant schedule, no gradient clipping
            lowerCAmelCase    		=			AdamW(params=[w] ,		lr=2e-1 ,		weight_decay=0.0					)
            for _ in range(100					):
                  lowerCAmelCase    		=			criterion(__lowerCamelCase ,		__lowerCamelCase					)
                  loss.backward()
                  optimizer.step()
                  w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
                  w.grad.zero_()
            self.assertListAlmostEqual(w.tolist() ,		[0.4, 0.2, -0.5] ,		tol=1e-2					)





      def        _snake_case     (			self					)						->	Tuple:
            lowerCAmelCase    		=			torch.tensor([0.1, -0.2, -0.1] ,		requires_grad=__lowerCamelCase					)
            lowerCAmelCase    		=			torch.tensor([0.4, 0.2, -0.5]					)
            lowerCAmelCase    		=			nn.MSELoss()
            # No warmup, constant schedule, no gradient clipping
            lowerCAmelCase    		=			Adafactor(
                params=[w] ,		lr=1e-2 ,		eps=(1e-30, 1e-3) ,		clip_threshold=1.0 ,		decay_rate=-0.8 ,		betaa=__lowerCamelCase ,		weight_decay=0.0 ,		relative_step=__lowerCamelCase ,		scale_parameter=__lowerCamelCase ,		warmup_init=__lowerCamelCase ,		)
            for _ in range(1_000					):
                  lowerCAmelCase    		=			criterion(__lowerCamelCase ,		__lowerCamelCase					)
                  loss.backward()
                  optimizer.step()
                  w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
                  w.grad.zero_()
            self.assertListAlmostEqual(w.tolist() ,		[0.4, 0.2, -0.5] ,		tol=1e-2					)







@require_torch
class 						lowercase   (   unittest.TestCase   ):
      _SCREAMING_SNAKE_CASE							      = nn.Linear(50	,    50   ) if is_torch_available() else None
      _SCREAMING_SNAKE_CASE							      = AdamW(m.parameters()	,    lr=10.0   ) if is_torch_available() else None
      _SCREAMING_SNAKE_CASE							      = 10





      def        _snake_case     (			self ,		lowercase ,		lowercase ,		lowercase ,		lowercase=None					)						->	Tuple:
            self.assertEqual(len(__lowerCamelCase					) ,		len(__lowerCamelCase					)					)
            for a, b in zip(__lowerCamelCase ,		__lowerCamelCase					):
                  self.assertAlmostEqual(__lowerCamelCase ,		__lowerCamelCase ,		delta=__lowerCamelCase ,		msg=__lowerCamelCase					)





      def        _snake_case     (			self					)						->	Any:
            lowerCAmelCase    		=			{"""num_warmup_steps""": 2, """num_training_steps""": 10}
            # schedulers doct format
            # function: (sched_args_dict, expected_learning_rates)
            lowerCAmelCase    		=			{
                get_constant_schedule: ({}, [10.0] * self.num_steps),
                get_constant_schedule_with_warmup: (
                    {"""num_warmup_steps""": 4},
                    [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
                ),
                get_linear_schedule_with_warmup: (
                    {**common_kwargs},
                    [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
                ),
                get_cosine_schedule_with_warmup: (
                    {**common_kwargs},
                    [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
                ),
                get_cosine_with_hard_restarts_schedule_with_warmup: (
                    {**common_kwargs, """num_cycles""": 2},
                    [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
                ),
                get_polynomial_decay_schedule_with_warmup: (
                    {**common_kwargs, """power""": 2.0, """lr_end""": 1e-7},
                    [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
                ),
                get_inverse_sqrt_schedule: (
                    {"""num_warmup_steps""": 2},
                    [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
                ),
            }

            for scheduler_func, data in scheds.items():
                  lowerCAmelCase		,						lowerCAmelCase    		=			data

                  lowerCAmelCase    		=			scheduler_func(self.optimizer ,		**__lowerCamelCase					)
                  self.assertEqual(len([scheduler.get_lr()[0]]					) ,		1					)
                  lowerCAmelCase    		=			unwrap_schedule(__lowerCamelCase ,		self.num_steps					)
                  self.assertListAlmostEqual(
                      __lowerCamelCase ,		__lowerCamelCase ,		tol=1e-2 ,		msg=f'failed for {scheduler_func} in normal scheduler' ,		)

                  lowerCAmelCase    		=			scheduler_func(self.optimizer ,		**__lowerCamelCase					)
                  if scheduler_func.__name__ != "get_constant_schedule":
                        LambdaScheduleWrapper.wrap_scheduler(__lowerCamelCase					)  # wrap to test picklability of the schedule
                  lowerCAmelCase    		=			unwrap_and_save_reload_schedule(__lowerCamelCase ,		self.num_steps					)
                  self.assertListEqual(__lowerCamelCase ,		__lowerCamelCase ,		msg=f'failed for {scheduler_func} in save and reload'					)







class 						lowercase   :





      def __init__(			self ,		lowercase					)						->	List[str]:
            lowerCAmelCase    		=			fn





      def __call__(			self ,		*lowercase ,		**lowercase					)						->	Optional[int]:
            return self.fn(*__lowerCamelCase ,		**__lowerCamelCase					)





      @classmethod
      def        _snake_case     (			self ,		lowercase					)						->	Dict:
            lowerCAmelCase    		=			list(map(self ,		scheduler.lr_lambdas					)					)













import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen, xsplitext
from ..table import array_cast
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
		from .features import FeatureType


SCREAMING_SNAKE_CASE_	,SCREAMING_SNAKE_CASE_	,SCREAMING_SNAKE_CASE_   :      Any			=   False, False, False

@dataclass
class 		snake_case_   :







	'''simple docstring'''
	__UpperCamelCase   =		None
	__UpperCamelCase   =		True
	__UpperCamelCase   =		True
	__UpperCamelCase   =		None
	# Automatically constructed
	__UpperCamelCase   =		"dict"
	__UpperCamelCase   =		pa.struct({'''bytes''': pa.binary(), '''path''': pa.string()}						)
	__UpperCamelCase   =		field(default='''Audio'''     ,     init=UpperCAmelCase_     ,     repr=UpperCAmelCase_						)




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

			'''simple docstring'''




			return self.pa_type




	def    UpperCAmelCase			(						self							:     List[Any]					,      __lowerCamelCase							:     Union[str, bytes, dict]   )       -> dict:

			'''simple docstring'''




			try:
					import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
			except ImportError as err:
					raise ImportError('To support encoding audio data, please install \'soundfile\'.'   ) from err
			if isinstance(__lowerCamelCase					,      __lowerCamelCase   ):
					return {"bytes": None, "path": value}
			elif isinstance(__lowerCamelCase					,      __lowerCamelCase   ):
					return {"bytes": value, "path": None}
			elif "array" in value:
					# convert the audio array to wav bytes
					__lowercase											=     BytesIO()
					sf.write(__lowerCamelCase					,      value['array']					,      value['sampling_rate']					,      format='wav'   )
					return {"bytes": buffer.getvalue(), "path": None}
			elif value.get('path'   ) is not None and os.path.isfile(value['path']   ):
					# we set "bytes": None to not duplicate the data if they're already available locally
					if value["path"].endswith('pcm'   ):
							# "PCM" only has raw audio bytes
							if value.get('sampling_rate'   ) is None:
									# At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
									raise KeyError('To use PCM files, please specify a \'sampling_rate\' in Audio object'   )
							if value.get('bytes'   ):
									# If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
									__lowercase											=     np.frombuffer(value['bytes']					,      dtype=np.intaa   ).astype(np.floataa   ) / 32_767
							else:
									__lowercase											=     np.memmap(value['path']					,      dtype='h'					,      mode='r'   ).astype(np.floataa   ) / 32_767

							__lowercase											=     BytesIO(bytes()   )
							sf.write(__lowerCamelCase					,      __lowerCamelCase					,      value['sampling_rate']					,      format='wav'   )
							return {"bytes": buffer.getvalue(), "path": None}
					else:
							return {"bytes": None, "path": value.get('path'   )}
			elif value.get('bytes'   ) is not None or value.get('path'   ) is not None:
					# store the audio bytes, and path is used to infer the audio format using the file extension
					return {"bytes": value.get('bytes'   ), "path": value.get('path'   )}
			else:
					raise ValueError(
					    F"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."   )




	def    UpperCAmelCase			(						self							:     Any					,      __lowerCamelCase							:     dict					,      __lowerCamelCase							:     Optional[Dict[str, Union[str, bool, None]]] = None   )       -> dict:

			'''simple docstring'''




			if not self.decode:
					raise RuntimeError('Decoding is disabled for this feature. Please use Audio(decode=True) instead.'   )

			__lowercase						,      __lowercase											=     (value['path'], BytesIO(value['bytes']   )) if value['bytes'] is not None else (value['path'], None)
			if path is None and file is None:
					raise ValueError(F"An audio sample should have one of 'path' or 'bytes' but both are None in {value}."   )

			try:
					import librosa
					import soundfile as sf
			except ImportError as err:
					raise ImportError('To support decoding audio files, please install \'librosa\' and \'soundfile\'.'   ) from err

			__lowercase											=     xsplitext(__lowerCamelCase   )[1][1:].lower() if path is not None else None
			if not config.IS_OPUS_SUPPORTED and audio_format == "opus":
					raise RuntimeError(
					    'Decoding \'opus\' files requires system library \'libsndfile\'>=1.0.31, '
					    'You can try to update `soundfile` python library: `pip install "soundfile>=0.12.1"`. '   )
			elif not config.IS_MP3_SUPPORTED and audio_format == "mp3":
					raise RuntimeError(
					    'Decoding \'mp3\' files requires system library \'libsndfile\'>=1.1.0, '
					    'You can try to update `soundfile` python library: `pip install "soundfile>=0.12.1"`. '   )

			if file is None:
					__lowercase											=     token_per_repo_id or {}
					__lowercase											=     path.split('::'   )[-1]
					try:
							__lowercase											=     string_to_dict(__lowerCamelCase					,      config.HUB_DATASETS_URL   )['repo_id']
							__lowercase											=     token_per_repo_id[repo_id]
					except (ValueError, KeyError):
							__lowercase											=     None

					with xopen(__lowerCamelCase					,      'rb'					,      use_auth_token=__lowerCamelCase   ) as f:
							__lowercase						,      __lowercase											=     sf.read(__lowerCamelCase   )

			else:
					__lowercase						,      __lowercase											=     sf.read(__lowerCamelCase   )

			__lowercase											=     array.T
			if self.mono:
					__lowercase											=     librosa.to_mono(__lowerCamelCase   )
			if self.sampling_rate and self.sampling_rate != sampling_rate:
					__lowercase											=     librosa.resample(__lowerCamelCase					,      orig_sr=__lowerCamelCase					,      target_sr=self.sampling_rate   )
					__lowercase											=     self.sampling_rate

			return {"path": path, "array": array, "sampling_rate": sampling_rate}




	def    UpperCAmelCase			(						self							:     int   )       -> Union["FeatureType", Dict[str, "FeatureType"]]:

			'''simple docstring'''




			from .features import Value

			if self.decode:
					raise ValueError('Cannot flatten a decoded Audio feature.'   )
			return {
			    "bytes": Value('binary'   ),
			    "path": Value('string'   ),
			}




	def    UpperCAmelCase			(						self							:     List[Any]					,      __lowerCamelCase							:     Union[pa.StringArray, pa.StructArray]   )       -> pa.StructArray:

			'''simple docstring'''




			if pa.types.is_string(storage.type   ):
					__lowercase											=     pa.array([None] * len(__lowerCamelCase   )					,      type=pa.binary()   )
					__lowercase											=     pa.StructArray.from_arrays([bytes_array, storage]					,      ['bytes', 'path']					,      mask=storage.is_null()   )
			elif pa.types.is_binary(storage.type   ):
					__lowercase											=     pa.array([None] * len(__lowerCamelCase   )					,      type=pa.string()   )
					__lowercase											=     pa.StructArray.from_arrays([storage, path_array]					,      ['bytes', 'path']					,      mask=storage.is_null()   )
			elif pa.types.is_struct(storage.type   ) and storage.type.get_all_field_indices('array'   ):
					__lowercase											=     pa.array([Audio().encode_example(__lowerCamelCase   ) if x is not None else None for x in storage.to_pylist()]   )
			elif pa.types.is_struct(storage.type   ):
					if storage.type.get_field_index('bytes'   ) >= 0:
							__lowercase											=     storage.field('bytes'   )
					else:
							__lowercase											=     pa.array([None] * len(__lowerCamelCase   )					,      type=pa.binary()   )
					if storage.type.get_field_index('path'   ) >= 0:
							__lowercase											=     storage.field('path'   )
					else:
							__lowercase											=     pa.array([None] * len(__lowerCamelCase   )					,      type=pa.string()   )
					__lowercase											=     pa.StructArray.from_arrays([bytes_array, path_array]					,      ['bytes', 'path']					,      mask=storage.is_null()   )
			return array_cast(__lowerCamelCase					,      self.pa_type   )






	def    UpperCAmelCase			(						self							:     Union[str, Any]					,      __lowerCamelCase							:     pa.StructArray   )       -> pa.StructArray:

			'''simple docstring'''




			@no_op_if_value_is_null
			def path_to_bytes(__lowerCamelCase							:     Any   ):
					with xopen(__lowerCamelCase					,      'rb'   ) as f:
							__lowercase											=     f.read()
					return bytes_

			__lowercase											=     pa.array(
			    [
			        (path_to_bytes(x['path']   ) if x['bytes'] is None else x['bytes']) if x is not None else None
			        for x in storage.to_pylist()
			    ]					,      type=pa.binary()					,      )
			__lowercase											=     pa.array(
			    [os.path.basename(__lowerCamelCase   ) if path is not None else None for path in storage.field('path'   ).to_pylist()]					,      type=pa.string()					,      )
			__lowercase											=     pa.StructArray.from_arrays([bytes_array, path_array]					,      ['bytes', 'path']					,      mask=bytes_array.is_null()   )
			return array_cast(__lowerCamelCase					,      self.pa_type   )




