




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  __magic_name__						(  __UpperCAmelCase):
    '''simple docstring'''







    def     _A							(							self:   Tuple		):
           SCREAMING_SNAKE_CASE_   	=    tempfile.mkdtemp()
           SCREAMING_SNAKE_CASE_   	=    8

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

           SCREAMING_SNAKE_CASE_   	=    os.path.join(self.tmpdirname		,					'''bart_tokenizer'''		)
           os.makedirs(_lowerCamelCase		,					exist_ok=_lowerCamelCase		)
           SCREAMING_SNAKE_CASE_   	=    os.path.join(_lowerCamelCase		,					BART_VOCAB_FILES_NAMES['''vocab_file''']		)
           SCREAMING_SNAKE_CASE_   	=    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     _A							(							self:   List[Any]		):
           return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname		,					'''dpr_tokenizer'''		)		)


    def     _A							(							self:   int		):
           return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname		,					'''dpr_tokenizer'''		)		)


    def     _A							(							self:   str		):
           return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname		,					'''bart_tokenizer'''		)		)


    def     _A							(							self:   Union[str, Any]		):
           shutil.rmtree(self.tmpdirname		)


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


    def     _A							(							self:   str		):
           SCREAMING_SNAKE_CASE_   	=    self.get_dummy_dataset()
           SCREAMING_SNAKE_CASE_   	=    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:
                  SCREAMING_SNAKE_CASE_   	=    dataset
                  SCREAMING_SNAKE_CASE_   	=    RagRetriever(
                      _lowerCamelCase		,					question_encoder_tokenizer=self.get_dpr_tokenizer()		,					generator_tokenizer=self.get_bart_tokenizer()		,					)
           return retriever


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


    def     _A							(							self:   Dict		):
           SCREAMING_SNAKE_CASE_   	=    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		)

           SCREAMING_SNAKE_CASE_   	=    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'''		)		)

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

           SCREAMING_SNAKE_CASE_   	=    RagConfig(
               retrieval_vector_size=self.retrieval_vector_size		,					question_encoder=DPRConfig().to_dict()		,					generator=BartConfig().to_dict()		,					index_name='''legacy'''		,					index_path=self.tmpdirname		,					)
           SCREAMING_SNAKE_CASE_   	=    RagRetriever(
               _lowerCamelCase		,					question_encoder_tokenizer=self.get_dpr_tokenizer()		,					generator_tokenizer=self.get_bart_tokenizer()		)
           return retriever


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


    def     _A							(							self:   List[Any]		):
           SCREAMING_SNAKE_CASE_   	=    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:
                         SCREAMING_SNAKE_CASE_   	=    self.get_dummy_dataset()
                         retriever.save_pretrained(_lowerCamelCase		)
                         SCREAMING_SNAKE_CASE_   	=    RagRetriever.from_pretrained(_lowerCamelCase		)
                         self.assertIsInstance(_lowerCamelCase		,					_lowerCamelCase		)
                         SCREAMING_SNAKE_CASE_   	=    np.array(
                             [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]		,					dtype=np.floataa		)
                         SCREAMING_SNAKE_CASE_   	=    retriever.retrieve(_lowerCamelCase		,					n_docs=1		)
                         self.assertTrue(out is not None		)


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


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


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


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


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


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


    @require_torch
    @require_tokenizers
    @require_sentencepiece
    def     _A							(							self:   Tuple		):
           import torch

           SCREAMING_SNAKE_CASE_   	=    1
           SCREAMING_SNAKE_CASE_   	=    self.get_dummy_canonical_hf_index_retriever()
           SCREAMING_SNAKE_CASE_   	=    [[5, 7], [10, 11]]
           SCREAMING_SNAKE_CASE_   	=    np.array(
               [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]		,					dtype=np.floataa		)
           SCREAMING_SNAKE_CASE_   	=    retriever(_lowerCamelCase		,					_lowerCamelCase		,					prefix=retriever.config.generator.prefix		,					n_docs=_lowerCamelCase		)
           SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    (
               out['''context_input_ids'''],
               out['''context_attention_mask'''],
               out['''retrieved_doc_embeds'''],
           )
           self.assertEqual(retrieved_doc_embeds.shape		,					(2, n_docs, self.retrieval_vector_size)		)
           self.assertIsInstance(_lowerCamelCase		,					_lowerCamelCase		)
           self.assertIsInstance(_lowerCamelCase		,					_lowerCamelCase		)
           self.assertIsInstance(_lowerCamelCase		,					np.ndarray		)

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





    @require_torch
    @require_tokenizers
    @require_sentencepiece
    def     _A							(							self:   Tuple		):
           SCREAMING_SNAKE_CASE_   	=    self.get_dpr_ctx_encoder_tokenizer()
           SCREAMING_SNAKE_CASE_   	=    1
           SCREAMING_SNAKE_CASE_   	=    self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCamelCase		)
           retriever.set_ctx_encoder_tokenizer(_lowerCamelCase		)

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

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








import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
			import jax






@require_flax
class  __magic_name__						(  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							str	      =					FlaxAutoencoderKL


				@property
				def     _A							(							self:   Any		):
											SCREAMING_SNAKE_CASE_   	=    4
											SCREAMING_SNAKE_CASE_   	=    3
											SCREAMING_SNAKE_CASE_   	=    (32, 32)

											SCREAMING_SNAKE_CASE_   	=    jax.random.PRNGKey(0		)
											SCREAMING_SNAKE_CASE_   	=    jax.random.uniform(_lowerCamelCase		,					((batch_size, num_channels) + sizes)		)

											return {"sample": image, "prng_key": prng_key}





				def     _A							(							self:   List[str]		):
											SCREAMING_SNAKE_CASE_   	=    {
											    '''block_out_channels''': [32, 64],
											    '''in_channels''': 3,
											    '''out_channels''': 3,
											    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
											    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
											    '''latent_channels''': 4,
											}
											SCREAMING_SNAKE_CASE_   	=    self.dummy_input
											return init_dict, inputs_dict


