





import argparse
import json
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


__A							:			Optional[int]         =			16
__A							:			Optional[Any]         =			32





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE = 16							,     _SCREAMING_SNAKE_CASE = "bert-base-cased"				)							-> int:




    """simple docstring"""




    _A     =							AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE				)
    _A     =							load_dataset('glue'							,     'mrpc'				)

    def tokenize_function(_SCREAMING_SNAKE_CASE				):
        # max_length=None => use the model max length (it's actually the default)
        _A     =							tokenizer(examples['sentence1']							,     examples['sentence2']							,     truncation=_SCREAMING_SNAKE_CASE							,     max_length=_SCREAMING_SNAKE_CASE				)
        return outputs

    # Apply the method we just defined to all the examples in all the splits of the dataset
    _A     =							datasets.map(
        _SCREAMING_SNAKE_CASE							,     batched=_SCREAMING_SNAKE_CASE							,     remove_columns=['idx', 'sentence1', 'sentence2']							,     load_from_cache_file=_SCREAMING_SNAKE_CASE				)

    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
    # transformers library
    _A     =							tokenized_datasets.rename_column('label'							,     'labels'				)

    def collate_fn(_SCREAMING_SNAKE_CASE				):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        if accelerator.distributed_type == DistributedType.TPU:
            return tokenizer.pad(_SCREAMING_SNAKE_CASE							,     padding='max_length'							,     max_length=128							,     return_tensors='pt'				)
        return tokenizer.pad(_SCREAMING_SNAKE_CASE							,     padding='longest'							,     return_tensors='pt'				)

    # Instantiate dataloaders.
    _A     =							DataLoader(
        tokenized_datasets['train']							,     shuffle=_SCREAMING_SNAKE_CASE							,     collate_fn=_SCREAMING_SNAKE_CASE							,     batch_size=_SCREAMING_SNAKE_CASE				)
    _A     =							DataLoader(
        tokenized_datasets['validation']							,     shuffle=_SCREAMING_SNAKE_CASE							,     collate_fn=_SCREAMING_SNAKE_CASE							,     batch_size=_SCREAMING_SNAKE_CASE				)

    return train_dataloader, eval_dataloader





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Union[str, Any]:




    """simple docstring"""




    model.eval()
    _A     =							0
    for step, batch in enumerate(_SCREAMING_SNAKE_CASE				):
        # We could avoid this line since we set the accelerator with `device_placement=True`.
        batch.to(accelerator.device				)
        with torch.no_grad():
            _A     =							model(**_SCREAMING_SNAKE_CASE				)
        _A     =							outputs.logits.argmax(dim=-1				)
        # It is slightly faster to call this once, than multiple times
        _A,			_A     =							accelerator.gather(
            (predictions, batch['labels'])				)  # If we are in a multiprocess environment, the last batch has duplicates
        if accelerator.use_distributed:
            if step == len(_SCREAMING_SNAKE_CASE				) - 1:
                _A     =							predictions[: len(eval_dataloader.dataset				) - samples_seen]
                _A     =							references[: len(eval_dataloader.dataset				) - samples_seen]
            else:
                samples_seen += references.shape[0]
        metric.add_batch(
            predictions=_SCREAMING_SNAKE_CASE							,     references=_SCREAMING_SNAKE_CASE							,     )

    _A     =							metric.compute()
    return eval_metric["accuracy"]





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Any:




    """simple docstring"""




    _A     =							Accelerator()

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    _A     =							config['lr']
    _A     =							int(config['num_epochs']				)
    _A     =							int(config['seed']				)
    _A     =							int(config['batch_size']				)
    _A     =							args.model_name_or_path

    set_seed(_SCREAMING_SNAKE_CASE				)
    _A,			_A     =							get_dataloaders(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    _A     =							AutoModelForSequenceClassification.from_pretrained(_SCREAMING_SNAKE_CASE							,     return_dict=_SCREAMING_SNAKE_CASE				)

    # Instantiate optimizer
    _A     =							(
        AdamW
        if accelerator.state.deepspeed_plugin is None
        or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    _A     =							optimizer_cls(params=model.parameters()							,     lr=_SCREAMING_SNAKE_CASE				)

    if accelerator.state.deepspeed_plugin is not None:
        _A     =							accelerator.state.deepspeed_plugin.deepspeed_config[
            'gradient_accumulation_steps'
        ]
    else:
        _A     =							1
    _A     =							(len(_SCREAMING_SNAKE_CASE				) * num_epochs) // gradient_accumulation_steps

    # Instantiate scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        _A     =							get_linear_schedule_with_warmup(
            optimizer=_SCREAMING_SNAKE_CASE							,     num_warmup_steps=0							,     num_training_steps=_SCREAMING_SNAKE_CASE							,     )
    else:
        _A     =							DummyScheduler(_SCREAMING_SNAKE_CASE							,     total_num_steps=_SCREAMING_SNAKE_CASE							,     warmup_num_steps=0				)

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    _A,			_A,			_A,			_A,			_A     =							accelerator.prepare(
        _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)

    # We need to keep track of how many total steps we have iterated over
    _A     =							0
    # We also need to keep track of the stating epoch so files are named properly
    _A     =							0
    _A     =							evaluate.load('glue'							,     'mrpc'				)
    _A     =							num_epochs

    if args.partial_train_epoch is not None:
        _A     =							args.partial_train_epoch

    if args.resume_from_checkpoint:
        accelerator.load_state(args.resume_from_checkpoint				)
        _A     =							args.resume_from_checkpoint.split('epoch_'				)[1]
        _A     =							''
        for char in epoch_string:
            if char.isdigit():
                state_epoch_num += char
            else:
                break
        _A     =							int(_SCREAMING_SNAKE_CASE				) + 1
        _A     =							evaluation_loop(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
        accelerator.print('resumed checkpoint performance:'							,     _SCREAMING_SNAKE_CASE				)
        accelerator.print('resumed checkpoint\'s scheduler\'s lr:'							,     lr_scheduler.get_lr()[0]				)
        accelerator.print('resumed optimizers\'s lr:'							,     optimizer.param_groups[0]['lr']				)
        with open(os.path.join(args.output_dir							,     F"state_{starting_epoch-1}.json"				)							,     'r'				) as f:
            _A     =							json.load(_SCREAMING_SNAKE_CASE				)
            assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
            assert (
                resumed_state["lr"] == lr_scheduler.get_lr()[0]
            ), "Scheduler learning rate mismatch, loading from checkpoint failed"
            assert (
                resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
            ), "Optimizer learning rate mismatch, loading from checkpoint failed"
            assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
            return

    # Now we train the model
    _A     =							{}
    for epoch in range(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        model.train()
        for step, batch in enumerate(_SCREAMING_SNAKE_CASE				):
            _A     =							model(**_SCREAMING_SNAKE_CASE				)
            _A     =							outputs.loss
            _A     =							loss / gradient_accumulation_steps
            accelerator.backward(_SCREAMING_SNAKE_CASE				)
            if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            overall_step += 1
        _A     =							F"epoch_{epoch}"
        _A     =							os.path.join(args.output_dir							,     _SCREAMING_SNAKE_CASE				)
        accelerator.save_state(_SCREAMING_SNAKE_CASE				)
        _A     =							evaluation_loop(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
        _A     =							accuracy
        _A     =							lr_scheduler.get_lr()[0]
        _A     =							optimizer.param_groups[0]['lr']
        _A     =							epoch
        _A     =							overall_step
        accelerator.print(F"epoch {epoch}:"							,     _SCREAMING_SNAKE_CASE				)

        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            with open(os.path.join(args.output_dir							,     F"state_{epoch}.json"				)							,     'w'				) as f:
                json.dump(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)





def 				__lowerCAmelCase(       )							-> Optional[Any]:




    """simple docstring"""




    _A     =							argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'				)
    parser.add_argument(
        '--model_name_or_path'							,     type=_SCREAMING_SNAKE_CASE							,     default='bert-base-cased'							,     help='Path to pretrained model or model identifier from huggingface.co/models.'							,     required=_SCREAMING_SNAKE_CASE							,     )
    parser.add_argument(
        '--output_dir'							,     type=_SCREAMING_SNAKE_CASE							,     default='.'							,     help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'							,     )
    parser.add_argument(
        '--resume_from_checkpoint'							,     type=_SCREAMING_SNAKE_CASE							,     default=_SCREAMING_SNAKE_CASE							,     help='If the training should continue from a checkpoint folder.'							,     )
    parser.add_argument(
        '--partial_train_epoch'							,     type=_SCREAMING_SNAKE_CASE							,     default=_SCREAMING_SNAKE_CASE							,     help='If passed, the training will stop after this number of epochs.'							,     )
    parser.add_argument(
        '--num_epochs'							,     type=_SCREAMING_SNAKE_CASE							,     default=2							,     help='Number of train epochs.'							,     )
    _A     =							parser.parse_args()
    _A     =							{'lr': 2e-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}

    training_function(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
 main()







import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import UperNetForSemanticSegmentation
		from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor



class    a_ :






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase=13				,			__UpperCamelCase=32				,			__UpperCamelCase=3				,			__UpperCamelCase=4				,			__UpperCamelCase=[10, 20, 30, 40]				,			__UpperCamelCase=[2, 2, 3, 2]				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=37				,			__UpperCamelCase="gelu"				,			__UpperCamelCase=10				,			__UpperCamelCase=0.02				,			__UpperCamelCase=["stage2", "stage3", "stage4"]				,			__UpperCamelCase=3				,			__UpperCamelCase=None				,			):
									_lowercase					       =					parent
									_lowercase					       =					batch_size
									_lowercase					       =					image_size
									_lowercase					       =					num_channels
									_lowercase					       =					num_stages
									_lowercase					       =					hidden_sizes
									_lowercase					       =					depths
									_lowercase					       =					is_training
									_lowercase					       =					use_labels
									_lowercase					       =					intermediate_size
									_lowercase					       =					hidden_act
									_lowercase					       =					type_sequence_label_size
									_lowercase					       =					initializer_range
									_lowercase					       =					out_features
									_lowercase					       =					num_labels
									_lowercase					       =					scope
									_lowercase					       =					num_stages






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

									_lowercase					       =					None
									if self.use_labels:
															_lowercase					       =					ids_tensor([self.batch_size]				,			self.type_sequence_label_size  )

									_lowercase					       =					self.get_config()

									return config, pixel_values, labels






			def    UpperCamelCase_				(      self  ):
									return ConvNextConfig(
									    num_channels=self.num_channels				,			num_stages=self.num_stages				,			hidden_sizes=self.hidden_sizes				,			depths=self.depths				,			is_training=self.is_training				,			intermediate_size=self.intermediate_size				,			hidden_act=self.hidden_act				,			out_features=self.out_features				,			)






			def    UpperCamelCase_				(      self  ):
									return UperNetConfig(
									    backbone_config=self.get_backbone_config()				,			hidden_size=512				,			pool_scales=[1, 2, 3, 6]				,			use_auxiliary_head=__UpperCamelCase				,			auxiliary_loss_weight=0.4				,			auxiliary_in_channels=40				,			auxiliary_channels=256				,			auxiliary_num_convs=1				,			auxiliary_concat_input=__UpperCamelCase				,			loss_ignore_index=255				,			num_labels=self.num_labels				,			)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									_lowercase					       =					UperNetForSemanticSegmentation(config=__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase  )
									self.parent.assertEqual(
									    result.logits.shape				,			(self.batch_size, self.num_labels, self.image_size, self.image_size)  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.prepare_config_and_inputs()
									(
									    (
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						
									)					       =					config_and_inputs
									_lowercase					       =					{"""pixel_values""": pixel_values}
									return config, inputs_dict







@require_torch
class    a_ (			_a					,			_a					,			unittest.TestCase    ):
			a	:			Optional[int]					   =							(UperNetForSemanticSegmentation,) if is_torch_available() else ()
			a	:			Optional[Any]					   =							{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
			a	:			Any					   =							False
			a	:			Union[str, Any]					   =							False
			a	:			Optional[Any]					   =							False
			a	:			int					   =							False
			a	:			Optional[int]					   =							False
			a	:			List[str]					   =							False






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					UperNetModelTester(self  )
									_lowercase					       =					ConfigTester(self				,			config_class=__UpperCamelCase				,			has_text_modality=__UpperCamelCase				,			hidden_size=37  )






			def    UpperCamelCase_				(      self  ):
									self.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.create_and_test_config_with_num_labels()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()






			def    UpperCamelCase_				(      self  ):
									return






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )
															_lowercase					       =					inspect.signature(model.forward  )
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															_lowercase					       =					[*signature.parameters.keys()]

															_lowercase					       =					["""pixel_values"""]
															self.assertListEqual(arg_names[:1]				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCamelCase  )






			@unittest.skip(reason="""UperNet does not use inputs_embeds"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not support input and output embeddings"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not have a base model"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not have a base model"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@require_torch_multi_gpu
			@unittest.skip(reason="""UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
			def    UpperCamelCase_				(      self  ):
									pass






			def    UpperCamelCase_				(      self  ):
									def check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
															_lowercase					       =					model_class(__UpperCamelCase  )
															model.to(__UpperCamelCase  )
															model.eval()

															with torch.no_grad():
																					_lowercase					       =					model(**self._prepare_for_class(__UpperCamelCase				,			__UpperCamelCase  )  )

															_lowercase					       =					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

															_lowercase					       =					self.model_tester.num_stages
															self.assertEqual(len(__UpperCamelCase  )				,			expected_num_stages + 1  )

															# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
															self.assertListEqual(
															    list(hidden_states[0].shape[-2:]  )				,			[self.model_tester.image_size // 4, self.model_tester.image_size // 4]				,			)

									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_lowercase					       =					True
															check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )

															# check that output_hidden_states also work using config
															del inputs_dict["output_hidden_states"]
															_lowercase					       =					True

															check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									_lowercase					       =					_config_zero_init(__UpperCamelCase  )
									_lowercase					       =					_config_zero_init(configs_no_init.backbone_config  )
									for model_class in self.all_model_classes:
															_lowercase					       =					model_class(config=__UpperCamelCase  )
															for name, param in model.named_parameters():
																					if param.requires_grad:
																											self.assertIn(
																											    ((param.data.mean() * 1E9).round() / 1E9).item()				,			[0.0, 1.0]				,			msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""				,			)






			@unittest.skip(reason="""UperNet does not have tied weights"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@slow
			def    UpperCamelCase_				(      self  ):
									for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained(__UpperCamelCase  )
															self.assertIsNotNone(__UpperCamelCase  )







def 	UpperCamelCase__  (     )  ->     str:
						_lowercase					       =					hf_hub_download(
						    repo_id="""hf-internal-testing/fixtures_ade20k"""		,   repo_type="""dataset"""		,   filename="""ADE_val_00000001.jpg"""       )
						_lowercase					       =					Image.open(SCREAMING_SNAKE_CASE_       ).convert("""RGB"""       )
						return image



@require_torch
@require_vision
@slow
class    a_ (			unittest.TestCase    ):






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""  )
									_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""  ).to(__UpperCamelCase  )

									_lowercase					       =					prepare_img()
									_lowercase					       =					processor(images=__UpperCamelCase				,			return_tensors="""pt"""  ).to(__UpperCamelCase  )

									with torch.no_grad():
															_lowercase					       =					model(**__UpperCamelCase  )

									_lowercase					       =					torch.Size((1, model.config.num_labels, 512, 512)  )
									self.assertEqual(outputs.logits.shape				,			__UpperCamelCase  )

									_lowercase					       =					torch.tensor(
									    [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]  ).to(__UpperCamelCase  )
									self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,			__UpperCamelCase				,			atol=1E-4  )  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""  )
									_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""  ).to(__UpperCamelCase  )

									_lowercase					       =					prepare_img()
									_lowercase					       =					processor(images=__UpperCamelCase				,			return_tensors="""pt"""  ).to(__UpperCamelCase  )

									with torch.no_grad():
															_lowercase					       =					model(**__UpperCamelCase  )

									_lowercase					       =					torch.Size((1, model.config.num_labels, 512, 512)  )
									self.assertEqual(outputs.logits.shape				,			__UpperCamelCase  )

									_lowercase					       =					torch.tensor(
									    [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]  ).to(__UpperCamelCase  )
									self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,			__UpperCamelCase				,			atol=1E-4  )  )