




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   A(						UpperCamelCase       ):



 '''simple docstring'''


 def     a__			(    self :  Any  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   tempfile.mkdtemp()
       lowerCamelCase_							=   8

       # DPR tok
       lowerCamelCase_							=   [
           '[UNK]',
           '[CLS]',
           '[SEP]',
           '[PAD]',
           '[MASK]',
           'want',
           '##want',
           '##ed',
           'wa',
           'un',
           'runn',
           '##ing',
           ',',
           'low',
           'lowest',
       ]
       lowerCamelCase_							=   os.path.join(self.tmpdirname						,							'dpr_tokenizer'  )
       os.makedirs(A_						,							exist_ok=A_  )
       lowerCamelCase_							=   os.path.join(A_						,							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_							=   [
           'l',
           'o',
           'w',
           'e',
           'r',
           's',
           't',
           'i',
           'd',
           'n',
           '\u0120',
           '\u0120l',
           '\u0120n',
           '\u0120lo',
           '\u0120low',
           'er',
           '\u0120lowest',
           '\u0120newer',
           '\u0120wider',
           '<unk>',
       ]
       lowerCamelCase_							=   dict(zip(A_						,							range(len(A_  )  )  )  )
       lowerCamelCase_							=   ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
       lowerCamelCase_							=   {'unk_token': '<unk>'}

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

 def     a__			(    self :  int  )       ->     DPRQuestionEncoderTokenizer:

       """simple docstring"""






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

 def     a__			(    self :  Tuple  )       ->     DPRContextEncoderTokenizer:

       """simple docstring"""






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

 def     a__			(    self :  List[str]  )       ->     BartTokenizer:

       """simple docstring"""






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

 def     a__			(    self :  Any  )       ->     Dict:

       """simple docstring"""






       shutil.rmtree(self.tmpdirname  )

 def     a__			(    self :  Optional[int]  )       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   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     a__			(    self :  List[Any]  )       ->     Tuple:

       """simple docstring"""






       lowerCamelCase_							=   self.get_dummy_dataset()
       lowerCamelCase_							=   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_							=   dataset
             lowerCamelCase_							=   RagRetriever(
                 A_						,							question_encoder_tokenizer=self.get_dpr_tokenizer()						,							generator_tokenizer=self.get_bart_tokenizer()						,							)
       return retriever

 def     a__			(    self :  List[Any]						,							A_ :  bool  )       ->     Tuple:

       """simple docstring"""






       lowerCamelCase_							=   self.get_dummy_dataset()
       lowerCamelCase_							=   RagConfig(
           retrieval_vector_size=self.retrieval_vector_size						,							question_encoder=DPRConfig().to_dict()						,							generator=BartConfig().to_dict()						,							index_name='custom'						,							)
       if from_disk:
             lowerCamelCase_							=   os.path.join(self.tmpdirname						,							'dataset'  )
             lowerCamelCase_							=   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_							=   RagRetriever(
                 A_						,							question_encoder_tokenizer=self.get_dpr_tokenizer()						,							generator_tokenizer=self.get_bart_tokenizer()						,							)
       else:
             lowerCamelCase_							=   RagRetriever(
                 A_						,							question_encoder_tokenizer=self.get_dpr_tokenizer()						,							generator_tokenizer=self.get_bart_tokenizer()						,							index=CustomHFIndex(config.retrieval_vector_size						,							A_  )						,							)
       return retriever

 def     a__			(    self :  Optional[int]  )       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   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_							=   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_							=   os.path.join(self.tmpdirname						,							'psgs_w100.tsv.pkl'  )
       lowerCamelCase_							=   {sample['id']: [sample['text'], sample['title']] for sample in dataset}
       pickle.dump(A_						,							open(A_						,							'wb'  )  )

       lowerCamelCase_							=   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_							=   RagRetriever(
           A_						,							question_encoder_tokenizer=self.get_dpr_tokenizer()						,							generator_tokenizer=self.get_bart_tokenizer()  )
       return retriever

 def     a__			(    self :  Any  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_canonical_hf_index_retriever()
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   retriever.retrieve(A_						,							n_docs=A_  )
       self.assertEqual(retrieved_doc_embeds.shape						,							(2, n_docs, self.retrieval_vector_size)  )
       self.assertEqual(len(A_  )						,							2  )
       self.assertEqual(sorted(doc_dicts[0]  )						,							['embeddings', 'id', 'text', 'title']  )
       self.assertEqual(len(doc_dicts[0]['id']  )						,							A_  )
       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     a__			(    self :  str  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   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_							=   self.get_dummy_dataset()
                   retriever.save_pretrained(A_  )
                   lowerCamelCase_							=   RagRetriever.from_pretrained(A_  )
                   self.assertIsInstance(A_						,							A_  )
                   lowerCamelCase_							=   np.array(
                       [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
                   lowerCamelCase_							=   retriever.retrieve(A_						,							n_docs=1  )
                   self.assertTrue(out is not None  )

 def     a__			(    self :  Optional[Any]  )       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_custom_hf_index_retriever(from_disk=A_  )
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   retriever.retrieve(A_						,							n_docs=A_  )
       self.assertEqual(retrieved_doc_embeds.shape						,							(2, n_docs, self.retrieval_vector_size)  )
       self.assertEqual(len(A_  )						,							2  )
       self.assertEqual(sorted(doc_dicts[0]  )						,							['embeddings', 'id', 'text', 'title']  )
       self.assertEqual(len(doc_dicts[0]['id']  )						,							A_  )
       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     a__			(    self :  str  )       ->     List[Any]:

       """simple docstring"""






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

 def     a__			(    self :  Dict  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_custom_hf_index_retriever(from_disk=A_  )
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   retriever.retrieve(A_						,							n_docs=A_  )
       self.assertEqual(retrieved_doc_embeds.shape						,							(2, n_docs, self.retrieval_vector_size)  )
       self.assertEqual(len(A_  )						,							2  )
       self.assertEqual(sorted(doc_dicts[0]  )						,							['embeddings', 'id', 'text', 'title']  )
       self.assertEqual(len(doc_dicts[0]['id']  )						,							A_  )
       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     a__			(    self :  List[Any]  )       ->     Optional[Any]:

       """simple docstring"""






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

 def     a__			(    self :  str  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_legacy_index_retriever()
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   retriever.retrieve(A_						,							n_docs=A_  )
       self.assertEqual(retrieved_doc_embeds.shape						,							(2, n_docs, self.retrieval_vector_size)  )
       self.assertEqual(len(A_  )						,							2  )
       self.assertEqual(sorted(doc_dicts[0]  )						,							['text', 'title']  )
       self.assertEqual(len(doc_dicts[0]['text']  )						,							A_  )
       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     a__			(    self :  Dict  )       ->     Any:

       """simple docstring"""






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

 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def     a__			(    self :  str  )       ->     Union[str, Any]:

       """simple docstring"""






       import torch

       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_canonical_hf_index_retriever()
       lowerCamelCase_							=   [[5, 7], [10, 11]]
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_							=   retriever(A_						,							A_						,							prefix=retriever.config.generator.prefix						,							n_docs=A_  )
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   (
           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(A_						,							A_  )
       self.assertIsInstance(A_						,							A_  )
       self.assertIsInstance(A_						,							np.ndarray  )

       lowerCamelCase_							=   retriever(
           A_						,							A_						,							prefix=retriever.config.generator.prefix						,							n_docs=A_						,							return_tensors='pt'						,							)
       lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   (  # 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(A_						,							torch.Tensor  )
       self.assertIsInstance(A_						,							torch.Tensor  )
       self.assertIsInstance(A_						,							torch.Tensor  )





 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def     a__			(    self :  Union[str, Any]  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   self.get_dpr_ctx_encoder_tokenizer()
       lowerCamelCase_							=   1
       lowerCamelCase_							=   self.get_dummy_custom_hf_index_retriever(from_disk=A_  )
       retriever.set_ctx_encoder_tokenizer(A_  )

       lowerCamelCase_							=   [[5, 7], [10, 11]]
       lowerCamelCase_							=   np.array(
           [np.ones(self.retrieval_vector_size  ), -np.ones(self.retrieval_vector_size  )]						,							dtype=np.floataa  )
       lowerCamelCase_							=   retriever(A_						,							A_						,							prefix=retriever.config.generator.prefix						,							n_docs=A_  )

       self.assertEqual(
           len(A_  )						,							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')  )						,							A_  )  # check for doc token related keys in dictionary.









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   A(						UpperCamelCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						42







class   A(						UpperCamelCase							,				UpperCamelCase       ):



 '''simple docstring'''


 @register_to_config
 def __init__(    self :  Tuple						,							A_ :  int = 32						,							A_ :  int = 64						,							A_ :  int = 20						,							A_ :  int = 768						,							A_ :  Optional[Any]=77						,							A_ :  Optional[int]=4						,							A_ :  float = 0.0						,							A_ :  str = "silu"						,							A_ :  Optional[str] = None						,							A_ :  Optional[str] = None						,							A_ :  Optional[str] = "linear"						,							A_ :  Optional[str] = "prd"						,							A_ :  Optional[int] = None						,							A_ :  Optional[int] = None						,							A_ :  Optional[int] = None						,							)       ->     List[Any]:

       """simple docstring"""






       super().__init__()
       lowerCamelCase_							=   num_attention_heads
       lowerCamelCase_							=   attention_head_dim
       lowerCamelCase_							=   num_attention_heads * attention_head_dim
       lowerCamelCase_							=   additional_embeddings

       lowerCamelCase_							=   time_embed_dim or inner_dim
       lowerCamelCase_							=   embedding_proj_dim or embedding_dim
       lowerCamelCase_							=   clip_embed_dim or embedding_dim

       lowerCamelCase_							=   Timesteps(A_						,							A_						,							0  )
       lowerCamelCase_							=   TimestepEmbedding(A_						,							A_						,							out_dim=A_						,							act_fn=A_  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       if embedding_proj_norm_type is None:
             lowerCamelCase_							=   None
       elif embedding_proj_norm_type == "layer":
             lowerCamelCase_							=   nn.LayerNorm(A_  )
       else:
             raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       if encoder_hid_proj_type is None:
             lowerCamelCase_							=   None
       elif encoder_hid_proj_type == "linear":
             lowerCamelCase_							=   nn.Linear(A_						,							A_  )
       else:
             raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""  )

       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							num_embeddings + additional_embeddings						,							A_  )  )

       if added_emb_type == "prd":
             lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							1						,							A_  )  )
       elif added_emb_type is None:
             lowerCamelCase_							=   None
       else:
             raise ValueError(
                 f"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""  )

       lowerCamelCase_							=   nn.ModuleList(
           [
               BasicTransformerBlock(
                   A_						,							A_						,							A_						,							dropout=A_						,							activation_fn='gelu'						,							attention_bias=A_						,							)
               for d in range(A_  )
           ]  )

       if norm_in_type == "layer":
             lowerCamelCase_							=   nn.LayerNorm(A_  )
       elif norm_in_type is None:
             lowerCamelCase_							=   None
       else:
             raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""  )

       lowerCamelCase_							=   nn.LayerNorm(A_  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       lowerCamelCase_							=   torch.full(
           [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]						,							-10000.0  )
       causal_attention_mask.triu_(1  )
       lowerCamelCase_							=   causal_attention_mask[None, ...]
       self.register_buffer('causal_attention_mask'						,							A_						,							persistent=A_  )

       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							A_  )  )
       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							A_  )  )

 @property
 # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
 def     a__			(    self :  str  )       ->     Dict[str, AttentionProcessor]:

       """simple docstring"""






       lowerCamelCase_							=   {}

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

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

             return processors

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

       return processors

 def     a__			(    self :  List[Any]						,							A_ :  Union[AttentionProcessor, Dict[str, AttentionProcessor]]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   len(self.attn_processors.keys()  )

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

       def fn_recursive_attn_processor(A_ :  str						,							A_ :  torch.nn.Module						,							A_ :  Union[str, Any]  ):
             if hasattr(A_						,							'set_processor'  ):
                   if not isinstance(A_						,							A_  ):
                         module.set_processor(A_  )
                   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}"""						,							A_						,							A_  )

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

 def     a__			(    self :  List[Any]  )       ->     List[Any]:

       """simple docstring"""






       self.set_attn_processor(AttnProcessor()  )

 def     a__			(    self :  Dict						,							A_ :  List[Any]						,							A_ :  Union[torch.Tensor, float, int]						,							A_ :  torch.FloatTensor						,							A_ :  Optional[torch.FloatTensor] = None						,							A_ :  Optional[torch.BoolTensor] = None						,							A_ :  bool = True						,							)       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   hidden_states.shape[0]

       lowerCamelCase_							=   timestep
       if not torch.is_tensor(A_  ):
             lowerCamelCase_							=   torch.tensor([timesteps]						,							dtype=torch.long						,							device=hidden_states.device  )
       elif torch.is_tensor(A_  ) and len(timesteps.shape  ) == 0:
             lowerCamelCase_							=   timesteps[None].to(hidden_states.device  )

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

       lowerCamelCase_							=   self.time_proj(A_  )

       # 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_							=   timesteps_projected.to(dtype=self.dtype  )
       lowerCamelCase_							=   self.time_embedding(A_  )

       if self.embedding_proj_norm is not None:
             lowerCamelCase_							=   self.embedding_proj_norm(A_  )

       lowerCamelCase_							=   self.embedding_proj(A_  )
       if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
             lowerCamelCase_							=   self.encoder_hidden_states_proj(A_  )
       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_							=   self.proj_in(A_  )

       lowerCamelCase_							=   self.positional_embedding.to(hidden_states.dtype  )

       lowerCamelCase_							=   []
       lowerCamelCase_							=   0

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

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

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

       lowerCamelCase_							=   additional_embeds + [
           proj_embeddings,
           time_embeddings[:, None, :],
           hidden_states,
       ]

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

       lowerCamelCase_							=   torch.cat(
           A_						,							dim=1						,							)

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

       lowerCamelCase_							=   hidden_states + positional_embeddings

       if attention_mask is not None:
             lowerCamelCase_							=   (1 - attention_mask.to(hidden_states.dtype  )) * -10000.0
             lowerCamelCase_							=   F.pad(A_						,							(0, self.additional_embeddings)						,							value=0.0  )
             lowerCamelCase_							=   (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype  )
             lowerCamelCase_							=   attention_mask.repeat_interleave(self.config.num_attention_heads						,							dim=0  )

       if self.norm_in is not None:
             lowerCamelCase_							=   self.norm_in(A_  )

       for block in self.transformer_blocks:
             lowerCamelCase_							=   block(A_						,							attention_mask=A_  )

       lowerCamelCase_							=   self.norm_out(A_  )

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

       lowerCamelCase_							=   self.proj_to_clip_embeddings(A_  )

       if not return_dict:
             return (predicted_image_embedding,)

       return PriorTransformerOutput(predicted_image_embedding=A_  )





 def     a__			(    self :  Tuple						,							A_ :  List[Any]  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   (prior_latents * self.clip_std) + self.clip_mean
       return prior_latents



