



import json
import os
import pickle
import subprocess
from hashlib import mda
from pathlib import Path
from tempfile import gettempdir
from textwrap import dedent
from types import CodeType, FunctionType
from unittest import TestCase
from unittest.mock import patch

import pytest
from multiprocess import Pool
from packaging import version

import datasets
from datasets.fingerprint import Hasher, fingerprint_transform
from datasets.table import InMemoryTable

from .utils import (
    require_regex,
    require_spacy,
    require_spacy_model,
    require_tiktoken,
    require_torch,
    require_transformers,
)


class 						lowerCamelCase__			:


						def __init__(			self    :	Tuple   ,							_a    :	str      ):
													a__:							Optional[int]									=foo







						def __call__(			self    :	Union[str, Any]      ):
													return self.foo


class 						lowerCamelCase__			(		datasets.Dataset ):


						@fingerprint_transform(inplace=_a      )
						def 	_lowerCamelCase      (			self    :	Optional[int]   ,							_a    :	List[Any]   ,							*_a    :	Union[str, Any]   ,							**_a    :	Optional[Any]      ):
													return DatasetChild(self.data   ,							fingerprint=_a      )







						@fingerprint_transform(inplace=_a      )
						def 	_lowerCamelCase      (			self    :	str   ,							_a    :	Optional[Any]   ,							*_a    :	Tuple   ,							**_a    :	int      ):
													return DatasetChild(self.data   ,							fingerprint=_a      )


class 						lowerCamelCase__			:


						def __init__(			self    :	List[Any]   ,							_a    :	Optional[int]      ):
													a__:							List[str]									=callable


						def __call__(			self    :	List[Any]   ,							*_a    :	Optional[int]   ,							**_a    :	str      ):
													if self.callable is not None:
																				return self.callable(*_a   ,							**_a      )







						def __getstate__(			self    :	int      ):
													raise pickle.PicklingError()


class 						lowerCamelCase__			(		_a ):


						@require_transformers
						@pytest.mark.integration
						def 	_lowerCamelCase      (			self    :	int      ):
													from transformers import AutoTokenizer

													def encode(_a    :	Dict      ):
																				return tokenizer(_a      )

													# TODO: add hash consistency tests across sessions
													a__:							str									=AutoTokenizer.from_pretrained("bert-base-uncased"      )
													a__:							Tuple									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Union[str, Any]									=mda(datasets.utils.py_utils.dumps(lambda _a      : tokenizer(_a      )      )      ).hexdigest()
													a__:							Optional[int]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							int									=AutoTokenizer.from_pretrained("bert-base-cased"      )
													a__:							List[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Optional[Any]									=mda(datasets.utils.py_utils.dumps(lambda _a      : tokenizer(_a      )      )      ).hexdigest()
													a__:							List[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Dict									=AutoTokenizer.from_pretrained("bert-base-uncased"      )
													a__:							Tuple									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Optional[int]									=mda(datasets.utils.py_utils.dumps(lambda _a      : tokenizer(_a      )      )      ).hexdigest()
													a__:							List[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						@require_transformers
						@pytest.mark.integration
						def 	_lowerCamelCase      (			self    :	Any      ):
													from transformers import AutoTokenizer

													a__:							Tuple									=AutoTokenizer.from_pretrained("gpt2"      )
													a__:							List[str]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													tokenizer("Hello world !"      )  # call once to change the tokenizer's cache
													a__:							Union[str, Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )







						@require_regex
						def 	_lowerCamelCase      (			self    :	List[str]      ):
													import regex

													a__:							Union[str, Any]									=regex.Regex("foo"      )
													a__:							List[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							str									=regex.Regex("bar"      )
													a__:							Optional[int]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Any									=regex.Regex("foo"      )
													a__:							Any									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


class 						lowerCamelCase__			(		_a ):


						def 	_lowerCamelCase      (			self    :	int      ):
													def func():
																				return foo

													a__:							Union[str, Any]									=[0]
													a__:							int									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Optional[int]									=[1]
													a__:							Dict									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Union[str, Any]									=[0]
													a__:							Any									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													def func():
																				pass

													a__:							Optional[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()

													def func():
																				pass

													a__:							str									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Any      ):
													a__:							str									=mda(datasets.utils.py_utils.dumps(Foo([0]      )      )      ).hexdigest()
													a__:							Dict									=mda(datasets.utils.py_utils.dumps(Foo([1]      )      )      ).hexdigest()
													a__:							str									=mda(datasets.utils.py_utils.dumps(Foo([0]      )      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	str      ):
													a__:							Tuple									=mda(datasets.utils.py_utils.dumps(Foo([0]      ).__call__      )      ).hexdigest()
													a__:							Tuple									=mda(datasets.utils.py_utils.dumps(Foo([1]      ).__call__      )      ).hexdigest()
													a__:							str									=mda(datasets.utils.py_utils.dumps(Foo([0]      ).__call__      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							Dict									=(
													    "co_argcount",
													    "co_kwonlyargcount",
													    "co_nlocals",
													    "co_stacksize",
													    "co_flags",
													    "co_code",
													    "co_consts",
													    "co_names",
													    "co_varnames",
													    "co_filename",
													    "co_name",
													    "co_firstlineno",
													    "co_lnotab",
													    "co_freevars",
													    "co_cellvars",
													)

													def _create_code_pyaa(*_a    :	List[Any]      ):
																				if hasattr(args[-3]   ,							"encode"      ):
																											a__:							Optional[Any]									=list(_a      )
																											a__:							Optional[int]									=args[-3].encode()  # co_lnotab
																											a__:							List[str]									=args[-1_0].encode()  # co_code
																				if hasattr(_a   ,							"co_posonlyargcount"      ):
																											if len(_a      ) == 1_6:
																																		return CodeType(*_a      )
																											elif len(_a      ) == 1_5:
																																		return CodeType(args[0]   ,							0   ,							*args[1:]      )
																											return CodeType(args[0]   ,							0   ,							0   ,							*args[1:]      )
																				elif hasattr(_a   ,							"co_kwonlyargcount"      ):
																											if len(_a      ) == 1_6:
																																		return CodeType(args[0]   ,							*args[2:]      )
																											elif len(_a      ) == 1_5:
																																		return CodeType(*_a      )
																											return CodeType(args[0]   ,							0   ,							*args[1:]      )
																				if len(_a      ) == 1_6:
																											return CodeType(args[0]   ,							*args[3:]      )
																				elif len(_a      ) == 1_5:
																											return CodeType(args[0]   ,							*args[2:]      )
																				return CodeType(*_a      )

													def create_ipython_func(_a    :	Optional[Any]   ,							_a    :	Tuple      ):
																				def func():
																											return returned_obj

																				a__:							Optional[int]									=func.__code__
																				# Use _create_code from dill in order to make it work for different python versions
																				if datasets.config.PY_VERSION < version.parse("3.8"      ):
																											a__:							Tuple									=_create_code_pyaa(
																											    *[getattr(_a   ,							_a      ) if k != "co_filename" else co_filename for k in code_args_pyaa]      )
																				else:
																											a__:							Tuple									=code.replace(co_filename=_a      )
																				return FunctionType(_a   ,							func.__globals__   ,							func.__name__   ,							func.__defaults__   ,							func.__closure__      )

													a__				,      a__:							List[Any]									="<ipython-input-2-e0383a102aae>", [0]
													a__:							int									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													a__				,      a__:							Any									="<ipython-input-2-e0383a102aae>", [1]
													a__:							str									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													a__				,      a__:							List[str]									="<ipython-input-5-713f6613acf3>", [0]
													a__:							Union[str, Any]									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )

													a__				,      a__:							Optional[Any]									=os.path.join(gettempdir()   ,							"ipykernel_12345"   ,							"321456789.py"      ), [0]
													a__:							Dict									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													a__				,      a__:							Union[str, Any]									=os.path.join(gettempdir()   ,							"ipykernel_12345"   ,							"321456789.py"      ), [1]
													a__:							Union[str, Any]									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													a__				,      a__:							Optional[Any]									=os.path.join(gettempdir()   ,							"ipykernel_12345"   ,							"654123987.py"      ), [0]
													a__:							Optional[Any]									=mda(datasets.utils.py_utils.dumps(create_ipython_func(_a   ,							_a      )      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )







						def 	_lowerCamelCase      (			self    :	int      ):
													a__				,      a__:							Optional[int]									=[0], [1]

													def func():
																				return foo, bar

													a__:							int									="__main__"

													def globalvars_mocka_side_effect(_a    :	int   ,							*_a    :	Optional[int]   ,							**_a    :	Union[str, Any]      ):
																				return {"foo": foo, "bar": bar}

													def globalvars_mocka_side_effect(_a    :	int   ,							*_a    :	List[str]   ,							**_a    :	List[Any]      ):
																				return {"bar": bar, "foo": foo}

													with patch("dill.detect.globalvars"   ,							side_effect=_a      ) as globalvars_mocka:
																				a__:							int									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
																				self.assertGreater(globalvars_mocka.call_count   ,							0      )
													with patch("dill.detect.globalvars"   ,							side_effect=_a      ) as globalvars_mocka:
																				a__:							Any									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
																				self.assertGreater(globalvars_mocka.call_count   ,							0      )
													self.assertEqual(_a   ,							_a      )


class 						lowerCamelCase__			(		_a ):


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							Tuple									=Hasher.hash("hello"      )
													a__:							Tuple									=Hasher.hash("hello"      )
													a__:							str									=Hasher.hash("there"      )
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													a__:							Optional[int]									=Hasher.hash(Foo("hello"      )      )
													a__:							List[str]									=Hasher.hash(Foo("hello"      )      )
													a__:							List[str]									=Hasher.hash(Foo("there"      )      )
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Dict      ):
													a__:							Tuple									=Hasher()
													for x in ["hello", Foo("hello"      )]:
																				hasher.update(_a      )
													a__:							Tuple									=hasher.hexdigest()
													a__:							str									=Hasher()
													for x in ["hello", Foo("hello"      )]:
																				hasher.update(_a      )
													a__:							Union[str, Any]									=hasher.hexdigest()
													a__:							str									=Hasher()
													for x in ["there", Foo("there"      )]:
																				hasher.update(_a      )
													a__:							Union[str, Any]									=hasher.hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Tuple      ):
													with self.assertRaises(pickle.PicklingError      ):
																				Hasher.hash(UnpicklableCallable(Foo("hello"      )      )      )


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							Optional[Any]									="abc"
													a__:							Union[str, Any]									=[string, string]  # two strings have the same ids
													a__:							Optional[Any]									=[string, string]
													a__:							Optional[Any]									=json.loads(F"[\"{string}\", \"{string}\"]"      )  # two strings have different ids
													self.assertIs(obja[0]   ,							_a      )
													self.assertIs(obja[0]   ,							obja[1]      )
													self.assertIs(obja[0]   ,							_a      )
													self.assertIs(obja[0]   ,							obja[1]      )
													self.assertIsNot(obja[0]   ,							_a      )
													self.assertIsNot(obja[0]   ,							obja[1]      )
													a__:							Any									=Hasher.hash(_a      )
													a__:							str									=Hasher.hash(_a      )
													a__:							Any									=Hasher.hash(_a      )
													self.assertEqual(_a   ,							_a      )
													self.assertEqual(_a   ,							_a      )


						@require_tiktoken
						def 	_lowerCamelCase      (			self    :	Union[str, Any]      ):
													import tiktoken

													a__:							Tuple									=tiktoken.get_encoding("gpt2"      )
													a__:							List[str]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Optional[Any]									=tiktoken.get_encoding("r50k_base"      )
													a__:							str									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Tuple									=tiktoken.get_encoding("gpt2"      )
													a__:							Union[str, Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )


						@require_torch
						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													import torch

													a__:							int									=torch.tensor([1.0]      )
													a__:							Optional[int]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							str									=torch.tensor([2.0]      )
													a__:							int									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Optional[Any]									=torch.tensor([1.0]      )
													a__:							Dict									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )







						@require_spacy
						@require_spacy_model("en_core_web_sm"      )
						@require_spacy_model("fr_core_news_sm"      )
						@pytest.mark.integration
						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													import spacy

													a__:							Union[str, Any]									=spacy.load("en_core_web_sm"      )
													a__:							Optional[Any]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							int									=spacy.load("fr_core_news_sm"      )
													a__:							List[str]									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													a__:							Union[str, Any]									=spacy.load("en_core_web_sm"      )
													a__:							str									=mda(datasets.utils.py_utils.dumps(_a      )      ).hexdigest()
													self.assertEqual(_a   ,							_a      )
													self.assertNotEqual(_a   ,							_a      )





@pytest.mark.integration
def     __lowerCamelCase							(   __magic_name__						:	Path    ):
							a__:							List[Any]									=tmp_path / "dir1"
							a__:							Dict									=tmp_path / "dir2"
							dira.mkdir()
							dira.mkdir()
							a__:							Optional[int]									="script.py"
							a__:							List[Any]									=dedent(
							    "\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    "    )
							a__:							Optional[int]									=dira / script_filename
							a__:							List[Any]									=dira / script_filename
							with script_patha.open("w"    ) as f:
														f.write(__magic_name__    )
							with script_patha.open("w"    ) as f:
														f.write(__magic_name__    )
							a__:							List[Any]									=subprocess.check_output(["python", str(__magic_name__    )]    )
							a__:							Dict									=subprocess.check_output(["python", str(__magic_name__    )]    )
							assert fingerprinta == fingerprinta





def     __lowerCamelCase							(   ):
							a__:							int									={"a": [0, 1, 2]}
							a__:							str									=DatasetChild(InMemoryTable.from_pydict(__magic_name__    )    )
							a__:							Any									=dataset.funca()._fingerprint
							assert expected_fingerprint == dataset.funca()._fingerprint
							assert expected_fingerprint != dataset.funca()._fingerprint

							with Pool(2    ) as p:
														assert expected_fingerprint == p.apply_async(dataset.funca    ).get()._fingerprint
														assert expected_fingerprint != p.apply_async(dataset.funca    ).get()._fingerprint





def     __lowerCamelCase							(   ):
							a__:							List[str]									={"a": [0, 1, 2]}


							class 						lowerCamelCase__			(		datasets.Dataset ):


													@fingerprint_transform(inplace=_a      )
													def 	_lowerCamelCase      (			self    :	int   ,							_a    :	Dict      ):
																				return DummyDatasetChild(self.data   ,							fingerprint=_a      )


							a__:							int									=DummyDatasetChild(InMemoryTable.from_pydict(__magic_name__    )    ).func()


							class 						lowerCamelCase__			(		datasets.Dataset ):


													@fingerprint_transform(inplace=_a   ,							version="1.0.0"      )
													def 	_lowerCamelCase      (			self    :	Optional[int]   ,							_a    :	Optional[Any]      ):
																				return DummyDatasetChild(self.data   ,							fingerprint=_a      )


							a__:							int									=DummyDatasetChild(InMemoryTable.from_pydict(__magic_name__    )    ).func()


							class 						lowerCamelCase__			(		datasets.Dataset ):


													@fingerprint_transform(inplace=_a   ,							version="2.0.0"      )
													def 	_lowerCamelCase      (			self    :	Union[str, Any]   ,							_a    :	int      ):
																				return DummyDatasetChild(self.data   ,							fingerprint=_a      )


							a__:							List[Any]									=DummyDatasetChild(InMemoryTable.from_pydict(__magic_name__    )    ).func()

							assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}    ) == 3







def     __lowerCamelCase							(   ):
							# AttributeError: module 'dill._dill' has no attribute 'stack'
							a__:							Dict									=Hasher()
							hasher.update(lambda __magic_name__    : x    )










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    :	Optional[Any]      ):
													a__:							Tuple									=False


						def 	_lowerCamelCase      (			self    :	Dict   ,							_a    :	Any   ,							_a    :	Union[str, Any]   ,							_a    :	List[str]   ,							_a    :	Optional[int]      ):
													if not self.initialized:
																				a__:							Any									=RagRetriever(
																				    _a   ,							question_encoder_tokenizer=_a   ,							generator_tokenizer=_a   ,							index=_a   ,							init_retrieval=_a   ,							)
																				a__:							str									=True


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													self.retriever.index.init_index()


						def 	_lowerCamelCase      (			self    :	Optional[int]      ):
													# delete the old self.retriever object before assigning the new index
													del self.retriever
													a__:							Dict									=False







						def 	_lowerCamelCase      (			self    :	str   ,							_a    :	List[Any]   ,							_a    :	Any      ):
													a__				,      a__:							List[str]									=self.retriever._main_retrieve(_a   ,							_a      )
													a__:							List[str]									=self.retriever.index.get_doc_dicts(_a      )
													return doc_ids, retrieved_doc_embeds, doc_dicts


class 						lowerCamelCase__			(		_a ):


						def __init__(			self    :	Optional[int]   ,							_a    :	List[Any]   ,							_a    :	Any   ,							_a    :	Optional[Any]   ,							_a    :	Union[str, Any]   ,							_a    :	str=None      ):
													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   ,							)

													a__:							Optional[Any]									=retrieval_workers
													a__:							str									=question_encoder_tokenizer
													a__:							Optional[int]									=generator_tokenizer
													if len(self.retrieval_workers      ) > 0:
																				ray.get(
																				    [
																				        worker.create_rag_retriever.remote(_a   ,							_a   ,							_a   ,							_a      )
																				        for worker in self.retrieval_workers
																				    ]      )


						def 	_lowerCamelCase      (			self    :	List[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 	_lowerCamelCase      (			self    :	List[Any]   ,							_a    :	Union[str, Any]   ,							_a    :	Union[str, Any]      ):
													if len(self.retrieval_workers      ) > 0:
																				# Select a random retrieval actor.
																				a__:							Tuple									=self.retrieval_workers[random.randint(0   ,							len(self.retrieval_workers      ) - 1      )]
																				a__				,      a__				,      a__:							Union[str, Any]									=ray.get(
																				    random_worker.retrieve.remote(_a   ,							_a      )      )
													else:
																				a__				,      a__:							List[Any]									=self._main_retrieve(_a   ,							_a      )
																				a__:							List[str]									=self.index.get_doc_dicts(_a      )
													return retrieved_doc_embeds, doc_ids, doc_dicts


						@classmethod
						def 	_lowerCamelCase      (			cls    :	str   ,							_a    :	Union[str, Any]   ,							_a    :	Any=None   ,							**_a    :	Dict      ):
													return super(_a   ,							cls      ).get_tokenizers(_a   ,							_a   ,							**_a      )


						@classmethod
						def 	_lowerCamelCase      (			cls    :	Dict   ,							_a    :	Optional[int]   ,							_a    :	Tuple   ,							_a    :	int=None   ,							**_a    :	List[str]      ):
													a__:							Union[str, Any]									=kwargs.pop("config"   ,							_a      ) or RagConfig.from_pretrained(_a   ,							**_a      )
													a__:							List[Any]									=RagTokenizer.from_pretrained(_a   ,							config=_a      )
													a__:							str									=rag_tokenizer.question_encoder
													a__:							Any									=rag_tokenizer.generator

													if indexed_dataset is not None:
																				a__:							Dict									="custom"
																				a__:							Dict									=CustomHFIndex(config.retrieval_vector_size   ,							_a      )
													else:
																				a__:							List[str]									=cls._build_index(_a      )

													return cls(
													    _a   ,							question_encoder_tokenizer=_a   ,							generator_tokenizer=_a   ,							retrieval_workers=_a   ,							index=_a   ,							)







						def 	_lowerCamelCase      (			self    :	Union[str, Any]      ):
													logger.info("re-loading the new dataset with embeddings"      )
													# access from the training loop

													ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]      )

													# build the index object again
													a__:							str									=self._build_index(self.config      )

													ray.get(
													    [
													        worker.create_rag_retriever.remote(
													            self.config   ,							self.question_encoder_tokenizer   ,							self.generator_tokenizer   ,							_a      )
													        for worker in self.retrieval_workers
													    ]      )





