


"""simple docstring"""

import logging
import random

import ray

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


UpperCAmelCase__  							=logging.getLogger(__name__)





class    lowerCamelCase__						:
      def __init__(		self	:	Dict							):






       '''simple docstring'''


       __lowercase									=	False
      def SCREAMING_SNAKE_CASE_	(		self	:	Any      ,       A_	:	List[Any]      ,       A_	:	List[Any]      ,       A_	:	int      ,       A_	:	str							):






       '''simple docstring'''


       if not self.initialized:
        __lowercase									=	RagRetriever(
            A_      ,       question_encoder_tokenizer=A_      ,       generator_tokenizer=A_      ,       index=A_      ,       init_retrieval=A_      ,       )
        __lowercase									=	True
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       self.retriever.index.init_index()






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]      ,       A_	:	int      ,       A_	:	Optional[int]							):






       '''simple docstring'''


       __lowercase							,		__lowercase									=	self.retriever._main_retrieve(A_      ,       A_							)
       return doc_ids, retrieved_doc_embeds






class    lowerCamelCase__						(		_a    ):
      def __init__(		self	:	Tuple      ,       A_	:	Optional[int]      ,       A_	:	Dict      ,       A_	:	Optional[Any]      ,       A_	:	List[str]      ,       A_	:	str=None							):






       '''simple docstring'''


       if index is not None and index.is_initialized() and len(A_							) > 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__(
           A_      ,       question_encoder_tokenizer=A_      ,       generator_tokenizer=A_      ,       index=A_      ,       init_retrieval=A_      ,       )
       __lowercase									=	retrieval_workers
       if len(self.retrieval_workers							) > 0:
        ray.get(
            [
                worker.create_rag_retriever.remote(A_      ,       A_      ,       A_      ,       A_							)
                for worker in self.retrieval_workers
            ]							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''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 SCREAMING_SNAKE_CASE_	(		self	:	List[str]      ,       A_	:	int      ,       A_	:	str							):






       '''simple docstring'''


       if len(self.retrieval_workers							) > 0:
        # Select a random retrieval actor.
        __lowercase									=	self.retrieval_workers[random.randint(0      ,       len(self.retrieval_workers							) - 1							)]
        __lowercase							,		__lowercase									=	ray.get(random_worker.retrieve.remote(A_      ,       A_							)							)
       else:
        __lowercase							,		__lowercase									=	self._main_retrieve(A_      ,       A_							)
       return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(A_							)
      @classmethod
      def SCREAMING_SNAKE_CASE_	(		cls	:	Dict      ,       A_	:	Union[str, Any]      ,       A_	:	Union[str, Any]=None      ,       **A_	:	str							):






       '''simple docstring'''


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






      @classmethod
      def SCREAMING_SNAKE_CASE_	(		cls	:	Any      ,       A_	:	str      ,       A_	:	str      ,       A_	:	int=None      ,       **A_	:	Optional[int]							):






       '''simple docstring'''


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




"""simple docstring"""

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def lowerCAmelCase_      (				UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Any ):






 """simple docstring"""







 assert isinstance(UpperCamelCase__				,   UpperCamelCase__ )
 assert dataset.num_rows == 4
 assert dataset.num_columns == 1
 assert dataset.column_names == ["text"]
 for feature, expected_dtype in expected_features.items():
  assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""				,   [False, True] )
def lowerCAmelCase_      (				UpperCamelCase__	: Dict				,   UpperCamelCase__	: Optional[int]				,   UpperCamelCase__	: List[str] ):






 """simple docstring"""







 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
  __lowercase									=	TextDatasetReader(UpperCamelCase__				,   cache_dir=UpperCamelCase__				,   keep_in_memory=UpperCamelCase__ ).read()
 _check_text_dataset(UpperCamelCase__				,   UpperCamelCase__ )


@pytest.mark.parametrize(
    """features"""				,   [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,   )
def lowerCAmelCase_      (				UpperCamelCase__	: List[Any]				,   UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Union[str, Any] ):






 """simple docstring"""







 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 __lowercase									=	features.copy() if features else default_expected_features
 __lowercase									=	(
     Features({feature: Value(UpperCamelCase__ ) for feature, dtype in features.items()} ) if features is not None else None
 )
 __lowercase									=	TextDatasetReader(UpperCamelCase__				,   features=UpperCamelCase__				,   cache_dir=UpperCamelCase__ ).read()
 _check_text_dataset(UpperCamelCase__				,   UpperCamelCase__ )


@pytest.mark.parametrize("""split"""				,   [None, NamedSplit("""train""" ), """train""", """test"""] )
def lowerCAmelCase_      (				UpperCamelCase__	: Union[str, Any]				,   UpperCamelCase__	: str				,   UpperCamelCase__	: Tuple ):






 """simple docstring"""







 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 __lowercase									=	TextDatasetReader(UpperCamelCase__				,   cache_dir=UpperCamelCase__				,   split=UpperCamelCase__ ).read()
 _check_text_dataset(UpperCamelCase__				,   UpperCamelCase__ )
 assert dataset.split == split if split else "train"


@pytest.mark.parametrize("""path_type"""				,   [str, list] )
def lowerCAmelCase_      (				UpperCamelCase__	: Union[str, Any]				,   UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Any ):






 """simple docstring"""







 if issubclass(UpperCamelCase__				,   UpperCamelCase__ ):
  __lowercase									=	text_path
 elif issubclass(UpperCamelCase__				,   UpperCamelCase__ ):
  __lowercase									=	[text_path]
 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 __lowercase									=	TextDatasetReader(UpperCamelCase__				,   cache_dir=UpperCamelCase__ ).read()
 _check_text_dataset(UpperCamelCase__				,   UpperCamelCase__ )


def lowerCAmelCase_      (				UpperCamelCase__	: str				,   UpperCamelCase__	: str				,   UpperCamelCase__	: Any=("train",) ):






 """simple docstring"""







 assert isinstance(UpperCamelCase__				,   UpperCamelCase__ )
 for split in splits:
  __lowercase									=	dataset_dict[split]
  assert dataset.num_rows == 4
  assert dataset.num_columns == 1
  assert dataset.column_names == ["text"]
  for feature, expected_dtype in expected_features.items():
   assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""				,   [False, True] )
def lowerCAmelCase_      (				UpperCamelCase__	: int				,   UpperCamelCase__	: List[str]				,   UpperCamelCase__	: List[Any] ):






 """simple docstring"""







 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
  __lowercase									=	TextDatasetReader({"""train""": text_path}				,   cache_dir=UpperCamelCase__				,   keep_in_memory=UpperCamelCase__ ).read()
 _check_text_datasetdict(UpperCamelCase__				,   UpperCamelCase__ )


@pytest.mark.parametrize(
    """features"""				,   [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,   )
def lowerCAmelCase_      (				UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: Dict				,   UpperCamelCase__	: Union[str, Any] ):






 """simple docstring"""







 __lowercase									=	tmp_path / """cache"""
 # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
 __lowercase									=	{"""text""": """string"""}
 __lowercase									=	features.copy() if features else default_expected_features
 __lowercase									=	(
     Features({feature: Value(UpperCamelCase__ ) for feature, dtype in features.items()} ) if features is not None else None
 )
 __lowercase									=	TextDatasetReader({"""train""": text_path}				,   features=UpperCamelCase__				,   cache_dir=UpperCamelCase__ ).read()
 _check_text_datasetdict(UpperCamelCase__				,   UpperCamelCase__ )


@pytest.mark.parametrize("""split"""				,   [None, NamedSplit("""train""" ), """train""", """test"""] )
def lowerCAmelCase_      (				UpperCamelCase__	: Dict				,   UpperCamelCase__	: int				,   UpperCamelCase__	: int ):






 """simple docstring"""







 if split:
  __lowercase									=	{split: text_path}
 else:
  __lowercase									=	"""train"""
  __lowercase									=	{"""train""": text_path, """test""": text_path}
 __lowercase									=	tmp_path / """cache"""
 __lowercase									=	{"""text""": """string"""}
 __lowercase									=	TextDatasetReader(UpperCamelCase__				,   cache_dir=UpperCamelCase__ ).read()
 _check_text_datasetdict(UpperCamelCase__				,   UpperCamelCase__				,   splits=list(path.keys() ) )
 assert all(dataset[split].split == split for split in path.keys() )
