
import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
    import jax

    from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
    from transformers.models.bert.modeling_flax_bert import FlaxBertModel
    from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel




@require_flax
class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     @slow
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           for model_name in ["bert-base-cased", "bert-large-uncased"]:
                 with self.subTest(__SCREAMING_SNAKE_CASE  ):
                       UpperCamelCase__    :   str					=				AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE  )
                       self.assertIsNotNone(__SCREAMING_SNAKE_CASE  )
                       self.assertIsInstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )

                       UpperCamelCase__    :   List[str]					=				FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE  )
                       self.assertIsNotNone(__SCREAMING_SNAKE_CASE  )
                       self.assertIsInstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






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




           """simple docstring"""







           for model_name in ["roberta-base", "roberta-large"]:
                 with self.subTest(__SCREAMING_SNAKE_CASE  ):
                       UpperCamelCase__    :   Optional[Any]					=				AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE  )
                       self.assertIsNotNone(__SCREAMING_SNAKE_CASE  )
                       self.assertIsInstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )

                       UpperCamelCase__    :   Optional[Any]					=				FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE  )
                       self.assertIsNotNone(__SCREAMING_SNAKE_CASE  )
                       self.assertIsInstance(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






     @slow
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           for model_name in ["bert-base-cased", "bert-large-uncased"]:
                 UpperCamelCase__    :   Dict					=				AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  )
                 UpperCamelCase__    :   str					=				FlaxBertModel.from_pretrained(__SCREAMING_SNAKE_CASE  )
                 UpperCamelCase__    :   Tuple					=				tokenizer('''Do you support jax jitted function?'''						,						return_tensors=TensorType.JAX  )

                 @jax.jit
                 def eval(**__SCREAMING_SNAKE_CASE  ):
                       return model(**__SCREAMING_SNAKE_CASE  )

                 eval(**__SCREAMING_SNAKE_CASE  ).block_until_ready()






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




           """simple docstring"""







           for model_name in ["roberta-base", "roberta-large"]:
                 UpperCamelCase__    :   List[str]					=				AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  )
                 UpperCamelCase__    :   Dict					=				FlaxRobertaModel.from_pretrained(__SCREAMING_SNAKE_CASE  )
                 UpperCamelCase__    :   str					=				tokenizer('''Do you support jax jitted function?'''						,						return_tensors=TensorType.JAX  )

                 @jax.jit
                 def eval(**__SCREAMING_SNAKE_CASE  ):
                       return model(**__SCREAMING_SNAKE_CASE  )

                 eval(**__SCREAMING_SNAKE_CASE  ).block_until_ready()






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







           with self.assertRaisesRegex(
               __SCREAMING_SNAKE_CASE						,						'''bert-base is not a local folder and is not a valid model identifier'''  ):
                 UpperCamelCase__    :   Tuple					=				FlaxAutoModel.from_pretrained('''bert-base'''  )






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




           """simple docstring"""







           with self.assertRaisesRegex(
               __SCREAMING_SNAKE_CASE						,						r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''  ):
                 UpperCamelCase__    :   str					=				FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE						,						revision='''aaaaaa'''  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







           with self.assertRaisesRegex(
               __SCREAMING_SNAKE_CASE						,						'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'''						,						):
                 UpperCamelCase__    :   Union[str, Any]					=				FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




           """simple docstring"""







           with self.assertRaisesRegex(__SCREAMING_SNAKE_CASE						,						'''Use `from_pt=True` to load this model'''  ):
                 UpperCamelCase__    :   int					=				FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''  )




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


