





import json
import os
import random
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import BertTokenizer, BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import FlavaImageProcessor, FlavaProcessor
  from transformers.models.flava.image_processing_flava import (
      FLAVA_CODEBOOK_MEAN,
      FLAVA_CODEBOOK_STD,
      FLAVA_IMAGE_MEAN,
      FLAVA_IMAGE_STD,
  )



@require_vision
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= tempfile.mkdtemp()

        # fmt: off
        SCREAMING_SNAKE_CASE__       							= ['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''want''', '''##want''', '''##ed''', '''wa''', '''un''', '''runn''', '''##ing''', ''',''', '''low''', '''lowest''']
        # fmt: on
        SCREAMING_SNAKE_CASE__       							= os.path.join(self.tmpdirname      ,							VOCAB_FILES_NAMES['''vocab_file''']    )

        with open(self.vocab_file      ,							'''w'''      ,							encoding='''utf-8'''    ) as fp:
            fp.write(''''''.join([x + '''\n''' for x in vocab_tokens]    )    )

        SCREAMING_SNAKE_CASE__       							= {
            '''image_mean''': FLAVA_IMAGE_MEAN,
            '''image_std''': FLAVA_IMAGE_STD,
            '''do_normalize''': True,
            '''do_resize''': True,
            '''size''': 2_24,
            '''do_center_crop''': True,
            '''crop_size''': 2_24,
            '''input_size_patches''': 14,
            '''total_mask_patches''': 75,
            '''mask_group_max_patches''': None,
            '''mask_group_min_patches''': 16,
            '''mask_group_min_aspect_ratio''': 0.3,
            '''mask_group_max_aspect_ratio''': None,
            '''codebook_do_resize''': True,
            '''codebook_size''': 1_12,
            '''codebook_do_center_crop''': True,
            '''codebook_crop_size''': 1_12,
            '''codebook_do_map_pixels''': True,
            '''codebook_do_normalize''': True,
            '''codebook_image_mean''': FLAVA_CODEBOOK_MEAN,
            '''codebook_image_std''': FLAVA_CODEBOOK_STD,
        }

        SCREAMING_SNAKE_CASE__       							= os.path.join(self.tmpdirname      ,							A_    )
        with open(self.image_processor_file      ,							'''w'''      ,							encoding='''utf-8'''    ) as fp:
            json.dump(A_      ,							A_    )



    def    lowercase_				(					self      ,							**A_    ):





        '''simple docstring'''

        return BertTokenizer.from_pretrained(self.tmpdirname      ,							**A_    )



    def    lowercase_				(					self      ,							**A_    ):





        '''simple docstring'''

        return BertTokenizerFast.from_pretrained(self.tmpdirname      ,							**A_    )



    def    lowercase_				(					self      ,							**A_    ):





        '''simple docstring'''

        return FlavaImageProcessor.from_pretrained(self.tmpdirname      ,							**A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        shutil.rmtree(self.tmpdirname    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= [np.random.randint(2_55      ,							size=(3, 30, 4_00)      ,							dtype=np.uinta    )]

        SCREAMING_SNAKE_CASE__       							= [Image.fromarray(np.moveaxis(A_      ,							0      ,							-1    )    ) for x in image_inputs]

        return image_inputs



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()
        SCREAMING_SNAKE_CASE__       							= self.get_rust_tokenizer()
        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )
        processor_slow.save_pretrained(self.tmpdirname    )
        SCREAMING_SNAKE_CASE__       							= FlavaProcessor.from_pretrained(self.tmpdirname      ,							use_fast=A_    )

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )
        processor_fast.save_pretrained(self.tmpdirname    )
        SCREAMING_SNAKE_CASE__       							= FlavaProcessor.from_pretrained(self.tmpdirname    )

        self.assertEqual(processor_slow.tokenizer.get_vocab()      ,							tokenizer_slow.get_vocab()    )
        self.assertEqual(processor_fast.tokenizer.get_vocab()      ,							tokenizer_fast.get_vocab()    )
        self.assertEqual(tokenizer_slow.get_vocab()      ,							tokenizer_fast.get_vocab()    )
        self.assertIsInstance(processor_slow.tokenizer      ,							A_    )
        self.assertIsInstance(processor_fast.tokenizer      ,							A_    )

        self.assertEqual(processor_slow.image_processor.to_json_string()      ,							image_processor.to_json_string()    )
        self.assertEqual(processor_fast.image_processor.to_json_string()      ,							image_processor.to_json_string()    )
        self.assertIsInstance(processor_slow.image_processor      ,							A_    )
        self.assertIsInstance(processor_fast.image_processor      ,							A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=self.get_tokenizer()      ,							image_processor=self.get_image_processor()    )
        processor.save_pretrained(self.tmpdirname    )

        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer(bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''    )
        SCREAMING_SNAKE_CASE__       							= self.get_image_processor(do_normalize=A_      ,							padding_value=1.0    )

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor.from_pretrained(
            self.tmpdirname      ,							bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''      ,							do_normalize=A_      ,							padding_value=1.0    )

        self.assertEqual(processor.tokenizer.get_vocab()      ,							tokenizer_add_kwargs.get_vocab()    )
        self.assertIsInstance(processor.tokenizer      ,							A_    )

        self.assertEqual(processor.image_processor.to_json_string()      ,							image_processor_add_kwargs.to_json_string()    )
        self.assertIsInstance(processor.image_processor      ,							A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )

        SCREAMING_SNAKE_CASE__       							= self.prepare_image_inputs()

        SCREAMING_SNAKE_CASE__       							= image_processor(A_      ,							return_tensors='''np'''    )
        SCREAMING_SNAKE_CASE__       							= processor(images=A_      ,							return_tensors='''np'''    )

        for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum()      ,							input_processor[key].sum()      ,							delta=1E-2    )

        # With rest of the args
        random.seed(12_34    )
        SCREAMING_SNAKE_CASE__       							= image_processor(
            A_      ,							return_image_mask=A_      ,							return_codebook_pixels=A_      ,							return_tensors='''np'''    )
        random.seed(12_34    )
        SCREAMING_SNAKE_CASE__       							= processor(
            images=A_      ,							return_image_mask=A_      ,							return_codebook_pixels=A_      ,							return_tensors='''np'''    )

        for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum()      ,							input_processor[key].sum()      ,							delta=1E-2    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )

        SCREAMING_SNAKE_CASE__       							= '''lower newer'''

        SCREAMING_SNAKE_CASE__       							= processor(text=A_    )

        SCREAMING_SNAKE_CASE__       							= tokenizer(A_    )

        for key in encoded_tok.keys():
            self.assertListEqual(encoded_tok[key]      ,							encoded_processor[key]    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )

        SCREAMING_SNAKE_CASE__       							= '''lower newer'''
        SCREAMING_SNAKE_CASE__       							= self.prepare_image_inputs()

        SCREAMING_SNAKE_CASE__       							= processor(text=A_      ,							images=A_    )

        self.assertListEqual(list(inputs.keys()    )      ,							['''input_ids''', '''token_type_ids''', '''attention_mask''', '''pixel_values''']    )

        # add extra args
        SCREAMING_SNAKE_CASE__       							= processor(text=A_      ,							images=A_      ,							return_codebook_pixels=A_      ,							return_image_mask=A_    )

        self.assertListEqual(
            list(inputs.keys()    )      ,							[
                '''input_ids''',
                '''token_type_ids''',
                '''attention_mask''',
                '''pixel_values''',
                '''codebook_pixel_values''',
                '''bool_masked_pos''',
            ]      ,							)

        # test if it raises when no input is passed
        with pytest.raises(A_    ):
            processor()



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )

        SCREAMING_SNAKE_CASE__       							= [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        SCREAMING_SNAKE_CASE__       							= processor.batch_decode(A_    )
        SCREAMING_SNAKE_CASE__       							= tokenizer.batch_decode(A_    )

        self.assertListEqual(A_      ,							A_    )




    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.get_image_processor()
        SCREAMING_SNAKE_CASE__       							= self.get_tokenizer()

        SCREAMING_SNAKE_CASE__       							= FlavaProcessor(tokenizer=A_      ,							image_processor=A_    )

        SCREAMING_SNAKE_CASE__       							= '''lower newer'''
        SCREAMING_SNAKE_CASE__       							= self.prepare_image_inputs()

        SCREAMING_SNAKE_CASE__       							= processor(text=A_      ,							images=A_    )

        self.assertListEqual(list(inputs.keys()    )      ,							processor.model_input_names    )







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True  )							->	Tuple:
    if not is_fpa_available():
        raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''  )
    for name, module in model.named_children():
        if isinstance(lowerCAmelCase_   ,       nn.Linear  ) and to_transformer_engine and _convert_linear:
            # Return early if the linear layer weights are not multiples of 16
            if any(p % 1_6 != 0 for p in module.weight.shape  ):
                return
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= te.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       nn.LayerNorm  ) and to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= te.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.Linear  ) and not to_transformer_engine and _convert_linear:
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= nn.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.LayerNorm  ) and not to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= nn.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            convert_model(
                lowerCAmelCase_   ,       to_transformer_engine=lowerCAmelCase_   ,       _convert_linear=_convert_linear   ,       _convert_ln=_convert_ln   ,       )



def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    if not is_fpa_available():
        raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''  )
    for m in model.modules():
        if isinstance(lowerCAmelCase_   ,       (te.LayerNorm, te.Linear)  ):
            return True
    return False
