





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 		UpperCamelCase__ (      lowerCAmelCase_   ):



     '''simple docstring'''

     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		List[Any]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Union[str, Any]       )     ->		DPRQuestionEncoderTokenizer:




           '''simple docstring'''


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






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		DPRContextEncoderTokenizer:




           '''simple docstring'''


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






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		BartTokenizer:




           '''simple docstring'''


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






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Dict:




           '''simple docstring'''


           shutil.rmtree(self.tmpdirname       )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		Optional[int]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Union[str, Any]       )     ->		str:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]						,lowerCamelCase__						:    bool       )     ->		Optional[Any]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Dict       )     ->		int:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		str:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		str:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		Optional[int]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Tuple       )     ->		Optional[int]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    List[Any]       )     ->		Any:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Any       )     ->		Any:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		int:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Union[str, Any]       )     ->		List[str]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Union[str, Any]       )     ->		List[Any]:




           '''simple docstring'''


           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 				SCREAMING_SNAKE_CASE__  (							self						:    Tuple       )     ->		Union[str, Any]:




           '''simple docstring'''


           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.









from argparse import ArgumentParser

from .env import EnvironmentCommand



def    __lowercase   (      ) ->   Dict:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						ArgumentParser("""Diffusers CLI tool""" ,						usage="""diffusers-cli <command> [<args>]"""	)
      SCREAMING_SNAKE_CASE		=						parser.add_subparsers(help="""diffusers-cli command helpers"""	)

      # Register commands
      EnvironmentCommand.register_subcommand(_SCREAMING_SNAKE_CASE	)

      # Let's go
      SCREAMING_SNAKE_CASE		=						parser.parse_args()

      if not hasattr(_SCREAMING_SNAKE_CASE ,						"""func"""	):
            parser.print_help()
            exit(1	)

      # Run
      SCREAMING_SNAKE_CASE		=						args.func(_SCREAMING_SNAKE_CASE	)
      service.run()


if __name__ == "__main__":
      main()


