




'''simple docstring'''



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


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

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

lowerCamelCase_									=		25_00_04
lowerCamelCase_									=		25_00_20







@require_sentencepiece
@require_tokenizers
class 						lowercase_       (					SCREAMING_SNAKE_CASE_	,    unittest.TestCase  ):


   """simple docstring"""

   lowerCamelCase_							    =		MBartaaTokenizer
   lowerCamelCase_							    =		MBartaaTokenizerFast
   lowerCamelCase_							    =		True
   lowerCamelCase_							    =		True



   def 				lowerCAmelCase_     (       self   :      Any       ):

      """simple docstring"""
      super().setUp()

      # We have a SentencePiece fixture for testing
      _SCREAMING_SNAKE_CASE			= MBartaaTokenizer(snake_case__    , src_lang="en_XX"    , tgt_lang="ro_RO"    , keep_accents=snake_case__       )
      tokenizer.save_pretrained(self.tmpdirname       )



   def 				lowerCAmelCase_     (       self   :      str       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= '<s>'
      _SCREAMING_SNAKE_CASE			= 0

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



   def 				lowerCAmelCase_     (       self   :      Optional[int]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= 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(snake_case__       )    , 1_0_5_4       )



   def 				lowerCAmelCase_     (       self   :      Dict       ):

      """simple docstring"""
      self.assertEqual(self.get_tokenizer().vocab_size    , 1_0_5_4       )



   def 				lowerCAmelCase_     (       self   :      Optional[int]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= MBartaaTokenizer(snake_case__    , src_lang="en_XX"    , tgt_lang="ro_RO"    , keep_accents=snake_case__       )

      _SCREAMING_SNAKE_CASE			= tokenizer.tokenize("This is a test"       )
      self.assertListEqual(snake_case__    , ["▁This", "▁is", "▁a", "▁t", "est"]       )

      self.assertListEqual(
          tokenizer.convert_tokens_to_ids(snake_case__       )    , [value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]    , )

      _SCREAMING_SNAKE_CASE			= tokenizer.tokenize("I was born in 92000, and this is falsé."       )
      self.assertListEqual(
          snake_case__    , [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", "é", "."]    , )
      _SCREAMING_SNAKE_CASE			= tokenizer.convert_tokens_to_ids(snake_case__       )
      self.assertListEqual(
          snake_case__    , [
              value + tokenizer.fairseq_offset
              for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
          ]    , )

      _SCREAMING_SNAKE_CASE			= tokenizer.convert_ids_to_tokens(snake_case__       )
      self.assertListEqual(
          snake_case__    , [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 				lowerCAmelCase_     (       self   :      Union[str, Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= {'input_ids': [[2_5_0_0_0_4, 1_1_0_6_2, 8_2_7_7_2, 7, 1_5, 8_2_7_7_2, 5_3_8, 5_1_5_2_9, 2_3_7, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 2_1_5_1_7_5, 1_3_1_4, 1_3_6, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 5_6_3_5_9, 4_2, 1_2_2_0_0_9, 9, 1_6_4_6_6, 1_6, 8_7_3_4_4, 4_5_3_7, 9, 4_7_1_7, 7_8_3_8_1, 6, 1_5_9_9_5_8, 7, 1_5, 2_4_4_8_0, 6_1_8, 4, 5_2_7, 2_2_6_9_3, 5_4_2_8, 4, 2_7_7_7, 2_4_4_8_0, 9_8_7_4, 4, 4_3_5_2_3, 5_9_4, 4, 8_0_3, 1_8_3_9_2, 3_3_1_8_9, 1_8, 4, 4_3_5_2_3, 2_4_4_4_7, 1_2_3_9_9, 1_0_0, 2_4_9_5_5, 8_3_6_5_8, 9_6_2_6, 1_4_4_0_5_7, 1_5, 8_3_9, 2_2_3_3_5, 1_6, 1_3_6, 2_4_9_5_5, 8_3_6_5_8, 8_3_4_7_9, 1_5, 3_9_1_0_2, 7_2_4, 1_6, 6_7_8, 6_4_5, 2_7_8_9, 1_3_2_8, 4_5_8_9, 4_2, 1_2_2_0_0_9, 1_1_5_7_7_4, 2_3, 8_0_5, 1_3_2_8, 4_6_8_7_6, 7, 1_3_6, 5_3_8_9_4, 1_9_4_0, 4_2_2_2_7, 4_1_1_5_9, 1_7_7_2_1, 8_2_3, 4_2_5, 4, 2_7_5_1_2, 9_8_7_2_2, 2_0_6, 1_3_6, 5_5_3_1, 4_9_7_0, 9_1_9, 1_7_3_3_6, 5, 2], [2_5_0_0_0_4, 2_0_0_8_0, 6_1_8, 8_3, 8_2_7_7_5, 4_7, 4_7_9, 9, 1_5_1_7, 7_3, 5_3_8_9_4, 3_3_3, 8_0_5_8_1, 1_1_0_1_1_7, 1_8_8_1_1, 5_2_5_6, 1_2_9_5, 5_1, 1_5_2_5_2_6, 2_9_7, 7_9_8_6, 3_9_0, 1_2_4_4_1_6, 5_3_8, 3_5_4_3_1, 2_1_4, 9_8, 1_5_0_4_4, 2_5_7_3_7, 1_3_6, 7_1_0_8, 4_3_7_0_1, 2_3, 7_5_6, 1_3_5_3_5_5, 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], [2_5_0_0_0_4, 5_8_1, 6_3_7_7_3, 1_1_9_4_5_5, 6, 1_4_7_7_9_7, 8_8_2_0_3, 7, 6_4_5, 7_0, 2_1, 3_2_8_5, 1_0_2_6_9, 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=snake_case__    , model_name="facebook/mbart-large-50"    , revision="d3913889c59cd5c9e456b269c376325eabad57e2"    , )



   def 				lowerCAmelCase_     (       self   :      Optional[Any]       ):

      """simple docstring"""

      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

      _SCREAMING_SNAKE_CASE			= (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})"""       ):
            _SCREAMING_SNAKE_CASE			= self.rust_tokenizer_class.from_pretrained(snake_case__    , **snake_case__       )
            _SCREAMING_SNAKE_CASE			= self.tokenizer_class.from_pretrained(snake_case__    , **snake_case__       )

            _SCREAMING_SNAKE_CASE			= tempfile.mkdtemp()

            _SCREAMING_SNAKE_CASE			= tokenizer_r.save_pretrained(snake_case__       )
            _SCREAMING_SNAKE_CASE			= tokenizer_p.save_pretrained(snake_case__       )

            # 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       )       )
            _SCREAMING_SNAKE_CASE			= tuple(f for f in tokenizer_r_files if "tokenizer.json" not in f       )
            self.assertSequenceEqual(snake_case__    , snake_case__       )

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

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

            shutil.rmtree(snake_case__       )

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

            _SCREAMING_SNAKE_CASE			= tokenizer_r.save_pretrained(snake_case__    , legacy_format=snake_case__       )
            _SCREAMING_SNAKE_CASE			= tokenizer_p.save_pretrained(snake_case__       )

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

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

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

            shutil.rmtree(snake_case__       )

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

            _SCREAMING_SNAKE_CASE			= tokenizer_r.save_pretrained(snake_case__    , legacy_format=snake_case__       )
            _SCREAMING_SNAKE_CASE			= tokenizer_p.save_pretrained(snake_case__       )

            # 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
            _SCREAMING_SNAKE_CASE			= tokenizer_r.from_pretrained(snake_case__       )
            _SCREAMING_SNAKE_CASE			= tokenizer_p.from_pretrained(snake_case__       )

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

            shutil.rmtree(snake_case__       )







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


   """simple docstring"""

   lowerCamelCase_							    =		'''facebook/mbart-large-50-one-to-many-mmt'''
   lowerCamelCase_							    =		[
       ''' 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.''',
   ]
   lowerCamelCase_							    =		[
       '''Ş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.''',
   ]
   lowerCamelCase_							    =		[EN_CODE, 8274, 127873, 25916, 7, 8622, 2071, 438, 67485, 53, 187895, 23, 51712, 2]



   @classmethod
   def 				lowerCAmelCase_     (       cls   :      int       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= MBartaaTokenizer.from_pretrained(
          cls.checkpoint_name    , src_lang="en_XX"    , tgt_lang="ro_RO"       )
      _SCREAMING_SNAKE_CASE			= 1
      return cls



   def 				lowerCAmelCase_     (       self   :      List[Any]       ):

      """simple docstring"""
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ar_AR"]    , 2_5_0_0_0_1       )
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["en_EN"]    , 2_5_0_0_0_4       )
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ro_RO"]    , 2_5_0_0_2_0       )
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["mr_IN"]    , 2_5_0_0_3_8       )



   def 				lowerCAmelCase_     (       self   :      Optional[Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.tokenizer.batch_encode_plus(self.src_text       ).input_ids[0]
      self.assertListEqual(self.expected_src_tokens    , snake_case__       )



   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      self.assertIn(snake_case__    , self.tokenizer.all_special_ids       )
      _SCREAMING_SNAKE_CASE			= [RO_CODE, 8_8_4, 9_0_1_9, 9_6, 9, 9_1_6, 8_6_7_9_2, 3_6, 1_8_7_4_3, 1_5_5_9_6, 5, 2]
      _SCREAMING_SNAKE_CASE			= self.tokenizer.decode(snake_case__    , skip_special_tokens=snake_case__       )
      _SCREAMING_SNAKE_CASE			= self.tokenizer.decode(generated_ids[1:]    , skip_special_tokens=snake_case__       )
      self.assertEqual(snake_case__    , snake_case__       )
      self.assertNotIn(self.tokenizer.eos_token    , snake_case__       )



   def 				lowerCAmelCase_     (       self   :      Union[str, Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= ['this is gunna be a long sentence ' * 2_0]
      assert isinstance(src_text[0]    , snake_case__       )
      _SCREAMING_SNAKE_CASE			= 1_0
      _SCREAMING_SNAKE_CASE			= self.tokenizer(snake_case__    , max_length=snake_case__    , truncation=snake_case__       ).input_ids[0]
      self.assertEqual(ids[0]    , snake_case__       )
      self.assertEqual(ids[-1]    , 2       )
      self.assertEqual(len(snake_case__       )    , snake_case__       )



   def 				lowerCAmelCase_     (       self   :      Any       ):

      """simple docstring"""
      self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "ar_AR"]       )    , [2_5_0_0_5_3, 2_5_0_0_0_1]       )



   def 				lowerCAmelCase_     (       self   :      Optional[int]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= tempfile.mkdtemp()
      _SCREAMING_SNAKE_CASE			= self.tokenizer.fairseq_tokens_to_ids
      self.tokenizer.save_pretrained(snake_case__       )
      _SCREAMING_SNAKE_CASE			= MBartaaTokenizer.from_pretrained(snake_case__       )
      self.assertDictEqual(new_tok.fairseq_tokens_to_ids    , snake_case__       )



   @require_torch
   def 				lowerCAmelCase_     (       self   :      Union[str, Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.tokenizer(self.src_text    , text_target=self.tgt_text    , padding=snake_case__    , return_tensors="pt"       )
      _SCREAMING_SNAKE_CASE			= 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 				lowerCAmelCase_     (       self   :      List[str]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.tokenizer(
          self.src_text    , text_target=self.tgt_text    , padding=snake_case__    , truncation=snake_case__    , max_length=len(self.expected_src_tokens       )    , return_tensors="pt"    , )
      _SCREAMING_SNAKE_CASE			= shift_tokens_right(batch["labels"]    , self.tokenizer.pad_token_id       )

      self.assertIsInstance(snake_case__    , snake_case__       )

      self.assertEqual((2, 1_4)    , batch.input_ids.shape       )
      self.assertEqual((2, 1_4)    , batch.attention_mask.shape       )
      _SCREAMING_SNAKE_CASE			= batch.input_ids.tolist()[0]
      self.assertListEqual(self.expected_src_tokens    , snake_case__       )
      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 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.tokenizer(self.src_text    , padding=snake_case__    , truncation=snake_case__    , max_length=3    , return_tensors="pt"       )
      _SCREAMING_SNAKE_CASE			= self.tokenizer(
          text_target=self.tgt_text    , padding=snake_case__    , truncation=snake_case__    , max_length=1_0    , return_tensors="pt"       )
      _SCREAMING_SNAKE_CASE			= targets['input_ids']
      _SCREAMING_SNAKE_CASE			= shift_tokens_right(snake_case__    , self.tokenizer.pad_token_id       )

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



   @require_torch
   def 				lowerCAmelCase_     (       self   :      Tuple       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.tokenizer._build_translation_inputs(
          "A test"    , return_tensors="pt"    , src_lang="en_XX"    , tgt_lang="ar_AR"       )

      self.assertEqual(
          nested_simplify(snake_case__       )    , {
              # en_XX, A, test, EOS
              "input_ids": [[2_5_0_0_0_4, 6_2, 3_0_3_4, 2]],
              "attention_mask": [[1, 1, 1, 1]],
              # ar_AR
              "forced_bos_token_id": 2_5_0_0_0_1,
          }    , )





'''simple docstring'''







from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block
@dataclass
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     __magic_name__           =     42
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     @register_to_config
     def __init__(			self   ,						snake_case__ = 6_5536   ,						snake_case__ = None   ,						snake_case__ = 2   ,						snake_case__ = 2   ,						snake_case__ = 0   ,						snake_case__ = "fourier"   ,						snake_case__ = True   ,						snake_case__ = False   ,						snake_case__ = 0.0   ,						snake_case__ = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")   ,						snake_case__ = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")   ,						snake_case__ = "UNetMidBlock1D"   ,						snake_case__ = None   ,						snake_case__ = (32, 32, 64)   ,						snake_case__ = None   ,						snake_case__ = 8   ,						snake_case__ = 1   ,						snake_case__ = False   ,						):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase     :				Optional[int]       =							sample_size

            # time
            if time_embedding_type == "fourier":
                   _lowerCAmelCase     :				Tuple       =							GaussianFourierProjection(
                       embedding_size=8   ,						set_W_to_weight=snake_case__   ,						log=snake_case__   ,						flip_sin_to_cos=snake_case__ )
                   _lowerCAmelCase     :				List[Any]       =							2 * block_out_channels[0]
            elif time_embedding_type == "positional":
                   _lowerCAmelCase     :				Optional[Any]       =							Timesteps(
                       block_out_channels[0]   ,						flip_sin_to_cos=snake_case__   ,						downscale_freq_shift=snake_case__ )
                   _lowerCAmelCase     :				Optional[int]       =							block_out_channels[0]

            if use_timestep_embedding:
                   _lowerCAmelCase     :				Any       =							block_out_channels[0] * 4
                   _lowerCAmelCase     :				List[str]       =							TimestepEmbedding(
                       in_channels=snake_case__   ,						time_embed_dim=snake_case__   ,						act_fn=snake_case__   ,						out_dim=block_out_channels[0]   ,						)

            _lowerCAmelCase     :				Tuple       =							nn.ModuleList([] )
            _lowerCAmelCase     :				Tuple       =							None
            _lowerCAmelCase     :				Optional[Any]       =							nn.ModuleList([] )
            _lowerCAmelCase     :				str       =							None

            # down
            _lowerCAmelCase     :				List[str]       =							in_channels
            for i, down_block_type in enumerate(snake_case__ ):
                   _lowerCAmelCase     :				str       =							output_channel
                   _lowerCAmelCase     :				Any       =							block_out_channels[i]

                   if i == 0:
                          input_channel += extra_in_channels

                   _lowerCAmelCase     :				Tuple       =							i == len(snake_case__ ) - 1

                   _lowerCAmelCase     :				List[str]       =							get_down_block(
                       snake_case__   ,						num_layers=snake_case__   ,						in_channels=snake_case__   ,						out_channels=snake_case__   ,						temb_channels=block_out_channels[0]   ,						add_downsample=not is_final_block or downsample_each_block   ,						)
                   self.down_blocks.append(snake_case__ )

            # mid
            _lowerCAmelCase     :				List[str]       =							get_mid_block(
                snake_case__   ,						in_channels=block_out_channels[-1]   ,						mid_channels=block_out_channels[-1]   ,						out_channels=block_out_channels[-1]   ,						embed_dim=block_out_channels[0]   ,						num_layers=snake_case__   ,						add_downsample=snake_case__   ,						)

            # up
            _lowerCAmelCase     :				Union[str, Any]       =							list(reversed(snake_case__ ) )
            _lowerCAmelCase     :				List[Any]       =							reversed_block_out_channels[0]
            if out_block_type is None:
                   _lowerCAmelCase     :				Tuple       =							out_channels
            else:
                   _lowerCAmelCase     :				Tuple       =							block_out_channels[0]

            for i, up_block_type in enumerate(snake_case__ ):
                   _lowerCAmelCase     :				int       =							output_channel
                   _lowerCAmelCase     :				Dict       =							(
                       reversed_block_out_channels[i + 1] if i < len(snake_case__ ) - 1 else final_upsample_channels
                   )

                   _lowerCAmelCase     :				Union[str, Any]       =							i == len(snake_case__ ) - 1

                   _lowerCAmelCase     :				Tuple       =							get_up_block(
                       snake_case__   ,						num_layers=snake_case__   ,						in_channels=snake_case__   ,						out_channels=snake_case__   ,						temb_channels=block_out_channels[0]   ,						add_upsample=not is_final_block   ,						)
                   self.up_blocks.append(snake_case__ )
                   _lowerCAmelCase     :				List[str]       =							output_channel

            # out
            _lowerCAmelCase     :				Tuple       =							norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4   ,						32 )
            _lowerCAmelCase     :				Any       =							get_out_block(
                out_block_type=snake_case__   ,						num_groups_out=snake_case__   ,						embed_dim=block_out_channels[0]   ,						out_channels=snake_case__   ,						act_fn=snake_case__   ,						fc_dim=block_out_channels[-1] // 4   ,						)







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__ = True   ,						):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							timestep
            if not torch.is_tensor(snake_case__ ):
                   _lowerCAmelCase     :				Optional[int]       =							torch.tensor([timesteps]   ,						dtype=torch.long   ,						device=sample.device )
            elif torch.is_tensor(snake_case__ ) and len(timesteps.shape ) == 0:
                   _lowerCAmelCase     :				List[Any]       =							timesteps[None].to(sample.device )

            _lowerCAmelCase     :				int       =							self.time_proj(snake_case__ )
            if self.config.use_timestep_embedding:
                   _lowerCAmelCase     :				Optional[int]       =							self.time_mlp(snake_case__ )
            else:
                   _lowerCAmelCase     :				List[str]       =							timestep_embed[..., None]
                   _lowerCAmelCase     :				Dict       =							timestep_embed.repeat([1, 1, sample.shape[2]] ).to(sample.dtype )
                   _lowerCAmelCase     :				Tuple       =							timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:]) )

            # 2. down
            _lowerCAmelCase     :				str       =							()
            for downsample_block in self.down_blocks:
                   _lowerCAmelCase      ,  _lowerCAmelCase     :				Union[str, Any]       =							downsample_block(hidden_states=snake_case__   ,						temb=snake_case__ )
                   down_block_res_samples += res_samples

            # 3. mid
            if self.mid_block:
                   _lowerCAmelCase     :				Any       =							self.mid_block(snake_case__   ,						snake_case__ )

            # 4. up
            for i, upsample_block in enumerate(self.up_blocks ):
                   _lowerCAmelCase     :				Union[str, Any]       =							down_block_res_samples[-1:]
                   _lowerCAmelCase     :				Union[str, Any]       =							down_block_res_samples[:-1]
                   _lowerCAmelCase     :				Union[str, Any]       =							upsample_block(snake_case__   ,						res_hidden_states_tuple=snake_case__   ,						temb=snake_case__ )

            # 5. post-process
            if self.out_block:
                   _lowerCAmelCase     :				Optional[Any]       =							self.out_block(snake_case__   ,						snake_case__ )

            if not return_dict:
                   return (sample,)

            return UNetaDOutput(sample=snake_case__ )

