



'''simple docstring'''



import argparse
import gc
import json
import os

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


lowerCAmelCase_  :							Optional[int]             =    16
lowerCAmelCase_  :							Optional[Any]             =    32
def        UpperCAmelCase (    A	:					Union[str, Any]					):
							return int(x / 2**20					)
class 						lowerCamelCase_				:







						def __enter__(	self :   Optional[Any]					):

													"""simple docstring"""
													gc.collect()
													torch.cuda.empty_cache()
													torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
													SCREAMING_SNAKE_CASE       :   Tuple         =      torch.cuda.memory_allocated()
													return self







						def __exit__(	self :   Any					,     *lowerCAmelCase__ :   str					):

													"""simple docstring"""
													gc.collect()
													torch.cuda.empty_cache()
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      torch.cuda.memory_allocated()
													SCREAMING_SNAKE_CASE       :   Tuple         =      torch.cuda.max_memory_allocated()
													SCREAMING_SNAKE_CASE       :   Tuple         =      bamb(self.end - self.begin					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      bamb(self.peak - self.begin					)
													# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


def        UpperCAmelCase (    A	:					Accelerator       ,		A	:					int = 16       ,		A	:					str = "bert-base-cased"       ,		A	:					int = 320       ,		A	:					int = 160       ,		):
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      AutoTokenizer.from_pretrained(A					)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      load_dataset(
							    '''glue'''       ,		'''mrpc'''       ,		split={'''train''': F"""train[:{n_train}]""", '''validation''': F"""validation[:{n_val}]"""}					)

							def tokenize_function(A	:					List[str]					):
														# max_length=None => use the model max length (it's actually the default)
														SCREAMING_SNAKE_CASE       :   Optional[Any]         =      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
							SCREAMING_SNAKE_CASE       :   Dict         =      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
							SCREAMING_SNAKE_CASE       :   Tuple         =      tokenized_datasets.rename_column('''label'''       ,		'''labels'''					)

							def collate_fn(A	:					Any					):
														# 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.
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      DataLoader(
							    tokenized_datasets['''train''']       ,		shuffle=A       ,		collate_fn=A       ,		batch_size=A					)
							SCREAMING_SNAKE_CASE       :   Dict         =      DataLoader(
							    tokenized_datasets['''validation''']       ,		shuffle=A       ,		collate_fn=A       ,		batch_size=A					)

							return train_dataloader, eval_dataloader
def        UpperCAmelCase (    A	:					Tuple       ,		A	:					Optional[int]					):
							# Initialize accelerator
							SCREAMING_SNAKE_CASE       :   List[Any]         =      Accelerator()

							# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
							SCREAMING_SNAKE_CASE       :   List[str]         =      config['''lr''']
							SCREAMING_SNAKE_CASE       :   List[str]         =      int(config['''num_epochs''']					)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      int(config['''seed''']					)
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      int(config['''batch_size''']					)
							SCREAMING_SNAKE_CASE       :   List[Any]         =      args.model_name_or_path

							set_seed(A					)
							SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   List[Any]         =      get_dataloaders(A       ,		A       ,		A       ,		args.n_train       ,		args.n_val					)

							# Instantiate the model (we build the model here so that the seed also control new weights initialization)
							SCREAMING_SNAKE_CASE       :   Optional[Any]         =      AutoModelForSequenceClassification.from_pretrained(A       ,		return_dict=A					)

							# Instantiate optimizer
							SCREAMING_SNAKE_CASE       :   List[Any]         =      (
							    AdamW
							    if accelerator.state.deepspeed_plugin is None
							    or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
							    else DummyOptim
							)
							SCREAMING_SNAKE_CASE       :   Tuple         =      optimizer_cls(params=model.parameters()       ,		lr=A					)

							if accelerator.state.deepspeed_plugin is not None:
														SCREAMING_SNAKE_CASE       :   Tuple         =      accelerator.state.deepspeed_plugin.deepspeed_config[
														    '''gradient_accumulation_steps'''
														]
							else:
														SCREAMING_SNAKE_CASE       :   str         =      1
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      (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
							):
														SCREAMING_SNAKE_CASE       :   Optional[int]         =      get_linear_schedule_with_warmup(
														    optimizer=A       ,		num_warmup_steps=0       ,		num_training_steps=A       ,		)
							else:
														SCREAMING_SNAKE_CASE       :   Optional[int]         =      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.
							SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   int         =      accelerator.prepare(
							    A       ,		A       ,		A       ,		A       ,		A					)

							# We need to keep track of how many total steps we have iterated over
							SCREAMING_SNAKE_CASE       :   int         =      0
							# We also need to keep track of the stating epoch so files are named properly
							SCREAMING_SNAKE_CASE       :   Dict         =      0

							# Now we train the model
							SCREAMING_SNAKE_CASE       :   int         =      {}
							for epoch in range(A       ,		A					):
														with TorchTracemalloc() as tracemalloc:
																					model.train()
																					for step, batch in enumerate(A					):
																												SCREAMING_SNAKE_CASE       :   Dict         =      model(**A					)
																												SCREAMING_SNAKE_CASE       :   Any         =      outputs.loss
																												SCREAMING_SNAKE_CASE       :   Any         =      loss / gradient_accumulation_steps
																												accelerator.backward(A					)
																												if step % gradient_accumulation_steps == 0:
																																			optimizer.step()
																																			lr_scheduler.step()
																																			optimizer.zero_grad()

																												overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
														accelerator.print('''Memory before entering the train : {}'''.format(bamb(tracemalloc.begin					)					)					)
														accelerator.print('''Memory consumed at the end of the train (end-begin): {}'''.format(tracemalloc.used					)					)
														accelerator.print('''Peak Memory consumed during the train (max-begin): {}'''.format(tracemalloc.peaked					)					)
														accelerator.print(
														    '''Total Peak Memory consumed during the train (max): {}'''.format(
														        tracemalloc.peaked + bamb(tracemalloc.begin					)					)					)
														SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      tracemalloc.peaked + bamb(tracemalloc.begin					)
														if args.peak_memory_upper_bound is not None:
																					assert (
																					    train_total_peak_memory[F"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
																					), "Peak memory usage exceeded the upper bound"

							accelerator.wait_for_everyone()
							if accelerator.is_main_process:
														with open(os.path.join(args.output_dir       ,		'''peak_memory_utilization.json'''					)       ,		'''w'''					) as f:
																					json.dump(A       ,		A					)
def        UpperCAmelCase (    ):
							SCREAMING_SNAKE_CASE       :   List[Any]         =      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(
							    '''--peak_memory_upper_bound'''       ,		type=A       ,		default=A       ,		help='''The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.'''       ,		)
							parser.add_argument(
							    '''--n_train'''       ,		type=A       ,		default=320       ,		help='''Number of training examples to use.'''       ,		)
							parser.add_argument(
							    '''--n_val'''       ,		type=A       ,		default=160       ,		help='''Number of validation examples to use.'''       ,		)
							parser.add_argument(
							    '''--num_epochs'''       ,		type=A       ,		default=1       ,		help='''Number of train epochs.'''       ,		)
							SCREAMING_SNAKE_CASE       :   List[Any]         =      parser.parse_args()
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      {'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
							training_function(A       ,		A					)


if __name__ == "__main__":
							main()









'''simple docstring'''



import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


lowerCAmelCase_  :							int             =    logging.get_logger(__name__)
def        UpperCAmelCase (    A	:					List[Any]       ,		A	:					Optional[int]					):
							SCREAMING_SNAKE_CASE       :   Any         =      nn.functional.normalize(A					)
							SCREAMING_SNAKE_CASE       :   List[str]         =      nn.functional.normalize(A					)
							return torch.mm(A       ,		normalized_text_embeds.t()					)






class 						lowerCamelCase_				(       snake_case_      ):
						_lowerCAmelCase  : Optional[Any]						  =					CLIPConfig

						_lowerCAmelCase  : List[str]						  =					['CLIPEncoderLayer']







						def __init__(	self :   int					,     lowerCAmelCase__ :   CLIPConfig					):

													"""simple docstring"""
													super().__init__(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      CLIPVisionModel(config.vision_config					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Linear(config.vision_config.hidden_size					,     config.projection_dim					,     bias=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      nn.Parameter(torch.ones(17					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      nn.Parameter(torch.ones(3					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      nn.Parameter(torch.ones(17					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Parameter(torch.ones(3					)					,     requires_grad=lowerCAmelCase__					)







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   List[str]					,     lowerCAmelCase__ :   Optional[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   Any         =      self.visual_projection(lowerCAmelCase__					)

													# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					).cpu().float().numpy()
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					).cpu().float().numpy()

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      []
													SCREAMING_SNAKE_CASE       :   int         =      image_embeds.shape[0]
													for i in range(lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   Any         =      {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

																				# increase this value to create a stronger `nfsw` filter
																				# at the cost of increasing the possibility of filtering benign images
																				SCREAMING_SNAKE_CASE       :   int         =      0.0

																				for concept_idx in range(len(special_cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[int]         =      special_cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   int         =      self.special_care_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["special_scores"][concept_idx] > 0:
																																		result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}					)
																																		SCREAMING_SNAKE_CASE       :   Tuple         =      0.01

																				for concept_idx in range(len(cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.concept_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Dict         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["concept_scores"][concept_idx] > 0:
																																		result_img["bad_concepts"].append(lowerCAmelCase__					)

																				result.append(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      [len(res['''bad_concepts''']					) > 0 for res in result]

													return images, has_nsfw_concepts







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   torch.FloatTensor					,     lowerCAmelCase__ :   torch.FloatTensor					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.visual_projection(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					)

													# increase this value to create a stronger `nsfw` filter
													# at the cost of increasing the possibility of filtering benign images
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      0.0

													SCREAMING_SNAKE_CASE       :   Any         =      special_cos_dist - self.special_care_embeds_weights + adjustment
													# special_scores = special_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      torch.any(special_scores > 0					,     dim=1					)
													SCREAMING_SNAKE_CASE       :   Dict         =      special_care * 0.01
													SCREAMING_SNAKE_CASE       :   List[str]         =      special_adjustment.unsqueeze(1					).expand(-1					,     cos_dist.shape[1]					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      (cos_dist - self.concept_embeds_weights) + special_adjustment
													# concept_scores = concept_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   Any         =      torch.any(concept_scores > 0					,     dim=1					)

													return images, has_nsfw_concepts




