



"""simple docstring"""





import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


lowerCamelCase      =		logging.getLogger(__name__)
class    lowercase__     :







       '''simple docstring'''


       def __init__(							self    :						int					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_    				=							False



       def   lowercase__       (							self    :						Union[str, Any]	,       _UpperCAmelCase    :						Union[str, Any]	,       _UpperCAmelCase    :						Tuple	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						str					)				->		Union[str, Any]:


             '''simple docstring'''




             if not self.initialized:
                   UpperCAmelCase_    				=							RagRetriever(
                       __snake_case	,       question_encoder_tokenizer=__snake_case	,       generator_tokenizer=__snake_case	,       index=__snake_case	,       init_retrieval=__snake_case	,       )
                   UpperCAmelCase_    				=							True



       def   lowercase__       (							self    :						str					)				->		Optional[int]:


             '''simple docstring'''




             self.retriever.index.init_index()







       def   lowercase__       (							self    :						Union[str, Any]	,       _UpperCAmelCase    :						List[Any]	,       _UpperCAmelCase    :						Union[str, Any]					)				->		List[str]:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.retriever._main_retrieve(__snake_case	,       __snake_case					)
             return doc_ids, retrieved_doc_embeds



class    lowercase__     (					__SCREAMING_SNAKE_CASE   ):







       '''simple docstring'''


       def __init__(							self    :						Any	,       _UpperCAmelCase    :						str	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						Optional[int]	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						Union[str, Any]=None					)				->		List[str]:


             '''simple docstring'''




             if index is not None and index.is_initialized() and len(__snake_case					) > 0:
                   raise ValueError(
                       "When using Ray for distributed fine-tuning, "
                       "you\'ll need to provide the paths instead, "
                       "as the dataset and the index are loaded "
                       "separately. More info in examples/rag/use_own_knowledge_dataset.py "					)
             super().__init__(
                 __snake_case	,       question_encoder_tokenizer=__snake_case	,       generator_tokenizer=__snake_case	,       index=__snake_case	,       init_retrieval=__snake_case	,       )
             UpperCAmelCase_    				=							retrieval_workers
             if len(self.retrieval_workers					) > 0:
                   ray.get(
                       [
                           worker.create_rag_retriever.remote(__snake_case	,       __snake_case	,       __snake_case	,       __snake_case					)
                           for worker in self.retrieval_workers
                       ]					)



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


             '''simple docstring'''




             logger.info("initializing retrieval"					)

             if len(self.retrieval_workers					) > 0:
                   ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]					)
             else:
                   # Non-distributed training. Load index into this same process.
                   self.index.init_index()



       def   lowercase__       (							self    :						List[str]	,       _UpperCAmelCase    :						List[Any]	,       _UpperCAmelCase    :						Optional[int]					)				->		Optional[Any]:


             '''simple docstring'''




             if len(self.retrieval_workers					) > 0:
                   # Select a random retrieval actor.
                   UpperCAmelCase_    				=							self.retrieval_workers[random.randint(0	,       len(self.retrieval_workers					) - 1					)]
                   UpperCAmelCase_    				=							ray.get(random_worker.retrieve.remote(__snake_case	,       __snake_case					)					)
             else:
                   UpperCAmelCase_    				=							self._main_retrieve(__snake_case	,       __snake_case					)
             return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__snake_case					)



       @classmethod
       def   lowercase__       (							cls    :						int	,       _UpperCAmelCase    :						Dict	,       _UpperCAmelCase    :						List[Any]=None	,       **_UpperCAmelCase    :						str					)				->		Optional[int]:


             '''simple docstring'''




             return super(__snake_case	,       cls					).get_tokenizers(__snake_case	,       __snake_case	,       **__snake_case					)







       @classmethod
       def   lowercase__       (							cls    :						Optional[Any]	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						List[str]	,       _UpperCAmelCase    :						Union[str, Any]=None	,       **_UpperCAmelCase    :						Tuple					)				->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_    				=							kwargs.pop("config"	,       __snake_case					) or RagConfig.from_pretrained(__snake_case	,       **__snake_case					)
             UpperCAmelCase_    				=							RagTokenizer.from_pretrained(__snake_case	,       config=__snake_case					)
             UpperCAmelCase_    				=							rag_tokenizer.question_encoder
             UpperCAmelCase_    				=							rag_tokenizer.generator
             if indexed_dataset is not None:
                   UpperCAmelCase_    				=							'''custom'''
                   UpperCAmelCase_    				=							CustomHFIndex(config.retrieval_vector_size	,       __snake_case					)
             else:
                   UpperCAmelCase_    				=							cls._build_index(__snake_case					)
             return cls(
                 __snake_case	,       question_encoder_tokenizer=__snake_case	,       generator_tokenizer=__snake_case	,       retrieval_workers=__snake_case	,       index=__snake_case	,       )




"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCamelCase      =		{
    """configuration_perceiver""": ["""PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PerceiverConfig""", """PerceiverOnnxConfig"""],
    """tokenization_perceiver""": ["""PerceiverTokenizer"""],
}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase      =		["""PerceiverFeatureExtractor"""]
  lowerCamelCase      =		["""PerceiverImageProcessor"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase      =		[
      """PERCEIVER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """PerceiverForImageClassificationConvProcessing""",
      """PerceiverForImageClassificationFourier""",
      """PerceiverForImageClassificationLearned""",
      """PerceiverForMaskedLM""",
      """PerceiverForMultimodalAutoencoding""",
      """PerceiverForOpticalFlow""",
      """PerceiverForSequenceClassification""",
      """PerceiverLayer""",
      """PerceiverModel""",
      """PerceiverPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_perceiver import PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP, PerceiverConfig, PerceiverOnnxConfig
  from .tokenization_perceiver import PerceiverTokenizer

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_perceiver import PerceiverFeatureExtractor
    from .image_processing_perceiver import PerceiverImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_perceiver import (
        PERCEIVER_PRETRAINED_MODEL_ARCHIVE_LIST,
        PerceiverForImageClassificationConvProcessing,
        PerceiverForImageClassificationFourier,
        PerceiverForImageClassificationLearned,
        PerceiverForMaskedLM,
        PerceiverForMultimodalAutoencoding,
        PerceiverForOpticalFlow,
        PerceiverForSequenceClassification,
        PerceiverLayer,
        PerceiverModel,
        PerceiverPreTrainedModel,
    )

else:
  import sys

  lowerCamelCase      =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




