




import json
import os
import shutil
import sys
import tempfile
import unittest
from unittest import TestCase
from unittest.mock import patch

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
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_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
 from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
 a_      						=   None

if is_ray_available():
 import ray  # noqa: E402 # isort:skip
 from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
 a_      						=   None
 a_      						=   None
 a_      						=   None




def     __SCREAMING_SNAKE_CASE      (     lowercase_					)				->		Optional[Any]:

    """simple docstring"""


    if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
        __UpperCamelCase     						=  unittest.skip('''test requires Datasets, Faiss, psutil'''					)(lowercase_					)
    return test_case


@require_distributed_retrieval
class        _lowerCamelCase  (	__UpperCAmelCase      ):

      """simple docstring"""

      def 					snake_case		(   self    :       List[str] ):
          __UpperCamelCase     						=  tempfile.mkdtemp()
          __UpperCamelCase     						=  8

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

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







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







      def 					snake_case		(   self    :       List[Any] ):
          shutil.rmtree(self.tmpdirname )







      def 					snake_case		(   self    :       Optional[Any] ):
          __UpperCamelCase     						=  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 					snake_case		(   self    :       Any       ,				snake_case    :       bool       ,				snake_case    :       Union[str, Any]=12345 ):
          __UpperCamelCase     						=  self.get_dummy_dataset()
          __UpperCamelCase     						=  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:
              __UpperCamelCase     						=  dataset
              __UpperCamelCase     						=  RagPyTorchDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				)
              if init_retrieval:
                  retriever.init_retrieval(_lowerCamelCase )
          return retriever







      def 					snake_case		(   self    :       int       ,				snake_case    :       bool ):
          # Have to run in local mode because sys.path modifications at top of
          # file are not propogated to remote workers.
          # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
          ray.init(local_mode=_lowerCamelCase )
          __UpperCamelCase     						=  RagConfig(
              retrieval_vector_size=self.retrieval_vector_size       ,				question_encoder=DPRConfig().to_dict()       ,				generator=BartConfig().to_dict()       ,				)
          __UpperCamelCase     						=  ray.remote(_lowerCamelCase )
          __UpperCamelCase     						=  [remote_cls.remote() for _ in range(1 )]
          with patch('''transformers.models.rag.retrieval_rag.load_dataset''' ) as mock_load_dataset:
              __UpperCamelCase     						=  self.get_dummy_dataset()
              __UpperCamelCase     						=  RagRayDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				retrieval_workers=_lowerCamelCase       ,				)
              if init_retrieval:
                  retriever.init_retrieval()
          return retriever







      def 					snake_case		(   self    :       Union[str, Any]       ,				snake_case    :       bool       ,				snake_case    :       bool       ,				snake_case    :       List[Any]=12345 ):
          __UpperCamelCase     						=  self.get_dummy_dataset()
          __UpperCamelCase     						=  RagConfig(
              retrieval_vector_size=self.retrieval_vector_size       ,				question_encoder=DPRConfig().to_dict()       ,				generator=BartConfig().to_dict()       ,				index_name='''custom'''       ,				)
          if from_disk:
              __UpperCamelCase     						=  os.path.join(self.tmpdirname       ,				'''dataset''' )
              __UpperCamelCase     						=  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
              __UpperCamelCase     						=  RagPyTorchDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				)
          else:
              __UpperCamelCase     						=  RagPyTorchDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				index=CustomHFIndex(config.retrieval_vector_size       ,				_lowerCamelCase )       ,				)
          if init_retrieval:
              retriever.init_retrieval(_lowerCamelCase )
          return retriever







      def 					snake_case		(   self    :       Tuple       ,				snake_case    :       bool       ,				snake_case    :       bool ):
          # Have to run in local mode because sys.path modifications at top of
          # file are not propogated to remote workers.
          # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
          ray.init(local_mode=_lowerCamelCase )
          __UpperCamelCase     						=  self.get_dummy_dataset()
          __UpperCamelCase     						=  RagConfig(
              retrieval_vector_size=self.retrieval_vector_size       ,				question_encoder=DPRConfig().to_dict()       ,				generator=BartConfig().to_dict()       ,				index_name='''custom'''       ,				)
          __UpperCamelCase     						=  ray.remote(_lowerCamelCase )
          __UpperCamelCase     						=  [remote_cls.remote() for _ in range(1 )]
          if from_disk:
              __UpperCamelCase     						=  os.path.join(self.tmpdirname       ,				'''dataset''' )
              __UpperCamelCase     						=  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
              __UpperCamelCase     						=  RagRayDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				retrieval_workers=_lowerCamelCase       ,				index=CustomHFIndex.load_from_disk(
                      vector_size=config.retrieval_vector_size       ,				dataset_path=config.passages_path       ,				index_path=config.index_path       ,				)       ,				)
          else:
              __UpperCamelCase     						=  RagRayDistributedRetriever(
                  _lowerCamelCase       ,				question_encoder_tokenizer=self.get_dpr_tokenizer()       ,				generator_tokenizer=self.get_bart_tokenizer()       ,				retrieval_workers=_lowerCamelCase       ,				index=CustomHFIndex(config.retrieval_vector_size       ,				_lowerCamelCase )       ,				)
          if init_retrieval:
              retriever.init_retrieval()
          return retriever







      def 					snake_case		(   self    :       Tuple       ,				snake_case    :       RagRetriever       ,				snake_case    :       np.array       ,				snake_case    :       int ):
          __UpperCamelCase     						=  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 					snake_case		(   self    :       int ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )

          self.distributed_retriever_check(
              self.get_dummy_pytorch_distributed_retriever(init_retrieval=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase )







      def 					snake_case		(   self    :       Dict ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )

          self.distributed_retriever_check(
              self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowerCamelCase       ,				from_disk=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase       ,				)







      def 					snake_case		(   self    :       Union[str, Any] ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )

          self.distributed_retriever_check(
              self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowerCamelCase       ,				from_disk=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase       ,				)







      @require_ray
      def 					snake_case		(   self    :       Union[str, Any] ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )

          self.distributed_retriever_check(
              self.get_dummy_ray_distributed_retriever(init_retrieval=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase )
          ray.shutdown()







      @require_ray
      def 					snake_case		(   self    :       Optional[int] ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )
          with self.assertRaises(_lowerCamelCase ):
              self.distributed_retriever_check(
                  self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowerCamelCase       ,				from_disk=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase       ,				)
          ray.shutdown()







      @require_ray
      def 					snake_case		(   self    :       List[Any] ):
          __UpperCamelCase     						=  1
          __UpperCamelCase     						=  np.array(
              [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]       ,				dtype=np.floataa )

          self.distributed_retriever_check(
              self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowerCamelCase       ,				from_disk=_lowerCamelCase )       ,				_lowerCamelCase       ,				_lowerCamelCase )
          ray.shutdown()













"""simple docstring"""





from ctypes import c_float, sizeof
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union


if TYPE_CHECKING:
    from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore





class 			snake_case		(   __UpperCAmelCase   ):
      lowerCamelCase__						= c_float

      @property
      def 						SCREAMING_SNAKE_CASE_			(			self   :Optional[Any]						):
       return sizeof(self.value						)







def 		lowerCAmelCase_	(  lowercase_   :							int ,      lowercase_   :							int ,      lowercase_   :							int = 0		):

 '''simple docstring'''






 if dimension <= 0:
  __SCREAMING_SNAKE_CASE :			str												=       fixed_dimension

 dimension -= num_token_to_add
 return dimension




def 		lowerCAmelCase_	(  lowercase_   :							int ,      lowercase_   :							ParameterFormat		):

 '''simple docstring'''






 return num_parameters * dtype.size






def 		lowerCAmelCase_	(  lowercase_   :							str		):

 '''simple docstring'''






 from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

 try:
  return AutoProcessor.from_pretrained(lowercase_		)
 except (ValueError, OSError, KeyError):
  __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			int												=       None, None
  try:
   __SCREAMING_SNAKE_CASE :			Optional[Any]												=       AutoTokenizer.from_pretrained(lowercase_		)
  except (OSError, KeyError):
   pass
  try:
   __SCREAMING_SNAKE_CASE :			Optional[Any]												=       AutoFeatureExtractor.from_pretrained(lowercase_		)
  except (OSError, KeyError):
   pass

  if tokenizer is not None and feature_extractor is not None:
   raise ValueError(
       F'''Couldn\'t auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor.'''		)
  elif tokenizer is None and feature_extractor is None:
   return None
  elif tokenizer is not None:
   return tokenizer
  else:
   return feature_extractor

