



import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil



class  __lowerCAmelCase     :




 def __init__(						self	,				__UpperCAmelCase = "cpu"	,				__UpperCAmelCase = "openai/clip-vit-large-patch14"				):



    '''simple docstring'''
    __lowerCamelCase													=    device
    __lowerCamelCase													=    CLIPTokenizerFast.from_pretrained(lowerCAmelCase_				)
    __lowerCamelCase													=    [0.48_145_466, 0.4_578_275, 0.40_821_073]
    __lowerCamelCase													=    [0.26_862_954, 0.26_130_258, 0.27_577_711]
    __lowerCamelCase													=    torchvision.transforms.Normalize(self.image_mean	,				self.image_std				)
    __lowerCamelCase													=    torchvision.transforms.Resize(224				)
    __lowerCamelCase													=    torchvision.transforms.CenterCrop(224				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.resize(lowerCAmelCase_				)
    __lowerCamelCase													=    self.center_crop(lowerCAmelCase_				)
    __lowerCamelCase													=    self.normalize(lowerCAmelCase_				)
    return images




 def __call__(						self	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				**__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.tokenizer(text=lowerCAmelCase_	,				**lowerCAmelCase_				)
    __lowerCamelCase													=    self.preprocess_img(lowerCAmelCase_				)
    __lowerCamelCase													=    {key: value.to(self.device				) for (key, value) in encoding.items()}
    return encoding



class  __lowerCAmelCase     (	nn.Module   ):




 def __init__(						self	,				__UpperCAmelCase=10	,				__UpperCAmelCase=0.01	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=False	,				__UpperCAmelCase=True	,				__UpperCAmelCase="image"	,				__UpperCAmelCase=True	,				__UpperCAmelCase=False	,				__UpperCAmelCase=False	,				__UpperCAmelCase=False	,				):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase													=    None
    __lowerCamelCase													=    device if device else get_device()
    if vqgan:
       __lowerCamelCase													=    vqgan
    else:
       __lowerCamelCase													=    load_vqgan(self.device	,				conf_path=lowerCAmelCase_	,				ckpt_path=lowerCAmelCase_				)
    self.vqgan.eval()
    if clip:
       __lowerCamelCase													=    clip
    else:
       __lowerCamelCase													=    CLIPModel.from_pretrained('''openai/clip-vit-base-patch32'''				)
    self.clip.to(self.device				)
    __lowerCamelCase													=    ProcessorGradientFlow(device=self.device				)

    __lowerCamelCase													=    iterations
    __lowerCamelCase													=    lr
    __lowerCamelCase													=    log
    __lowerCamelCase													=    make_grid
    __lowerCamelCase													=    return_val
    __lowerCamelCase													=    quantize
    __lowerCamelCase													=    self.vqgan.decoder.z_shape




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=5	,				__UpperCAmelCase=True				):



    '''simple docstring'''
    __lowerCamelCase													=    []
    if output_path is None:
       __lowerCamelCase													=    '''./animation.gif'''
    if input_path is None:
       __lowerCamelCase													=    self.save_path
    __lowerCamelCase													=    sorted(glob(input_path + '''/*'''				)				)
    if not len(lowerCAmelCase_				):
       raise ValueError(
           '''No images found in save path, aborting (did you pass save_intermediate=True to the generate'''
           ''' function?)'''				)
    if len(lowerCAmelCase_				) == 1:
       print('''Only one image found in save path, (did you pass save_intermediate=True to the generate function?)'''				)
    __lowerCamelCase													=    total_duration / len(lowerCAmelCase_				)
    __lowerCamelCase													=    [frame_duration] * len(lowerCAmelCase_				)
    if extend_frames:
       __lowerCamelCase													=    1.5
       __lowerCamelCase													=    3
    for file_name in paths:
       if file_name.endswith('''.png'''				):
          images.append(imageio.imread(lowerCAmelCase_				)				)
    imageio.mimsave(lowerCAmelCase_	,				lowerCAmelCase_	,				duration=lowerCAmelCase_				)
    print(F"""gif saved to {output_path}"""				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None				):



    '''simple docstring'''







    if not (path or img):
       raise ValueError('''Input either path or tensor'''				)
    if img is not None:
       raise NotImplementedError
    __lowerCamelCase													=    preprocess(Image.open(lowerCAmelCase_				)	,				target_image_size=256				).to(self.device				)
    __lowerCamelCase													=    preprocess_vqgan(lowerCAmelCase_				)
    __lowerCamelCase							,*__lowerCamelCase													=    self.vqgan.encode(lowerCAmelCase_				)
    return z




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.latent.detach().requires_grad_()
    __lowerCamelCase													=    base_latent + transform_vector
    if self.quantize:
       __lowerCamelCase							,*__lowerCamelCase													=    self.vqgan.quantize(lowerCAmelCase_				)
    else:
       __lowerCamelCase													=    trans_latent
    return self.vqgan.decode(lowerCAmelCase_				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase=None				):



    '''simple docstring'''
    __lowerCamelCase													=    self.clip_preprocessor(text=lowerCAmelCase_	,				images=lowerCAmelCase_	,				return_tensors='''pt'''	,				padding=lowerCAmelCase_				)
    __lowerCamelCase													=    self.clip(**lowerCAmelCase_				)
    __lowerCamelCase													=    clip_outputs.logits_per_image
    if weights is not None:
       __lowerCamelCase													=    similarity_logits * weights
    return similarity_logits.sum()




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self._get_clip_similarity(pos_prompts['''prompts''']	,				lowerCAmelCase_	,				weights=(1 / pos_prompts['''weights'''])				)
    if neg_prompts:
       __lowerCamelCase													=    self._get_clip_similarity(neg_prompts['''prompts''']	,				lowerCAmelCase_	,				weights=neg_prompts['''weights''']				)
    else:
       __lowerCamelCase													=    torch.tensor([1]	,				device=self.device				)
    __lowerCamelCase													=    -torch.log(lowerCAmelCase_				) + torch.log(lowerCAmelCase_				)
    return loss




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    torch.randn_like(self.latent	,				requires_grad=lowerCAmelCase_	,				device=self.device				)
    __lowerCamelCase													=    torch.optim.Adam([vector]	,				lr=self.lr				)

    for i in range(self.iterations				):
       optim.zero_grad()
       __lowerCamelCase													=    self._add_vector(lowerCAmelCase_				)
       __lowerCamelCase													=    loop_post_process(lowerCAmelCase_				)
       __lowerCamelCase													=    self._get_CLIP_loss(lowerCAmelCase_	,				lowerCAmelCase_	,				lowerCAmelCase_				)
       print('''CLIP loss'''	,				lowerCAmelCase_				)
       if self.log:
          wandb.log({'''CLIP Loss''': clip_loss}				)
       clip_loss.backward(retain_graph=lowerCAmelCase_				)
       optim.step()
       if self.return_val == "image":
          yield custom_to_pil(transformed_img[0]				)
       else:
          yield vector




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    wandb.init(reinit=lowerCAmelCase_	,				project='''face-editor'''				)
    wandb.config.update({'''Positive Prompts''': positive_prompts}				)
    wandb.config.update({'''Negative Prompts''': negative_prompts}				)
    wandb.config.update({'''lr''': self.lr, '''iterations''': self.iterations}				)
    if image_path:
       __lowerCamelCase													=    Image.open(lowerCAmelCase_				)
       __lowerCamelCase													=    image.resize((256, 256)				)
       wandb.log('''Original Image'''	,				wandb.Image(lowerCAmelCase_				)				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''







    if not prompts:
       return []
    __lowerCamelCase													=    []
    __lowerCamelCase													=    []
    if isinstance(lowerCAmelCase_	,				lowerCAmelCase_				):
       __lowerCamelCase													=    [prompt.strip() for prompt in prompts.split('''|'''				)]
    for prompt in prompts:
       if isinstance(lowerCAmelCase_	,				(tuple, list)				):
          __lowerCamelCase													=    prompt[0]
          __lowerCamelCase													=    float(prompt[1]				)
       elif ":" in prompt:
          __lowerCamelCase							,__lowerCamelCase													=    prompt.split(''':'''				)
          __lowerCamelCase													=    float(lowerCAmelCase_				)
       else:
          __lowerCamelCase													=    prompt
          __lowerCamelCase													=    1.0
       processed_prompts.append(lowerCAmelCase_				)
       weights.append(lowerCAmelCase_				)
    return {
        "prompts": processed_prompts,
        "weights": torch.tensor(lowerCAmelCase_	,				device=self.device				),
    }




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase=None	,				__UpperCAmelCase=None	,				__UpperCAmelCase=True	,				__UpperCAmelCase=False	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=None	,				):



    '''simple docstring'''







    if image_path:
       __lowerCamelCase													=    self._get_latent(lowerCAmelCase_				)
    else:
       __lowerCamelCase													=    torch.randn(self.latent_dim	,				device=self.device				)
    if self.log:
       self._init_logging(lowerCAmelCase_	,				lowerCAmelCase_	,				lowerCAmelCase_				)

    assert pos_prompts, "You must provide at least one positive prompt."
    __lowerCamelCase													=    self.process_prompts(lowerCAmelCase_				)
    __lowerCamelCase													=    self.process_prompts(lowerCAmelCase_				)
    if save_final and save_path is None:
       __lowerCamelCase													=    os.path.join('''./outputs/'''	,				'''_'''.join(pos_prompts['''prompts''']				)				)
    if not os.path.exists(lowerCAmelCase_				):
       os.makedirs(lowerCAmelCase_				)
    else:
       __lowerCamelCase													=    save_path + '''_''' + get_timestamp()
       os.makedirs(lowerCAmelCase_				)
    __lowerCamelCase													=    save_path

    __lowerCamelCase													=    self.vqgan.decode(self.latent				)[0]
    if show_intermediate:
       print('''Original Image'''				)
       show_pil(custom_to_pil(lowerCAmelCase_				)				)

    __lowerCamelCase													=    loop_post_process(lowerCAmelCase_				)
    for iter, transformed_img in enumerate(self._optimize_CLIP(lowerCAmelCase_	,				lowerCAmelCase_	,				lowerCAmelCase_				)				):
       if show_intermediate:
          show_pil(lowerCAmelCase_				)
       if save_intermediate:
          transformed_img.save(os.path.join(self.save_path	,				F"""iter_{iter:03d}.png"""				)				)
       if self.log:
          wandb.log({'''Image''': wandb.Image(lowerCAmelCase_				)}				)
    if show_final:
       show_pil(lowerCAmelCase_				)
    if save_final:
       transformed_img.save(os.path.join(self.save_path	,				F"""iter_{iter:03d}_final.png"""				)				)











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  __lowerCAmelCase     (	lowerCAmelCase__   ):




 def        lowerCamelCase      (						self				):



    '''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(__UpperCAmelCase	,				exist_ok=__UpperCAmelCase				)
    __lowerCamelCase													=    os.path.join(__UpperCAmelCase	,				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(__UpperCAmelCase	,				range(len(__UpperCAmelCase				)				)				)				)
    __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(__UpperCAmelCase	,				exist_ok=__UpperCAmelCase				)
    __lowerCamelCase													=    os.path.join(__UpperCAmelCase	,				BART_VOCAB_FILES_NAMES['''vocab_file''']				)
    __lowerCamelCase													=    os.path.join(__UpperCAmelCase	,				BART_VOCAB_FILES_NAMES['''merges_file''']				)
    with open(self.vocab_file	,				'''w'''	,				encoding='''utf-8'''				) as fp:
       fp.write(json.dumps(__UpperCAmelCase				) + '''\n'''				)
    with open(self.merges_file	,				'''w'''	,				encoding='''utf-8'''				) as fp:
       fp.write('''\n'''.join(__UpperCAmelCase				)				)




 def        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    shutil.rmtree(self.tmpdirname				)




 def        lowerCamelCase      (						self				):



    '''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        lowerCamelCase      (						self				):



    '''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(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				)
    return retriever




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''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(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				)
    else:
       __lowerCamelCase													=    RagRetriever(
           __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()	,				index=CustomHFIndex(config.retrieval_vector_size	,				__UpperCAmelCase				)	,				)
    return retriever




 def        lowerCamelCase      (						self				):



    '''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(__UpperCAmelCase	,				open(__UpperCAmelCase	,				'''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(
        __UpperCAmelCase	,				question_encoder_tokenizer=self.get_dpr_tokenizer()	,				generator_tokenizer=self.get_bart_tokenizer()				)
    return retriever




 def        lowerCamelCase      (						self				):



    '''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(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    __lowerCamelCase													=    np.array(
        [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]	,				dtype=np.floataa				)
    __lowerCamelCase							,__lowerCamelCase							,__lowerCamelCase													=    retriever.retrieve(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''embeddings''', '''id''', '''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''id''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



    '''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(__UpperCAmelCase	,				n_docs=__UpperCAmelCase				)
    self.assertEqual(retrieved_doc_embeds.shape	,				(2, n_docs, self.retrieval_vector_size)				)
    self.assertEqual(len(__UpperCAmelCase				)	,				2				)
    self.assertEqual(sorted(doc_dicts[0]				)	,				['''text''', '''title''']				)
    self.assertEqual(len(doc_dicts[0]['''text''']				)	,				__UpperCAmelCase				)
    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        lowerCamelCase      (						self				):



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




 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        lowerCamelCase      (						self				):



    '''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(__UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase				)
    __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(__UpperCAmelCase	,				__UpperCAmelCase				)
    self.assertIsInstance(__UpperCAmelCase	,				__UpperCAmelCase				)
    self.assertIsInstance(__UpperCAmelCase	,				np.ndarray				)

    __lowerCamelCase													=    retriever(
        __UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase	,				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(__UpperCAmelCase	,				torch.Tensor				)
    self.assertIsInstance(__UpperCAmelCase	,				torch.Tensor				)
    self.assertIsInstance(__UpperCAmelCase	,				torch.Tensor				)




 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dpr_ctx_encoder_tokenizer()
    __lowerCamelCase													=    1
    __lowerCamelCase													=    self.get_dummy_custom_hf_index_retriever(from_disk=__UpperCAmelCase				)
    retriever.set_ctx_encoder_tokenizer(__UpperCAmelCase				)

    __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(__UpperCAmelCase	,				__UpperCAmelCase	,				prefix=retriever.config.generator.prefix	,				n_docs=__UpperCAmelCase				)

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






