





import json
import os
import shutil
import tempfile
from unittest import TestCase

from transformers import BartTokenizer, BartTokenizerFast, DPRQuestionEncoderTokenizer, DPRQuestionEncoderTokenizerFast
from transformers.models.bart.configuration_bart import BartConfig
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.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_tokenizers, require_torch, slow
from transformers.utils import is_datasets_available, is_faiss_available, is_torch_available


if is_torch_available() and is_datasets_available() and is_faiss_available():
  from transformers.models.rag.configuration_rag import RagConfig
  from transformers.models.rag.tokenization_rag import RagTokenizer







@require_faiss
@require_torch
class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       def 							UpperCAmelCase__			(  self			:   Any		):
         __snake_case:						str			  =						tempfile.mkdtemp()
         __snake_case:						Union[str, Any]			  =						8

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

         __snake_case:						List[Any]			  =						os.path.join(self.tmpdirname							,    """bart_tokenizer"""		)
         os.makedirs(lowerCamelCase_							,    exist_ok=lowerCamelCase_		)
         __snake_case:						Union[str, Any]			  =						os.path.join(lowerCamelCase_							,    BART_VOCAB_FILES_NAMES["""vocab_file"""]		)
         __snake_case:						Optional[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 							UpperCAmelCase__			(  self			:   Optional[int]		):
         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,    """dpr_tokenizer"""		)		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname							,    """bart_tokenizer"""		)		)




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         shutil.rmtree(self.tmpdirname		)




       @require_tokenizers
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Dict			  =						os.path.join(self.tmpdirname							,    """rag_tokenizer"""		)
         __snake_case:						List[str]			  =						RagConfig(question_encoder=DPRConfig().to_dict()							,    generator=BartConfig().to_dict()		)
         __snake_case:						Optional[Any]			  =						RagTokenizer(question_encoder=self.get_dpr_tokenizer()							,    generator=self.get_bart_tokenizer()		)
         rag_config.save_pretrained(lowerCamelCase_		)
         rag_tokenizer.save_pretrained(lowerCamelCase_		)
         __snake_case:						Any			  =						RagTokenizer.from_pretrained(lowerCamelCase_							,    config=lowerCamelCase_		)
         self.assertIsInstance(new_rag_tokenizer.question_encoder							,    lowerCamelCase_		)
         self.assertEqual(new_rag_tokenizer.question_encoder.get_vocab()							,    rag_tokenizer.question_encoder.get_vocab()		)
         self.assertIsInstance(new_rag_tokenizer.generator							,    lowerCamelCase_		)
         self.assertEqual(new_rag_tokenizer.generator.get_vocab()							,    rag_tokenizer.generator.get_vocab()		)




       @slow
       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case:						Optional[int]			  =						RagTokenizer.from_pretrained("""facebook/rag-token-nq"""		)
         __snake_case:						Tuple			  =						[
             """who got the first nobel prize in physics""",
             """when is the next deadpool movie being released""",
             """which mode is used for short wave broadcast service""",
             """who is the owner of reading football club""",
             """when is the next scandal episode coming out""",
             """when is the last time the philadelphia won the superbowl""",
             """what is the most current adobe flash player version""",
             """how many episodes are there in dragon ball z""",
             """what is the first step in the evolution of the eye""",
             """where is gall bladder situated in human body""",
             """what is the main mineral in lithium batteries""",
             """who is the president of usa right now""",
             """where do the greasers live in the outsiders""",
             """panda is a national animal of which country""",
             """what is the name of manchester united stadium""",
         ]
         __snake_case:						Optional[int]			  =						tokenizer(lowerCamelCase_		)
         self.assertIsNotNone(lowerCamelCase_		)






       @slow
       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case:						Optional[Any]			  =						RagTokenizer.from_pretrained("""facebook/rag-sequence-nq"""		)
         __snake_case:						int			  =						[
             """who got the first nobel prize in physics""",
             """when is the next deadpool movie being released""",
             """which mode is used for short wave broadcast service""",
             """who is the owner of reading football club""",
             """when is the next scandal episode coming out""",
             """when is the last time the philadelphia won the superbowl""",
             """what is the most current adobe flash player version""",
             """how many episodes are there in dragon ball z""",
             """what is the first step in the evolution of the eye""",
             """where is gall bladder situated in human body""",
             """what is the main mineral in lithium batteries""",
             """who is the president of usa right now""",
             """where do the greasers live in the outsiders""",
             """panda is a national animal of which country""",
             """what is the name of manchester united stadium""",
         ]
         __snake_case:						Tuple			  =						tokenizer(lowerCamelCase_		)
         self.assertIsNotNone(lowerCamelCase_		)













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

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder







@dataclass
class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					42







class    __snake_case	(							__lowerCamelCase       ,      __lowerCamelCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					True




       @register_to_config
       def __init__(  self			:   Union[str, Any]							,    A			:   int = 3							,    A			:   int = 3							,    A			:   Tuple[str] = ("DownEncoderBlock2D",)							,    A			:   Tuple[str] = ("UpDecoderBlock2D",)							,    A			:   Tuple[int] = (64,)							,    A			:   int = 1							,    A			:   str = "silu"							,    A			:   int = 4							,    A			:   int = 32							,    A			:   int = 32							,    A			:   float = 0.1_8215							,    ):
         super().__init__()

         # pass init params to Encoder
         __snake_case:						Any			  =						Encoder(
             in_channels=A							,    out_channels=A							,    down_block_types=A							,    block_out_channels=A							,    layers_per_block=A							,    act_fn=A							,    norm_num_groups=A							,    double_z=A							,    )

         # pass init params to Decoder
         __snake_case:						int			  =						Decoder(
             in_channels=A							,    out_channels=A							,    up_block_types=A							,    block_out_channels=A							,    layers_per_block=A							,    norm_num_groups=A							,    act_fn=A							,    )

         __snake_case:						Dict			  =						nn.Convad(2 * latent_channels							,    2 * latent_channels							,    1		)
         __snake_case:						int			  =						nn.Convad(A							,    A							,    1		)

         __snake_case:						List[str]			  =						False
         __snake_case:						Optional[int]			  =						False

         # only relevant if vae tiling is enabled
         __snake_case:						Any			  =						self.config.sample_size
         __snake_case:						int			  =						(
             self.config.sample_size[0]
             if isinstance(self.config.sample_size							,    (list, tuple)		)
             else self.config.sample_size
         )
         __snake_case:						Union[str, Any]			  =						int(sample_size / (2 ** (len(self.config.block_out_channels		) - 1))		)
         __snake_case:						Optional[int]			  =						0.25




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   List[str]							,    A			:   Optional[Any]=False		):
         if isinstance(A							,    (Encoder, Decoder)		):
           __snake_case:						str			  =						value




       def 							UpperCAmelCase__			(  self			:   str							,    A			:   bool = True		):
         __snake_case:						Union[str, Any]			  =						use_tiling




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         self.enable_tiling(A		)




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						List[str]			  =						True




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						List[str]			  =						False




       @property
       # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Any			  =						{}

         def fn_recursive_add_processors(A			:   str							,    A			:   torch.nn.Module							,    A			:   Dict[str, AttentionProcessor]		):
           if hasattr(A							,    """set_processor"""		):
             __snake_case:						List[Any]			  =						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 							UpperCAmelCase__			(  self			:   Optional[int]							,    A			:   Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):
         __snake_case:						Any			  =						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			:   Optional[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 							UpperCAmelCase__			(  self			:   List[str]		):
         self.set_attn_processor(AttnProcessor()		)




       @apply_forward_hook
       def 							UpperCAmelCase__			(  self			:   Optional[Any]							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
           return self.tiled_encode(A							,    return_dict=A		)

         if self.use_slicing and x.shape[0] > 1:
           __snake_case:						List[Any]			  =						[self.encoder(A		) for x_slice in x.split(1		)]
           __snake_case:						Optional[Any]			  =						torch.cat(A		)
         else:
           __snake_case:						str			  =						self.encoder(A		)

         __snake_case:						Any			  =						self.quant_conv(A		)
         __snake_case:						Tuple			  =						DiagonalGaussianDistribution(A		)

         if not return_dict:
           return (posterior,)

         return AutoencoderKLOutput(latent_dist=A		)




       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
           return self.tiled_decode(A							,    return_dict=A		)

         __snake_case:						Optional[int]			  =						self.post_quant_conv(A		)
         __snake_case:						Union[str, Any]			  =						self.decoder(A		)

         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)




       @apply_forward_hook
       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         if self.use_slicing and z.shape[0] > 1:
           __snake_case:						Union[str, Any]			  =						[self._decode(A		).sample for z_slice in z.split(1		)]
           __snake_case:						List[str]			  =						torch.cat(A		)
         else:
           __snake_case:						int			  =						self._decode(A		).sample

         if not return_dict:
           return (decoded,)

         return DecoderOutput(sample=A		)




       def 							UpperCAmelCase__			(  self			:   Any							,    A			:   Tuple							,    A			:   int							,    A			:   List[Any]		):
         __snake_case:						int			  =						min(a.shape[2]							,    b.shape[2]							,    A		)
         for y in range(A		):
           __snake_case:						Dict			  =						a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
         return b




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   Optional[Any]							,    A			:   List[str]							,    A			:   List[str]		):
         __snake_case:						Dict			  =						min(a.shape[3]							,    b.shape[3]							,    A		)
         for x in range(A		):
           __snake_case:						Tuple			  =						a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
         return b




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         __snake_case:						List[str]			  =						int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)		)
         __snake_case:						Dict			  =						int(self.tile_latent_min_size * self.tile_overlap_factor		)
         __snake_case:						Dict			  =						self.tile_latent_min_size - blend_extent

         # Split the image into 512x512 tiles and encode them separately.
         __snake_case:						Optional[int]			  =						[]
         for i in range(0							,    x.shape[2]							,    A		):
           __snake_case:						Optional[int]			  =						[]
           for j in range(0							,    x.shape[3]							,    A		):
             __snake_case:						int			  =						x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
             __snake_case:						Tuple			  =						self.encoder(A		)
             __snake_case:						Dict			  =						self.quant_conv(A		)
             row.append(A		)
           rows.append(A		)
         __snake_case:						Tuple			  =						[]
         for i, row in enumerate(A		):
           __snake_case:						str			  =						[]
           for j, tile in enumerate(A		):
             # blend the above tile and the left tile
             # to the current tile and add the current tile to the result row
             if i > 0:
               __snake_case:						Optional[Any]			  =						self.blend_v(rows[i - 1][j]							,    A							,    A		)
             if j > 0:
               __snake_case:						Tuple			  =						self.blend_h(row[j - 1]							,    A							,    A		)
             result_row.append(tile[:, :, :row_limit, :row_limit]		)
           result_rows.append(torch.cat(A							,    dim=3		)		)

         __snake_case:						Tuple			  =						torch.cat(A							,    dim=2		)
         __snake_case:						Optional[int]			  =						DiagonalGaussianDistribution(A		)

         if not return_dict:
           return (posterior,)

         return AutoencoderKLOutput(latent_dist=A		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   torch.FloatTensor							,    A			:   bool = True		):
         __snake_case:						Optional[Any]			  =						int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)		)
         __snake_case:						str			  =						int(self.tile_sample_min_size * self.tile_overlap_factor		)
         __snake_case:						int			  =						self.tile_sample_min_size - blend_extent

         # Split z into overlapping 64x64 tiles and decode them separately.
         # The tiles have an overlap to avoid seams between tiles.
         __snake_case:						List[Any]			  =						[]
         for i in range(0							,    z.shape[2]							,    A		):
           __snake_case:						Optional[Any]			  =						[]
           for j in range(0							,    z.shape[3]							,    A		):
             __snake_case:						Dict			  =						z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
             __snake_case:						Any			  =						self.post_quant_conv(A		)
             __snake_case:						Optional[Any]			  =						self.decoder(A		)
             row.append(A		)
           rows.append(A		)
         __snake_case:						Optional[Any]			  =						[]
         for i, row in enumerate(A		):
           __snake_case:						Optional[Any]			  =						[]
           for j, tile in enumerate(A		):
             # blend the above tile and the left tile
             # to the current tile and add the current tile to the result row
             if i > 0:
               __snake_case:						Tuple			  =						self.blend_v(rows[i - 1][j]							,    A							,    A		)
             if j > 0:
               __snake_case:						List[str]			  =						self.blend_h(row[j - 1]							,    A							,    A		)
             result_row.append(tile[:, :, :row_limit, :row_limit]		)
           result_rows.append(torch.cat(A							,    dim=3		)		)

         __snake_case:						Dict			  =						torch.cat(A							,    dim=2		)
         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)






       def 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   torch.FloatTensor							,    A			:   bool = False							,    A			:   bool = True							,    A			:   Optional[torch.Generator] = None							,    ):
         __snake_case:						Optional[Any]			  =						sample
         __snake_case:						Union[str, Any]			  =						self.encode(A		).latent_dist
         if sample_posterior:
           __snake_case:						Optional[Any]			  =						posterior.sample(generator=A		)
         else:
           __snake_case:						Dict			  =						posterior.mode()
         __snake_case:						Any			  =						self.decode(A		).sample

         if not return_dict:
           return (dec,)

         return DecoderOutput(sample=A		)






