


"""simple docstring"""






import importlib
import os
import platform
from pathlib import Path

from packaging import version

from .utils.logging import get_logger


SCREAMING_SNAKE_CASE__				:  Optional[int]						= get_logger(__name__)

# Datasets
SCREAMING_SNAKE_CASE__				:  int						= "https://s3.amazonaws.com/datasets.huggingface.co/datasets/datasets"
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= "https://cdn-datasets.huggingface.co/datasets/datasets"
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= "https://raw.githubusercontent.com/huggingface/datasets/{revision}/datasets/{path}/{name}"

# Metrics
SCREAMING_SNAKE_CASE__				:  Optional[int]						= "https://s3.amazonaws.com/datasets.huggingface.co/datasets/metrics"
SCREAMING_SNAKE_CASE__				:  Optional[int]						= "https://cdn-datasets.huggingface.co/datasets/metric"
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= "https://raw.githubusercontent.com/huggingface/datasets/{revision}/metrics/{path}/{name}"

# Hub
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.environ.get("HF_ENDPOINT", "https://huggingface.co")
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= HF_ENDPOINT + "/datasets/{repo_id}/resolve/{revision}/{path}"
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= "main"

SCREAMING_SNAKE_CASE__				:  List[str]						= version.parse(platform.python_version())

if PY_VERSION < version.parse("3.8"):
 import importlib_metadata
else:
 import importlib.metadata as importlib_metadata

# General environment variables accepted values for booleans
SCREAMING_SNAKE_CASE__				:  List[str]						= {"1", "ON", "YES", "TRUE"}
SCREAMING_SNAKE_CASE__				:  List[Any]						= ENV_VARS_TRUE_VALUES.union({"AUTO"})


# Imports
SCREAMING_SNAKE_CASE__				:  int						= version.parse(importlib_metadata.version("dill"))
SCREAMING_SNAKE_CASE__				:  str						= version.parse(importlib_metadata.version("pandas"))
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= version.parse(importlib_metadata.version("pyarrow"))

SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.environ.get("USE_TF", "AUTO").upper()
SCREAMING_SNAKE_CASE__				:  Tuple						= os.environ.get("USE_TORCH", "AUTO").upper()
SCREAMING_SNAKE_CASE__				:  Tuple						= os.environ.get("USE_JAX", "AUTO").upper()

SCREAMING_SNAKE_CASE__				:  Any						= "N/A"
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= False

if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
 SCREAMING_SNAKE_CASE__				:  Optional[int]						= importlib.util.find_spec("torch") is not None
 if TORCH_AVAILABLE:
  try:
   SCREAMING_SNAKE_CASE__				:  Tuple						= version.parse(importlib_metadata.version("torch"))
   logger.info(F'PyTorch version {TORCH_VERSION} available.')
  except importlib_metadata.PackageNotFoundError:
   pass
else:
 logger.info("Disabling PyTorch because USE_TF is set")

SCREAMING_SNAKE_CASE__				:  List[Any]						= "N/A"
SCREAMING_SNAKE_CASE__				:  Dict						= False

if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES:
 SCREAMING_SNAKE_CASE__				:  List[Any]						= importlib.util.find_spec("tensorflow") is not None
 if TF_AVAILABLE:
  # For the metadata, we have to look for both tensorflow and tensorflow-cpu
  for package in [
      "tensorflow",
      "tensorflow-cpu",
      "tensorflow-gpu",
      "tf-nightly",
      "tf-nightly-cpu",
      "tf-nightly-gpu",
      "intel-tensorflow",
      "tensorflow-rocm",
      "tensorflow-macos",
  ]:
   try:
    SCREAMING_SNAKE_CASE__				:  Any						= version.parse(importlib_metadata.version(package))
   except importlib_metadata.PackageNotFoundError:
    continue
   else:
    break
  else:
   SCREAMING_SNAKE_CASE__				:  Optional[Any]						= False
 if TF_AVAILABLE:
  if TF_VERSION.major < 2:
   logger.info(F'TensorFlow found but with version {TF_VERSION}. `datasets` requires version 2 minimum.')
   SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= False
  else:
   logger.info(F'TensorFlow version {TF_VERSION} available.')
else:
 logger.info("Disabling Tensorflow because USE_TORCH is set")


SCREAMING_SNAKE_CASE__				:  Any						= "N/A"
SCREAMING_SNAKE_CASE__				:  Optional[int]						= False

if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES:
 SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= importlib.util.find_spec("jax") is not None and importlib.util.find_spec("jaxlib") is not None
 if JAX_AVAILABLE:
  try:
   SCREAMING_SNAKE_CASE__				:  int						= version.parse(importlib_metadata.version("jax"))
   logger.info(F'JAX version {JAX_VERSION} available.')
  except importlib_metadata.PackageNotFoundError:
   pass
else:
 logger.info("Disabling JAX because USE_JAX is set to False")


SCREAMING_SNAKE_CASE__				:  Optional[Any]						= os.environ.get("USE_BEAM", "AUTO").upper()
SCREAMING_SNAKE_CASE__				:  Dict						= "N/A"
SCREAMING_SNAKE_CASE__				:  Optional[int]						= False
if USE_BEAM in ENV_VARS_TRUE_AND_AUTO_VALUES:
 try:
  SCREAMING_SNAKE_CASE__				:  Tuple						= version.parse(importlib_metadata.version("apache_beam"))
  SCREAMING_SNAKE_CASE__				:  Optional[Any]						= True
  logger.info(F'Apache Beam version {BEAM_VERSION} available.')
 except importlib_metadata.PackageNotFoundError:
  pass
else:
 logger.info("Disabling Apache Beam because USE_BEAM is set to False")


# Optional tools for data loading
SCREAMING_SNAKE_CASE__				:  List[str]						= importlib.util.find_spec("sqlalchemy") is not None

# Optional tools for feature decoding
SCREAMING_SNAKE_CASE__				:  Optional[int]						= importlib.util.find_spec("PIL") is not None
SCREAMING_SNAKE_CASE__				:  Any						= importlib.util.find_spec("soundfile") is not None and version.parse(
    importlib.import_module("soundfile").__libsndfile_version__
) >= version.parse("1.0.31")
SCREAMING_SNAKE_CASE__				:  List[str]						= importlib.util.find_spec("soundfile") is not None and version.parse(
    importlib.import_module("soundfile").__libsndfile_version__
) >= version.parse("1.1.0")

# Optional compression tools
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= importlib.util.find_spec("rarfile") is not None
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= importlib.util.find_spec("zstandard") is not None
SCREAMING_SNAKE_CASE__				:  List[str]						= importlib.util.find_spec("lz4") is not None
SCREAMING_SNAKE_CASE__				:  Dict						= importlib.util.find_spec("py7zr") is not None


# Cache location
SCREAMING_SNAKE_CASE__				:  int						= "~/.cache"
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= os.getenv("XDG_CACHE_HOME", DEFAULT_XDG_CACHE_HOME)
SCREAMING_SNAKE_CASE__				:  List[str]						= os.path.join(XDG_CACHE_HOME, "huggingface")
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.path.expanduser(os.getenv("HF_HOME", DEFAULT_HF_CACHE_HOME))

SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.path.join(HF_CACHE_HOME, "datasets")
SCREAMING_SNAKE_CASE__				:  List[str]						= Path(os.getenv("HF_DATASETS_CACHE", DEFAULT_HF_DATASETS_CACHE))

SCREAMING_SNAKE_CASE__				:  Dict						= os.path.join(HF_CACHE_HOME, "metrics")
SCREAMING_SNAKE_CASE__				:  List[str]						= Path(os.getenv("HF_METRICS_CACHE", DEFAULT_HF_METRICS_CACHE))

SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.path.join(HF_CACHE_HOME, "modules")
SCREAMING_SNAKE_CASE__				:  Dict						= Path(os.getenv("HF_MODULES_CACHE", DEFAULT_HF_MODULES_CACHE))

SCREAMING_SNAKE_CASE__				:  Tuple						= "downloads"
SCREAMING_SNAKE_CASE__				:  Tuple						= os.path.join(HF_DATASETS_CACHE, DOWNLOADED_DATASETS_DIR)
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= Path(os.getenv("HF_DATASETS_DOWNLOADED_DATASETS_PATH", DEFAULT_DOWNLOADED_DATASETS_PATH))

SCREAMING_SNAKE_CASE__				:  Tuple						= "extracted"
SCREAMING_SNAKE_CASE__				:  Tuple						= os.path.join(DEFAULT_DOWNLOADED_DATASETS_PATH, EXTRACTED_DATASETS_DIR)
SCREAMING_SNAKE_CASE__				:  Any						= Path(os.getenv("HF_DATASETS_EXTRACTED_DATASETS_PATH", DEFAULT_EXTRACTED_DATASETS_PATH))

# Download count for the website
SCREAMING_SNAKE_CASE__				:  List[str]						= (
    os.environ.get("HF_UPDATE_DOWNLOAD_COUNTS", "AUTO").upper() in ENV_VARS_TRUE_AND_AUTO_VALUES
)

# Batch size constants. For more info, see:
# https://github.com/apache/arrow/blob/master/docs/source/cpp/arrays.rst#size-limitations-and-recommendations)
SCREAMING_SNAKE_CASE__				:  Optional[int]						= 1000

# Size of the preloaded record batch in `Dataset.__iter__`
SCREAMING_SNAKE_CASE__				:  List[str]						= 10

# Pickling tables works only for small tables (<4GiB)
# For big tables, we write them on disk instead
SCREAMING_SNAKE_CASE__				:  Optional[int]						= 4 << 30

# Max shard size in bytes (e.g. to shard parquet datasets in push_to_hub or download_and_prepare)
SCREAMING_SNAKE_CASE__				:  List[Any]						= "500MB"

# Parquet configuration
SCREAMING_SNAKE_CASE__				:  int						= 100
SCREAMING_SNAKE_CASE__				:  int						= 100
SCREAMING_SNAKE_CASE__				:  List[Any]						= 100

# Offline mode
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= os.environ.get("HF_DATASETS_OFFLINE", "AUTO").upper() in ENV_VARS_TRUE_VALUES

# In-memory
SCREAMING_SNAKE_CASE__				:  Dict						= 0  # Disabled
SCREAMING_SNAKE_CASE__				:  Optional[int]						= float(os.environ.get("HF_DATASETS_IN_MEMORY_MAX_SIZE", DEFAULT_IN_MEMORY_MAX_SIZE))

# File names
SCREAMING_SNAKE_CASE__				:  Any						= "dataset.arrow"
SCREAMING_SNAKE_CASE__				:  List[Any]						= "indices.arrow"
SCREAMING_SNAKE_CASE__				:  int						= "state.json"
SCREAMING_SNAKE_CASE__				:  Union[str, Any]						= "dataset_info.json"
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= "dataset_infos.json"
SCREAMING_SNAKE_CASE__				:  Optional[int]						= "LICENSE"
SCREAMING_SNAKE_CASE__				:  str						= "metric_info.json"
SCREAMING_SNAKE_CASE__				:  Any						= "dataset_dict.json"

SCREAMING_SNAKE_CASE__				:  Dict						= "datasets_modules"

SCREAMING_SNAKE_CASE__				:  Optional[int]						= 255

# Streaming
SCREAMING_SNAKE_CASE__				:  List[str]						= 20
SCREAMING_SNAKE_CASE__				:  List[Any]						= 5

# Datasets without script
SCREAMING_SNAKE_CASE__				:  Dict						= 200
SCREAMING_SNAKE_CASE__				:  Optional[int]						= 10
SCREAMING_SNAKE_CASE__				:  Optional[Any]						= 200

# Progress bars
SCREAMING_SNAKE_CASE__				:  List[Any]						= 0.0_5  # 20 progress updates per sec









"""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__				:  Optional[Any]						= logging.getLogger(__name__)


class 	A_       :





  """simple docstring"""

  def __init__(						self						)				->					Dict:
     a   :	List[str]						=    False
  def 							lowercase_     (						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase						)				->					List[str]:
     if not self.initialized:
        a   :	Dict						=    RagRetriever(
            __UpperCAmelCase  ,					question_encoder_tokenizer=__UpperCAmelCase  ,					generator_tokenizer=__UpperCAmelCase  ,					index=__UpperCAmelCase  ,					init_retrieval=__UpperCAmelCase  ,					)
        a   :	Dict						=    True
  def 							lowercase_     (						self						)				->					List[Any]:
     self.retriever.index.init_index()
  def 							lowercase_     (						self						)				->					List[Any]:
     # delete the old self.retriever object before assigning the new index
     del self.retriever
     a   :	Dict						=    False






  def 							lowercase_     (						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase						)				->					Optional[int]:
     a	, a   :	List[str]						=    self.retriever._main_retrieve(__UpperCAmelCase  ,					__UpperCAmelCase						)
     a   :	List[Any]						=    self.retriever.index.get_doc_dicts(__UpperCAmelCase						)
     return doc_ids, retrieved_doc_embeds, doc_dicts



class 	A_       (     _UpperCAmelCase   ):





  """simple docstring"""

  def __init__(						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase=None						)				->					Tuple:
     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  ,					)

     a   :	Tuple						=    retrieval_workers
     a   :	List[str]						=    question_encoder_tokenizer
     a   :	Union[str, Any]						=    generator_tokenizer
     if len(self.retrieval_workers						) > 0:
        ray.get(
            [
                worker.create_rag_retriever.remote(__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase						)
                for worker in self.retrieval_workers
            ]						)
  def 							lowercase_     (						self						)				->					Tuple:
     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  ,					__UpperCAmelCase  ,					__UpperCAmelCase						)				->					Tuple:
     if len(self.retrieval_workers						) > 0:
        # Select a random retrieval actor.
        a   :	Dict						=    self.retrieval_workers[random.randint(0  ,					len(self.retrieval_workers						) - 1						)]
        a	, a	, a   :	Any						=    ray.get(
            random_worker.retrieve.remote(__UpperCAmelCase  ,					__UpperCAmelCase						)						)
     else:
        a	, a   :	str						=    self._main_retrieve(__UpperCAmelCase  ,					__UpperCAmelCase						)
        a   :	str						=    self.index.get_doc_dicts(__UpperCAmelCase						)
     return retrieved_doc_embeds, doc_ids, doc_dicts
  @classmethod
  def 							lowercase_     (						cls  ,					__UpperCAmelCase  ,					__UpperCAmelCase=None  ,					**__UpperCAmelCase						)				->					int:
     return super(__UpperCAmelCase  ,					cls						).get_tokenizers(__UpperCAmelCase  ,					__UpperCAmelCase  ,					**__UpperCAmelCase						)
  @classmethod
  def 							lowercase_     (						cls  ,					__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase=None  ,					**__UpperCAmelCase						)				->					str:
     a   :	Union[str, Any]						=    kwargs.pop('config'  ,					__UpperCAmelCase						) or RagConfig.from_pretrained(__UpperCAmelCase  ,					**__UpperCAmelCase						)
     a   :	Optional[Any]						=    RagTokenizer.from_pretrained(__UpperCAmelCase  ,					config=__UpperCAmelCase						)
     a   :	List[str]						=    rag_tokenizer.question_encoder
     a   :	int						=    rag_tokenizer.generator

     if indexed_dataset is not None:
        a   :	Optional[int]						=    'custom'
        a   :	str						=    CustomHFIndex(config.retrieval_vector_size  ,					__UpperCAmelCase						)
     else:
        a   :	Dict						=    cls._build_index(__UpperCAmelCase						)

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






  def 							lowercase_     (						self						)				->					List[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   :	Optional[Any]						=    self._build_index(self.config						)

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





