


import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartaaTokenizer, MBartaaTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)

from ...test_tokenization_common import TokenizerTesterMixin


_A			=						get_tests_dir('fixtures/test_sentencepiece.model')

if is_torch_available():
   from transformers.models.mbart.modeling_mbart import shift_tokens_right

_A			=						25_0004
_A			=						25_0020






@require_sentencepiece
@require_tokenizers
class    UpperCAmelCase__	( A_       ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  MBartaaTokenizer
    UpperCAmelCase__						:		Optional[int]              =  MBartaaTokenizerFast
    UpperCAmelCase__						:		Tuple              =  True
    UpperCAmelCase__						:		List[Any]              =  True






    def _a (      self   )  ->	List[Any]:
         super().setUp()

         # We have a SentencePiece fixture for testing
         __UpperCamelCase   							=MBartaaTokenizer(A_ ,       src_lang='en_XX' ,       tgt_lang='ro_RO' ,       keep_accents=A_   )
         tokenizer.save_pretrained(self.tmpdirname   )






    def _a (      self   )  ->	List[Any]:
         __UpperCamelCase   							='<s>'
         __UpperCamelCase   							=0

         self.assertEqual(self.get_tokenizer()._convert_token_to_id(A_   ) ,       A_   )
         self.assertEqual(self.get_tokenizer()._convert_id_to_token(A_   ) ,       A_   )






    def _a (      self   )  ->	str:
         __UpperCamelCase   							=list(self.get_tokenizer().get_vocab().keys()   )

         self.assertEqual(vocab_keys[0] ,       '<s>'   )
         self.assertEqual(vocab_keys[1] ,       '<pad>'   )
         self.assertEqual(vocab_keys[-1] ,       '<mask>'   )
         self.assertEqual(len(A_   ) ,       1054   )






    def _a (      self   )  ->	Any:
         self.assertEqual(self.get_tokenizer().vocab_size ,       1054   )






    def _a (      self   )  ->	Optional[int]:
         __UpperCamelCase   							=MBartaaTokenizer(A_ ,       src_lang='en_XX' ,       tgt_lang='ro_RO' ,       keep_accents=A_   )

         __UpperCamelCase   							=tokenizer.tokenize('This is a test'   )
         self.assertListEqual(A_ ,       ['▁This', '▁is', '▁a', '▁t', 'est']   )

         self.assertListEqual(
             tokenizer.convert_tokens_to_ids(A_   ) ,       [value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]] ,       )

         __UpperCamelCase   							=tokenizer.tokenize('I was born in 92000, and this is falsé.'   )
         self.assertListEqual(
             A_ ,       [SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '9', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', 'é', '.'] ,       )
         __UpperCamelCase   							=tokenizer.convert_tokens_to_ids(A_   )
         self.assertListEqual(
             A_ ,       [
                 value + tokenizer.fairseq_offset
                 for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
             ] ,       )

         __UpperCamelCase   							=tokenizer.convert_ids_to_tokens(A_   )
         self.assertListEqual(
             A_ ,       [SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '<unk>', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', '<unk>', '.'] ,       )






    @slow
    def _a (      self   )  ->	int:
         # fmt: off
         __UpperCamelCase   							={'input_ids': [[250004, 11062, 82772, 7, 15, 82772, 538, 51529, 237, 17198, 1290, 206, 9, 215175, 1314, 136, 17198, 1290, 206, 9, 56359, 42, 122009, 9, 16466, 16, 87344, 4537, 9, 4717, 78381, 6, 159958, 7, 15, 24480, 618, 4, 527, 22693, 5428, 4, 2777, 24480, 9874, 4, 43523, 594, 4, 803, 18392, 33189, 18, 4, 43523, 24447, 12399, 100, 24955, 83658, 9626, 144057, 15, 839, 22335, 16, 136, 24955, 83658, 83479, 15, 39102, 724, 16, 678, 645, 2789, 1328, 4589, 42, 122009, 115774, 23, 805, 1328, 46876, 7, 136, 53894, 1940, 42227, 41159, 17721, 823, 425, 4, 27512, 98722, 206, 136, 5531, 4970, 919, 17336, 5, 2], [250004, 20080, 618, 83, 82775, 47, 479, 9, 1517, 73, 53894, 333, 80581, 110117, 18811, 5256, 1295, 51, 152526, 297, 7986, 390, 124416, 538, 35431, 214, 98, 15044, 25737, 136, 7108, 43701, 23, 756, 135355, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [250004, 581, 63773, 119455, 6, 147797, 88203, 7, 645, 70, 21, 3285, 10269, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
         # fmt: on

         self.tokenizer_integration_test_util(
             expected_encoding=A_ ,       model_name='facebook/mbart-large-50' ,       revision='d3913889c59cd5c9e456b269c376325eabad57e2' ,       )






    def _a (      self   )  ->	Union[str, Any]:
         if not self.test_slow_tokenizer:
              # as we don't have a slow version, we can't compare the outputs between slow and fast versions
              return

         __UpperCamelCase   							=(self.rust_tokenizer_class, 'hf-internal-testing/tiny-random-mbart50', {})
         for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'   ):
                   __UpperCamelCase   							=self.rust_tokenizer_class.from_pretrained(A_ ,       **A_   )
                   __UpperCamelCase   							=self.tokenizer_class.from_pretrained(A_ ,       **A_   )

                   __UpperCamelCase   							=tempfile.mkdtemp()

                   __UpperCamelCase   							=tokenizer_r.save_pretrained(A_   )
                   __UpperCamelCase   							=tokenizer_p.save_pretrained(A_   )

                   # Checks it save with the same files + the tokenizer.json file for the fast one
                   self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files   )   )
                   __UpperCamelCase   							=tuple(f for f in tokenizer_r_files if 'tokenizer.json' not in f   )
                   self.assertSequenceEqual(A_ ,       A_   )

                   # Checks everything loads correctly in the same way
                   __UpperCamelCase   							=tokenizer_r.from_pretrained(A_   )
                   __UpperCamelCase   							=tokenizer_p.from_pretrained(A_   )

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(A_ ,       A_   )   )
                        # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
                        # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

                   shutil.rmtree(A_   )

                   # Save tokenizer rust, legacy_format=True
                   __UpperCamelCase   							=tempfile.mkdtemp()

                   __UpperCamelCase   							=tokenizer_r.save_pretrained(A_ ,       legacy_format=A_   )
                   __UpperCamelCase   							=tokenizer_p.save_pretrained(A_   )

                   # Checks it save with the same files
                   self.assertSequenceEqual(A_ ,       A_   )

                   # Checks everything loads correctly in the same way
                   __UpperCamelCase   							=tokenizer_r.from_pretrained(A_   )
                   __UpperCamelCase   							=tokenizer_p.from_pretrained(A_   )

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(A_ ,       A_   )   )

                   shutil.rmtree(A_   )

                   # Save tokenizer rust, legacy_format=False
                   __UpperCamelCase   							=tempfile.mkdtemp()

                   __UpperCamelCase   							=tokenizer_r.save_pretrained(A_ ,       legacy_format=A_   )
                   __UpperCamelCase   							=tokenizer_p.save_pretrained(A_   )

                   # Checks it saved the tokenizer.json file
                   self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files   )   )

                   # Checks everything loads correctly in the same way
                   __UpperCamelCase   							=tokenizer_r.from_pretrained(A_   )
                   __UpperCamelCase   							=tokenizer_p.from_pretrained(A_   )

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(A_ ,       A_   )   )

                   shutil.rmtree(A_   )






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



    """simple docstring"""






    UpperCAmelCase__						:		Any              =  "facebook/mbart-large-50-one-to-many-mmt"
    UpperCAmelCase__						:		Tuple              =  [
        " UN Chief Says There Is No Military Solution in Syria",
        " Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.",
    ]
    UpperCAmelCase__						:		Tuple              =  [
        "Şeful ONU declară că nu există o soluţie militară în Siria",
        "Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei"
        " pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor"
        " face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.",
    ]
    UpperCAmelCase__						:		List[str]              =  [EN_CODE, 8_2_7_4, 1_2_7_8_7_3, 2_5_9_1_6, 7, 8_6_2_2, 2_0_7_1, 4_3_8, 6_7_4_8_5, 5_3, 1_8_7_8_9_5, 2_3, 5_1_7_1_2, 2]






    @classmethod
    def _a (      cls   )  ->	Optional[Any]:
         __UpperCamelCase   							=MBartaaTokenizer.from_pretrained(
             cls.checkpoint_name ,       src_lang='en_XX' ,       tgt_lang='ro_RO'   )
         __UpperCamelCase   							=1
         return cls






    def _a (      self   )  ->	Optional[int]:
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ar_AR'] ,       250001   )
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['en_EN'] ,       250004   )
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ro_RO'] ,       250020   )
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['mr_IN'] ,       250038   )






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=self.tokenizer.batch_encode_plus(self.src_text   ).input_ids[0]
         self.assertListEqual(self.expected_src_tokens ,       A_   )






    def _a (      self   )  ->	List[str]:
         self.assertIn(A_ ,       self.tokenizer.all_special_ids   )
         __UpperCamelCase   							=[RO_CODE, 884, 9019, 96, 9, 916, 86792, 36, 18743, 15596, 5, 2]
         __UpperCamelCase   							=self.tokenizer.decode(A_ ,       skip_special_tokens=A_   )
         __UpperCamelCase   							=self.tokenizer.decode(generated_ids[1:] ,       skip_special_tokens=A_   )
         self.assertEqual(A_ ,       A_   )
         self.assertNotIn(self.tokenizer.eos_token ,       A_   )






    def _a (      self   )  ->	Any:
         __UpperCamelCase   							=['this is gunna be a long sentence ' * 20]
         assert isinstance(src_text[0] ,       A_   )
         __UpperCamelCase   							=10
         __UpperCamelCase   							=self.tokenizer(A_ ,       max_length=A_ ,       truncation=A_   ).input_ids[0]
         self.assertEqual(ids[0] ,       A_   )
         self.assertEqual(ids[-1] ,       2   )
         self.assertEqual(len(A_   ) ,       A_   )






    def _a (      self   )  ->	int:
         self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']   ) ,       [250053, 250001]   )






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							=tempfile.mkdtemp()
         __UpperCamelCase   							=self.tokenizer.fairseq_tokens_to_ids
         self.tokenizer.save_pretrained(A_   )
         __UpperCamelCase   							=MBartaaTokenizer.from_pretrained(A_   )
         self.assertDictEqual(new_tok.fairseq_tokens_to_ids ,       A_   )






    @require_torch
    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=self.tokenizer(self.src_text ,       text_target=self.tgt_text ,       padding=A_ ,       return_tensors='pt'   )
         __UpperCamelCase   							=shift_tokens_right(batch['labels'] ,       self.tokenizer.pad_token_id   )

         # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
         assert batch.input_ids[1][0] == EN_CODE
         assert batch.input_ids[1][-1] == 2
         assert batch.labels[1][0] == RO_CODE
         assert batch.labels[1][-1] == 2
         assert batch.decoder_input_ids[1][:2].tolist() == [2, RO_CODE]






    @require_torch
    def _a (      self   )  ->	List[Any]:
         __UpperCamelCase   							=self.tokenizer(
             self.src_text ,       text_target=self.tgt_text ,       padding=A_ ,       truncation=A_ ,       max_length=len(self.expected_src_tokens   ) ,       return_tensors='pt' ,       )
         __UpperCamelCase   							=shift_tokens_right(batch['labels'] ,       self.tokenizer.pad_token_id   )

         self.assertIsInstance(A_ ,       A_   )

         self.assertEqual((2, 14) ,       batch.input_ids.shape   )
         self.assertEqual((2, 14) ,       batch.attention_mask.shape   )
         __UpperCamelCase   							=batch.input_ids.tolist()[0]
         self.assertListEqual(self.expected_src_tokens ,       A_   )
         self.assertEqual(2 ,       batch.decoder_input_ids[0, 0]   )  # decoder_start_token_id
         # Test that special tokens are reset
         self.assertEqual(self.tokenizer.prefix_tokens ,       [EN_CODE]   )
         self.assertEqual(self.tokenizer.suffix_tokens ,       [self.tokenizer.eos_token_id]   )






    def _a (      self   )  ->	Optional[int]:
         __UpperCamelCase   							=self.tokenizer(self.src_text ,       padding=A_ ,       truncation=A_ ,       max_length=3 ,       return_tensors='pt'   )
         __UpperCamelCase   							=self.tokenizer(
             text_target=self.tgt_text ,       padding=A_ ,       truncation=A_ ,       max_length=10 ,       return_tensors='pt'   )
         __UpperCamelCase   							=targets['input_ids']
         __UpperCamelCase   							=shift_tokens_right(A_ ,       self.tokenizer.pad_token_id   )

         self.assertEqual(batch.input_ids.shape[1] ,       3   )
         self.assertEqual(batch.decoder_input_ids.shape[1] ,       10   )






    @require_torch
    def _a (      self   )  ->	Any:
         __UpperCamelCase   							=self.tokenizer._build_translation_inputs(
             'A test' ,       return_tensors='pt' ,       src_lang='en_XX' ,       tgt_lang='ar_AR'   )

         self.assertEqual(
             nested_simplify(A_   ) ,       {
                 # en_XX, A, test, EOS
                 'input_ids': [[250004, 62, 3034, 2]],
                 'attention_mask': [[1, 1, 1, 1]],
                 # ar_AR
                 'forced_bos_token_id': 250001,
             } ,       )









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


_A			=						{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 ,       A_   )  ->	Optional[int]:
         super().__init__()
         __UpperCamelCase   							=torchvision.models.resnetaaa(pretrained=A_   )
         __UpperCamelCase   							=list(model.children()   )[:-2]
         __UpperCamelCase   							=nn.Sequential(*A_   )
         __UpperCamelCase   							=nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds]   )






    def _a (      self ,       A_   )  ->	int:
         # Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx2048
         __UpperCamelCase   							=self.pool(self.model(A_   )   )
         __UpperCamelCase   							=torch.flatten(A_ ,       start_dim=2   )
         __UpperCamelCase   							=out.transpose(1 ,       2   ).contiguous()
         return out  # BxNx2048






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	List[str]:
         __UpperCamelCase   							=[json.loads(A_   ) for l in open(A_   )]
         __UpperCamelCase   							=os.path.dirname(A_   )
         __UpperCamelCase   							=tokenizer
         __UpperCamelCase   							=labels
         __UpperCamelCase   							=len(A_   )
         __UpperCamelCase   							=max_seq_length

         __UpperCamelCase   							=transforms






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






    def __getitem__(      self ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=torch.LongTensor(self.tokenizer.encode(self.data[index]['text'] ,       add_special_tokens=A_   )   )
         __UpperCamelCase							,				__UpperCamelCase							,				__UpperCamelCase   							=sentence[0], sentence[1:-1], sentence[-1]
         __UpperCamelCase   							=sentence[: self.max_seq_length]

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

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

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






    def _a (      self   )  ->	List[str]:
         __UpperCamelCase   							=Counter()
         for row in self.data:
              label_freqs.update(row['label']   )
         return label_freqs

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Union[str, Any]    ):
     __UpperCamelCase   							=[len(row['sentence']    ) for row in batch]
     __UpperCamelCase							,				__UpperCamelCase   							=len(SCREAMING_SNAKE_CASE__    ), max(SCREAMING_SNAKE_CASE__    )

     __UpperCamelCase   							=torch.zeros(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					dtype=torch.long    )
     __UpperCamelCase   							=torch.zeros(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					dtype=torch.long    )

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

     __UpperCamelCase   							=torch.stack([row['image'] for row in batch]    )
     __UpperCamelCase   							=torch.stack([row['label'] for row in batch]    )
     __UpperCamelCase   							=torch.stack([row['image_start_token'] for row in batch]    )
     __UpperCamelCase   							=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 						_UpperCAmelCase    (      ):
     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 						_UpperCAmelCase    (      ):
     return transforms.Compose(
         [
             transforms.Resize(2_56    ),
             transforms.CenterCrop(2_24    ),
             transforms.ToTensor(),
             transforms.Normalize(
                 mean=[0.46777044, 0.44531429, 0.40661017]      ,					std=[0.12221994, 0.12145835, 0.14380469]      ,					),
         ]    )





