



import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
  from transformers import (
      FlaxBertModel,
      FlaxCLIPVisionModel,
      FlaxVisionTextDualEncoderModel,
      FlaxViTModel,
      VisionTextDualEncoderConfig,
      VisionTextDualEncoderProcessor,
  )
  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )


if is_torch_available():
  import torch

  from transformers import VisionTextDualEncoderModel

if is_vision_available():
  from PIL import Image





def 		a__      (		_UpperCamelCase      :			Optional[int]				):
   if isinstance(_UpperCamelCase			,collections.abc.Iterable				):
      return x
   return (x, x)



@require_flax
class 			__lowerCAmelCase  :




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    pass




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    pass




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    pass




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     np.abs((a - b)					).max()
    self.assertLessEqual(__UpperCAmelCase						,					__UpperCAmelCase						,					F"""Difference between torch and flax is {diff} (>= {tol})."""					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None						,					**__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel(__UpperCAmelCase					)

    __lowerCamelCase    				=     model(input_ids=__UpperCAmelCase						,					pixel_values=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase					)

    self.assertEqual(output['''text_embeds'''].shape						,					(input_ids.shape[0], config.projection_dim)					)
    self.assertEqual(output['''image_embeds'''].shape						,					(pixel_values.shape[0], config.projection_dim)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None						,					**__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.get_vision_text_model(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     {'''vision_model''': vision_model, '''text_model''': text_model}
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCAmelCase					)

    __lowerCamelCase    				=     model(input_ids=__UpperCAmelCase						,					pixel_values=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase					)

    self.assertEqual(output['''text_embeds'''].shape						,					(input_ids.shape[0], model.config.projection_dim)					)
    self.assertEqual(output['''image_embeds'''].shape						,					(pixel_values.shape[0], model.config.projection_dim)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None						,					**__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.get_vision_text_model(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     {'''vision_model''': vision_model, '''text_model''': text_model}
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCAmelCase					)

    __lowerCamelCase    				=     model(input_ids=__UpperCAmelCase						,					pixel_values=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase					)
    __lowerCamelCase    				=     output[0]

    with tempfile.TemporaryDirectory() as tmpdirname:
       model.save_pretrained(__UpperCAmelCase					)
       __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase					)

       __lowerCamelCase    				=     model(input_ids=__UpperCAmelCase						,					pixel_values=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase					)
       __lowerCamelCase    				=     after_output[0]
       __lowerCamelCase    				=     np.amax(np.abs(out_a - out_a					)					)
       self.assertLessEqual(__UpperCAmelCase						,					1E-3					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None						,					**__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.get_vision_text_model(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     {'''vision_model''': vision_model, '''text_model''': text_model}
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCAmelCase					)

    __lowerCamelCase    				=     model(
        input_ids=__UpperCAmelCase						,					pixel_values=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					output_attentions=__UpperCAmelCase					)

    __lowerCamelCase    				=     output.vision_model_output.attentions
    self.assertEqual(len(__UpperCAmelCase					)						,					vision_config.num_hidden_layers					)

    # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
    __lowerCamelCase    				=     to_atuple(vision_model.config.image_size					)
    __lowerCamelCase    				=     to_atuple(vision_model.config.patch_size					)
    __lowerCamelCase    				=     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
    __lowerCamelCase    				=     num_patches + 1
    self.assertEqual(vision_attentions[0].shape[-3:]						,					(vision_config.num_attention_heads, seq_len, seq_len)					)

    __lowerCamelCase    				=     output.text_model_output.attentions
    self.assertEqual(len(__UpperCAmelCase					)						,					text_config.num_hidden_layers					)

    self.assertEqual(
        text_attentions[0].shape[-3:]						,					(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])						,					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    pt_model.to(__UpperCAmelCase					)
    pt_model.eval()

    # prepare inputs
    __lowerCamelCase    				=     inputs_dict
    __lowerCamelCase    				=     {k: torch.tensor(v.tolist()					) for k, v in flax_inputs.items()}

    with torch.no_grad():
       __lowerCamelCase    				=     pt_model(**__UpperCAmelCase					).to_tuple()

    __lowerCamelCase    				=     fx_model(**__UpperCAmelCase					).to_tuple()
    self.assertEqual(len(__UpperCAmelCase					)						,					len(__UpperCAmelCase					)						,					'''Output lengths differ between Flax and PyTorch'''					)
    for fx_output, pt_output in zip(fx_outputs[:4]						,					pt_outputs[:4]					):
       self.assert_almost_equals(__UpperCAmelCase						,					pt_output.numpy()						,					4E-2					)

    # PT -> Flax
    with tempfile.TemporaryDirectory() as tmpdirname:
       pt_model.save_pretrained(__UpperCAmelCase					)
       __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase						,					from_pt=__UpperCAmelCase					)

    __lowerCamelCase    				=     fx_model_loaded(**__UpperCAmelCase					).to_tuple()
    self.assertEqual(len(__UpperCAmelCase					)						,					len(__UpperCAmelCase					)						,					'''Output lengths differ between Flax and PyTorch'''					)
    for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]						,					pt_outputs[:4]					):
       self.assert_almost_equals(__UpperCAmelCase						,					pt_output.numpy()						,					4E-2					)

    # Flax -> PT
    with tempfile.TemporaryDirectory() as tmpdirname:
       fx_model.save_pretrained(__UpperCAmelCase					)
       __lowerCamelCase    				=     VisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase						,					from_flax=__UpperCAmelCase					)

    pt_model_loaded.to(__UpperCAmelCase					)
    pt_model_loaded.eval()

    with torch.no_grad():
       __lowerCamelCase    				=     pt_model_loaded(**__UpperCAmelCase					).to_tuple()

    self.assertEqual(len(__UpperCAmelCase					)						,					len(__UpperCAmelCase					)						,					'''Output lengths differ between Flax and PyTorch'''					)
    for fx_output, pt_output_loaded in zip(fx_outputs[:4]						,					pt_outputs_loaded[:4]					):
       self.assert_almost_equals(__UpperCAmelCase						,					pt_output_loaded.numpy()						,					4E-2					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     VisionTextDualEncoderModel(__UpperCAmelCase					)
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel(__UpperCAmelCase					)

    __lowerCamelCase    				=     convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     fx_state

    self.check_pt_flax_equivalence(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     VisionTextDualEncoderModel(__UpperCAmelCase					)
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel(__UpperCAmelCase					)

    __lowerCamelCase    				=     load_flax_weights_in_pytorch_model(__UpperCAmelCase						,					fx_model.params					)

    self.check_pt_flax_equivalence(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    self.check_model_from_pretrained_configs(**__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    self.check_vision_text_dual_encoder_from_pretrained(**__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    self.check_save_load(**__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    self.check_vision_text_output_attention(**__UpperCAmelCase					)




 @is_pt_flax_cross_test
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    __lowerCamelCase    				=     config_inputs_dict.pop('''vision_config'''					)
    __lowerCamelCase    				=     config_inputs_dict.pop('''text_config'''					)

    __lowerCamelCase    				=     config_inputs_dict

    self.check_equivalence_pt_to_flax(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)
    self.check_equivalence_flax_to_pt(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.get_pretrained_model_and_inputs()

    __lowerCamelCase    				=     model_a(**__UpperCAmelCase					)
    __lowerCamelCase    				=     outputs[0]

    with tempfile.TemporaryDirectory() as tmp_dirname:
       model_a.save_pretrained(__UpperCAmelCase					)
       __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase					)

       __lowerCamelCase    				=     model_a(**__UpperCAmelCase					)
       __lowerCamelCase    				=     after_outputs[0]
       __lowerCamelCase    				=     np.amax(np.abs(out_a - out_a					)					)
       self.assertLessEqual(__UpperCAmelCase						,					1E-5					)



@require_flax
class 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
        '''hf-internal-testing/tiny-random-vit'''						,					'''hf-internal-testing/tiny-bert'''						,					vision_from_pt=__UpperCAmelCase						,					text_from_pt=__UpperCAmelCase						,					)
    __lowerCamelCase    				=     13
    __lowerCamelCase    				=     floats_tensor(
        [
            batch_size,
            model.config.vision_config.num_channels,
            model.config.vision_config.image_size,
            model.config.vision_config.image_size,
        ]					)
    __lowerCamelCase    				=     ids_tensor([batch_size, 4]						,					model.config.text_config.vocab_size					)
    __lowerCamelCase    				=     random_attention_mask([batch_size, 4]					)
    __lowerCamelCase    				=     {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

    return model, inputs




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxViTModel(__UpperCAmelCase					)
    __lowerCamelCase    				=     FlaxBertModel(__UpperCAmelCase					)
    return vision_model, text_model




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxViTModelTester(self					)
    __lowerCamelCase    				=     FlaxBertModelTester(self					)
    __lowerCamelCase    				=     vit_model_tester.prepare_config_and_inputs()
    __lowerCamelCase    				=     bert_model_tester.prepare_config_and_inputs()

    __lowerCamelCase				,__lowerCamelCase    				=     vision_config_and_inputs

    __lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase    				=     text_config_and_inputs

    # make sure that cross attention layers are added
    return {
        "text_config": text_config,
        "vision_config": vision_config,
        "pixel_values": pixel_values,
        "attention_mask": attention_mask,
        "input_ids": input_ids,
        "token_type_ids": token_type_ids,
    }



@require_torch
class 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
        '''hf-internal-testing/tiny-random-clip'''						,					'''hf-internal-testing/tiny-bert'''						,					vision_from_pt=__UpperCAmelCase						,					text_from_pt=__UpperCAmelCase						,					)
    __lowerCamelCase    				=     13
    __lowerCamelCase    				=     floats_tensor(
        [
            batch_size,
            model.config.vision_config.num_channels,
            model.config.vision_config.image_size,
            model.config.vision_config.image_size,
        ]					)
    __lowerCamelCase    				=     ids_tensor([batch_size, 4]						,					model.config.text_config.vocab_size					)
    __lowerCamelCase    				=     random_attention_mask([batch_size, 4]					)
    __lowerCamelCase    				=     {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

    return model, inputs




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxCLIPVisionModel(__UpperCAmelCase					)
    __lowerCamelCase    				=     FlaxBertModel(__UpperCAmelCase					)
    return vision_model, text_model




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxCLIPVisionModelTester(self					)
    __lowerCamelCase    				=     FlaxBertModelTester(self					)
    __lowerCamelCase    				=     clip_model_tester.prepare_config_and_inputs()
    __lowerCamelCase    				=     bert_model_tester.prepare_config_and_inputs()

    __lowerCamelCase				,__lowerCamelCase    				=     vision_config_and_inputs

    __lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase    				=     text_config_and_inputs

    # make sure that cross attention layers are added
    return {
        "text_config": text_config,
        "vision_config": vision_config,
        "pixel_values": pixel_values,
        "attention_mask": attention_mask,
        "input_ids": input_ids,
        "token_type_ids": token_type_ids,
    }



@require_flax
@require_vision
class 			__lowerCAmelCase  (						unittest.TestCase ):




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''						,					logit_scale_init_value=1.0					)
    __lowerCamelCase    				=     VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''					)

    __lowerCamelCase    				=     Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
    __lowerCamelCase    				=     processor(
        text=['''una foto di un gatto''', '''una foto di un cane''']						,					images=__UpperCAmelCase						,					padding=__UpperCAmelCase						,					return_tensors='''np'''					)

    __lowerCamelCase    				=     model(**__UpperCAmelCase					)

    # verify the logits
    self.assertEqual(outputs.logits_per_image.shape						,					(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])					)
    self.assertEqual(
        outputs.logits_per_text.shape						,					(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])						,					)

    __lowerCamelCase    				=     np.array([[1.2_284_727, 0.3_104_122]]					)

    self.assertTrue(np.allclose(outputs.logits_per_image						,					__UpperCAmelCase						,					atol=1E-3					)					)











import json
import os
import unittest

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES, XLMTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin



class 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):
 lowerCAmelCase__							       =	XLMTokenizer
 lowerCAmelCase__							       =	False




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    super().setUp()

    # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
    __lowerCamelCase    				=     [
        '''l''',
        '''o''',
        '''w''',
        '''e''',
        '''r''',
        '''s''',
        '''t''',
        '''i''',
        '''d''',
        '''n''',
        '''w</w>''',
        '''r</w>''',
        '''t</w>''',
        '''lo''',
        '''low''',
        '''er</w>''',
        '''low</w>''',
        '''lowest</w>''',
        '''newer</w>''',
        '''wider</w>''',
        '''<unk>''',
    ]
    __lowerCamelCase    				=     dict(zip(__UpperCAmelCase						,					range(len(__UpperCAmelCase					)					)					)					)
    __lowerCamelCase    				=     ['''l o 123''', '''lo w 1456''', '''e r</w> 1789''', '''''']

    __lowerCamelCase    				=     os.path.join(self.tmpdirname						,					VOCAB_FILES_NAMES['''vocab_file''']					)
    __lowerCamelCase    				=     os.path.join(self.tmpdirname						,					VOCAB_FILES_NAMES['''merges_file''']					)
    with open(self.vocab_file						,					'''w'''					) as fp:
       fp.write(json.dumps(__UpperCAmelCase					)					)
    with open(self.merges_file						,					'''w'''					) as fp:
       fp.write('''\n'''.join(__UpperCAmelCase					)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     '''lower newer'''
    __lowerCamelCase    				=     '''lower newer'''
    return input_text, output_text




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     XLMTokenizer(self.vocab_file						,					self.merges_file					)

    __lowerCamelCase    				=     '''lower'''
    __lowerCamelCase    				=     ['''low''', '''er</w>''']
    __lowerCamelCase    				=     tokenizer.tokenize(__UpperCAmelCase					)
    self.assertListEqual(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     tokens + ['''<unk>''']
    __lowerCamelCase    				=     [14, 15, 20]
    self.assertListEqual(tokenizer.convert_tokens_to_ids(__UpperCAmelCase					)						,					__UpperCAmelCase					)




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     XLMTokenizer.from_pretrained('''xlm-mlm-en-2048'''					)

    __lowerCamelCase    				=     tokenizer.encode('''sequence builders'''						,					add_special_tokens=__UpperCAmelCase					)
    __lowerCamelCase    				=     tokenizer.encode('''multi-sequence build'''						,					add_special_tokens=__UpperCAmelCase					)

    __lowerCamelCase    				=     tokenizer.build_inputs_with_special_tokens(__UpperCAmelCase					)
    __lowerCamelCase    				=     tokenizer.build_inputs_with_special_tokens(__UpperCAmelCase						,					__UpperCAmelCase					)

    assert encoded_sentence == [0] + text + [1]
    assert encoded_pair == [0] + text + [1] + text_a + [1]






