



'''simple docstring'''





import logging
import random

import ray

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


SCREAMING_SNAKE_CASE__      					=     logging.getLogger(__name__)


class 				snake_case :



       def __init__(				self      )  ->							str:
           lowercase__		       =						False



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							int:
           if not self.initialized:
               lowercase__		       =						RagRetriever(
                   UpperCAmelCase_     ,question_encoder_tokenizer=UpperCAmelCase_     ,generator_tokenizer=UpperCAmelCase_     ,index=UpperCAmelCase_     ,init_retrieval=UpperCAmelCase_     ,)
               lowercase__		       =						True



       def 							_a      (				self      )  ->							Any:
           self.retriever.index.init_index()







       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Dict:
           lowercase__       ,       lowercase__		       =						self.retriever._main_retrieve(UpperCAmelCase_     ,UpperCAmelCase_      )
           return doc_ids, retrieved_doc_embeds





class 				snake_case (UpperCamelCase   ):



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



       def 							_a      (				self      )  ->							int:
           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 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[int]:
           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(UpperCAmelCase_     ,UpperCAmelCase_      )      )
           else:
               lowercase__       ,       lowercase__		       =						self._main_retrieve(UpperCAmelCase_     ,UpperCAmelCase_      )
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(UpperCAmelCase_      )



       @classmethod
       def 							_a      (				cls     ,UpperCAmelCase_     ,UpperCAmelCase_=None     ,**UpperCAmelCase_      )  ->							int:
           return super(UpperCAmelCase_     ,cls      ).get_tokenizers(UpperCAmelCase_     ,UpperCAmelCase_     ,**UpperCAmelCase_      )







       @classmethod
       def 							_a      (				cls     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_=None     ,**UpperCAmelCase_      )  ->							Any:
           lowercase__		       =						kwargs.pop("config"     ,UpperCAmelCase_      ) or RagConfig.from_pretrained(UpperCAmelCase_     ,**UpperCAmelCase_      )
           lowercase__		       =						RagTokenizer.from_pretrained(UpperCAmelCase_     ,config=UpperCAmelCase_      )
           lowercase__		       =						rag_tokenizer.question_encoder
           lowercase__		       =						rag_tokenizer.generator
           if indexed_dataset is not None:
               lowercase__		       =						"custom"
               lowercase__		       =						CustomHFIndex(config.retrieval_vector_size     ,UpperCAmelCase_      )
           else:
               lowercase__		       =						cls._build_index(UpperCAmelCase_      )
           return cls(
               UpperCAmelCase_     ,question_encoder_tokenizer=UpperCAmelCase_     ,generator_tokenizer=UpperCAmelCase_     ,retrieval_workers=UpperCAmelCase_     ,index=UpperCAmelCase_     ,)









'''simple docstring'''





import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


SCREAMING_SNAKE_CASE__      					=     "http://www.mocksite.com/file1.txt"
SCREAMING_SNAKE_CASE__      					=     "\"text\": [\"foo\", \"foo\"]"
SCREAMING_SNAKE_CASE__      					=     "6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8"


class 				snake_case :
       lowerCAmelCase__  :Union[str, Any]						   =					200
       lowerCAmelCase__  :Tuple						   =					{"Content-Length": "100"}
       lowerCAmelCase__  :Any						   =					{}



       def 							_a      (				self     ,**UpperCAmelCase_      )  ->							List[Any]:
           return [bytes(UpperCAmelCase_     ,"utf-8"      )]





def        lowerCamelCase  (				*_snake_case  :						int     ,**_snake_case  :						int       ):





    '''simple docstring'''


    return MockResponse()




@pytest.mark.parametrize("urls_type"     ,[str, list, dict]       )
def        lowerCamelCase  (				_snake_case  :						str     ,_snake_case  :						List[Any]     ,_snake_case  :						List[Any]       ):





    '''simple docstring'''


    import requests

    monkeypatch.setattr(_snake_case     ,"request"     ,_snake_case       )

    lowercase__		       =						URL
    if issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						url
    elif issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						[url]
    elif issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						{"train": url}
    lowercase__		       =						"dummy"
    lowercase__		       =						"downloads"
    lowercase__		       =						tmp_path
    lowercase__		       =						DownloadConfig(
        cache_dir=os.path.join(_snake_case     ,_snake_case       )     ,use_etag=_snake_case     ,)
    lowercase__		       =						DownloadManager(dataset_name=_snake_case     ,download_config=_snake_case       )
    lowercase__		       =						dl_manager.download(_snake_case       )
    lowercase__		       =						urls
    for downloaded_paths in [downloaded_paths]:
        if isinstance(_snake_case     ,_snake_case       ):
            lowercase__		       =						[downloaded_paths]
            lowercase__		       =						[urls]
        elif isinstance(_snake_case     ,_snake_case       ):
            assert "train" in downloaded_paths.keys()
            lowercase__		       =						downloaded_paths.values()
            lowercase__		       =						urls.values()
        assert downloaded_paths
        for downloaded_path, input_url in zip(_snake_case     ,_snake_case       ):
            assert downloaded_path == dl_manager.downloaded_paths[input_url]
            lowercase__		       =						Path(_snake_case       )
            lowercase__		       =						downloaded_path.parts
            assert parts[-1] == HASH
            assert parts[-2] == cache_subdir
            assert downloaded_path.exists()
            lowercase__		       =						downloaded_path.read_text()
            assert content == CONTENT
            lowercase__		       =						downloaded_path.with_suffix(".json"       )
            assert metadata_downloaded_path.exists()
            lowercase__		       =						json.loads(metadata_downloaded_path.read_text()       )
            assert metadata_content == {"url": URL, "etag": None}




@pytest.mark.parametrize("paths_type"     ,[str, list, dict]       )
def        lowerCamelCase  (				_snake_case  :						Optional[Any]     ,_snake_case  :						List[str]     ,_snake_case  :						List[str]       ):





    '''simple docstring'''


    lowercase__		       =						str(_snake_case       )
    if issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						filename
    elif issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						[filename]
    elif issubclass(_snake_case     ,_snake_case       ):
        lowercase__		       =						{"train": filename}
    lowercase__		       =						"dummy"
    lowercase__		       =						xz_file.parent
    lowercase__		       =						"extracted"
    lowercase__		       =						DownloadConfig(
        cache_dir=_snake_case     ,use_etag=_snake_case     ,)
    lowercase__		       =						DownloadManager(dataset_name=_snake_case     ,download_config=_snake_case       )
    lowercase__		       =						dl_manager.extract(_snake_case       )
    lowercase__		       =						paths
    for extracted_paths in [extracted_paths]:
        if isinstance(_snake_case     ,_snake_case       ):
            lowercase__		       =						[extracted_paths]
            lowercase__		       =						[paths]
        elif isinstance(_snake_case     ,_snake_case       ):
            assert "train" in extracted_paths.keys()
            lowercase__		       =						extracted_paths.values()
            lowercase__		       =						paths.values()
        assert extracted_paths
        for extracted_path, input_path in zip(_snake_case     ,_snake_case       ):
            assert extracted_path == dl_manager.extracted_paths[input_path]
            lowercase__		       =						Path(_snake_case       )
            lowercase__		       =						extracted_path.parts
            assert parts[-1] == hash_url_to_filename(_snake_case     ,etag=_snake_case       )
            assert parts[-2] == extracted_subdir
            assert extracted_path.exists()
            lowercase__		       =						extracted_path.read_text()
            lowercase__		       =						text_file.read_text()
            assert extracted_file_content == expected_file_content




def        lowerCamelCase  (				_snake_case  :						Tuple     ,_snake_case  :						List[str]       ):





    '''simple docstring'''


    assert path.endswith(".jsonl"       )
    for num_items, line in enumerate(_snake_case     ,start=1       ):
        lowercase__		       =						json.loads(line.decode("utf-8"       )       )
        assert item.keys() == {"col_1", "col_2", "col_3"}
    assert num_items == 4




@pytest.mark.parametrize("archive_jsonl"     ,["tar_jsonl_path", "zip_jsonl_path"]       )
def        lowerCamelCase  (				_snake_case  :						Any     ,_snake_case  :						Union[str, Any]       ):





    '''simple docstring'''


    lowercase__		       =						request.getfixturevalue(_snake_case       )
    lowercase__		       =						DownloadManager()
    for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(_snake_case       )     ,start=1       ):
        _test_jsonl(_snake_case     ,_snake_case       )
    assert num_jsonl == 2




@pytest.mark.parametrize("archive_nested_jsonl"     ,["tar_nested_jsonl_path", "zip_nested_jsonl_path"]       )
def        lowerCamelCase  (				_snake_case  :						Optional[Any]     ,_snake_case  :						Dict       ):





    '''simple docstring'''


    lowercase__		       =						request.getfixturevalue(_snake_case       )
    lowercase__		       =						DownloadManager()
    for num_tar, (path, file) in enumerate(dl_manager.iter_archive(_snake_case       )     ,start=1       ):
        for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(_snake_case       )     ,start=1       ):
            _test_jsonl(_snake_case     ,_snake_case       )
    assert num_tar == 1
    assert num_jsonl == 2




def        lowerCamelCase  (				_snake_case  :						int       ):





    '''simple docstring'''


    lowercase__		       =						DownloadManager()
    for num_file, file in enumerate(dl_manager.iter_files(_snake_case       )     ,start=1       ):
        assert os.path.basename(_snake_case       ) == ("test.txt" if num_file == 1 else "train.txt")
    assert num_file == 2




