





"""simple docstring"""



import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           MODEL_FOR_PRETRAINING_MAPPING,
           MegatronBertForCausalLM,
           MegatronBertForMaskedLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
           MegatronBertModel,
       )






class 						__snake_case   :





 """simple docstring"""


 def __init__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase=13   ,		__lowerCamelCase=7   ,		__lowerCamelCase=True   ,		__lowerCamelCase=True   ,		__lowerCamelCase=True   ,		__lowerCamelCase=True   ,		__lowerCamelCase=99   ,		__lowerCamelCase=64   ,		__lowerCamelCase=32   ,		__lowerCamelCase=5   ,		__lowerCamelCase=4   ,		__lowerCamelCase=37   ,		__lowerCamelCase="gelu"   ,		__lowerCamelCase=0.1   ,		__lowerCamelCase=0.1   ,		__lowerCamelCase=512   ,		__lowerCamelCase=16   ,		__lowerCamelCase=2   ,		__lowerCamelCase=0.0_2   ,		__lowerCamelCase=3   ,		__lowerCamelCase=4   ,		__lowerCamelCase=None   ,		):



     '''simple docstring'''




     __A       :		str							=	parent
     __A       :		str							=	batch_size
     __A       :		Tuple							=	seq_length
     __A       :		Optional[int]							=	is_training
     __A       :		int							=	use_input_mask
     __A       :		List[str]							=	use_token_type_ids
     __A       :		str							=	use_labels
     __A       :		Union[str, Any]							=	vocab_size
     __A       :		str							=	hidden_size
     __A       :		Tuple							=	embedding_size
     __A       :		Dict							=	num_hidden_layers
     __A       :		Dict							=	num_attention_heads
     __A       :		Any							=	intermediate_size
     __A       :		Union[str, Any]							=	hidden_act
     __A       :		List[str]							=	hidden_dropout_prob
     __A       :		Tuple							=	attention_probs_dropout_prob
     __A       :		Optional[int]							=	max_position_embeddings
     __A       :		Optional[int]							=	type_vocab_size
     __A       :		Optional[int]							=	type_sequence_label_size
     __A       :		Optional[int]							=	initializer_range
     __A       :		Dict							=	num_labels
     __A       :		Tuple							=	num_choices
     __A       :		Any							=	scope

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Tuple							=	ids_tensor([self.batch_size, self.seq_length]   ,		self.vocab_size    )

     __A       :		Optional[Any]							=	None
     if self.use_input_mask:
         __A       :		List[str]							=	random_attention_mask([self.batch_size, self.seq_length]    )

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

     __A       :		Union[str, Any]							=	None
     __A       :		Union[str, Any]							=	None
     __A       :		Dict							=	None
     if self.use_labels:
         __A       :		Any							=	ids_tensor([self.batch_size]   ,		self.type_sequence_label_size    )
         __A       :		List[Any]							=	ids_tensor([self.batch_size, self.seq_length]   ,		self.num_labels    )
         __A       :		Tuple							=	ids_tensor([self.batch_size]   ,		self.num_choices    )

     __A       :		List[Any]							=	self.get_config()

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

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     return MegatronBertConfig(
         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   ,		embedding_size=self.embedding_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   ,		is_decoder=lowerCamelCase_   ,		initializer_range=self.initializer_range   ,		)

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		List[Any]							=	MegatronBertModel(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		List[Any]							=	model(lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_    )
     __A       :		List[str]							=	model(lowerCamelCase_   ,		token_type_ids=lowerCamelCase_    )
     __A       :		Tuple							=	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    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		str							=	MegatronBertForMaskedLM(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		List[str]							=	model(lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_    )
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.vocab_size)    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		List[str]							=	MegatronBertForCausalLM(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		List[Any]							=	model(lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_    )
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.vocab_size)    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Any							=	MegatronBertForNextSentencePrediction(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		Dict							=	model(
         lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_   ,		)
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, 2)    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		int							=	MegatronBertForPreTraining(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		str							=	model(
         lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_   ,		next_sentence_label=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    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Optional[int]							=	MegatronBertForQuestionAnswering(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		Optional[int]							=	model(
         lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		start_positions=lowerCamelCase_   ,		end_positions=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    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Optional[Any]							=	self.num_labels
     __A       :		str							=	MegatronBertForSequenceClassification(lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		Any							=	model(lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_    )
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.num_labels)    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		str							=	self.num_labels
     __A       :		Union[str, Any]							=	MegatronBertForTokenClassification(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		Optional[Any]							=	model(lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_    )
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.num_labels)    )

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Optional[Any]							=	self.num_choices
     __A       :		List[Any]							=	MegatronBertForMultipleChoice(config=lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.eval()
     __A       :		List[str]							=	input_ids.unsqueeze(1    ).expand(-1   ,		self.num_choices   ,		-1    ).contiguous()
     __A       :		Union[str, Any]							=	token_type_ids.unsqueeze(1    ).expand(-1   ,		self.num_choices   ,		-1    ).contiguous()
     __A       :		int							=	input_mask.unsqueeze(1    ).expand(-1   ,		self.num_choices   ,		-1    ).contiguous()
     __A       :		List[Any]							=	model(
         lowerCamelCase_   ,		attention_mask=lowerCamelCase_   ,		token_type_ids=lowerCamelCase_   ,		labels=lowerCamelCase_   ,		)
     self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.num_choices)    )







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		str							=	self.prepare_config_and_inputs()
     (
         __A							
     )       :		List[Any]							=	config_and_inputs
     __A       :		Union[str, Any]							=	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
     return config, inputs_dict






@require_torch
class 						__snake_case   (						lowercase_				,				lowercase_				,				unittest.TestCase	):





 """simple docstring"""


 _lowerCamelCase								=	(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 _lowerCamelCase								=	(
     {
         """feature-extraction""": MegatronBertModel,
         """fill-mask""": MegatronBertForMaskedLM,
         """question-answering""": MegatronBertForQuestionAnswering,
         """text-classification""": MegatronBertForSequenceClassification,
         """text-generation""": MegatronBertForCausalLM,
         """token-classification""": MegatronBertForTokenClassification,
         """zero-shot""": MegatronBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 _lowerCamelCase								=	True
 # test_resize_embeddings = False
 _lowerCamelCase								=	False

 def    UpperCamelCase__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase=False    ):



     '''simple docstring'''




     __A       :		Optional[Any]							=	super()._prepare_for_class(lowerCamelCase_   ,		lowerCamelCase_   ,		return_labels=lowerCamelCase_    )

     if return_labels:
         if model_class in get_values(lowerCamelCase_    ):
             __A       :		Dict							=	torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.seq_length)   ,		dtype=torch.long   ,		device=lowerCamelCase_    )
             __A       :		Union[str, Any]							=	torch.zeros(
                 self.model_tester.batch_size   ,		dtype=torch.long   ,		device=lowerCamelCase_    )
     return inputs_dict

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Any							=	MegatronBertModelTester(self    )
     __A       :		Dict							=	ConfigTester(self   ,		config_class=lowerCamelCase_   ,		hidden_size=37    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     self.config_tester.run_common_tests()

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Optional[Any]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Union[str, Any]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Optional[int]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Optional[Any]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Dict							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Dict							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase_    )

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		List[Any]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase_    )







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Optional[int]							=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase_    )







def 						__lowercase	(  snake_case_    :    Any		)  ->Optional[int]:



    '''simple docstring'''







    return torch.tensor(
        lowerCamelCase_							,dtype=torch.long							,device=lowerCamelCase_							,)


a_					=     1e-4






@require_torch
@require_sentencepiece
@require_tokenizers
class 						__snake_case   (						unittest.TestCase	):





 """simple docstring"""


 @slow
 @unittest.skip('''Model is not available.'''    )
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		List[str]							=	"""nvidia/megatron-bert-uncased-345m"""
     if "MYDIR" in os.environ:
         __A       :		str							=	os.path.join(os.environ['''MYDIR''']   ,		lowerCamelCase_    )
     __A       :		List[Any]							=	MegatronBertModel.from_pretrained(lowerCamelCase_    )
     model.to(lowerCamelCase_    )
     model.half()
     __A       :		Optional[Any]							=	_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]    )
     with torch.no_grad():
         __A       :		List[str]							=	model(lowerCamelCase_    )[0]
     __A       :		Any							=	torch.Size((1, 9, 1024)    )
     self.assertEqual(output.shape   ,		lowerCamelCase_    )

     __A       :		int							=	[-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
     for ii in range(3    ):
         for jj in range(3    ):
             __A       :		Dict							=	output[0, ii, jj]
             __A       :		Union[str, Any]							=	expected[3 * ii + jj]
             __A       :		Optional[int]							=	"""ii={} jj={} a={} b={}""".format(lowerCamelCase_   ,		lowerCamelCase_   ,		lowerCamelCase_   ,		lowerCamelCase_    )
             self.assertTrue(math.isclose(lowerCamelCase_   ,		lowerCamelCase_   ,		rel_tol=lowerCamelCase_   ,		abs_tol=lowerCamelCase_    )   ,		msg=lowerCamelCase_    )











'''simple docstring'''




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			):




					"""simple docstring"""


					def __init__(						self				:			Dict      ,						lowerCamelCase_				:			Union[str, Any]						):

							'''simple docstring'''



							super().__init__()
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torchvision.models.resnetaaa(pretrained=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Any		=			list(model.children()						)[:-2]
							SCREAMING_SNAKE_CASE :				Any		=			nn.Sequential(*lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds]						)



					def lowerCamelCase_			(						self				:			List[str]      ,						lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			self.pool(self.model(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				int		=			torch.flatten(lowerCamelCase_      ,						start_dim=2						)
							SCREAMING_SNAKE_CASE :				List[Any]		=			out.transpose(1      ,						2						).contiguous()
							return out  # BxNx2048







class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					def __init__(						self				:			str      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			Tuple						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			[json.loads(lowerCamelCase_						) for l in open(lowerCamelCase_						)]
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			os.path.dirname(lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Tuple		=			tokenizer
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			labels
							SCREAMING_SNAKE_CASE :				str		=			len(lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				int		=			max_seq_length

							SCREAMING_SNAKE_CASE :				List[Any]		=			transforms


					def __len__(						self				:			Any						):

							'''simple docstring'''



							return len(self.data						)


					def __getitem__(						self				:			Optional[int]      ,						lowerCamelCase_				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			torch.LongTensor(self.tokenizer.encode(self.data[index]["""text"""]      ,						add_special_tokens=lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE :				Tuple		=			sentence[0], sentence[1:-1], sentence[-1]
							SCREAMING_SNAKE_CASE :				int		=			sentence[: self.max_seq_length]

							SCREAMING_SNAKE_CASE :				List[Any]		=			torch.zeros(self.n_classes						)
							SCREAMING_SNAKE_CASE :				List[str]		=			1

							SCREAMING_SNAKE_CASE :				int		=			Image.open(os.path.join(self.data_dir      ,						self.data[index]["""img"""]						)						).convert("""RGB"""						)
							SCREAMING_SNAKE_CASE :				str		=			self.transforms(lowerCamelCase_						)

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



					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Any		=			Counter()
							for row in self.data:
									label_freqs.update(row["""label"""]						)
							return label_freqs







def 		__A		(  lowerCamelCase_	):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE :				Tuple		=			[len(row["""sentence"""]	) for row in batch]
		SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE :				Any		=			len(lowerCamelCase_	), max(lowerCamelCase_	)

		SCREAMING_SNAKE_CASE :				int		=			torch.zeros(lowerCamelCase_			,			lowerCamelCase_			,			dtype=torch.long	)
		SCREAMING_SNAKE_CASE :				Any		=			torch.zeros(lowerCamelCase_			,			lowerCamelCase_			,			dtype=torch.long	)

		for i_batch, (input_row, length) in enumerate(zip(lowerCamelCase_			,			lowerCamelCase_	)	):
				SCREAMING_SNAKE_CASE :				Union[str, Any]		=			input_row["""sentence"""]
				SCREAMING_SNAKE_CASE :				Union[str, Any]		=			1

		SCREAMING_SNAKE_CASE :				List[str]		=			torch.stack([row["""image"""] for row in batch]	)
		SCREAMING_SNAKE_CASE :				int		=			torch.stack([row["""label"""] for row in batch]	)
		SCREAMING_SNAKE_CASE :				str		=			torch.stack([row["""image_start_token"""] for row in batch]	)
		SCREAMING_SNAKE_CASE :				Any		=			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 		__A		(  ):


		"""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 		__A		(  ):


		"""simple docstring"""



		return transforms.Compose(
		    [
		        transforms.Resize(2_56	),
		        transforms.CenterCrop(2_24	),
		        transforms.ToTensor(),
		        transforms.Normalize(
		            mean=[0.46_777_044, 0.44_531_429, 0.40_661_017]			,			std=[0.12_221_994, 0.12_145_835, 0.14_380_469]			,			),
		    ]	)



