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


UpperCamelCase__    							= 1_6
UpperCamelCase__    							= 3_2




def        lowerCAmelCase_			(					__A,  __A = 16,  __A = "bert-base-cased"							)					->    Tuple:







						'''simple docstring'''


						UpperCAmelCase__								=  AutoTokenizer.from_pretrained(__A							)
						UpperCAmelCase__								=  load_dataset("glue",  "mrpc"							)

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

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

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

						def collate_fn(__A							):
												# 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(__A,  padding="max_length",  max_length=128,  return_tensors="pt"							)
												return tokenizer.pad(__A,  padding="longest",  return_tensors="pt"							)

						# Instantiate dataloaders.
						UpperCAmelCase__								=  DataLoader(
						    tokenized_datasets["train"],  shuffle=__A,  collate_fn=__A,  batch_size=__A							)
						UpperCAmelCase__								=  DataLoader(
						    tokenized_datasets["validation"],  shuffle=__A,  collate_fn=__A,  batch_size=__A							)

						return train_dataloader, eval_dataloader




def        lowerCAmelCase_			(					__A,  __A,  __A,  __A							)					->    int:







						'''simple docstring'''


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

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




def        lowerCAmelCase_			(					__A,  __A							)					->    str:







						'''simple docstring'''


						UpperCAmelCase__								=  Accelerator()

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

						set_seed(__A							)
						UpperCAmelCase__		,				UpperCAmelCase__								=  get_dataloaders(__A,  __A,  __A							)

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

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

						if accelerator.state.deepspeed_plugin is not None:
												UpperCAmelCase__								=  accelerator.state.deepspeed_plugin.deepspeed_config[
												    "gradient_accumulation_steps"
												]
						else:
												UpperCAmelCase__								=  1
						UpperCAmelCase__								=  (len(__A							) * num_epochs) // gradient_accumulation_steps

						# Instantiate scheduler
						if (
						    accelerator.state.deepspeed_plugin is None
						    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
						):
												UpperCAmelCase__								=  get_linear_schedule_with_warmup(
												    optimizer=__A,  num_warmup_steps=0,  num_training_steps=__A,  )
						else:
												UpperCAmelCase__								=  DummyScheduler(__A,  total_num_steps=__A,  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.
						UpperCAmelCase__		,				UpperCAmelCase__		,				UpperCAmelCase__		,				UpperCAmelCase__		,				UpperCAmelCase__								=  accelerator.prepare(
						    __A,  __A,  __A,  __A,  __A							)

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

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

						if args.resume_from_checkpoint:
												accelerator.load_state(args.resume_from_checkpoint							)
												UpperCAmelCase__								=  args.resume_from_checkpoint.split("epoch_"							)[1]
												UpperCAmelCase__								=  ""
												for char in epoch_string:
																		if char.isdigit():
																								state_epoch_num += char
																		else:
																								break
												UpperCAmelCase__								=  int(__A							) + 1
												UpperCAmelCase__								=  evaluation_loop(__A,  __A,  __A,  __A							)
												accelerator.print("resumed checkpoint performance:",  __A							)
												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:
																		UpperCAmelCase__								=  json.load(__A							)
																		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
						UpperCAmelCase__								=  {}
						for epoch in range(__A,  __A							):
												model.train()
												for step, batch in enumerate(__A							):
																		UpperCAmelCase__								=  model(**__A							)
																		UpperCAmelCase__								=  outputs.loss
																		UpperCAmelCase__								=  loss / gradient_accumulation_steps
																		accelerator.backward(__A							)
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

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

												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(__A,  __A							)




def        lowerCAmelCase_			(					)					->    Dict:







						'''simple docstring'''


						UpperCAmelCase__								=  argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."							)
						parser.add_argument(
						    "--model_name_or_path",  type=__A,  default="bert-base-cased",  help="Path to pretrained model or model identifier from huggingface.co/models.",  required=__A,  )
						parser.add_argument(
						    "--output_dir",  type=__A,  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=__A,  default=__A,  help="If the training should continue from a checkpoint folder.",  )
						parser.add_argument(
						    "--partial_train_epoch",  type=__A,  default=__A,  help="If passed, the training will stop after this number of epochs.",  )
						parser.add_argument(
						    "--num_epochs",  type=__A,  default=2,  help="Number of train epochs.",  )
						UpperCAmelCase__								=  parser.parse_args()
						UpperCAmelCase__								=  {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}

						training_function(__A,  __A							)


if __name__ == "__main__":
						main()




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers.models.distilbert.modeling_tf_distilbert import (
          TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFDistilBertForMaskedLM,
          TFDistilBertForMultipleChoice,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertModel,
      )





class 					A							:


   def __init__(self					:						Tuple	,  __UpperCAmelCase					:						str	,  )						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  1_3
         UpperCAmelCase__								=  7
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  False
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  9_9
         UpperCAmelCase__								=  3_2
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  3_7
         UpperCAmelCase__								=  "gelu"
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  5_1_2
         UpperCAmelCase__								=  1_6
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  0.02
         UpperCAmelCase__								=  3
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  None


   def       lowercase_				(self					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size		)

         UpperCAmelCase__								=  None
         if self.use_input_mask:
               UpperCAmelCase__								=  random_attention_mask([self.batch_size, self.seq_length]		)

         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         if self.use_labels:
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.type_sequence_label_size		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.num_labels		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.num_choices		)

         UpperCAmelCase__								=  DistilBertConfig(
             vocab_size=self.vocab_size	,  dim=self.hidden_size	,  n_layers=self.num_hidden_layers	,  n_heads=self.num_attention_heads	,  hidden_dim=self.intermediate_size	,  hidden_act=self.hidden_act	,  dropout=self.hidden_dropout_prob	,  attention_dropout=self.attention_probs_dropout_prob	,  max_position_embeddings=self.max_position_embeddings	,  initializer_range=self.initializer_range	,  )

         return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Tuple		)						->    Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         UpperCAmelCase__								=  [input_ids, input_mask]

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size)		)


   def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Dict		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForMaskedLM(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						List[str]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForQuestionAnswering(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.start_logits.shape	,  (self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape	,  (self.batch_size, self.seq_length)		)


   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForSequenceClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels)		)


   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Any		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_choices
         UpperCAmelCase__								=  TFDistilBertForMultipleChoice(__UpperCAmelCase		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_choices)		)


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Optional[int]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForTokenClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.num_labels)		)





   def       lowercase_				(self					:						Any		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.prepare_config_and_inputs()
         ((UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__))								=  config_and_inputs
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_tf
class 					A							(    UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase					):
   __UpperCAmelCase			:			Union[str, Any]              =      (
       (
           TFDistilBertModel,
           TFDistilBertForMaskedLM,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertForMultipleChoice,
       )
       if is_tf_available()
       else None
   )
   __UpperCAmelCase			:			Optional[int]              =      (
       {
           'feature-extraction': TFDistilBertModel,
           'fill-mask': TFDistilBertForMaskedLM,
           'question-answering': TFDistilBertForQuestionAnswering,
           'text-classification': TFDistilBertForSequenceClassification,
           'token-classification': TFDistilBertForTokenClassification,
           'zero-shot': TFDistilBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase			:			Tuple              =      False
   __UpperCAmelCase			:			str              =      False


   def       lowercase_				(self					:						Union[str, Any]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self	,  config_class=__UpperCAmelCase	,  dim=3_7		)


   def       lowercase_				(self					:						Any		)						->    int:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def       lowercase_				(self					:						int		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Optional[int]		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Any		)						->    Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*__UpperCAmelCase		)


   def       lowercase_				(self					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Dict		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*__UpperCAmelCase		)


   def       lowercase_				(self					:						str		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*__UpperCAmelCase		)





   @slow
   def       lowercase_				(self					:						Optional[int]		)						->    Optional[int]:
         """simple docstring"""



         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
               UpperCAmelCase__								=  TFDistilBertModel.from_pretrained(__UpperCAmelCase		)
               self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class 					A							(    unittest.TestCase					):


   @slow
   def       lowercase_				(self					:						List[Any]		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel.from_pretrained("distilbert-base-uncased"		)
         UpperCAmelCase__								=  tf.constant([[0, 1, 2, 3, 4, 5]]		)
         UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]

         UpperCAmelCase__								=  [1, 6, 7_6_8]
         self.assertEqual(output.shape	,  __UpperCAmelCase		)

         UpperCAmelCase__								=  tf.constant(
             [
                 [
                     [0.19261885, -0.13732955, 0.4119799],
                     [0.22150156, -0.07422661, 0.39037204],
                     [0.22756018, -0.0896414, 0.3701467],
                 ]
             ]		)
         tf.debugging.assert_near(output[:, :3, :3]	,  __UpperCAmelCase	,  atol=1E-4		)



