
import inspect
import os
import sys
import unittest

import accelerate
from accelerate.test_utils import execute_subprocess_async, require_tpu




class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				inspect.getfile(accelerate.test_utils  )
           UpperCamelCase__    :   Any					=				os.path.sep.join(mod_file.split(os.path.sep  )[:-1] + ['''scripts''', '''test_script.py''']  )
           UpperCamelCase__    :   Tuple					=				os.path.sep.join(inspect.getfile(self.__class__  ).split(os.path.sep  )[:-1]  )






     @require_tpu
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[Any]:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				F'''
            {self.test_dir}/xla_spawn.py
            --num_cores 8
            {self.test_file_path}
        '''.split()
           UpperCamelCase__    :   List[str]					=				[sys.executable] + distributed_args
           execute_subprocess_async(__SCREAMING_SNAKE_CASE						,						env=os.environ.copy()  )




import torch
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel




class 				_lowerCamelCase (      UpperCamelCase_ ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  '''M-CLIP'''






     def __init__(							self						,						__SCREAMING_SNAKE_CASE=1_0_2_4						,						__SCREAMING_SNAKE_CASE=7_6_8						,						**__SCREAMING_SNAKE_CASE  )  ->   List[Any]:




           """simple docstring"""







           UpperCamelCase__    :   Dict					=				transformerDimSize
           UpperCamelCase__    :   Any					=				imageDimSize
           super().__init__(**__SCREAMING_SNAKE_CASE  )




class 				_lowerCamelCase (      UpperCamelCase_ ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  MCLIPConfig






     def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )  ->   Optional[int]:




           """simple docstring"""







           super().__init__(__SCREAMING_SNAKE_CASE						,						*__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Union[str, Any]					=				XLMRobertaModel(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Tuple					=				torch.nn.Linear(
               in_features=config.transformerDimensions						,						out_features=config.numDims  )






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   Tuple					=				self.transformer(input_ids=__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE  )[0]
           UpperCamelCase__    :   Tuple					=				(embs * attention_mask.unsqueeze(2  )).sum(dim=1  ) / attention_mask.sum(dim=1  )[:, None]
           return self.LinearTransformation(__SCREAMING_SNAKE_CASE  ), embs


