"""simple docstring"""




import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available






def   UpperCAmelCase					(	*UpperCamelCase__				):






							"""simple docstring"""





							if not isinstance(UpperCamelCase__       ,    UpperCamelCase__				):
														A__	     =							list(UpperCamelCase__				)
							for i in range(len(UpperCamelCase__				)				):
														A__	     =							None
							gc.collect()
							if is_xpu_available():
														torch.xpu.empty_cache()
							elif is_npu_available():
														torch.npu.empty_cache()
							else:
														torch.cuda.empty_cache()
							return objects






def   UpperCAmelCase					(	UpperCamelCase__				):






							"""simple docstring"""





							A__	     =							[
							    'CUDA out of memory.',  # CUDA OOM
							    'cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.',  # CUDNN SNAFU
							    'DefaultCPUAllocator: can\'t allocate memory',  # CPU OOM
							]
							if isinstance(UpperCamelCase__       ,    UpperCamelCase__				) and len(exception.args				) == 1:
														return any(err in exception.args[0] for err in _statements				)
							return False






def   UpperCAmelCase					(	UpperCamelCase__ = None       ,    UpperCamelCase__ = 128				):






							"""simple docstring"""





							if function is None:
														return functools.partial(UpperCamelCase__       ,    starting_batch_size=UpperCamelCase__				)

							A__	     =							starting_batch_size

							def decorator(*UpperCamelCase__       ,    **UpperCamelCase__				):
														nonlocal batch_size
														gc.collect()
														if is_xpu_available():
																					torch.xpu.empty_cache()
														elif is_npu_available():
																					torch.npu.empty_cache()
														else:
																					torch.cuda.empty_cache()
														A__	     =							list(inspect.signature(UpperCamelCase__				).parameters.keys()				)
														# Guard against user error
														if len(UpperCamelCase__				) < (len(UpperCamelCase__				) + 1):
																					A__	     =							', '.join([F'''{arg}={value}''' for arg, value in zip(params[1:]       ,    args[1:]				)]				)
																					raise TypeError(
																					    F'''Batch size was passed into `{function.__name__}` as the first argument when called.'''
																					    F'''Remove this as the decorator already does so: `{function.__name__}({arg_str})`'''				)
														while True:
																					if batch_size == 0:
																												raise RuntimeError('No executable batch size found, reached zero.'				)
																					try:
																												return function(UpperCamelCase__       ,    *UpperCamelCase__       ,    **UpperCamelCase__				)
																					except Exception as e:
																												if should_reduce_batch_size(UpperCamelCase__				):
																																			gc.collect()
																																			if is_xpu_available():
																																										torch.xpu.empty_cache()
																																			elif is_npu_available():
																																										torch.npu.empty_cache()
																																			else:
																																										torch.cuda.empty_cache()
																																			batch_size //= 2
																												else:
																																			raise

							return decorator






"""simple docstring"""




import torch
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						Union[str, Any]													=		'M-CLIP'





						def __init__(       self		,__UpperCAmelCase=10_24		,__UpperCAmelCase=7_68		,**__UpperCAmelCase     )					->				List[str]:
													A__	     =							transformerDimSize
													A__	     =							imageDimSize
													super().__init__(**__UpperCAmelCase     )






class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						Union[str, Any]													=		MCLIPConfig





						def __init__(       self		,__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase     )					->				List[Any]:
													super().__init__(__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase     )
													A__	     =							XLMRobertaModel(__UpperCAmelCase     )
													A__	     =							torch.nn.Linear(
													    in_features=config.transformerDimensions		,out_features=config.numDims     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase     )					->				Tuple:
													A__	     =							self.transformer(input_ids=__UpperCAmelCase		,attention_mask=__UpperCAmelCase     )[0]
													A__	     =							(embs * attention_mask.unsqueeze(2     )).sum(dim=1     ) / attention_mask.sum(dim=1     )[:, None]
													return self.LinearTransformation(__UpperCAmelCase     ), embs





