






'''simple docstring'''







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block





@dataclass
class 	SCREAMING_SNAKE_CASE					(					__lowerCamelCase       ):
			snake_case__  =				42





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			int  ,						__SCREAMING_SNAKE_CASE      :			str=3  ,						__SCREAMING_SNAKE_CASE      :			str=3  ,						__SCREAMING_SNAKE_CASE      :			str=("DownEncoderBlock2D",)  ,						__SCREAMING_SNAKE_CASE      :			Any=(64,)  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=2  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=32  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]="silu"  ,						__SCREAMING_SNAKE_CASE      :			int=True  ,						)			->			List[Any]:
							super().__init__()
							a_		:  Tuple											=      layers_per_block

							a_		:  Any											=      torch.nn.Convad(
							    SCREAMING_SNAKE_CASE_  ,						block_out_channels[0]  ,						kernel_size=3  ,						stride=1  ,						padding=1  ,						)

							a_		:  Any											=      None
							a_		:  Union[str, Any]											=      nn.ModuleList([]					)

							# down
							a_		:  Tuple											=      block_out_channels[0]
							for i, down_block_type in enumerate(SCREAMING_SNAKE_CASE_					):
											a_		:  List[Any]											=      output_channel
											a_		:  Dict											=      block_out_channels[i]
											a_		:  int											=      i == len(SCREAMING_SNAKE_CASE_					) - 1

											a_		:  Any											=      get_down_block(
											    SCREAMING_SNAKE_CASE_  ,						num_layers=self.layers_per_block  ,						in_channels=SCREAMING_SNAKE_CASE_  ,						out_channels=SCREAMING_SNAKE_CASE_  ,						add_downsample=not is_final_block  ,						resnet_eps=1e-6  ,						downsample_padding=0  ,						resnet_act_fn=SCREAMING_SNAKE_CASE_  ,						resnet_groups=SCREAMING_SNAKE_CASE_  ,						attention_head_dim=SCREAMING_SNAKE_CASE_  ,						temb_channels=SCREAMING_SNAKE_CASE_  ,						)
											self.down_blocks.append(SCREAMING_SNAKE_CASE_					)

							# mid
							a_		:  str											=      UNetMidBlockaD(
							    in_channels=block_out_channels[-1]  ,						resnet_eps=1e-6  ,						resnet_act_fn=SCREAMING_SNAKE_CASE_  ,						output_scale_factor=1  ,						resnet_time_scale_shift='''default'''  ,						attention_head_dim=block_out_channels[-1]  ,						resnet_groups=SCREAMING_SNAKE_CASE_  ,						temb_channels=SCREAMING_SNAKE_CASE_  ,						)

							# out
							a_		:  Tuple											=      nn.GroupNorm(num_channels=block_out_channels[-1]  ,						num_groups=SCREAMING_SNAKE_CASE_  ,						eps=1e-6					)
							a_		:  Any											=      nn.SiLU()

							a_		:  Dict											=      2 * out_channels if double_z else out_channels
							a_		:  Any											=      nn.Convad(block_out_channels[-1]  ,						SCREAMING_SNAKE_CASE_  ,						3  ,						padding=1					)

							a_		:  int											=      False







			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			Tuple					)			->			Any:
							a_		:  Union[str, Any]											=      x
							a_		:  Optional[Any]											=      self.conv_in(SCREAMING_SNAKE_CASE_					)

							if self.training and self.gradient_checkpointing:

											def create_custom_forward(__SCREAMING_SNAKE_CASE      :			Union[str, Any]					):
															def custom_forward(*__SCREAMING_SNAKE_CASE      :			Optional[int]					):
																			return module(*SCREAMING_SNAKE_CASE_					)

															return custom_forward

											# down
											if is_torch_version('''>='''  ,						'''1.11.0'''					):
															for down_block in self.down_blocks:
																			a_		:  List[Any]											=      torch.utils.checkpoint.checkpoint(
																			    create_custom_forward(SCREAMING_SNAKE_CASE_					)  ,						SCREAMING_SNAKE_CASE_  ,						use_reentrant=SCREAMING_SNAKE_CASE_					)
															# middle
															a_		:  str											=      torch.utils.checkpoint.checkpoint(
															    create_custom_forward(self.mid_block					)  ,						SCREAMING_SNAKE_CASE_  ,						use_reentrant=SCREAMING_SNAKE_CASE_					)
											else:
															for down_block in self.down_blocks:
																			a_		:  Dict											=      torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE_					)  ,						SCREAMING_SNAKE_CASE_					)
															# middle
															a_		:  Any											=      torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block					)  ,						SCREAMING_SNAKE_CASE_					)

							else:
											# down
											for down_block in self.down_blocks:
															a_		:  List[str]											=      down_block(SCREAMING_SNAKE_CASE_					)

											# middle
											a_		:  Union[str, Any]											=      self.mid_block(SCREAMING_SNAKE_CASE_					)

							# post-process
							a_		:  Any											=      self.conv_norm_out(SCREAMING_SNAKE_CASE_					)
							a_		:  Union[str, Any]											=      self.conv_act(SCREAMING_SNAKE_CASE_					)
							a_		:  Any											=      self.conv_out(SCREAMING_SNAKE_CASE_					)

							return sample





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			Dict=3  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=3  ,						__SCREAMING_SNAKE_CASE      :			str=("UpDecoderBlock2D",)  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=(64,)  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=2  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=32  ,						__SCREAMING_SNAKE_CASE      :			Tuple="silu"  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]="group"  ,						)			->			str:
							super().__init__()
							a_		:  List[Any]											=      layers_per_block

							a_		:  Union[str, Any]											=      nn.Convad(
							    SCREAMING_SNAKE_CASE_  ,						block_out_channels[-1]  ,						kernel_size=3  ,						stride=1  ,						padding=1  ,						)

							a_		:  str											=      None
							a_		:  Dict											=      nn.ModuleList([]					)

							a_		:  Union[str, Any]											=      in_channels if norm_type == '''spatial''' else None

							# mid
							a_		:  Optional[int]											=      UNetMidBlockaD(
							    in_channels=block_out_channels[-1]  ,						resnet_eps=1e-6  ,						resnet_act_fn=SCREAMING_SNAKE_CASE_  ,						output_scale_factor=1  ,						resnet_time_scale_shift='''default''' if norm_type == '''group''' else norm_type  ,						attention_head_dim=block_out_channels[-1]  ,						resnet_groups=SCREAMING_SNAKE_CASE_  ,						temb_channels=SCREAMING_SNAKE_CASE_  ,						)

							# up
							a_		:  Tuple											=      list(reversed(SCREAMING_SNAKE_CASE_					)					)
							a_		:  List[str]											=      reversed_block_out_channels[0]
							for i, up_block_type in enumerate(SCREAMING_SNAKE_CASE_					):
											a_		:  List[Any]											=      output_channel
											a_		:  Dict											=      reversed_block_out_channels[i]

											a_		:  List[str]											=      i == len(SCREAMING_SNAKE_CASE_					) - 1

											a_		:  List[str]											=      get_up_block(
											    SCREAMING_SNAKE_CASE_  ,						num_layers=self.layers_per_block + 1  ,						in_channels=SCREAMING_SNAKE_CASE_  ,						out_channels=SCREAMING_SNAKE_CASE_  ,						prev_output_channel=SCREAMING_SNAKE_CASE_  ,						add_upsample=not is_final_block  ,						resnet_eps=1e-6  ,						resnet_act_fn=SCREAMING_SNAKE_CASE_  ,						resnet_groups=SCREAMING_SNAKE_CASE_  ,						attention_head_dim=SCREAMING_SNAKE_CASE_  ,						temb_channels=SCREAMING_SNAKE_CASE_  ,						resnet_time_scale_shift=SCREAMING_SNAKE_CASE_  ,						)
											self.up_blocks.append(SCREAMING_SNAKE_CASE_					)
											a_		:  str											=      output_channel

							# out
							if norm_type == "spatial":
											a_		:  Dict											=      SpatialNorm(block_out_channels[0]  ,						SCREAMING_SNAKE_CASE_					)
							else:
											a_		:  Tuple											=      nn.GroupNorm(num_channels=block_out_channels[0]  ,						num_groups=SCREAMING_SNAKE_CASE_  ,						eps=1e-6					)
							a_		:  Tuple											=      nn.SiLU()
							a_		:  Any											=      nn.Convad(block_out_channels[0]  ,						SCREAMING_SNAKE_CASE_  ,						3  ,						padding=1					)

							a_		:  int											=      False







			def 			SCREAMING_SNAKE_CASE      ( self      :			str  ,						__SCREAMING_SNAKE_CASE      :			str  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=None					)			->			str:
							a_		:  str											=      z
							a_		:  List[Any]											=      self.conv_in(SCREAMING_SNAKE_CASE_					)

							a_		:  Any											=      next(iter(self.up_blocks.parameters()					)					).dtype
							if self.training and self.gradient_checkpointing:

											def create_custom_forward(__SCREAMING_SNAKE_CASE      :			Tuple					):
															def custom_forward(*__SCREAMING_SNAKE_CASE      :			List[Any]					):
																			return module(*SCREAMING_SNAKE_CASE_					)

															return custom_forward

											if is_torch_version('''>='''  ,						'''1.11.0'''					):
															# middle
															a_		:  Optional[int]											=      torch.utils.checkpoint.checkpoint(
															    create_custom_forward(self.mid_block					)  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						use_reentrant=SCREAMING_SNAKE_CASE_					)
															a_		:  Union[str, Any]											=      sample.to(SCREAMING_SNAKE_CASE_					)

															# up
															for up_block in self.up_blocks:
																			a_		:  List[Any]											=      torch.utils.checkpoint.checkpoint(
																			    create_custom_forward(SCREAMING_SNAKE_CASE_					)  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						use_reentrant=SCREAMING_SNAKE_CASE_					)
											else:
															# middle
															a_		:  List[str]											=      torch.utils.checkpoint.checkpoint(
															    create_custom_forward(self.mid_block					)  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_					)
															a_		:  int											=      sample.to(SCREAMING_SNAKE_CASE_					)

															# up
															for up_block in self.up_blocks:
																			a_		:  List[str]											=      torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE_					)  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_					)
							else:
											# middle
											a_		:  Tuple											=      self.mid_block(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_					)
											a_		:  Optional[int]											=      sample.to(SCREAMING_SNAKE_CASE_					)

											# up
											for up_block in self.up_blocks:
															a_		:  Optional[int]											=      up_block(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_					)

        # post-process
							if latent_embeds is None:
											a_		:  str											=      self.conv_norm_out(SCREAMING_SNAKE_CASE_					)
							else:
											a_		:  Tuple											=      self.conv_norm_out(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_					)
							a_		:  Optional[int]											=      self.conv_act(SCREAMING_SNAKE_CASE_					)
							a_		:  Optional[int]											=      self.conv_out(SCREAMING_SNAKE_CASE_					)

							return sample





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=None  ,						__SCREAMING_SNAKE_CASE      :			int="random"  ,						__SCREAMING_SNAKE_CASE      :			List[str]=False  ,						__SCREAMING_SNAKE_CASE      :			Dict=True					)			->			Tuple:
							super().__init__()
							a_		:  int											=      n_e
							a_		:  Dict											=      vq_embed_dim
							a_		:  Any											=      beta
							a_		:  List[Any]											=      legacy

							a_		:  Dict											=      nn.Embedding(self.n_e  ,						self.vq_embed_dim					)
							self.embedding.weight.data.uniform_(-1.0 / self.n_e  ,						1.0 / self.n_e					)

							a_		:  List[Any]											=      remap
							if self.remap is not None:
											self.register_buffer('''used'''  ,						torch.tensor(np.load(self.remap					)					)					)
											a_		:  Dict											=      self.used.shape[0]
											a_		:  List[str]											=      unknown_index  # "random" or "extra" or integer
											if self.unknown_index == "extra":
															a_		:  List[Any]											=      self.re_embed
															a_		:  Optional[int]											=      self.re_embed + 1
											print(
											    f'Remapping {self.n_e} indices to {self.re_embed} indices. '
											    f'Using {self.unknown_index} for unknown indices.'					)
							else:
											a_		:  Any											=      n_e

							a_		:  List[Any]											=      sane_index_shape






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			List[Any]					)			->			Any:
							a_		:  List[Any]											=      inds.shape
							assert len(SCREAMING_SNAKE_CASE_					) > 1
							a_		:  Dict											=      inds.reshape(ishape[0]  ,						-1					)
							a_		:  List[Any]											=      self.used.to(SCREAMING_SNAKE_CASE_					)
							a_		:  int											=      (inds[:, :, None] == used[None, None, ...]).long()
							a_		:  Union[str, Any]											=      match.argmax(-1					)
							a_		:  List[str]											=      match.sum(2					) < 1
							if self.unknown_index == "random":
											a_		:  Optional[int]											=      torch.randint(0  ,						self.re_embed  ,						size=new[unknown].shape					).to(device=new.device					)
							else:
											a_		:  Dict											=      self.unknown_index
							return new.reshape(SCREAMING_SNAKE_CASE_					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			Dict					)			->			str:
							a_		:  List[Any]											=      inds.shape
							assert len(SCREAMING_SNAKE_CASE_					) > 1
							a_		:  Optional[Any]											=      inds.reshape(ishape[0]  ,						-1					)
							a_		:  Any											=      self.used.to(SCREAMING_SNAKE_CASE_					)
							if self.re_embed > self.used.shape[0]:  # extra token
											a_		:  List[Any]											=      0  # simply set to zero
							a_		:  List[Any]											=      torch.gather(used[None, :][inds.shape[0] * [0], :]  ,						1  ,						SCREAMING_SNAKE_CASE_					)
							return back.reshape(SCREAMING_SNAKE_CASE_					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]					)			->			List[str]:
							# reshape z -> (batch, height, width, channel) and flatten
							a_		:  List[str]											=      z.permute(0  ,						2  ,						3  ,						1					).contiguous()
							a_		:  Tuple											=      z.view(-1  ,						self.vq_embed_dim					)

							# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
							a_		:  Union[str, Any]											=      torch.argmin(torch.cdist(SCREAMING_SNAKE_CASE_  ,						self.embedding.weight					)  ,						dim=1					)

							a_		:  Union[str, Any]											=      self.embedding(SCREAMING_SNAKE_CASE_					).view(z.shape					)
							a_		:  List[Any]											=      None
							a_		:  List[Any]											=      None

							# compute loss for embedding
							if not self.legacy:
											a_		:  List[str]											=      self.beta * torch.mean((z_q.detach() - z) ** 2					) + torch.mean((z_q - z.detach()) ** 2					)
							else:
											a_		:  Optional[int]											=      torch.mean((z_q.detach() - z) ** 2					) + self.beta * torch.mean((z_q - z.detach()) ** 2					)

							# preserve gradients
							a_		:  str											=      z + (z_q - z).detach()

							# reshape back to match original input shape
							a_		:  Optional[Any]											=      z_q.permute(0  ,						3  ,						1  ,						2					).contiguous()

							if self.remap is not None:
											a_		:  Union[str, Any]											=      min_encoding_indices.reshape(z.shape[0]  ,						-1					)  # add batch axis
											a_		:  List[str]											=      self.remap_to_used(SCREAMING_SNAKE_CASE_					)
											a_		:  Dict											=      min_encoding_indices.reshape(-1  ,						1					)  # flatten

							if self.sane_index_shape:
											a_		:  Optional[int]											=      min_encoding_indices.reshape(z_q.shape[0]  ,						z_q.shape[2]  ,						z_q.shape[3]					)

							return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			str  ,						__SCREAMING_SNAKE_CASE      :			str					)			->			Optional[Any]:
							# shape specifying (batch, height, width, channel)
							if self.remap is not None:
											a_		:  str											=      indices.reshape(shape[0]  ,						-1					)  # add batch axis
											a_		:  Tuple											=      self.unmap_to_all(SCREAMING_SNAKE_CASE_					)
											a_		:  Tuple											=      indices.reshape(-1					)  # flatten again

							# get quantized latent vectors
							a_		:  str											=      self.embedding(SCREAMING_SNAKE_CASE_					)

							if shape is not None:
											a_		:  Optional[int]											=      z_q.view(SCREAMING_SNAKE_CASE_					)
											# reshape back to match original input shape
											a_		:  Optional[Any]											=      z_q.permute(0  ,						3  ,						1  ,						2					).contiguous()

							return z_q





class 	SCREAMING_SNAKE_CASE					(					__lowerCamelCase       ):






			def __init__( self      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=False					)			->			Any:
							a_		:  Dict											=      parameters
							a_				,			a_		:  int											=      torch.chunk(SCREAMING_SNAKE_CASE_  ,						2  ,						dim=1					)
							a_		:  str											=      torch.clamp(self.logvar  ,						-30.0  ,						20.0					)
							a_		:  Dict											=      deterministic
							a_		:  Optional[Any]											=      torch.exp(0.5 * self.logvar					)
							a_		:  Any											=      torch.exp(self.logvar					)
							if self.deterministic:
											a_		:  List[str]											=      torch.zeros_like(
											    self.mean  ,						device=self.parameters.device  ,						dtype=self.parameters.dtype					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			Any = None					)			->			Tuple:
							# make sure sample is on the same device as the parameters and has same dtype
							a_		:  Any											=      randn_tensor(
							    self.mean.shape  ,						generator=SCREAMING_SNAKE_CASE_  ,						device=self.parameters.device  ,						dtype=self.parameters.dtype					)
							a_		:  Any											=      self.mean + self.std * sample
							return x






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]  ,						__SCREAMING_SNAKE_CASE      :			str=None					)			->			Optional[Any]:
							if self.deterministic:
											return torch.Tensor([0.0]					)
							else:
											if other is None:
															return 0.5 * torch.sum(torch.pow(self.mean  ,						2					) + self.var - 1.0 - self.logvar  ,						dim=[1, 2, 3]					)
											else:
															return 0.5 * torch.sum(
															    torch.pow(self.mean - other.mean  ,						2					) / other.var
															    + self.var / other.var
															    - 1.0
															    - self.logvar
															    + other.logvar  ,						dim=[1, 2, 3]  ,						)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			Dict=[1, 2, 3]					)			->			int:
							if self.deterministic:
											return torch.Tensor([0.0]					)
							a_		:  Optional[int]											=      np.log(2.0 * np.pi					)
							return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean  ,						2					) / self.var  ,						dim=SCREAMING_SNAKE_CASE_					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			Any:
							return self.mean












'''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 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Dict:
							a_		:  Union[str, Any]											=      tempfile.mkdtemp()
							a_		:  Union[str, Any]											=      8

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

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






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			DPRQuestionEncoderTokenizer:
							return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,						'''dpr_tokenizer'''					)					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			DPRContextEncoderTokenizer:
							return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,						'''dpr_tokenizer'''					)					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			BartTokenizer:
							return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,						'''bart_tokenizer'''					)					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			Union[str, Any]:
							shutil.rmtree(self.tmpdirname					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Union[str, Any]:
							a_		:  str											=      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      :			List[str]					)			->			Tuple:
							a_		:  List[str]											=      self.get_dummy_dataset()
							a_		:  Tuple											=      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:
											a_		:  Tuple											=      dataset
											a_		:  Any											=      RagRetriever(
											    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()  ,						)
							return retriever






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			bool					)			->			Dict:
							a_		:  Dict											=      self.get_dummy_dataset()
							a_		:  Dict											=      RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size  ,						question_encoder=DPRConfig().to_dict()  ,						generator=BartConfig().to_dict()  ,						index_name='''custom'''  ,						)
							if from_disk:
											a_		:  Optional[int]											=      os.path.join(self.tmpdirname  ,						'''dataset'''					)
											a_		:  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
											a_		:  int											=      RagRetriever(
											    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()  ,						)
							else:
											a_		:  Optional[Any]											=      RagRetriever(
											    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()  ,						index=CustomHFIndex(config.retrieval_vector_size  ,						__SCREAMING_SNAKE_CASE					)  ,						)
							return retriever






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			Optional[int]:
							a_		:  str											=      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					)

							a_		:  Optional[int]											=      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'''					)					)

							a_		:  Union[str, Any]											=      os.path.join(self.tmpdirname  ,						'''psgs_w100.tsv.pkl'''					)
							a_		:  Dict											=      {sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
							pickle.dump(__SCREAMING_SNAKE_CASE  ,						open(__SCREAMING_SNAKE_CASE  ,						'''wb'''					)					)

							a_		:  Optional[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  ,						)
							a_		:  int											=      RagRetriever(
							    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()					)
							return retriever






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Dict:
							a_		:  Optional[Any]											=      1
							a_		:  Dict											=      self.get_dummy_canonical_hf_index_retriever()
							a_		:  Tuple											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  str											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=__SCREAMING_SNAKE_CASE					)
							self.assertEqual(retrieved_doc_embeds.shape  ,						(2, n_docs, self.retrieval_vector_size)					)
							self.assertEqual(len(__SCREAMING_SNAKE_CASE					)  ,						2					)
							self.assertEqual(sorted(doc_dicts[0]					)  ,						['''embeddings''', '''id''', '''text''', '''title''']					)
							self.assertEqual(len(doc_dicts[0]['''id''']					)  ,						__SCREAMING_SNAKE_CASE					)
							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      :			Tuple					)			->			Dict:
							a_		:  str											=      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:
															a_		:  List[str]											=      self.get_dummy_dataset()
															retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
															a_		:  Optional[Any]											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
															self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
															a_		:  Union[str, Any]											=      np.array(
															    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
															a_		:  List[str]											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
															self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			int:
							a_		:  Union[str, Any]											=      1
							a_		:  Optional[Any]											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							a_		:  List[Any]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  Any											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=__SCREAMING_SNAKE_CASE					)
							self.assertEqual(retrieved_doc_embeds.shape  ,						(2, n_docs, self.retrieval_vector_size)					)
							self.assertEqual(len(__SCREAMING_SNAKE_CASE					)  ,						2					)
							self.assertEqual(sorted(doc_dicts[0]					)  ,						['''embeddings''', '''id''', '''text''', '''title''']					)
							self.assertEqual(len(doc_dicts[0]['''id''']					)  ,						__SCREAMING_SNAKE_CASE					)
							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					)			->			List[str]:
							a_		:  Dict											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							with tempfile.TemporaryDirectory() as tmp_dirname:
											retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  List[str]											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  Tuple											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Dict											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Optional[int]:
							a_		:  Union[str, Any]											=      1
							a_		:  str											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							a_		:  Optional[int]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  Tuple											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=__SCREAMING_SNAKE_CASE					)
							self.assertEqual(retrieved_doc_embeds.shape  ,						(2, n_docs, self.retrieval_vector_size)					)
							self.assertEqual(len(__SCREAMING_SNAKE_CASE					)  ,						2					)
							self.assertEqual(sorted(doc_dicts[0]					)  ,						['''embeddings''', '''id''', '''text''', '''title''']					)
							self.assertEqual(len(doc_dicts[0]['''id''']					)  ,						__SCREAMING_SNAKE_CASE					)
							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      :			int					)			->			Dict:
							a_		:  List[str]											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							with tempfile.TemporaryDirectory() as tmp_dirname:
											retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  Any											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  Tuple											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Dict											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			Any:
							a_		:  str											=      1
							a_		:  Tuple											=      self.get_dummy_legacy_index_retriever()
							a_		:  Union[str, Any]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  Any											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=__SCREAMING_SNAKE_CASE					)
							self.assertEqual(retrieved_doc_embeds.shape  ,						(2, n_docs, self.retrieval_vector_size)					)
							self.assertEqual(len(__SCREAMING_SNAKE_CASE					)  ,						2					)
							self.assertEqual(sorted(doc_dicts[0]					)  ,						['''text''', '''title''']					)
							self.assertEqual(len(doc_dicts[0]['''text''']					)  ,						__SCREAMING_SNAKE_CASE					)
							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      :			int					)			->			Optional[int]:
							a_		:  List[str]											=      self.get_dummy_legacy_index_retriever()
							with tempfile.TemporaryDirectory() as tmp_dirname:
											retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  Any											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  List[str]											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Optional[Any]											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			@require_torch
			@require_tokenizers
			@require_sentencepiece
			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			List[str]:
							import torch

							a_		:  Any											=      1
							a_		:  List[Any]											=      self.get_dummy_canonical_hf_index_retriever()
							a_		:  Union[str, Any]											=      [[5, 7], [10, 11]]
							a_		:  Optional[int]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_		:  str											=      retriever(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						prefix=retriever.config.generator.prefix  ,						n_docs=__SCREAMING_SNAKE_CASE					)
							a_				,			a_				,			a_		:  List[str]											=      (
							    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(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						np.ndarray					)

							a_		:  Any											=      retriever(
							    __SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						prefix=retriever.config.generator.prefix  ,						n_docs=__SCREAMING_SNAKE_CASE  ,						return_tensors='''pt'''  ,						)
							a_				,			a_				,			a_				,			a_		:  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(__SCREAMING_SNAKE_CASE  ,						torch.Tensor					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						torch.Tensor					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						torch.Tensor					)







			@require_torch
			@require_tokenizers
			@require_sentencepiece
			def 			SCREAMING_SNAKE_CASE      ( self      :			Union[str, Any]					)			->			Any:
							a_		:  str											=      self.get_dpr_ctx_encoder_tokenizer()
							a_		:  Tuple											=      1
							a_		:  Any											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							retriever.set_ctx_encoder_tokenizer(__SCREAMING_SNAKE_CASE					)

							a_		:  Dict											=      [[5, 7], [10, 11]]
							a_		:  List[str]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_		:  List[Any]											=      retriever(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						prefix=retriever.config.generator.prefix  ,						n_docs=__SCREAMING_SNAKE_CASE					)

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




