




import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A							):







	'''simple docstring'''





	if isinstance(UpperCamelCase_     ,	torch.Tensor							):
		return image
	elif isinstance(UpperCamelCase_     ,	PIL.Image.Image							):
		snake_case_		    =     [image]

	if isinstance(image[0]     ,	PIL.Image.Image							):
		snake_case_		    =     [np.array(i.resize((w, h)     ,	resample=PIL_INTERPOLATION["lanczos"]							)							)[None, :] for i in image]
		snake_case_		    =     np.concatenate(UpperCamelCase_     ,	axis=0							)
		snake_case_		    =     np.array(UpperCamelCase_							).astype(np.floataa							) / 2_55.0
		snake_case_		    =     image.transpose(0     ,	3     ,	1     ,	2							)
		snake_case_		    =     2.0 * image - 1.0
		snake_case_		    =     torch.from_numpy(UpperCamelCase_							)
	elif isinstance(image[0]     ,	torch.Tensor							):
		snake_case_		    =     torch.cat(UpperCamelCase_     ,	dim=0							)
	return image


def 						lowerCamelCase__    (		_A     ,	_A     ,	_A     ,	_A=0.99_95							):







	'''simple docstring'''





	if not isinstance(UpperCamelCase_     ,	np.ndarray							):
		snake_case_		    =     True
		snake_case_		    =     va.device
		snake_case_		    =     va.cpu().numpy()
		snake_case_		    =     va.cpu().numpy()

	snake_case_		    =     np.sum(va * va / (np.linalg.norm(UpperCamelCase_							) * np.linalg.norm(UpperCamelCase_							))							)
	if np.abs(UpperCamelCase_							) > DOT_THRESHOLD:
		snake_case_		    =     (1 - t) * va + t * va
	else:
		snake_case_		    =     np.arccos(UpperCamelCase_							)
		snake_case_		    =     np.sin(UpperCamelCase_							)
		snake_case_		    =     theta_a * t
		snake_case_		    =     np.sin(UpperCamelCase_							)
		snake_case_		    =     np.sin(theta_a - theta_t							) / sin_theta_a
		snake_case_		    =     sin_theta_t / sin_theta_a
		snake_case_		    =     sa * va + sa * va

	if inputs_are_torch:
		snake_case_		    =     torch.from_numpy(UpperCamelCase_							).to(UpperCamelCase_							)

	return va


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	snake_case_		    =     F.normalize(UpperCamelCase_     ,	dim=-1							)
	snake_case_		    =     F.normalize(UpperCamelCase_     ,	dim=-1							)
	return (x - y).norm(dim=-1							).div(2							).arcsin().pow(2							).mul(2							)


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	for param in model.parameters():
		snake_case_		    =     value






class 		UpperCAmelCase   (				_UpperCAmelCase			):







						'''simple docstring'''





						def __init__(						self							:     Optional[Any]       ,	__lowercase							:     AutoencoderKL       ,	__lowercase							:     CLIPTextModel       ,	__lowercase							:     CLIPModel       ,	__lowercase							:     CLIPTokenizer       ,	__lowercase							:     UNetaDConditionModel       ,	__lowercase							:     Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]       ,	__lowercase							:     CLIPFeatureExtractor       ,	__lowercase							:     str=None       ,	__lowercase							:     List[Any]=None       ,	__lowercase							:     List[str]=None       ,	):






							"""simple docstring"""






							super().__init__()
							self.register_modules(
							    vae=UpperCAmelCase_       ,	text_encoder=UpperCAmelCase_       ,	clip_model=UpperCAmelCase_       ,	tokenizer=UpperCAmelCase_       ,	unet=UpperCAmelCase_       ,	scheduler=UpperCAmelCase_       ,	feature_extractor=UpperCAmelCase_       ,	coca_model=UpperCAmelCase_       ,	coca_tokenizer=UpperCAmelCase_       ,	coca_transform=UpperCAmelCase_       ,	)
							snake_case_		    =     (
							    feature_extractor.size
							    if isinstance(feature_extractor.size       ,	UpperCAmelCase_						)
							    else feature_extractor.size["shortest_edge"]
							)
							snake_case_		    =     transforms.Normalize(mean=feature_extractor.image_mean       ,	std=feature_extractor.image_std						)
							set_requires_grad(self.text_encoder       ,	UpperCAmelCase_						)
							set_requires_grad(self.clip_model       ,	UpperCAmelCase_						)






						def        snake_case__						(						self							:     Union[str, Any]       ,	__lowercase							:     Optional[Union[str, int]] = "auto"						):






							"""simple docstring"""






							if slice_size == "auto":
								# half the attention head size is usually a good trade-off between
								# speed and memory
								snake_case_		    =     self.unet.config.attention_head_dim // 2
							self.unet.set_attention_slice(UpperCAmelCase_						)






						def        snake_case__						(						self							:     Any						):






							"""simple docstring"""






							self.enable_attention_slicing(UpperCAmelCase_						)






						def        snake_case__						(						self							:     List[str]						):






							"""simple docstring"""






							set_requires_grad(self.vae       ,	UpperCAmelCase_						)






						def        snake_case__						(						self							:     str						):






							"""simple docstring"""






							set_requires_grad(self.vae       ,	UpperCAmelCase_						)






						def        snake_case__						(						self							:     Optional[int]						):






							"""simple docstring"""






							set_requires_grad(self.unet       ,	UpperCAmelCase_						)






						def        snake_case__						(						self							:     Any						):






							"""simple docstring"""






							set_requires_grad(self.unet       ,	UpperCAmelCase_						)






						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Any       ,	__lowercase							:     Any       ,	__lowercase							:     Union[str, Any]						):






							"""simple docstring"""






							snake_case_		    =     min(int(num_inference_steps * strength						)       ,	UpperCAmelCase_						)

							snake_case_		    =     max(num_inference_steps - init_timestep       ,	0						)
							snake_case_		    =     self.scheduler.timesteps[t_start:]

							return timesteps, num_inference_steps - t_start






						def        snake_case__						(						self							:     Any       ,	__lowercase							:     int       ,	__lowercase							:     List[str]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Tuple       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Dict=None						):






							"""simple docstring"""






							if not isinstance(UpperCAmelCase_       ,	torch.Tensor						):
								raise ValueError(f"`image` has to be of type `torch.Tensor` but is {type(UpperCAmelCase_						)}"						)

							snake_case_		    =     image.to(device=UpperCAmelCase_       ,	dtype=UpperCAmelCase_						)

							if isinstance(UpperCAmelCase_       ,	UpperCAmelCase_						):
								snake_case_		    =     [
								    self.vae.encode(image[i : i + 1]						).latent_dist.sample(generator[i]						) for i in range(UpperCAmelCase_						)
								]
								snake_case_		    =     torch.cat(UpperCAmelCase_       ,	dim=0						)
							else:
								snake_case_		    =     self.vae.encode(UpperCAmelCase_						).latent_dist.sample(UpperCAmelCase_						)

							# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							snake_case_		    =     0.1_8215 * init_latents
							snake_case_		    =     init_latents.repeat_interleave(UpperCAmelCase_       ,	dim=0						)

							snake_case_		    =     randn_tensor(init_latents.shape       ,	generator=UpperCAmelCase_       ,	device=UpperCAmelCase_       ,	dtype=UpperCAmelCase_						)

							# get latents
							snake_case_		    =     self.scheduler.add_noise(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)
							snake_case_		    =     init_latents

							return latents






						def        snake_case__						(						self							:     Optional[int]       ,	__lowercase							:     str						):






							"""simple docstring"""






							snake_case_		    =     self.coca_transform(UpperCAmelCase_						).unsqueeze(0						)
							with torch.no_grad(), torch.cuda.amp.autocast():
								snake_case_		    =     self.coca_model.generate(transformed_image.to(device=self.device       ,	dtype=self.coca_model.dtype						)						)
							snake_case_		    =     self.coca_tokenizer.decode(generated[0].cpu().numpy()						)
							return generated.split("<end_of_text>"						)[0].replace("<start_of_text>"       ,	""						).rstrip(" .,"						)






						def        snake_case__						(						self							:     Union[str, Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     List[str]						):






							"""simple docstring"""






							snake_case_		    =     self.feature_extractor.preprocess(UpperCAmelCase_						)
							snake_case_		    =     torch.from_numpy(clip_image_input["pixel_values"][0]						).unsqueeze(0						).to(self.device						).half()
							snake_case_		    =     self.clip_model.get_image_features(UpperCAmelCase_						)
							snake_case_		    =     image_embeddings_clip / image_embeddings_clip.norm(p=2       ,	dim=-1       ,	keepdim=UpperCAmelCase_						)
							snake_case_		    =     image_embeddings_clip.repeat_interleave(UpperCAmelCase_       ,	dim=0						)
							return image_embeddings_clip






						@torch.enable_grad()
						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     str       ,	__lowercase							:     List[Any]       ,	__lowercase							:     str       ,	__lowercase							:     List[Any]       ,	__lowercase							:     str       ,	__lowercase							:     Optional[Any]       ,	):






							"""simple docstring"""






							snake_case_		    =     latents.detach().requires_grad_()

							snake_case_		    =     self.scheduler.scale_model_input(UpperCAmelCase_       ,	UpperCAmelCase_						)

							# predict the noise residual
							snake_case_		    =     self.unet(UpperCAmelCase_       ,	UpperCAmelCase_       ,	encoder_hidden_states=UpperCAmelCase_						).sample

							if isinstance(self.scheduler       ,	(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)						):
								snake_case_		    =     self.scheduler.alphas_cumprod[timestep]
								snake_case_		    =     1 - alpha_prod_t
								# compute predicted original sample from predicted noise also called
								# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
								snake_case_		    =     (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

								snake_case_		    =     torch.sqrt(UpperCAmelCase_						)
								snake_case_		    =     pred_original_sample * (fac) + latents * (1 - fac)
							elif isinstance(self.scheduler       ,	UpperCAmelCase_						):
								snake_case_		    =     self.scheduler.sigmas[index]
								snake_case_		    =     latents - sigma * noise_pred
							else:
								raise ValueError(f"scheduler type {type(self.scheduler						)} not supported"						)

							# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							snake_case_		    =     1 / 0.1_8215 * sample
							snake_case_		    =     self.vae.decode(UpperCAmelCase_						).sample
							snake_case_		    =     (image / 2 + 0.5).clamp(0       ,	1						)

							snake_case_		    =     transforms.Resize(self.feature_extractor_size						)(UpperCAmelCase_						)
							snake_case_		    =     self.normalize(UpperCAmelCase_						).to(latents.dtype						)

							snake_case_		    =     self.clip_model.get_image_features(UpperCAmelCase_						)
							snake_case_		    =     image_embeddings_clip / image_embeddings_clip.norm(p=2       ,	dim=-1       ,	keepdim=UpperCAmelCase_						)

							snake_case_		    =     spherical_dist_loss(UpperCAmelCase_       ,	UpperCAmelCase_						).mean() * clip_guidance_scale

							snake_case_		    =     -torch.autograd.grad(UpperCAmelCase_       ,	UpperCAmelCase_						)[0]

							if isinstance(self.scheduler       ,	UpperCAmelCase_						):
								snake_case_		    =     latents.detach() + grads * (sigma**2)
								snake_case_		    =     noise_pred_original
							else:
								snake_case_		    =     noise_pred_original - torch.sqrt(UpperCAmelCase_						) * grads
							return noise_pred, latents






						@torch.no_grad()
						def __call__(						self							:     Optional[int]       ,	__lowercase							:     Union[torch.FloatTensor, PIL.Image.Image]       ,	__lowercase							:     Union[torch.FloatTensor, PIL.Image.Image]       ,	__lowercase							:     Optional[str] = None       ,	__lowercase							:     Optional[str] = None       ,	__lowercase							:     Optional[int] = 5_12       ,	__lowercase							:     Optional[int] = 5_12       ,	__lowercase							:     float = 0.6       ,	__lowercase							:     Optional[int] = 50       ,	__lowercase							:     Optional[float] = 7.5       ,	__lowercase							:     Optional[int] = 1       ,	__lowercase							:     float = 0.0       ,	__lowercase							:     Optional[float] = 1_00       ,	__lowercase							:     Optional[torch.Generator] = None       ,	__lowercase							:     Optional[str] = "pil"       ,	__lowercase							:     bool = True       ,	__lowercase							:     float = 0.8       ,	__lowercase							:     float = 0.1       ,	__lowercase							:     float = 0.1       ,	):






							"""simple docstring"""






							if isinstance(UpperCAmelCase_       ,	UpperCAmelCase_						) and len(UpperCAmelCase_						) != batch_size:
								raise ValueError(f"You have passed {batch_size} batch_size, but only {len(UpperCAmelCase_						)} generators."						)

							if height % 8 != 0 or width % 8 != 0:
								raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}."						)

							if isinstance(UpperCAmelCase_       ,	torch.Generator						) and batch_size > 1:
								snake_case_		    =     [generator] + [None] * (batch_size - 1)

							snake_case_		    =     [
							    ("model", self.coca_model is None),
							    ("tokenizer", self.coca_tokenizer is None),
							    ("transform", self.coca_transform is None),
							]
							snake_case_		    =     [x[0] for x in coca_is_none if x[1]]
							snake_case_		    =     ", ".join(UpperCAmelCase_						)
							# generate prompts with coca model if prompt is None
							if content_prompt is None:
								if len(UpperCAmelCase_						):
									raise ValueError(
									    f"Content prompt is None and CoCa [{coca_is_none_str}] is None."
									    f"Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."						)
								snake_case_		    =     self.get_image_description(UpperCAmelCase_						)
							if style_prompt is None:
								if len(UpperCAmelCase_						):
									raise ValueError(
									    f"Style prompt is None and CoCa [{coca_is_none_str}] is None."
									    f" Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."						)
								snake_case_		    =     self.get_image_description(UpperCAmelCase_						)

							# get prompt text embeddings for content and style
							snake_case_		    =     self.tokenizer(
							    UpperCAmelCase_       ,	padding="max_length"       ,	max_length=self.tokenizer.model_max_length       ,	truncation=UpperCAmelCase_       ,	return_tensors="pt"       ,	)
							snake_case_		    =     self.text_encoder(content_text_input.input_ids.to(self.device						)						)[0]

							snake_case_		    =     self.tokenizer(
							    UpperCAmelCase_       ,	padding="max_length"       ,	max_length=self.tokenizer.model_max_length       ,	truncation=UpperCAmelCase_       ,	return_tensors="pt"       ,	)
							snake_case_		    =     self.text_encoder(style_text_input.input_ids.to(self.device						)						)[0]

							snake_case_		    =     slerp(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)

							# duplicate text embeddings for each generation per prompt
							snake_case_		    =     text_embeddings.repeat_interleave(UpperCAmelCase_       ,	dim=0						)

							# set timesteps
							snake_case_		    =     "offset" in set(inspect.signature(self.scheduler.set_timesteps						).parameters.keys()						)
							snake_case_		    =     {}
							if accepts_offset:
								snake_case_		    =     1

							self.scheduler.set_timesteps(UpperCAmelCase_       ,	**UpperCAmelCase_						)
							# Some schedulers like PNDM have timesteps as arrays
							# It's more optimized to move all timesteps to correct device beforehand
							self.scheduler.timesteps.to(self.device						)

							snake_case_    ,      snake_case_		    =     self.get_timesteps(UpperCAmelCase_       ,	UpperCAmelCase_       ,	self.device						)
							snake_case_		    =     timesteps[:1].repeat(UpperCAmelCase_						)

							# Preprocess image
							snake_case_		    =     preprocess(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)
							snake_case_		    =     self.prepare_latents(
							    UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	text_embeddings.dtype       ,	self.device       ,	UpperCAmelCase_						)

							snake_case_		    =     preprocess(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)
							snake_case_		    =     self.prepare_latents(
							    UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	text_embeddings.dtype       ,	self.device       ,	UpperCAmelCase_						)

							snake_case_		    =     slerp(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)

							if clip_guidance_scale > 0:
								snake_case_		    =     self.get_clip_image_embeddings(UpperCAmelCase_       ,	UpperCAmelCase_						)
								snake_case_		    =     self.get_clip_image_embeddings(UpperCAmelCase_       ,	UpperCAmelCase_						)
								snake_case_		    =     slerp(
								    UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_						)

							# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
							# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
							# corresponds to doing no classifier free guidance.
							snake_case_		    =     guidance_scale > 1.0
							# get unconditional embeddings for classifier free guidance
							if do_classifier_free_guidance:
								snake_case_		    =     content_text_input.input_ids.shape[-1]
								snake_case_		    =     self.tokenizer([""]       ,	padding="max_length"       ,	max_length=UpperCAmelCase_       ,	return_tensors="pt"						)
								snake_case_		    =     self.text_encoder(uncond_input.input_ids.to(self.device						)						)[0]
								# duplicate unconditional embeddings for each generation per prompt
								snake_case_		    =     uncond_embeddings.repeat_interleave(UpperCAmelCase_       ,	dim=0						)

								# For classifier free guidance, we need to do two forward passes.
								# Here we concatenate the unconditional and text embeddings into a single batch
								# to avoid doing two forward passes
								snake_case_		    =     torch.cat([uncond_embeddings, text_embeddings]						)

							# get the initial random noise unless the user supplied it

							# Unlike in other pipelines, latents need to be generated in the target device
							# for 1-to-1 results reproducibility with the CompVis implementation.
							# However this currently doesn't work in `mps`.
							snake_case_		    =     (batch_size, self.unet.config.in_channels, height // 8, width // 8)
							snake_case_		    =     text_embeddings.dtype
							if latents is None:
								if self.device.type == "mps":
									# randn does not work reproducibly on mps
									snake_case_		    =     torch.randn(UpperCAmelCase_       ,	generator=UpperCAmelCase_       ,	device="cpu"       ,	dtype=UpperCAmelCase_						).to(
									    self.device						)
								else:
									snake_case_		    =     torch.randn(UpperCAmelCase_       ,	generator=UpperCAmelCase_       ,	device=self.device       ,	dtype=UpperCAmelCase_						)
							else:
								if latents.shape != latents_shape:
									raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"						)
								snake_case_		    =     latents.to(self.device						)

							# scale the initial noise by the standard deviation required by the scheduler
							snake_case_		    =     latents * self.scheduler.init_noise_sigma

							# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
							# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
							# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
							# and should be between [0, 1]
							snake_case_		    =     "eta" in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
							snake_case_		    =     {}
							if accepts_eta:
								snake_case_		    =     eta

							# check if the scheduler accepts generator
							snake_case_		    =     "generator" in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
							if accepts_generator:
								snake_case_		    =     generator

							with self.progress_bar(total=UpperCAmelCase_						):
								for i, t in enumerate(UpperCAmelCase_						):
									# expand the latents if we are doing classifier free guidance
									snake_case_		    =     torch.cat([latents] * 2						) if do_classifier_free_guidance else latents
									snake_case_		    =     self.scheduler.scale_model_input(UpperCAmelCase_       ,	UpperCAmelCase_						)

									# predict the noise residual
									snake_case_		    =     self.unet(UpperCAmelCase_       ,	UpperCAmelCase_       ,	encoder_hidden_states=UpperCAmelCase_						).sample

									# perform classifier free guidance
									if do_classifier_free_guidance:
										snake_case_    ,      snake_case_		    =     noise_pred.chunk(2						)
										snake_case_		    =     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

									# perform clip guidance
									if clip_guidance_scale > 0:
										snake_case_		    =     (
										    text_embeddings.chunk(2						)[1] if do_classifier_free_guidance else text_embeddings
										)
										snake_case_    ,      snake_case_		    =     self.cond_fn(
										    UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	)

									# compute the previous noisy sample x_t -> x_t-1
									snake_case_		    =     self.scheduler.step(UpperCAmelCase_       ,	UpperCAmelCase_       ,	UpperCAmelCase_       ,	**UpperCAmelCase_						).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							snake_case_		    =     1 / 0.1_8215 * latents
							snake_case_		    =     self.vae.decode(UpperCAmelCase_						).sample

							snake_case_		    =     (image / 2 + 0.5).clamp(0       ,	1						)
							snake_case_		    =     image.cpu().permute(0       ,	2       ,	3       ,	1						).numpy()

							if output_type == "pil":
								snake_case_		    =     self.numpy_to_pil(UpperCAmelCase_						)

							if not return_dict:
								return (image, None)

							return StableDiffusionPipelineOutput(images=UpperCAmelCase_       ,	nsfw_content_detected=UpperCAmelCase_						)


from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization 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_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
				from fairscale.optim import OSS


__snake_case								=  logging.get_logger(__name__)

__snake_case								=  {
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    """constant""": get_constant_schedule,
    """constant_w_warmup""": get_constant_schedule_with_warmup,
}




class 						lowercase__  (						_UpperCAmelCase  ):
	def __init__(			self							:   str       ,   UpperCAmelCase_							:   int=None       ,   UpperCAmelCase_							:   List[str]=None       ,   *UpperCAmelCase_							:   Union[str, Any]       ,   **UpperCAmelCase_							:   Optional[Any]							):
						super().__init__(*UpperCAmelCase_       ,   **UpperCAmelCase_							)

						if config is None:
											assert isinstance(self.model       ,   UpperCAmelCase_							), (
											    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
											    F' {self.model.__class__}'
											)
											SCREAMING_SNAKE_CASE__													= self.model.config
						else:
											SCREAMING_SNAKE_CASE__													= config

						SCREAMING_SNAKE_CASE__													= data_args
						SCREAMING_SNAKE_CASE__													= self.config.tgt_vocab_size if isinstance(self.config       ,   UpperCAmelCase_							) else self.config.vocab_size

						if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
											assert self.config.pad_token_id is not None, (
											    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
											    " calculation or doing label smoothing."
											)

						if self.config.pad_token_id is None and self.config.eos_token_id is not None:
											logger.warning(
											    F'The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'
											    ' padding..'							)

						if self.args.label_smoothing == 0:
											SCREAMING_SNAKE_CASE__													= torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id							)
						else:
											# dynamically import label_smoothed_nll_loss
											from utils import label_smoothed_nll_loss

											SCREAMING_SNAKE_CASE__													= label_smoothed_nll_loss
	def      A_				(			self							:   Tuple       ,   UpperCAmelCase_							:   int							):
						if self.optimizer is None:
											SCREAMING_SNAKE_CASE__													= ['bias', 'LayerNorm.weight']
											SCREAMING_SNAKE_CASE__													= [
											    {
											        'params': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay							)],
											        'weight_decay': self.args.weight_decay,
											    },
											    {
											        'params': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay							)],
											        'weight_decay': 0.0,
											    },
											]
											SCREAMING_SNAKE_CASE__													= Adafactor if self.args.adafactor else AdamW
											if self.args.adafactor:
																SCREAMING_SNAKE_CASE__													= Adafactor
																SCREAMING_SNAKE_CASE__													= {'scale_parameter': False, 'relative_step': False}
											else:
																SCREAMING_SNAKE_CASE__													= AdamW
																SCREAMING_SNAKE_CASE__													= {
																    'betas': (self.args.adam_betaa, self.args.adam_betaa),
																    'eps': self.args.adam_epsilon,
																}
											SCREAMING_SNAKE_CASE__													= self.args.learning_rate
											if self.sharded_ddp:
																SCREAMING_SNAKE_CASE__													= OSS(
																    params=UpperCAmelCase_       ,   optim=UpperCAmelCase_       ,   **UpperCAmelCase_       ,   )
											else:
																SCREAMING_SNAKE_CASE__													= optimizer_cls(UpperCAmelCase_       ,   **UpperCAmelCase_							)

						if self.lr_scheduler is None:
											SCREAMING_SNAKE_CASE__													= self._get_lr_scheduler(UpperCAmelCase_							)
						else:  # ignoring --lr_scheduler
											logger.warning('scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'							)
	def      A_				(			self							:   str       ,   UpperCAmelCase_							:   List[Any]							):
						SCREAMING_SNAKE_CASE__													= arg_to_scheduler[self.args.lr_scheduler]
						if self.args.lr_scheduler == "constant":
											SCREAMING_SNAKE_CASE__													= schedule_func(self.optimizer							)
						elif self.args.lr_scheduler == "constant_w_warmup":
											SCREAMING_SNAKE_CASE__													= schedule_func(self.optimizer       ,   num_warmup_steps=self.args.warmup_steps							)
						else:
											SCREAMING_SNAKE_CASE__													= schedule_func(
											    self.optimizer       ,   num_warmup_steps=self.args.warmup_steps       ,   num_training_steps=UpperCAmelCase_							)
						return scheduler
	def      A_				(			self							:   List[str]							):
						if isinstance(self.train_dataset       ,   torch.utils.data.IterableDataset							):
											return None
						elif is_torch_tpu_available():
											return get_tpu_sampler(self.train_dataset							)
						else:
											if self.args.sortish_sampler:
																self.train_dataset.make_sortish_sampler(
																    self.args.per_device_train_batch_size       ,   distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)       ,   )

											return (
											    RandomSampler(self.train_dataset							)
											    if self.args.local_rank == -1
											    else DistributedSampler(self.train_dataset							)
											)
	def      A_				(			self							:   Any       ,   UpperCAmelCase_							:   Tuple       ,   UpperCAmelCase_							:   Optional[Any]       ,   UpperCAmelCase_							:   Dict							):
						if self.args.label_smoothing == 0:
											if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
																# force training to ignore pad token
																SCREAMING_SNAKE_CASE__													= model(**UpperCAmelCase_       ,   use_cache=UpperCAmelCase_							)[0]
																SCREAMING_SNAKE_CASE__													= self.loss_fn(logits.view(-1       ,   logits.shape[-1]							)       ,   labels.view(-1							)							)
											else:
																# compute usual loss via models
																SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= model(**UpperCAmelCase_       ,   labels=UpperCAmelCase_       ,   use_cache=UpperCAmelCase_							)[:2]
						else:
											# compute label smoothed loss
											SCREAMING_SNAKE_CASE__													= model(**UpperCAmelCase_       ,   use_cache=UpperCAmelCase_							)[0]
											SCREAMING_SNAKE_CASE__													= torch.nn.functional.log_softmax(UpperCAmelCase_       ,   dim=-1							)
											SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= self.loss_fn(UpperCAmelCase_       ,   UpperCAmelCase_       ,   self.args.label_smoothing       ,   ignore_index=self.config.pad_token_id							)
						return loss, logits
	def      A_				(			self							:   Union[str, Any]       ,   UpperCAmelCase_							:   str       ,   UpperCAmelCase_							:   List[Any]							):
						SCREAMING_SNAKE_CASE__													= inputs.pop('labels'							)
						SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= self._compute_loss(UpperCAmelCase_       ,   UpperCAmelCase_       ,   UpperCAmelCase_							)
						return loss
	def      A_				(			self							:   List[str]       ,   UpperCAmelCase_							:   nn.Module       ,   UpperCAmelCase_							:   Dict[str, Union[torch.Tensor, Any]]       ,   UpperCAmelCase_							:   bool       ,   UpperCAmelCase_							:   Optional[List[str]] = None       ,   ):
						SCREAMING_SNAKE_CASE__													= self._prepare_inputs(UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= {
						    'max_length': self.data_args.val_max_target_length
						    if self.data_args is not None
						    else self.config.max_length,
						    'num_beams': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
						}

						if self.args.predict_with_generate and not self.args.prediction_loss_only:
											SCREAMING_SNAKE_CASE__													= self.model.generate(
											    inputs['input_ids']       ,   attention_mask=inputs['attention_mask']       ,   **UpperCAmelCase_       ,   )
											# in case the batch is shorter than max length, the output should be padded
											if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																SCREAMING_SNAKE_CASE__													= self._pad_tensors_to_max_len(UpperCAmelCase_       ,   gen_kwargs['max_length']							)

						SCREAMING_SNAKE_CASE__													= inputs.pop('labels'							)
						with torch.no_grad():
											# compute loss on predict data
											SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= self._compute_loss(UpperCAmelCase_       ,   UpperCAmelCase_       ,   UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= loss.mean().detach()
						if self.args.prediction_loss_only:
											return (loss, None, None)

						SCREAMING_SNAKE_CASE__													= generated_tokens if self.args.predict_with_generate else logits

						if labels.shape[-1] < gen_kwargs["max_length"]:
											SCREAMING_SNAKE_CASE__													= self._pad_tensors_to_max_len(UpperCAmelCase_       ,   gen_kwargs['max_length']							)

						return (loss, logits, labels)





	def      A_				(			self							:   Union[str, Any]       ,   UpperCAmelCase_							:   Tuple       ,   UpperCAmelCase_							:   List[str]							):
						# If PAD token is not defined at least EOS token has to be defined
						SCREAMING_SNAKE_CASE__													= self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

						if pad_token_id is None:
											raise ValueError(
											    'Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'
											    F' padded to `max_length`={max_length}'							)

						SCREAMING_SNAKE_CASE__													= pad_token_id * torch.ones(
						    (tensor.shape[0], max_length)       ,   dtype=tensor.dtype       ,   device=tensor.device							)
						SCREAMING_SNAKE_CASE__													= tensor
						return padded_tensor





