



import os
import sys
import warnings
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen
from ..table import array_cast
from ..utils.file_utils import is_local_path
from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
		import PIL.Image

		from .features import FeatureType


snake_case    :						Optional[List[str]]  =						None
snake_case    :						Any  =						'<' if sys.byteorder == 'little' else '>'
# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
snake_case    :						List[Any]  =						[
    np.dtype('|b1'),
    np.dtype('|u1'),
    np.dtype('<u2'),
    np.dtype('>u2'),
    np.dtype('<i2'),
    np.dtype('>i2'),
    np.dtype('<u4'),
    np.dtype('>u4'),
    np.dtype('<i4'),
    np.dtype('>i4'),
    np.dtype('<f4'),
    np.dtype('>f4'),
    np.dtype('<f8'),
    np.dtype('>f8'),
]




@dataclass
class 							lowerCAmelCase__  :
					__A     :    bool           =						True
					__A     :    Optional[str]           =						None
					# Automatically constructed
					__A     :    ClassVar[str]           =						"PIL.Image.Image"
					__A     :    ClassVar[Any]           =						pa.struct({'bytes': pa.binary(), 'path': pa.string()}						)
					__A     :    str           =						field(default='Image'					,							init=UpperCamelCase					,							repr=UpperCamelCase						)







					def __call__(       self      :     List[str]):
											return self.pa_type







					def  _lowercase			(       self      :     Dict			,		_A      :     Union[str, bytes, dict, np.ndarray, "PIL.Image.Image"]):
											if config.PIL_AVAILABLE:
																	import PIL.Image
											else:
																	raise ImportError("To support encoding images, please install 'Pillow'.")

											if isinstance(_A			,		_A):
																	A__	:							Dict						    = np.array(_A)

											if isinstance(_A			,		_A):
																	return {"path": value, "bytes": None}
											elif isinstance(_A			,		_A):
																	return {"path": None, "bytes": value}
											elif isinstance(_A			,		np.ndarray):
																	# convert the image array to PNG/TIFF bytes
																	return encode_np_array(_A)
											elif isinstance(_A			,		PIL.Image.Image):
																	# convert the PIL image to bytes (default format is PNG/TIFF)
																	return encode_pil_image(_A)
											elif value.get("path") is not None and os.path.isfile(value["path"]):
																	# we set "bytes": None to not duplicate the data if they're already available locally
																	return {"bytes": None, "path": value.get("path")}
											elif value.get("bytes") is not None or value.get("path") is not None:
																	# store the image bytes, and path is used to infer the image format using the file extension
																	return {"bytes": value.get("bytes"), "path": value.get("path")}
											else:
																	raise ValueError(
																	    F'An image sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.')







					def  _lowercase			(       self      :     str			,		_A      :     dict			,		_A      :     Any=None):
											if not self.decode:
																	raise RuntimeError("Decoding is disabled for this feature. Please use Image(decode=True) instead.")

											if config.PIL_AVAILABLE:
																	import PIL.Image
											else:
																	raise ImportError("To support decoding images, please install 'Pillow'.")

											if token_per_repo_id is None:
																	A__	:							Any						    = {}

											A__    ,       A__	:							List[str]						    = value["path"], value["bytes"]
											if bytes_ is None:
																	if path is None:
																							raise ValueError(F'An image should have one of \'path\' or \'bytes\' but both are None in {value}.')
																	else:
																							if is_local_path(_A):
																													A__	:							Optional[int]						    = PIL.Image.open(_A)
																							else:
																													A__	:							Union[str, Any]						    = path.split("::")[-1]
																													try:
																																			A__	:							List[str]						    = string_to_dict(_A			,		config.HUB_DATASETS_URL)["repo_id"]
																																			A__	:							List[Any]						    = token_per_repo_id.get(_A)
																													except ValueError:
																																			A__	:							Optional[int]						    = None
																													with xopen(_A			,		"rb"			,		use_auth_token=_A) as f:
																																			A__	:							Tuple						    = BytesIO(f.read())
																													A__	:							List[str]						    = PIL.Image.open(bytes_)
											else:
																	A__	:							Optional[Any]						    = PIL.Image.open(BytesIO(bytes_))
											image.load()  # to avoid "Too many open files" errors
											return image







					def  _lowercase			(       self      :     Optional[Any]):
											from .features import Value

											return (
											    self
											    if self.decode
											    else {
											        "bytes": Value("binary"),
											        "path": Value("string"),
											    }
											)







					def  _lowercase			(       self      :     List[Any]			,		_A      :     Union[pa.StringArray, pa.StructArray, pa.ListArray]):
											if pa.types.is_string(storage.type):
																	A__	:							Any						    = pa.array([None] * len(_A)			,		type=pa.binary())
																	A__	:							Any						    = pa.StructArray.from_arrays([bytes_array, storage]			,		["bytes", "path"]			,		mask=storage.is_null())
											elif pa.types.is_binary(storage.type):
																	A__	:							List[str]						    = pa.array([None] * len(_A)			,		type=pa.string())
																	A__	:							Dict						    = pa.StructArray.from_arrays([storage, path_array]			,		["bytes", "path"]			,		mask=storage.is_null())
											elif pa.types.is_struct(storage.type):
																	if storage.type.get_field_index("bytes") >= 0:
																							A__	:							Optional[Any]						    = storage.field("bytes")
																	else:
																							A__	:							Tuple						    = pa.array([None] * len(_A)			,		type=pa.binary())
																	if storage.type.get_field_index("path") >= 0:
																							A__	:							List[str]						    = storage.field("path")
																	else:
																							A__	:							List[str]						    = pa.array([None] * len(_A)			,		type=pa.string())
																	A__	:							int						    = pa.StructArray.from_arrays([bytes_array, path_array]			,		["bytes", "path"]			,		mask=storage.is_null())
											elif pa.types.is_list(storage.type):
																	A__	:							List[Any]						    = pa.array(
																	    [encode_np_array(np.array(_A))["bytes"] if arr is not None else None for arr in storage.to_pylist()]			,		type=pa.binary()			,		)
																	A__	:							Union[str, Any]						    = pa.array([None] * len(_A)			,		type=pa.string())
																	A__	:							Optional[int]						    = pa.StructArray.from_arrays(
																	    [bytes_array, path_array]			,		["bytes", "path"]			,		mask=bytes_array.is_null())
											return array_cast(_A			,		self.pa_type)







					def  _lowercase			(       self      :     List[Any]			,		_A      :     pa.StructArray):

											@no_op_if_value_is_null
											def path_to_bytes(_A      :     Dict):
																	with xopen(_A			,		"rb") as f:
																							A__	:							int						    = f.read()
																	return bytes_

											A__	:							Union[str, Any]						    = pa.array(
											    [
											        (path_to_bytes(x["path"]) if x["bytes"] is None else x["bytes"]) if x is not None else None
											        for x in storage.to_pylist()
											    ]			,		type=pa.binary()			,		)
											A__	:							str						    = pa.array(
											    [os.path.basename(_A) if path is not None else None for path in storage.field("path").to_pylist()]			,		type=pa.string()			,		)
											A__	:							str						    = pa.StructArray.from_arrays([bytes_array, path_array]			,		["bytes", "path"]			,		mask=bytes_array.is_null())
											return array_cast(_A			,		self.pa_type)




def 	snake_case__       (					)		->		List[str]:







						"""simple docstring"""

						if config.PIL_AVAILABLE:
												import PIL.Image
						else:
												raise ImportError("To support encoding images, please install 'Pillow'."       )

						global _IMAGE_COMPRESSION_FORMATS
						if _IMAGE_COMPRESSION_FORMATS is None:
												PIL.Image.init()
												A__	:							Any						    = list(set(PIL.Image.OPEN.keys()       ) & set(PIL.Image.SAVE.keys()       )       )
						return _IMAGE_COMPRESSION_FORMATS
def 	snake_case__       (					__lowercase       )		->		bytes:







						"""simple docstring"""

						A__	:							Dict						    = BytesIO()
						if image.format in list_image_compression_formats():
												A__	:							List[Any]						    = image.format
						else:
												A__	:							List[str]						    = "PNG" if image.mode in ["1", "L", "LA", "RGB", "RGBA"] else "TIFF"
						image.save(__lowercase							,     format=__lowercase       )
						return buffer.getvalue()
def 	snake_case__       (					__lowercase       )		->		dict:







						"""simple docstring"""

						if hasattr(__lowercase							,     "filename"       ) and image.filename != "":
												return {"path": image.filename, "bytes": None}
						else:
												return {"path": None, "bytes": image_to_bytes(__lowercase       )}
def 	snake_case__       (					__lowercase       )		->		dict:







						"""simple docstring"""

						if config.PIL_AVAILABLE:
												import PIL.Image
						else:
												raise ImportError("To support encoding images, please install 'Pillow'."       )

						A__	:							Optional[int]						    = array.dtype
						A__	:							List[str]						    = dtype.byteorder if dtype.byteorder != "=" else _NATIVE_BYTEORDER
						A__	:							Dict						    = dtype.kind
						A__	:							Union[str, Any]						    = dtype.itemsize

						A__	:							int						    = None

						# Multi-channel array case (only np.dtype("|u1") is allowed)
						if array.shape[2:]:
												A__	:							Union[str, Any]						    = np.dtype("|u1"       )
												if dtype_kind not in ["u", "i"]:
																		raise TypeError(
																		    F'Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays.'       )
												if dtype is not dest_dtype:
																		warnings.warn(F'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''       )
    # Exact match
						elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
												A__	:							str						    = dtype
						else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
												while dtype_itemsize >= 1:
																		A__	:							Optional[Any]						    = dtype_byteorder + dtype_kind + str(__lowercase       )
																		A__	:							Union[str, Any]						    = np.dtype(__lowercase       )
																		if dest_dtype in _VALID_IMAGE_ARRAY_DTPYES:
																								warnings.warn(F'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''       )
																								break
																		else:
																								dtype_itemsize //= 2

						if dest_dtype is None:
												raise TypeError(
												    F'Cannot convert dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}'       )
						A__	:							Optional[int]						    = PIL.Image.fromarray(array.astype(__lowercase       )       )
						return {"path": None, "bytes": image_to_bytes(__lowercase       )}




def 	snake_case__       (					__lowercase       )		->		List[dict]:







						"""simple docstring"""

						if config.PIL_AVAILABLE:
												import PIL.Image
						else:
												raise ImportError("To support encoding images, please install 'Pillow'."       )

						if objs:
												A__    ,       A__	:							List[str]						    = first_non_null_value(__lowercase       )
												if isinstance(__lowercase							,     __lowercase       ):
																		return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
												if isinstance(__lowercase							,     np.ndarray       ):
																		A__	:							List[str]						    = no_op_if_value_is_null(__lowercase       )
																		return [obj_to_image_dict_func(__lowercase       ) for obj in objs]
												elif isinstance(__lowercase							,     PIL.Image.Image       ):
																		A__	:							Tuple						    = no_op_if_value_is_null(__lowercase       )
																		return [obj_to_image_dict_func(__lowercase       ) for obj in objs]
												else:
																		return objs
						else:
												return objs




import logging
import random

import ray

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


snake_case    :						Any  =						logging.getLogger(__name__)




class 							lowerCAmelCase__  :







					def __init__(       self      :     Union[str, Any]):
											A__	:							Dict						    = False







					def  _lowercase			(       self      :     Optional[int]			,		_A      :     int			,		_A      :     Optional[int]			,		_A      :     List[str]			,		_A      :     Optional[Any]):
											if not self.initialized:
																	A__	:							List[Any]						    = RagRetriever(
																	    _A			,		question_encoder_tokenizer=_A			,		generator_tokenizer=_A			,		index=_A			,		init_retrieval=_A			,		)
																	A__	:							Dict						    = True







					def  _lowercase			(       self      :     Union[str, Any]):
											self.retriever.index.init_index()







					def  _lowercase			(       self      :     List[str]			,		_A      :     Optional[int]			,		_A      :     List[str]):
											A__    ,       A__	:							str						    = self.retriever._main_retrieve(_A			,		_A)
											return doc_ids, retrieved_doc_embeds




class 							lowerCAmelCase__  (	UpperCamelCase						):







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







					@classmethod
					def  _lowercase			(       cls      :     int			,		_A      :     Optional[int]			,		_A      :     List[str]=None			,		**_A      :     Optional[Any]):
											return super(_A			,		cls).get_tokenizers(_A			,		_A			,		**_A)







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