






"""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     UpperCamelCase_ :





	def __init__(   self      :				List[Any]		)		->    Dict:
				UpperCAmelCase_  :				int      =      False





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				Optional[int]			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				List[Any]			,			lowerCAmelCase_      :				Union[str, Any]		)		->    Tuple:
				if not self.initialized:
							UpperCAmelCase_  :				Tuple      =      RagRetriever(
							    _lowerCAmelCase			,			question_encoder_tokenizer=_lowerCAmelCase			,			generator_tokenizer=_lowerCAmelCase			,			index=_lowerCAmelCase			,			init_retrieval=_lowerCAmelCase			,			)
							UpperCAmelCase_  :				Optional[int]      =      True





	def   _SCREAMING_SNAKE_CASE   (   self      :				Tuple		)		->    Dict:
				self.retriever.index.init_index()





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any			,			lowerCAmelCase_      :				List[str]			,			lowerCAmelCase_      :				Dict		)		->    int:
				UpperCAmelCase_  ,			UpperCAmelCase_  :				int      =      self.retriever._main_retrieve(_lowerCAmelCase			,			_lowerCAmelCase		)
				return doc_ids, retrieved_doc_embeds







class     UpperCamelCase_ (__A  ):





	def __init__(   self      :				Dict			,			lowerCAmelCase_      :				Tuple			,			lowerCAmelCase_      :				List[str]			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				Optional[int]=None		)		->    List[str]:
				if index is not None and index.is_initialized() and len(_lowerCAmelCase		) > 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__(
				    _lowerCAmelCase			,			question_encoder_tokenizer=_lowerCAmelCase			,			generator_tokenizer=_lowerCAmelCase			,			index=_lowerCAmelCase			,			init_retrieval=_lowerCAmelCase			,			)
				UpperCAmelCase_  :				Any      =      retrieval_workers
				if len(self.retrieval_workers		) > 0:
							ray.get(
							    [
							        worker.create_rag_retriever.remote(_lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase		)
							        for worker in self.retrieval_workers
							    ]		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    Optional[Any]:
				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   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				Optional[int]		)		->    Union[str, Any]:
				if len(self.retrieval_workers		) > 0:
							# Select a random retrieval actor.
							UpperCAmelCase_  :				Dict      =      self.retrieval_workers[random.randint(0			,			len(self.retrieval_workers		) - 1		)]
							UpperCAmelCase_  ,			UpperCAmelCase_  :				Dict      =      ray.get(random_worker.retrieve.remote(_lowerCAmelCase			,			_lowerCAmelCase		)		)
				else:
							UpperCAmelCase_  ,			UpperCAmelCase_  :				List[str]      =      self._main_retrieve(_lowerCAmelCase			,			_lowerCAmelCase		)
				return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_lowerCAmelCase		)





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				str			,			lowerCAmelCase_      :				Optional[Any]			,			lowerCAmelCase_      :				List[str]=None			,			**lowerCAmelCase_      :				Union[str, Any]		)		->    List[str]:
				return super(_lowerCAmelCase			,			cls		).get_tokenizers(_lowerCAmelCase			,			_lowerCAmelCase			,			**_lowerCAmelCase		)





	@classmethod
	def   _SCREAMING_SNAKE_CASE   (   cls      :				List[Any]			,			lowerCAmelCase_      :				List[str]			,			lowerCAmelCase_      :				Optional[int]			,			lowerCAmelCase_      :				List[str]=None			,			**lowerCAmelCase_      :				Union[str, Any]		)		->    Optional[int]:
				UpperCAmelCase_  :				Tuple      =      kwargs.pop("config"			,			_lowerCAmelCase		) or RagConfig.from_pretrained(_lowerCAmelCase			,			**_lowerCAmelCase		)
				UpperCAmelCase_  :				Optional[Any]      =      RagTokenizer.from_pretrained(_lowerCAmelCase			,			config=_lowerCAmelCase		)
				UpperCAmelCase_  :				List[Any]      =      rag_tokenizer.question_encoder
				UpperCAmelCase_  :				List[Any]      =      rag_tokenizer.generator
				if indexed_dataset is not None:
							UpperCAmelCase_  :				List[str]      =      "custom"
							UpperCAmelCase_  :				Union[str, Any]      =      CustomHFIndex(config.retrieval_vector_size			,			_lowerCAmelCase		)
				else:
							UpperCAmelCase_  :				Dict      =      cls._build_index(_lowerCAmelCase		)
				return cls(
				    _lowerCAmelCase			,			question_encoder_tokenizer=_lowerCAmelCase			,			generator_tokenizer=_lowerCAmelCase			,			retrieval_workers=_lowerCAmelCase			,			index=_lowerCAmelCase			,			)











"""simple docstring"""
from typing import TYPE_CHECKING

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


lowerCamelCase_       			=   {
    '''configuration_clip''': [
        '''CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''CLIPConfig''',
        '''CLIPOnnxConfig''',
        '''CLIPTextConfig''',
        '''CLIPVisionConfig''',
    ],
    '''processing_clip''': ['''CLIPProcessor'''],
    '''tokenization_clip''': ['''CLIPTokenizer'''],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''CLIPTokenizerFast''']

try:
			if not is_vision_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''CLIPFeatureExtractor''']
			lowerCamelCase_       			=   ['''CLIPImageProcessor''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''CLIPModel''',
			    '''CLIPPreTrainedModel''',
			    '''CLIPTextModel''',
			    '''CLIPTextModelWithProjection''',
			    '''CLIPVisionModel''',
			    '''CLIPVisionModelWithProjection''',
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''TFCLIPModel''',
			    '''TFCLIPPreTrainedModel''',
			    '''TFCLIPTextModel''',
			    '''TFCLIPVisionModel''',
			]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''FlaxCLIPModel''',
			    '''FlaxCLIPPreTrainedModel''',
			    '''FlaxCLIPTextModel''',
			    '''FlaxCLIPTextPreTrainedModel''',
			    '''FlaxCLIPVisionModel''',
			    '''FlaxCLIPVisionPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_clip import (
			    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    CLIPConfig,
			    CLIPOnnxConfig,
			    CLIPTextConfig,
			    CLIPVisionConfig,
			)
			from .processing_clip import CLIPProcessor
			from .tokenization_clip import CLIPTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_clip_fast import CLIPTokenizerFast

			try:
						if not is_vision_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .feature_extraction_clip import CLIPFeatureExtractor
						from .image_processing_clip import CLIPImageProcessor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_clip import (
						    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
						    CLIPModel,
						    CLIPPreTrainedModel,
						    CLIPTextModel,
						    CLIPTextModelWithProjection,
						    CLIPVisionModel,
						    CLIPVisionModelWithProjection,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_clip import (
						    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFCLIPModel,
						    TFCLIPPreTrainedModel,
						    TFCLIPTextModel,
						    TFCLIPVisionModel,
						)

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_clip import (
						    FlaxCLIPModel,
						    FlaxCLIPPreTrainedModel,
						    FlaxCLIPTextModel,
						    FlaxCLIPTextPreTrainedModel,
						    FlaxCLIPVisionModel,
						    FlaxCLIPVisionPreTrainedModel,
						)


else:
			import sys

			lowerCamelCase_       			=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



