




'''simple docstring'''







import argparse
import json
import os
import sys
import tempfile
import unittest
from argparse import Namespace
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import List, Literal, Optional

import yaml

from transformers import HfArgumentParser, TrainingArguments
from transformers.hf_argparser import make_choice_type_function, string_to_bool


# Since Python 3.10, we can use the builtin `|` operator for Union types
# See PEP 604: https://peps.python.org/pep-0604
lowercase      :   Tuple      =       sys.version_info >= (3, 10)



def        lowerCamelCase__				(      __lowercase=None     ,   __lowercase=None				):
     return field(default_factory=lambda: default     ,   metadata=__lowercase				)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     int
     lowerCAmelCase_  :     float
     lowerCAmelCase_  :     str
     lowerCAmelCase_  :     bool







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     int          =					42
     lowerCAmelCase_  :     str          =					field(default="""toto""",		metadata={"""help""": """help message"""}	)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     bool          =					False
     lowerCAmelCase_  :     bool          =					True
     lowerCAmelCase_  :     Optional[bool]          =					None







class 					_a (a__	):





     '''simple docstring'''



     lowerCAmelCase_  :     int          =					"""titi"""
     lowerCAmelCase_  :     Optional[int]          =					"""toto"""







class 					_a (a__	):





     '''simple docstring'''



     lowerCAmelCase_  :     Any          =					"""titi"""
     lowerCAmelCase_  :     Dict          =					"""toto"""
     lowerCAmelCase_  :     int          =					42







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     BasicEnum          =					"toto"







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      Union[str, Any]        =			BasicEnum(self.foo							)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     MixedTypeEnum          =					"toto"







     def 			snake_case_      (				self							)      ->				Tuple:
          snake_case			:      List[str]        =			MixedTypeEnum(self.foo							)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     Optional[int]          =					None
     lowerCAmelCase_  :     Optional[float]          =					field(default=a__,		metadata={"""help""": """help message"""}	)
     lowerCAmelCase_  :     Optional[str]          =					None
     lowerCAmelCase_  :     Optional[List[str]]          =					list_field(default=[]	)
     lowerCAmelCase_  :     Optional[List[int]]          =					list_field(default=[]	)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     List[int]          =					list_field(default=[]	)
     lowerCAmelCase_  :     List[int]          =					list_field(default=[1, 2, 3]	)
     lowerCAmelCase_  :     List[str]          =					list_field(default=["""Hallo""", """Bonjour""", """Hello"""]	)
     lowerCAmelCase_  :     List[float]          =					list_field(default=[0.1, 0.2, 0.3]	)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     List[int]          =					field()
     lowerCAmelCase_  :     str          =					field()
     lowerCAmelCase_  :     BasicEnum          =					field()







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      Optional[int]        =			BasicEnum(self.required_enum							)







@dataclass
class 					_a :





     '''simple docstring'''



     lowerCAmelCase_  :     int
     lowerCAmelCase_  :     "BasicEnum"          =					field()
     lowerCAmelCase_  :     "Optional[bool]"          =					None
     lowerCAmelCase_  :     "str"          =					field(default="""toto""",		metadata={"""help""": """help message"""}	)
     lowerCAmelCase_  :     "List[str]"          =					list_field(default=["""Hallo""", """Bonjour""", """Hello"""]	)


if is_python_no_less_than_3_10:







      @dataclass
      class 					_a :





           '''simple docstring'''



           lowerCAmelCase_  :     bool          =					False
           lowerCAmelCase_  :     bool          =					True
           lowerCAmelCase_  :     bool | None          =					None







      @dataclass
      class 					_a :





           '''simple docstring'''



           lowerCAmelCase_  :     int | None          =					None
           lowerCAmelCase_  :     float | None          =					field(default=a__,		metadata={"""help""": """help message"""}	)
           lowerCAmelCase_  :     str | None          =					None
           lowerCAmelCase_  :     list[str] | None          =					list_field(default=[]	)
           lowerCAmelCase_  :     list[int] | None          =					list_field(default=[]	)







class 					_a (unittest.TestCase	):





     '''simple docstring'''



     def 			snake_case_      (				self							,__a							,__a							)      ->				Optional[int]:
          self.assertEqual(len(a._actions							)							,len(b._actions							)							)
          for x, y in zip(a._actions							,b._actions							):
               snake_case			:      Optional[int]        =			{k: v for k, v in vars(__a							).items() if k != """container"""}
               snake_case			:      Tuple        =			{k: v for k, v in vars(__a							).items() if k != """container"""}

               # Choices with mixed type have custom function as "type"
               # So we need to compare results directly for equality
               if xx.get("""choices"""							,__a							) and yy.get("""choices"""							,__a							):
                    for expected_choice in yy["choices"] + xx["choices"]:
                         self.assertEqual(xx["""type"""](__a							)							,yy["""type"""](__a							)							)
                    del xx["type"], yy["type"]

               self.assertEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      str        =			HfArgumentParser(__a							)

          snake_case			:      Dict        =			argparse.ArgumentParser()
          expected.add_argument("""--foo"""							,type=__a							,required=__a							)
          expected.add_argument("""--bar"""							,type=__a							,required=__a							)
          expected.add_argument("""--baz"""							,type=__a							,required=__a							)
          expected.add_argument("""--flag"""							,type=__a							,default=__a							,const=__a							,nargs="""?"""							)
          self.argparsersEqual(__a							,__a							)

          snake_case			:      Optional[int]        =			["""--foo""", """1""", """--baz""", """quux""", """--bar""", """0.5"""]
          ((snake_case) ,		)			:      Dict        =			parser.parse_args_into_dataclasses(__a							,look_for_args_file=__a							)
          self.assertFalse(example.flag							)







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      Optional[int]        =			HfArgumentParser(__a							)

          snake_case			:      str        =			argparse.ArgumentParser()
          expected.add_argument("""--foo"""							,default=42							,type=__a							)
          expected.add_argument("""--baz"""							,default="""toto"""							,type=__a							,help="""help message"""							)
          self.argparsersEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				Optional[Any]:
          snake_case			:      Optional[int]        =			argparse.ArgumentParser()
          expected.add_argument("""--foo"""							,type=__a							,default=__a							,const=__a							,nargs="""?"""							)
          expected.add_argument("""--baz"""							,type=__a							,default=__a							,const=__a							,nargs="""?"""							)
          # A boolean no_* argument always has to come after its "default: True" regular counter-part
          # and its default must be set to False
          expected.add_argument("""--no_baz"""							,action="""store_false"""							,default=__a							,dest="""baz"""							)
          expected.add_argument("""--opt"""							,type=__a							,default=__a							)

          snake_case			:      str        =			[WithDefaultBoolExample]
          if is_python_no_less_than_3_10:
               dataclass_types.append(__a							)

          for dataclass_type in dataclass_types:
               snake_case			:      Optional[int]        =			HfArgumentParser(__a							)
               self.argparsersEqual(__a							,__a							)

               snake_case			:      List[str]        =			parser.parse_args([]							)
               self.assertEqual(__a							,Namespace(foo=__a							,baz=__a							,opt=__a							)							)

               snake_case			:      str        =			parser.parse_args(["""--foo""", """--no_baz"""]							)
               self.assertEqual(__a							,Namespace(foo=__a							,baz=__a							,opt=__a							)							)

               snake_case			:      Union[str, Any]        =			parser.parse_args(["""--foo""", """--baz"""]							)
               self.assertEqual(__a							,Namespace(foo=__a							,baz=__a							,opt=__a							)							)

               snake_case			:      Union[str, Any]        =			parser.parse_args(["""--foo""", """True""", """--baz""", """True""", """--opt""", """True"""]							)
               self.assertEqual(__a							,Namespace(foo=__a							,baz=__a							,opt=__a							)							)

               snake_case			:      List[str]        =			parser.parse_args(["""--foo""", """False""", """--baz""", """False""", """--opt""", """False"""]							)
               self.assertEqual(__a							,Namespace(foo=__a							,baz=__a							,opt=__a							)							)







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      Any        =			HfArgumentParser(__a							)

          snake_case			:      Union[str, Any]        =			argparse.ArgumentParser()
          expected.add_argument(
              """--foo"""							,default="""toto"""							,choices=["""titi""", """toto""", 42]							,type=make_choice_type_function(["""titi""", """toto""", 42]							)							,)
          self.argparsersEqual(__a							,__a							)

          snake_case			:      Any        =			parser.parse_args([]							)
          self.assertEqual(args.foo							,"""toto"""							)
          snake_case			:      str        =			parser.parse_args_into_dataclasses([]							)[0]
          self.assertEqual(enum_ex.foo							,MixedTypeEnum.toto							)

          snake_case			:      Optional[Any]        =			parser.parse_args(["""--foo""", """titi"""]							)
          self.assertEqual(args.foo							,"""titi"""							)
          snake_case			:      Optional[Any]        =			parser.parse_args_into_dataclasses(["""--foo""", """titi"""]							)[0]
          self.assertEqual(enum_ex.foo							,MixedTypeEnum.titi							)

          snake_case			:      Any        =			parser.parse_args(["""--foo""", """42"""]							)
          self.assertEqual(args.foo							,42							)
          snake_case			:      Any        =			parser.parse_args_into_dataclasses(["""--foo""", """42"""]							)[0]
          self.assertEqual(enum_ex.foo							,MixedTypeEnum.fourtytwo							)







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







          @dataclass
          class 					_a :





               '''simple docstring'''



               lowerCAmelCase_  :     Literal["titi", "toto", 42]          =					"toto"
          snake_case			:      Tuple        =			HfArgumentParser(__a							)

          snake_case			:      List[Any]        =			argparse.ArgumentParser()
          expected.add_argument(
              """--foo"""							,default="""toto"""							,choices=("""titi""", """toto""", 42)							,type=make_choice_type_function(["""titi""", """toto""", 42]							)							,)
          self.argparsersEqual(__a							,__a							)

          snake_case			:      Optional[int]        =			parser.parse_args([]							)
          self.assertEqual(args.foo							,"""toto"""							)

          snake_case			:      int        =			parser.parse_args(["""--foo""", """titi"""]							)
          self.assertEqual(args.foo							,"""titi"""							)

          snake_case			:      Optional[int]        =			parser.parse_args(["""--foo""", """42"""]							)
          self.assertEqual(args.foo							,42							)







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      Any        =			HfArgumentParser(__a							)

          snake_case			:      Dict        =			argparse.ArgumentParser()
          expected.add_argument("""--foo_int"""							,nargs="""+"""							,default=[]							,type=__a							)
          expected.add_argument("""--bar_int"""							,nargs="""+"""							,default=[1, 2, 3]							,type=__a							)
          expected.add_argument("""--foo_str"""							,nargs="""+"""							,default=["""Hallo""", """Bonjour""", """Hello"""]							,type=__a							)
          expected.add_argument("""--foo_float"""							,nargs="""+"""							,default=[0.1, 0.2, 0.3]							,type=__a							)

          self.argparsersEqual(__a							,__a							)

          snake_case			:      int        =			parser.parse_args([]							)
          self.assertEqual(
              __a							,Namespace(foo_int=[]							,bar_int=[1, 2, 3]							,foo_str=["""Hallo""", """Bonjour""", """Hello"""]							,foo_float=[0.1, 0.2, 0.3]							)							,)

          snake_case			:      Optional[int]        =			parser.parse_args("""--foo_int 1 --bar_int 2 3 --foo_str a b c --foo_float 0.1 0.7""".split()							)
          self.assertEqual(__a							,Namespace(foo_int=[1]							,bar_int=[2, 3]							,foo_str=["""a""", """b""", """c"""]							,foo_float=[0.1, 0.7]							)							)







     def 			snake_case_      (				self							)      ->				Optional[int]:
          snake_case			:      List[Any]        =			argparse.ArgumentParser()
          expected.add_argument("""--foo"""							,default=__a							,type=__a							)
          expected.add_argument("""--bar"""							,default=__a							,type=__a							,help="""help message"""							)
          expected.add_argument("""--baz"""							,default=__a							,type=__a							)
          expected.add_argument("""--ces"""							,nargs="""+"""							,default=[]							,type=__a							)
          expected.add_argument("""--des"""							,nargs="""+"""							,default=[]							,type=__a							)

          snake_case			:      Optional[Any]        =			[OptionalExample]
          if is_python_no_less_than_3_10:
               dataclass_types.append(__a							)

          for dataclass_type in dataclass_types:
               snake_case			:      Any        =			HfArgumentParser(__a							)

               self.argparsersEqual(__a							,__a							)

               snake_case			:      Any        =			parser.parse_args([]							)
               self.assertEqual(__a							,Namespace(foo=__a							,bar=__a							,baz=__a							,ces=[]							,des=[]							)							)

               snake_case			:      Any        =			parser.parse_args("""--foo 12 --bar 3.14 --baz 42 --ces a b c --des 1 2 3""".split()							)
               self.assertEqual(__a							,Namespace(foo=12							,bar=3.14							,baz="""42"""							,ces=["""a""", """b""", """c"""]							,des=[1, 2, 3]							)							)







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      Optional[Any]        =			HfArgumentParser(__a							)

          snake_case			:      List[Any]        =			argparse.ArgumentParser()
          expected.add_argument("""--required_list"""							,nargs="""+"""							,type=__a							,required=__a							)
          expected.add_argument("""--required_str"""							,type=__a							,required=__a							)
          expected.add_argument(
              """--required_enum"""							,type=make_choice_type_function(["""titi""", """toto"""]							)							,choices=["""titi""", """toto"""]							,required=__a							,)
          self.argparsersEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      Optional[int]        =			HfArgumentParser(__a							)

          snake_case			:      Any        =			argparse.ArgumentParser()
          expected.add_argument("""--foo"""							,type=__a							,required=__a							)
          expected.add_argument(
              """--required_enum"""							,type=make_choice_type_function(["""titi""", """toto"""]							)							,choices=["""titi""", """toto"""]							,required=__a							,)
          expected.add_argument("""--opt"""							,type=__a							,default=__a							)
          expected.add_argument("""--baz"""							,default="""toto"""							,type=__a							,help="""help message"""							)
          expected.add_argument("""--foo_str"""							,nargs="""+"""							,default=["""Hallo""", """Bonjour""", """Hello"""]							,type=__a							)
          self.argparsersEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				Tuple:
          snake_case			:      str        =			HfArgumentParser(__a							)

          snake_case			:      Any        =			{
              """foo""": 12,
              """bar""": 3.14,
              """baz""": """42""",
              """flag""": True,
          }

          snake_case			:      List[str]        =			parser.parse_dict(__a							)[0]
          snake_case			:      str        =			BasicExample(**__a							)
          self.assertEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      Tuple        =			HfArgumentParser(__a							)

          snake_case			:      int        =			{
              """foo""": 12,
              """bar""": 3.14,
              """baz""": """42""",
              """flag""": True,
              """extra""": 42,
          }

          self.assertRaises(__a							,parser.parse_dict							,__a							,allow_extra_keys=__a							)







     def 			snake_case_      (				self							)      ->				Optional[int]:
          snake_case			:      List[str]        =			HfArgumentParser(__a							)

          snake_case			:      Any        =			{
              """foo""": 12,
              """bar""": 3.14,
              """baz""": """42""",
              """flag""": True,
          }
          with tempfile.TemporaryDirectory() as tmp_dir:
               snake_case			:      Dict        =			os.path.join(__a							,"""temp_json"""							)
               os.mkdir(__a							)
               with open(temp_local_path + """.json"""							,"""w+"""							) as f:
                    json.dump(__a							,__a							)
               snake_case			:      Any        =			parser.parse_yaml_file(Path(temp_local_path + """.json"""							)							)[0]

          snake_case			:      List[str]        =			BasicExample(**__a							)
          self.assertEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      Union[str, Any]        =			HfArgumentParser(__a							)

          snake_case			:      Any        =			{
              """foo""": 12,
              """bar""": 3.14,
              """baz""": """42""",
              """flag""": True,
          }
          with tempfile.TemporaryDirectory() as tmp_dir:
               snake_case			:      str        =			os.path.join(__a							,"""temp_yaml"""							)
               os.mkdir(__a							)
               with open(temp_local_path + """.yaml"""							,"""w+"""							) as f:
                    yaml.dump(__a							,__a							)
               snake_case			:      Union[str, Any]        =			parser.parse_yaml_file(Path(temp_local_path + """.yaml"""							)							)[0]
          snake_case			:      Tuple        =			BasicExample(**__a							)
          self.assertEqual(__a							,__a							)







     def 			snake_case_      (				self							)      ->				List[str]:
          snake_case			:      Union[str, Any]        =			HfArgumentParser(__a							)
          self.assertIsNotNone(__a							)









'''simple docstring'''







# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("""ignore""", category=UserWarning, module="""torch.optim.lr_scheduler""")







class 					_a :





     '''simple docstring'''



     def __init__(				self							,__a							,__a							,__a = True							,__a = False							)      ->				Dict:
          snake_case			:      List[Any]        =			scheduler
          snake_case			:      List[Any]        =			optimizers if isinstance(__a							,(list, tuple)							) else [optimizers]
          snake_case			:      List[str]        =			split_batches
          snake_case			:      List[str]        =			step_with_optimizer
          snake_case			:      int        =			GradientState()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          if not self.step_with_optimizer:
               # No link between scheduler and optimizer -> just step
               self.scheduler.step(*__a							,**__a							)
               return

          # Otherwise, first make sure the optimizer was stepped.
          if not self.gradient_state.sync_gradients:
               if self.gradient_state.adjust_scheduler:
                    self.scheduler._step_count += 1
               return

          for opt in self.optimizers:
               if opt.step_was_skipped:
                    return
          if self.split_batches:
               # Split batches -> the training dataloader batch size is not changed so one step per training step
               self.scheduler.step(*__a							,**__a							)
          else:
               # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
               # num_processes steps per training step
               snake_case			:      List[str]        =			AcceleratorState().num_processes
               for _ in range(__a							):
                    # Special case when using OneCycle and `drop_last` was not used
                    if hasattr(self.scheduler							,"""total_steps"""							):
                         if self.scheduler._step_count <= self.scheduler.total_steps:
                              self.scheduler.step(*__a							,**__a							)
                    else:
                         self.scheduler.step(*__a							,**__a							)







     def 			snake_case_      (				self							)      ->				Dict:
          return self.scheduler.get_last_lr()







     def 			snake_case_      (				self							)      ->				int:
          return self.scheduler.state_dict()







     def 			snake_case_      (				self							,__a							)      ->				Any:
          self.scheduler.load_state_dict(__a							)







     def 			snake_case_      (				self							)      ->				Any:
          return self.scheduler.get_lr()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          return self.scheduler.print_lr(*__a							,**__a							)



