



"""simple docstring"""






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
   import faiss






@require_faiss
class 	__snake_case	(     _lowercase):




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			tempfile.mkdtemp()
         _lowerCamelCase						:    int         =			8

         # DPR tok
         _lowerCamelCase						:    Union[str, Any]         =			[
             '''[UNK]''',
             '''[CLS]''',
             '''[SEP]''',
             '''[PAD]''',
             '''[MASK]''',
             '''want''',
             '''##want''',
             '''##ed''',
             '''wa''',
             '''un''',
             '''runn''',
             '''##ing''',
             ''',''',
             '''low''',
             '''lowest''',
         ]
         _lowerCamelCase						:    int         =			os.path.join(self.tmpdirname						, '''dpr_tokenizer'''		)
         os.makedirs(__lowerCAmelCase						, exist_ok=__lowerCAmelCase		)
         _lowerCamelCase						:    str         =			os.path.join(__lowerCAmelCase						, DPR_VOCAB_FILES_NAMES['''vocab_file''']		)
         with open(self.vocab_file						, '''w'''						, encoding='''utf-8'''		) as vocab_writer:
                vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]		)		)

         # BART tok
         _lowerCamelCase						:    Optional[Any]         =			[
             '''l''',
             '''o''',
             '''w''',
             '''e''',
             '''r''',
             '''s''',
             '''t''',
             '''i''',
             '''d''',
             '''n''',
             '''\u0120''',
             '''\u0120l''',
             '''\u0120n''',
             '''\u0120lo''',
             '''\u0120low''',
             '''er''',
             '''\u0120lowest''',
             '''\u0120newer''',
             '''\u0120wider''',
             '''<unk>''',
         ]
         _lowerCamelCase						:    List[str]         =			dict(zip(__lowerCAmelCase						, range(len(__lowerCAmelCase		)		)		)		)
         _lowerCamelCase						:    Any         =			['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
         _lowerCamelCase						:    Dict         =			{'''unk_token''': '''<unk>'''}

         _lowerCamelCase						:    Dict         =			os.path.join(self.tmpdirname						, '''bart_tokenizer'''		)
         os.makedirs(__lowerCAmelCase						, exist_ok=__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			os.path.join(__lowerCAmelCase						, BART_VOCAB_FILES_NAMES['''vocab_file''']		)
         _lowerCamelCase						:    Any         =			os.path.join(__lowerCAmelCase						, BART_VOCAB_FILES_NAMES['''merges_file''']		)
         with open(self.vocab_file						, '''w'''						, encoding='''utf-8'''		) as fp:
                fp.write(json.dumps(__lowerCAmelCase		) + '''\n'''		)
         with open(self.merges_file						, '''w'''						, encoding='''utf-8'''		) as fp:
                fp.write('''\n'''.join(__lowerCAmelCase		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname						, '''dpr_tokenizer'''		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname						, '''dpr_tokenizer'''		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname						, '''bart_tokenizer'''		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         shutil.rmtree(self.tmpdirname		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			Dataset.from_dict(
             {
                 '''id''': ['''0''', '''1'''],
                 '''text''': ['''foo''', '''bar'''],
                 '''title''': ['''Foo''', '''Bar'''],
                 '''embeddings''': [np.ones(self.retrieval_vector_size		), 2 * np.ones(self.retrieval_vector_size		)],
             }		)
         dataset.add_faiss_index('''embeddings'''						, string_factory='''Flat'''						, metric_type=faiss.METRIC_INNER_PRODUCT		)
         return dataset




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			self.get_dummy_dataset()
         _lowerCamelCase						:    Union[str, Any]         =			RagConfig(
             retrieval_vector_size=self.retrieval_vector_size						, question_encoder=DPRConfig().to_dict()						, generator=BartConfig().to_dict()						, )
         with patch('''transformers.models.rag.retrieval_rag.load_dataset'''		) as mock_load_dataset:
                _lowerCamelCase						:    Optional[int]         =			dataset
                _lowerCamelCase						:    Any         =			RagRetriever(
                    __lowerCAmelCase						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()						, )
         return retriever




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						bool		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			self.get_dummy_dataset()
         _lowerCamelCase						:    Optional[int]         =			RagConfig(
             retrieval_vector_size=self.retrieval_vector_size						, question_encoder=DPRConfig().to_dict()						, generator=BartConfig().to_dict()						, index_name='''custom'''						, )
         if from_disk:
                _lowerCamelCase						:    Optional[Any]         =			os.path.join(self.tmpdirname						, '''dataset'''		)
                _lowerCamelCase						:    str         =			os.path.join(self.tmpdirname						, '''index.faiss'''		)
                dataset.get_index('''embeddings'''		).save(os.path.join(self.tmpdirname						, '''index.faiss'''		)		)
                dataset.drop_index('''embeddings'''		)
                dataset.save_to_disk(os.path.join(self.tmpdirname						, '''dataset'''		)		)
                del dataset
                _lowerCamelCase						:    Optional[int]         =			RagRetriever(
                    __lowerCAmelCase						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()						, )
         else:
                _lowerCamelCase						:    Union[str, Any]         =			RagRetriever(
                    __lowerCAmelCase						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()						, index=CustomHFIndex(config.retrieval_vector_size						, __lowerCAmelCase		)						, )
         return retriever




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			Dataset.from_dict(
             {
                 '''id''': ['''0''', '''1'''],
                 '''text''': ['''foo''', '''bar'''],
                 '''title''': ['''Foo''', '''Bar'''],
                 '''embeddings''': [np.ones(self.retrieval_vector_size + 1		), 2 * np.ones(self.retrieval_vector_size + 1		)],
             }		)
         dataset.add_faiss_index('''embeddings'''						, string_factory='''Flat'''						, metric_type=faiss.METRIC_INNER_PRODUCT		)

         _lowerCamelCase						:    Any         =			os.path.join(self.tmpdirname						, '''hf_bert_base.hnswSQ8_correct_phi_128.c_index'''		)
         dataset.save_faiss_index('''embeddings'''						, index_file_name + '''.index.dpr'''		)
         pickle.dump(dataset['''id''']						, open(index_file_name + '''.index_meta.dpr'''						, '''wb'''		)		)

         _lowerCamelCase						:    Optional[int]         =			os.path.join(self.tmpdirname						, '''psgs_w100.tsv.pkl'''		)
         _lowerCamelCase						:    List[Any]         =			{sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
         pickle.dump(__lowerCAmelCase						, open(__lowerCAmelCase						, '''wb'''		)		)

         _lowerCamelCase						:    Union[str, Any]         =			RagConfig(
             retrieval_vector_size=self.retrieval_vector_size						, question_encoder=DPRConfig().to_dict()						, generator=BartConfig().to_dict()						, index_name='''legacy'''						, index_path=self.tmpdirname						, )
         _lowerCamelCase						:    str         =			RagRetriever(
             __lowerCAmelCase						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()		)
         return retriever




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






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			1
         _lowerCamelCase						:    Any         =			self.get_dummy_canonical_hf_index_retriever()
         _lowerCamelCase						:    str         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    int         =			retriever.retrieve(__lowerCAmelCase						, n_docs=__lowerCAmelCase		)
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertEqual(len(__lowerCAmelCase		)						, 2		)
         self.assertEqual(sorted(doc_dicts[0]		)						, ['''embeddings''', '''id''', '''text''', '''title''']		)
         self.assertEqual(len(doc_dicts[0]['''id''']		)						, __lowerCAmelCase		)
         self.assertEqual(doc_dicts[0]['''id'''][0]						, '''1'''		)  # max inner product is reached with second doc
         self.assertEqual(doc_dicts[1]['''id'''][0]						, '''0'''		)  # max inner product is reached with first doc
         self.assertListEqual(doc_ids.tolist()						, [[1], [0]]		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			self.get_dummy_canonical_hf_index_retriever()
         with tempfile.TemporaryDirectory() as tmp_dirname:
                with patch('''transformers.models.rag.retrieval_rag.load_dataset'''		) as mock_load_dataset:
                       _lowerCamelCase						:    Dict         =			self.get_dummy_dataset()
                       retriever.save_pretrained(__lowerCAmelCase		)
                       _lowerCamelCase						:    Union[str, Any]         =			RagRetriever.from_pretrained(__lowerCAmelCase		)
                       self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
                       _lowerCamelCase						:    int         =			np.array(
                           [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
                       _lowerCamelCase						:    Dict         =			retriever.retrieve(__lowerCAmelCase						, n_docs=1		)
                       self.assertTrue(out is not None		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			1
         _lowerCamelCase						:    Dict         =			self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    str         =			retriever.retrieve(__lowerCAmelCase						, n_docs=__lowerCAmelCase		)
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertEqual(len(__lowerCAmelCase		)						, 2		)
         self.assertEqual(sorted(doc_dicts[0]		)						, ['''embeddings''', '''id''', '''text''', '''title''']		)
         self.assertEqual(len(doc_dicts[0]['''id''']		)						, __lowerCAmelCase		)
         self.assertEqual(doc_dicts[0]['''id'''][0]						, '''1'''		)  # max inner product is reached with second doc
         self.assertEqual(doc_dicts[1]['''id'''][0]						, '''0'''		)  # max inner product is reached with first doc
         self.assertListEqual(doc_ids.tolist()						, [[1], [0]]		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCAmelCase		)
         with tempfile.TemporaryDirectory() as tmp_dirname:
                retriever.save_pretrained(__lowerCAmelCase		)
                _lowerCamelCase						:    Dict         =			RagRetriever.from_pretrained(__lowerCAmelCase		)
                self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    Dict         =			np.array(
                    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
                _lowerCamelCase						:    List[Any]         =			retriever.retrieve(__lowerCAmelCase						, n_docs=1		)
                self.assertTrue(out is not None		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			1
         _lowerCamelCase						:    Dict         =			self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    Optional[int]         =			retriever.retrieve(__lowerCAmelCase						, n_docs=__lowerCAmelCase		)
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertEqual(len(__lowerCAmelCase		)						, 2		)
         self.assertEqual(sorted(doc_dicts[0]		)						, ['''embeddings''', '''id''', '''text''', '''title''']		)
         self.assertEqual(len(doc_dicts[0]['''id''']		)						, __lowerCAmelCase		)
         self.assertEqual(doc_dicts[0]['''id'''][0]						, '''1'''		)  # max inner product is reached with second doc
         self.assertEqual(doc_dicts[1]['''id'''][0]						, '''0'''		)  # max inner product is reached with first doc
         self.assertListEqual(doc_ids.tolist()						, [[1], [0]]		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCAmelCase		)
         with tempfile.TemporaryDirectory() as tmp_dirname:
                retriever.save_pretrained(__lowerCAmelCase		)
                _lowerCamelCase						:    Union[str, Any]         =			RagRetriever.from_pretrained(__lowerCAmelCase		)
                self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    List[Any]         =			np.array(
                    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
                _lowerCamelCase						:    Dict         =			retriever.retrieve(__lowerCAmelCase						, n_docs=1		)
                self.assertTrue(out is not None		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			1
         _lowerCamelCase						:    int         =			self.get_dummy_legacy_index_retriever()
         _lowerCamelCase						:    List[Any]         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			retriever.retrieve(__lowerCAmelCase						, n_docs=__lowerCAmelCase		)
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertEqual(len(__lowerCAmelCase		)						, 2		)
         self.assertEqual(sorted(doc_dicts[0]		)						, ['''text''', '''title''']		)
         self.assertEqual(len(doc_dicts[0]['''text''']		)						, __lowerCAmelCase		)
         self.assertEqual(doc_dicts[0]['''text'''][0]						, '''bar'''		)  # max inner product is reached with second doc
         self.assertEqual(doc_dicts[1]['''text'''][0]						, '''foo'''		)  # max inner product is reached with first doc
         self.assertListEqual(doc_ids.tolist()						, [[1], [0]]		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.get_dummy_legacy_index_retriever()
         with tempfile.TemporaryDirectory() as tmp_dirname:
                retriever.save_pretrained(__lowerCAmelCase		)
                _lowerCamelCase						:    List[str]         =			RagRetriever.from_pretrained(__lowerCAmelCase		)
                self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    Optional[Any]         =			np.array(
                    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
                _lowerCamelCase						:    Any         =			retriever.retrieve(__lowerCAmelCase						, n_docs=1		)
                self.assertTrue(out is not None		)




  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         import torch

         _lowerCamelCase						:    Optional[int]         =			1
         _lowerCamelCase						:    Optional[Any]         =			self.get_dummy_canonical_hf_index_retriever()
         _lowerCamelCase						:    List[str]         =			[[5, 7], [1_0, 1_1]]
         _lowerCamelCase						:    str         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase						:    Tuple         =			retriever(__lowerCAmelCase						, __lowerCAmelCase						, prefix=retriever.config.generator.prefix						, n_docs=__lowerCAmelCase		)
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    Optional[int]         =			(
             out['''context_input_ids'''],
             out['''context_attention_mask'''],
             out['''retrieved_doc_embeds'''],
         )
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
         self.assertIsInstance(__lowerCAmelCase						, __lowerCAmelCase		)
         self.assertIsInstance(__lowerCAmelCase						, np.ndarray		)

         _lowerCamelCase						:    Optional[int]         =			retriever(
             __lowerCAmelCase						, __lowerCAmelCase						, prefix=retriever.config.generator.prefix						, n_docs=__lowerCAmelCase						, return_tensors='''pt'''						, )
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    str         =			(  # noqa: F841
             out['''context_input_ids'''],
             out['''context_attention_mask'''],
             out['''retrieved_doc_embeds'''],
             out['''doc_ids'''],
         )
         self.assertEqual(retrieved_doc_embeds.shape						, (2, n_docs, self.retrieval_vector_size)		)
         self.assertIsInstance(__lowerCAmelCase						, torch.Tensor		)
         self.assertIsInstance(__lowerCAmelCase						, torch.Tensor		)
         self.assertIsInstance(__lowerCAmelCase						, torch.Tensor		)




  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			self.get_dpr_ctx_encoder_tokenizer()
         _lowerCamelCase						:    Optional[Any]         =			1
         _lowerCamelCase						:    List[str]         =			self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCAmelCase		)
         retriever.set_ctx_encoder_tokenizer(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			[[5, 7], [1_0, 1_1]]
         _lowerCamelCase						:    str         =			np.array(
             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]						, dtype=np.floataa		)
         _lowerCamelCase						:    Optional[int]         =			retriever(__lowerCAmelCase						, __lowerCAmelCase						, prefix=retriever.config.generator.prefix						, n_docs=__lowerCAmelCase		)

         self.assertEqual(
             len(__lowerCAmelCase		)						, 6		)  # check whether the retriever output consist of 6 attributes including tokenized docs
         self.assertEqual(
             all(k in out for k in ('''tokenized_doc_ids''', '''tokenized_doc_attention_mask''')		)						, __lowerCAmelCase		)  # check for doc token related keys in dictionary.







"""simple docstring"""






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

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin






@dataclass
class 	__snake_case	(     _lowercase):
  snake_case__    :       torch.FloatTensor






class 	__snake_case	(     _lowercase  ,					_lowercase):




  @register_to_config
  def __init__(			self					:						int						, __lowerCAmelCase					:						int = 3_2						, __lowerCAmelCase					:						int = 6_4						, __lowerCAmelCase					:						int = 2_0						, __lowerCAmelCase					:						int = 7_6_8						, __lowerCAmelCase					:						Tuple=7_7						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						float = 0.0						, __lowerCAmelCase					:						str = "silu"						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = "linear"						, __lowerCAmelCase					:						Optional[str] = "prd"						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, ):






         """simple docstring"""






         super().__init__()
         _lowerCamelCase						:    Any         =			num_attention_heads
         _lowerCamelCase						:    Optional[int]         =			attention_head_dim
         _lowerCamelCase						:    Dict         =			num_attention_heads * attention_head_dim
         _lowerCamelCase						:    Optional[Any]         =			additional_embeddings

         _lowerCamelCase						:    List[Any]         =			time_embed_dim or inner_dim
         _lowerCamelCase						:    Optional[Any]         =			embedding_proj_dim or embedding_dim
         _lowerCamelCase						:    Tuple         =			clip_embed_dim or embedding_dim

         _lowerCamelCase						:    int         =			Timesteps(__lowerCAmelCase						, __lowerCAmelCase						, 0		)
         _lowerCamelCase						:    str         =			TimestepEmbedding(__lowerCAmelCase						, __lowerCAmelCase						, out_dim=__lowerCAmelCase						, act_fn=__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if embedding_proj_norm_type is None:
                _lowerCamelCase						:    Dict         =			None
         elif embedding_proj_norm_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''		)

         _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if encoder_hid_proj_type is None:
                _lowerCamelCase						:    Any         =			None
         elif encoder_hid_proj_type == "linear":
                _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''		)

         _lowerCamelCase						:    Optional[int]         =			nn.Parameter(torch.zeros(1						, num_embeddings + additional_embeddings						, __lowerCAmelCase		)		)

         if added_emb_type == "prd":
                _lowerCamelCase						:    Optional[Any]         =			nn.Parameter(torch.zeros(1						, 1						, __lowerCAmelCase		)		)
         elif added_emb_type is None:
                _lowerCamelCase						:    Union[str, Any]         =			None
         else:
                raise ValueError(
                    f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''		)

         _lowerCamelCase						:    Tuple         =			nn.ModuleList(
             [
                 BasicTransformerBlock(
                     __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, dropout=__lowerCAmelCase						, activation_fn='''gelu'''						, attention_bias=__lowerCAmelCase						, )
                 for d in range(__lowerCAmelCase		)
             ]		)

         if norm_in_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         elif norm_in_type is None:
                _lowerCamelCase						:    Any         =			None
         else:
                raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''		)

         _lowerCamelCase						:    str         =			nn.LayerNorm(__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    Optional[int]         =			torch.full(
             [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]						, -1_00_00.0		)
         causal_attention_mask.triu_(1		)
         _lowerCamelCase						:    Dict         =			causal_attention_mask[None, ...]
         self.register_buffer('''causal_attention_mask'''						, __lowerCAmelCase						, persistent=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)
         _lowerCamelCase						:    Tuple         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)




  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			{}

         def fn_recursive_add_processors(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Dict[str, AttentionProcessor]		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       _lowerCamelCase						:    str         =			module.processor

                for sub_name, child in module.named_children():
                       fn_recursive_add_processors(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

                return processors

         for name, module in self.named_children():
                fn_recursive_add_processors(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

         return processors




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			len(self.attn_processors.keys()		)

         if isinstance(__lowerCAmelCase						, __lowerCAmelCase		) and len(__lowerCAmelCase		) != count:
                raise ValueError(
                    f'''A dict of processors was passed, but the number of processors {len(__lowerCAmelCase		)} does not match the'''
                    f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''		)

         def fn_recursive_attn_processor(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Tuple		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       if not isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                              module.set_processor(__lowerCAmelCase		)
                       else:
                              module.set_processor(processor.pop(f'''{name}.processor'''		)		)

                for sub_name, child in module.named_children():
                       fn_recursive_attn_processor(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

         for name, module in self.named_children():
                fn_recursive_attn_processor(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         self.set_attn_processor(AttnProcessor()		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Union[torch.Tensor, float, int]						, __lowerCAmelCase					:						torch.FloatTensor						, __lowerCAmelCase					:						Optional[torch.FloatTensor] = None						, __lowerCAmelCase					:						Optional[torch.BoolTensor] = None						, __lowerCAmelCase					:						bool = True						, ):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			hidden_states.shape[0]

         _lowerCamelCase						:    Union[str, Any]         =			timestep
         if not torch.is_tensor(__lowerCAmelCase		):
                _lowerCamelCase						:    Dict         =			torch.tensor([timesteps]						, dtype=torch.long						, device=hidden_states.device		)
         elif torch.is_tensor(__lowerCAmelCase		) and len(timesteps.shape		) == 0:
                _lowerCamelCase						:    str         =			timesteps[None].to(hidden_states.device		)

         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         _lowerCamelCase						:    Tuple         =			timesteps * torch.ones(__lowerCAmelCase						, dtype=timesteps.dtype						, device=timesteps.device		)

         _lowerCamelCase						:    Tuple         =			self.time_proj(__lowerCAmelCase		)

         # timesteps does not contain any weights and will always return f32 tensors
         # but time_embedding might be fp16, so we need to cast here.
         _lowerCamelCase						:    Union[str, Any]         =			timesteps_projected.to(dtype=self.dtype		)
         _lowerCamelCase						:    Any         =			self.time_embedding(__lowerCAmelCase		)

         if self.embedding_proj_norm is not None:
                _lowerCamelCase						:    str         =			self.embedding_proj_norm(__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.embedding_proj(__lowerCAmelCase		)
         if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
                _lowerCamelCase						:    Optional[Any]         =			self.encoder_hidden_states_proj(__lowerCAmelCase		)
         elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
                raise ValueError('''`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'''		)

         _lowerCamelCase						:    str         =			self.proj_in(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			self.positional_embedding.to(hidden_states.dtype		)

         _lowerCamelCase						:    Tuple         =			[]
         _lowerCamelCase						:    Dict         =			0

         if encoder_hidden_states is not None:
                additional_embeds.append(__lowerCAmelCase		)
                additional_embeddings_len += encoder_hidden_states.shape[1]

         if len(proj_embeddings.shape		) == 2:
                _lowerCamelCase						:    str         =			proj_embeddings[:, None, :]

         if len(hidden_states.shape		) == 2:
                _lowerCamelCase						:    Dict         =			hidden_states[:, None, :]

         _lowerCamelCase						:    int         =			additional_embeds + [
             proj_embeddings,
             time_embeddings[:, None, :],
             hidden_states,
         ]

         if self.prd_embedding is not None:
                _lowerCamelCase						:    int         =			self.prd_embedding.to(hidden_states.dtype		).expand(__lowerCAmelCase						, -1						, -1		)
                additional_embeds.append(__lowerCAmelCase		)

         _lowerCamelCase						:    str         =			torch.cat(
             __lowerCAmelCase						, dim=1						, )

         # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
         _lowerCamelCase						:    Dict         =			additional_embeddings_len + proj_embeddings.shape[1] + 1
         if positional_embeddings.shape[1] < hidden_states.shape[1]:
                _lowerCamelCase						:    List[Any]         =			F.pad(
                    __lowerCAmelCase						, (
                        0,
                        0,
                        additional_embeddings_len,
                        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                    )						, value=0.0						, )

         _lowerCamelCase						:    Any         =			hidden_states + positional_embeddings

         if attention_mask is not None:
                _lowerCamelCase						:    Union[str, Any]         =			(1 - attention_mask.to(hidden_states.dtype		)) * -1_00_00.0
                _lowerCamelCase						:    Tuple         =			F.pad(__lowerCAmelCase						, (0, self.additional_embeddings)						, value=0.0		)
                _lowerCamelCase						:    Any         =			(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype		)
                _lowerCamelCase						:    int         =			attention_mask.repeat_interleave(self.config.num_attention_heads						, dim=0		)

         if self.norm_in is not None:
                _lowerCamelCase						:    List[str]         =			self.norm_in(__lowerCAmelCase		)

         for block in self.transformer_blocks:
                _lowerCamelCase						:    List[str]         =			block(__lowerCAmelCase						, attention_mask=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.norm_out(__lowerCAmelCase		)

         if self.prd_embedding is not None:
                _lowerCamelCase						:    Any         =			hidden_states[:, -1]
         else:
                _lowerCamelCase						:    Dict         =			hidden_states[:, additional_embeddings_len:]

         _lowerCamelCase						:    str         =			self.proj_to_clip_embeddings(__lowerCAmelCase		)

         if not return_dict:
                return (predicted_image_embedding,)

         return PriorTransformerOutput(predicted_image_embedding=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			(prior_latents * self.clip_std) + self.clip_mean
         return prior_latents


