





from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TF_MODEL_FOR_PRETRAINING_MAPPING,
					    TFMobileBertForMaskedLM,
					    TFMobileBertForMultipleChoice,
					    TFMobileBertForNextSentencePrediction,
					    TFMobileBertForPreTraining,
					    TFMobileBertForQuestionAnswering,
					    TFMobileBertForSequenceClassification,
					    TFMobileBertForTokenClassification,
					    TFMobileBertModel,
					)
@require_tf
class 			UpperCamelCase_		(  _lowerCamelCase							,  _lowerCamelCase							,  unittest.TestCase					):
					lowerCAmelCase_							=	(
					    (
					        TFMobileBertModel,
					        TFMobileBertForMaskedLM,
					        TFMobileBertForNextSentencePrediction,
					        TFMobileBertForPreTraining,
					        TFMobileBertForQuestionAnswering,
					        TFMobileBertForSequenceClassification,
					        TFMobileBertForTokenClassification,
					        TFMobileBertForMultipleChoice,
					    )
					    if is_tf_available()
					    else ()
					)
					lowerCAmelCase_							=	(
					    {
					        '''feature-extraction''': TFMobileBertModel,
					        '''fill-mask''': TFMobileBertForMaskedLM,
					        '''question-answering''': TFMobileBertForQuestionAnswering,
					        '''text-classification''': TFMobileBertForSequenceClassification,
					        '''token-classification''': TFMobileBertForTokenClassification,
					        '''zero-shot''': TFMobileBertForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	False




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=False  )		->							Tuple:
										_snake_case									=				super()._prepare_for_class(lowerCAmelCase_				,						lowerCAmelCase_				,						return_labels=lowerCAmelCase_  )

										if return_labels:
															if model_class in get_values(lowerCAmelCase_  ):
																				_snake_case									=				tf.zeros(self.model_tester.batch_size				,						dtype=tf.intaa  )

										return inputs_dict
					class 			UpperCamelCase_		(  _lowerCamelCase					):




										def __init__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_=13				,						lowerCAmelCase_=7				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=99				,						lowerCAmelCase_=32				,						lowerCAmelCase_=32				,						lowerCAmelCase_=2				,						lowerCAmelCase_=4				,						lowerCAmelCase_=37				,						lowerCAmelCase_="gelu"				,						lowerCAmelCase_=0.1				,						lowerCAmelCase_=0.1				,						lowerCAmelCase_=512				,						lowerCAmelCase_=16				,						lowerCAmelCase_=2				,						lowerCAmelCase_=0.02				,						lowerCAmelCase_=3				,						lowerCAmelCase_=4				,						lowerCAmelCase_=None				,						)		->							str:
															_snake_case									=				parent
															_snake_case									=				batch_size
															_snake_case									=				seq_length
															_snake_case									=				is_training
															_snake_case									=				use_input_mask
															_snake_case									=				use_token_type_ids
															_snake_case									=				use_labels
															_snake_case									=				vocab_size
															_snake_case									=				hidden_size
															_snake_case									=				num_hidden_layers
															_snake_case									=				num_attention_heads
															_snake_case									=				intermediate_size
															_snake_case									=				hidden_act
															_snake_case									=				hidden_dropout_prob
															_snake_case									=				attention_probs_dropout_prob
															_snake_case									=				max_position_embeddings
															_snake_case									=				type_vocab_size
															_snake_case									=				type_sequence_label_size
															_snake_case									=				initializer_range
															_snake_case									=				num_labels
															_snake_case									=				num_choices
															_snake_case									=				scope
															_snake_case									=				embedding_size




										def 	lowerCAmelCase			(		self  )		->							str:
															_snake_case									=				ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size  )

															_snake_case									=				None
															if self.use_input_mask:
																				_snake_case									=				random_attention_mask([self.batch_size, self.seq_length]  )

															_snake_case									=				None
															if self.use_token_type_ids:
																				_snake_case									=				ids_tensor([self.batch_size, self.seq_length]				,						self.type_vocab_size  )

															_snake_case									=				None
															_snake_case									=				None
															_snake_case									=				None
															if self.use_labels:
																				_snake_case									=				ids_tensor([self.batch_size]				,						self.type_sequence_label_size  )
																				_snake_case									=				ids_tensor([self.batch_size, self.seq_length]				,						self.num_labels  )
																				_snake_case									=				ids_tensor([self.batch_size]				,						self.num_choices  )

															_snake_case									=				MobileBertConfig(
															    vocab_size=self.vocab_size				,						hidden_size=self.hidden_size				,						num_hidden_layers=self.num_hidden_layers				,						num_attention_heads=self.num_attention_heads				,						intermediate_size=self.intermediate_size				,						hidden_act=self.hidden_act				,						hidden_dropout_prob=self.hidden_dropout_prob				,						attention_probs_dropout_prob=self.attention_probs_dropout_prob				,						max_position_embeddings=self.max_position_embeddings				,						type_vocab_size=self.type_vocab_size				,						initializer_range=self.initializer_range				,						embedding_size=self.embedding_size				,						)

															return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							List[Any]:
															_snake_case									=				TFMobileBertModel(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )

															_snake_case									=				[input_ids, input_mask]
															_snake_case									=				model(lowerCAmelCase_  )

															_snake_case									=				model(lowerCAmelCase_  )

															self.parent.assertEqual(
															    result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)  )
															self.parent.assertEqual(result.pooler_output.shape				,						(self.batch_size, self.hidden_size)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							List[str]:
															_snake_case									=				TFMobileBertForMaskedLM(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							str:
															_snake_case									=				TFMobileBertForNextSentencePrediction(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.logits.shape				,						(self.batch_size, 2)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Optional[Any]:
															_snake_case									=				TFMobileBertForPreTraining(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(
															    result.prediction_logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)  )
															self.parent.assertEqual(result.seq_relationship_logits.shape				,						(self.batch_size, 2)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Tuple:
															_snake_case									=				self.num_labels
															_snake_case									=				TFMobileBertForSequenceClassification(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.num_labels)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							int:
															_snake_case									=				self.num_choices
															_snake_case									=				TFMobileBertForMultipleChoice(config=lowerCAmelCase_  )
															_snake_case									=				tf.tile(tf.expand_dims(lowerCAmelCase_				,						1  )				,						(1, self.num_choices, 1)  )
															_snake_case									=				tf.tile(tf.expand_dims(lowerCAmelCase_				,						1  )				,						(1, self.num_choices, 1)  )
															_snake_case									=				tf.tile(tf.expand_dims(lowerCAmelCase_				,						1  )				,						(1, self.num_choices, 1)  )
															_snake_case									=				{
															    'input_ids': multiple_choice_inputs_ids,
															    'attention_mask': multiple_choice_input_mask,
															    'token_type_ids': multiple_choice_token_type_ids,
															}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.num_choices)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Dict:
															_snake_case									=				self.num_labels
															_snake_case									=				TFMobileBertForTokenClassification(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.num_labels)  )




										def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							List[Any]:
															_snake_case									=				TFMobileBertForQuestionAnswering(config=lowerCAmelCase_  )
															_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_snake_case									=				model(lowerCAmelCase_  )
															self.parent.assertEqual(result.start_logits.shape				,						(self.batch_size, self.seq_length)  )
															self.parent.assertEqual(result.end_logits.shape				,						(self.batch_size, self.seq_length)  )






										def 	lowerCAmelCase			(		self  )		->							str:
															_snake_case									=				self.prepare_config_and_inputs()
															(
															    (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     (
															    _snake_case
															)							,     
															)									=				config_and_inputs
															_snake_case									=				{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
															return config, inputs_dict





					def 	lowerCAmelCase			(		self  )		->							Any:
										_snake_case									=				TFMobileBertModelTest.TFMobileBertModelTester(self  )
										_snake_case									=				ConfigTester(self				,						config_class=lowerCAmelCase_				,						hidden_size=37  )




					def 	lowerCAmelCase			(		self  )		->							Tuple:
										self.config_tester.run_common_tests()




					def 	lowerCAmelCase			(		self  )		->							Any:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_model(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_masked_lm(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_multiple_choice(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_pretraining(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_question_answering(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_sequence_classification(*lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							Any:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_token_classification(*lowerCAmelCase_  )






					@slow
					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["google/mobilebert-uncased"]:
															_snake_case									=				TFMobileBertModel.from_pretrained(lowerCAmelCase_  )
															self.assertIsNotNone(lowerCAmelCase_  )





@require_tf
class 			UpperCamelCase_		(  unittest.TestCase					):




					@slow
					def 	lowerCAmelCase			(		self  )		->							Union[str, Any]:
										_snake_case									=				TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'  )
										_snake_case									=				tf.constant([[0, 1, 2, 3, 4, 5]]  )
										_snake_case									=				model(lowerCAmelCase_  )[0]

										_snake_case									=				[1, 6, 3_0522]
										self.assertEqual(output.shape				,						lowerCAmelCase_  )

										_snake_case									=				tf.constant(
										    [
										        [
										            [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
										            [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
										            [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
										        ]
										    ]  )
										tf.debugging.assert_near(output[:, :3, :3]				,						lowerCAmelCase_				,						atol=1E-4  )










import json
import os
from collections import Counter

import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torch.utils.data import Dataset


UpperCAmelCase_      							=     {1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}
class 			UpperCamelCase_		(  nn.Module					):




					def __init__(		self				,						lowerCAmelCase_  )		->							int:
										super().__init__()
										_snake_case									=				torchvision.models.resnetaaa(pretrained=lowerCAmelCase_  )
										_snake_case									=				list(model.children()  )[:-2]
										_snake_case									=				nn.Sequential(*lowerCAmelCase_  )
										_snake_case									=				nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds]  )






					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_  )		->							str:
										# Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx2048
										_snake_case									=				self.pool(self.model(lowerCAmelCase_  )  )
										_snake_case									=				torch.flatten(lowerCAmelCase_				,						start_dim=2  )
										_snake_case									=				out.transpose(1				,						2  ).contiguous()
										return out  # BxNx2048





class 			UpperCamelCase_		(  _lowerCamelCase					):




					def __init__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Optional[Any]:
										_snake_case									=				[json.loads(lowerCAmelCase_  ) for l in open(lowerCAmelCase_  )]
										_snake_case									=				os.path.dirname(lowerCAmelCase_  )
										_snake_case									=				tokenizer
										_snake_case									=				labels
										_snake_case									=				len(lowerCAmelCase_  )
										_snake_case									=				max_seq_length

										_snake_case									=				transforms




					def __len__(		self  )		->							Any:
										return len(self.data  )




					def __getitem__(		self				,						lowerCAmelCase_  )		->							Optional[int]:
										_snake_case									=				torch.LongTensor(self.tokenizer.encode(self.data[index]['text']				,						add_special_tokens=lowerCAmelCase_  )  )
										_snake_case							,     _snake_case							,     _snake_case									=				sentence[0], sentence[1:-1], sentence[-1]
										_snake_case									=				sentence[: self.max_seq_length]

										_snake_case									=				torch.zeros(self.n_classes  )
										_snake_case									=				1

										_snake_case									=				Image.open(os.path.join(self.data_dir				,						self.data[index]['img']  )  ).convert('RGB'  )
										_snake_case									=				self.transforms(lowerCAmelCase_  )

										return {
										    "image_start_token": start_token,
										    "image_end_token": end_token,
										    "sentence": sentence,
										    "image": image,
										    "label": label,
										}






					def 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				Counter()
										for row in self.data:
															label_freqs.update(row['label']  )
										return label_freqs





def 						lowerCamelCase__     (     UpperCamelCase__			:       str						)			-> Dict:



					'''simple docstring'''




					_snake_case									=				[len(row['sentence']						) for row in batch]
					_snake_case							,     _snake_case									=				len(UpperCamelCase__						), max(UpperCamelCase__						)

					_snake_case									=				torch.zeros(UpperCamelCase__       ,    UpperCamelCase__       ,    dtype=torch.long						)
					_snake_case									=				torch.zeros(UpperCamelCase__       ,    UpperCamelCase__       ,    dtype=torch.long						)

					for i_batch, (input_row, length) in enumerate(zip(UpperCamelCase__       ,    UpperCamelCase__						)						):
										_snake_case									=				input_row['sentence']
										_snake_case									=				1

					_snake_case									=				torch.stack([row['image'] for row in batch]						)
					_snake_case									=				torch.stack([row['label'] for row in batch]						)
					_snake_case									=				torch.stack([row['image_start_token'] for row in batch]						)
					_snake_case									=				torch.stack([row['image_end_token'] for row in batch]						)

					return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor





def 						lowerCamelCase__     (     )			-> str:



					'''simple docstring'''




					return [
					    "Crime",
					    "Drama",
					    "Thriller",
					    "Action",
					    "Comedy",
					    "Romance",
					    "Documentary",
					    "Short",
					    "Mystery",
					    "History",
					    "Family",
					    "Adventure",
					    "Fantasy",
					    "Sci-Fi",
					    "Western",
					    "Horror",
					    "Sport",
					    "War",
					    "Music",
					    "Musical",
					    "Animation",
					    "Biography",
					    "Film-Noir",
					]





def 						lowerCamelCase__     (     )			-> Tuple:



					'''simple docstring'''




					return transforms.Compose(
					    [
					        transforms.Resize(256						),
					        transforms.CenterCrop(224						),
					        transforms.ToTensor(),
					        transforms.Normalize(
					            mean=[0.4677_7044, 0.4453_1429, 0.4066_1017]       ,    std=[0.1222_1994, 0.1214_5835, 0.1438_0469]       ,    ),
					    ]						)



