


'''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       a				(      __a			, __a=10       )  ->     Tuple:



					'''simple docstring'''






					UpperCamelCase__       :Dict							 =			[]
					for _ in range(__a       ):
										lrs.append(scheduler.get_lr()[0]       )
										scheduler.step()
					return lrs


def       a				(      __a			, __a=10       )  ->     List[Any]:



					'''simple docstring'''






					UpperCamelCase__       :Union[str, Any]							 =			[]
					for step in range(__a       ):
										lrs.append(scheduler.get_lr()[0]       )
										scheduler.step()
										if step == num_steps // 2:
															with tempfile.TemporaryDirectory() as tmpdirname:
																				UpperCamelCase__       :List[str]							 =			os.path.join(__a			, '''schedule.bin'''       )
																				torch.save(scheduler.state_dict()			, __a       )

																				UpperCamelCase__       :Any							 =			torch.load(__a       )
																				scheduler.load_state_dict(__a       )
					return lrs
@require_torch
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							self.assertEqual(len(UpperCamelCase_							)   ,       len(UpperCamelCase_							)							)
							for a, b in zip(UpperCamelCase_   ,       UpperCamelCase_							):
												self.assertAlmostEqual(UpperCamelCase_   ,       UpperCamelCase_   ,       delta=UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Dict							 =			torch.tensor([0.1, -0.2, -0.1]   ,       requires_grad=UpperCamelCase_							)
							UpperCamelCase__       :int							 =			torch.tensor([0.4, 0.2, -0.5]							)
							UpperCamelCase__       :Optional[int]							 =			nn.MSELoss()
							# No warmup, constant schedule, no gradient clipping
							UpperCamelCase__       :Union[str, Any]							 =			AdamW(params=[w]   ,       lr=2e-1   ,       weight_decay=0.0							)
							for _ in range(100							):
												UpperCamelCase__       :Optional[Any]							 =			criterion(UpperCamelCase_   ,       UpperCamelCase_							)
												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    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Dict							 =			torch.tensor([0.1, -0.2, -0.1]   ,       requires_grad=UpperCamelCase_							)
							UpperCamelCase__       :Any							 =			torch.tensor([0.4, 0.2, -0.5]							)
							UpperCamelCase__       :Any							 =			nn.MSELoss()
							# No warmup, constant schedule, no gradient clipping
							UpperCamelCase__       :Union[str, Any]							 =			Adafactor(
							    params=[w]   ,       lr=1e-2   ,       eps=(1e-30, 1e-3)   ,       clip_threshold=1.0   ,       decay_rate=-0.8   ,       betaa=UpperCamelCase_   ,       weight_decay=0.0   ,       relative_step=UpperCamelCase_   ,       scale_parameter=UpperCamelCase_   ,       warmup_init=UpperCamelCase_   ,       )
							for _ in range(1000							):
												UpperCamelCase__       :Dict							 =			criterion(UpperCamelCase_   ,       UpperCamelCase_							)
												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 ):

		"""simple docstring"""

		_a   			=    nn.Linear(50     ,		50 ) if is_torch_available() else None
		_a   			=    AdamW(m.parameters()     ,		lr=1_0.0 ) if is_torch_available() else None
		_a   			=    10




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_=None							):

							'''simple docstring'''


							self.assertEqual(len(UpperCamelCase_							)   ,       len(UpperCamelCase_							)							)
							for a, b in zip(UpperCamelCase_   ,       UpperCamelCase_							):
												self.assertAlmostEqual(UpperCamelCase_   ,       UpperCamelCase_   ,       delta=UpperCamelCase_   ,       msg=UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			{'''num_warmup_steps''': 2, '''num_training_steps''': 10}
							# schedulers doct format
							# function: (sched_args_dict, expected_learning_rates)
							UpperCamelCase__       :str							 =			{
							    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():
												UpperCamelCase__ ,    UpperCamelCase__       :Tuple							 =			data

												UpperCamelCase__       :Tuple							 =			scheduler_func(self.optimizer   ,       **UpperCamelCase_							)
												self.assertEqual(len([scheduler.get_lr()[0]]							)   ,       1							)
												UpperCamelCase__       :List[Any]							 =			unwrap_schedule(UpperCamelCase_   ,       self.num_steps							)
												self.assertListAlmostEqual(
												    UpperCamelCase_   ,       UpperCamelCase_   ,       tol=1e-2   ,       msg=F'''failed for {scheduler_func} in normal scheduler'''   ,       )

												UpperCamelCase__       :Dict							 =			scheduler_func(self.optimizer   ,       **UpperCamelCase_							)
												if scheduler_func.__name__ != "get_constant_schedule":
																	LambdaScheduleWrapper.wrap_scheduler(UpperCamelCase_							)  # wrap to test picklability of the schedule
												UpperCamelCase__       :List[str]							 =			unwrap_and_save_reload_schedule(UpperCamelCase_   ,       self.num_steps							)
												self.assertListEqual(UpperCamelCase_   ,       UpperCamelCase_   ,       msg=F'''failed for {scheduler_func} in save and reload'''							)



class        lowercase     :

		"""simple docstring"""

		def __init__(		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			fn




		def __call__(		self   ,       *UpperCamelCase_   ,       **UpperCamelCase_							):

							'''simple docstring'''


							return self.fn(*UpperCamelCase_   ,       **UpperCamelCase_							)




		@classmethod
		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			list(map(self   ,       scheduler.lr_lambdas							)							)



'''simple docstring'''
import argparse
from collections import defaultdict

import yaml


__snake_case											=					'''docs/source/en/_toctree.yml'''


def       a				(      __a       )  ->     int:



					'''simple docstring'''






					UpperCamelCase__       :int							 =			defaultdict(__a       )
					UpperCamelCase__       :int							 =			[]
					UpperCamelCase__       :int							 =			[]
					for doc in doc_list:
										if "local" in doc:
															counts[doc["local"]] += 1

										if doc["title"].lower() == "overview":
															overview_doc.append({'''local''': doc['''local'''], '''title''': doc['''title''']}       )
										else:
															new_doc_list.append(__a       )

					UpperCamelCase__       :Union[str, Any]							 =			new_doc_list
					UpperCamelCase__       :Tuple							 =			[key for key, value in counts.items() if value > 1]

					UpperCamelCase__       :Union[str, Any]							 =			[]
					for duplicate_key in duplicates:
										UpperCamelCase__       :Dict							 =			list({doc['''title'''] for doc in doc_list if doc['''local'''] == duplicate_key}       )
										if len(__a       ) > 1:
															raise ValueError(
															    f'''{duplicate_key} is present several times in the documentation table of content at '''
															    '''`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the '''
															    '''others.'''       )
										# Only add this once
										new_doc.append({'''local''': duplicate_key, '''title''': titles[0]}       )

					# Add none duplicate-keys
					new_doc.extend([doc for doc in doc_list if '''local''' not in counts or counts[doc['''local''']] == 1]       )
					UpperCamelCase__       :Union[str, Any]							 =			sorted(__a			, key=lambda __a       : s["title"].lower()       )

					# "overview" gets special treatment and is always first
					if len(__a       ) > 1:
										raise ValueError('''{doc_list} has two \'overview\' docs which is not allowed.'''       )

					overview_doc.extend(__a       )

					# Sort
					return overview_doc


def       a				(      __a=False       )  ->     Any:



					'''simple docstring'''






					with open(__a			, encoding='''utf-8'''       ) as f:
										UpperCamelCase__       :Any							 =			yaml.safe_load(f.read()       )

					# Get to the API doc
					UpperCamelCase__       :str							 =			0
					while content[api_idx]["title"] != "API":
										api_idx += 1
					UpperCamelCase__       :str							 =			content[api_idx]['''sections''']

					# Then to the model doc
					UpperCamelCase__       :Optional[Any]							 =			0
					while api_doc[scheduler_idx]["title"] != "Schedulers":
										scheduler_idx += 1

					UpperCamelCase__       :List[Any]							 =			api_doc[scheduler_idx]['''sections''']
					UpperCamelCase__       :Union[str, Any]							 =			clean_doc_toc(__a       )

					UpperCamelCase__       :List[Any]							 =			False
					if new_scheduler_doc != scheduler_doc:
										UpperCamelCase__       :Optional[int]							 =			True
										if overwrite:
															UpperCamelCase__       :Dict							 =			new_scheduler_doc

					if diff:
										if overwrite:
															UpperCamelCase__       :Any							 =			api_doc
															with open(__a			, '''w'''			, encoding='''utf-8'''       ) as f:
																				f.write(yaml.dump(__a			, allow_unicode=__a       )       )
										else:
															raise ValueError(
															    '''The model doc part of the table of content is not properly sorted, run `make style` to fix this.'''       )


def       a				(      __a=False       )  ->     Optional[Any]:



					'''simple docstring'''






					with open(__a			, encoding='''utf-8'''       ) as f:
										UpperCamelCase__       :str							 =			yaml.safe_load(f.read()       )

					# Get to the API doc
					UpperCamelCase__       :Optional[Any]							 =			0
					while content[api_idx]["title"] != "API":
										api_idx += 1
					UpperCamelCase__       :Any							 =			content[api_idx]['''sections''']

					# Then to the model doc
					UpperCamelCase__       :str							 =			0
					while api_doc[pipeline_idx]["title"] != "Pipelines":
										pipeline_idx += 1

					UpperCamelCase__       :Any							 =			False
					UpperCamelCase__       :Union[str, Any]							 =			api_doc[pipeline_idx]['''sections''']
					UpperCamelCase__       :Tuple							 =			[]

					# sort sub pipeline docs
					for pipeline_doc in pipeline_docs:
										if "section" in pipeline_doc:
															UpperCamelCase__       :Dict							 =			pipeline_doc['''section''']
															UpperCamelCase__       :Optional[Any]							 =			clean_doc_toc(__a       )
															if overwrite:
																				UpperCamelCase__       :Optional[int]							 =			new_sub_pipeline_doc
										new_pipeline_docs.append(__a       )

					# sort overall pipeline doc
					UpperCamelCase__       :Optional[Any]							 =			clean_doc_toc(__a       )

					if new_pipeline_docs != pipeline_docs:
										UpperCamelCase__       :int							 =			True
										if overwrite:
															UpperCamelCase__       :Union[str, Any]							 =			new_pipeline_docs

					if diff:
										if overwrite:
															UpperCamelCase__       :Dict							 =			api_doc
															with open(__a			, '''w'''			, encoding='''utf-8'''       ) as f:
																				f.write(yaml.dump(__a			, allow_unicode=__a       )       )
										else:
															raise ValueError(
															    '''The model doc part of the table of content is not properly sorted, run `make style` to fix this.'''       )


if __name__ == "__main__":
		__snake_case											=					argparse.ArgumentParser()
		parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
		__snake_case											=					parser.parse_args()

		check_scheduler_doc(args.fix_and_overwrite)
		check_pipeline_doc(args.fix_and_overwrite)