
'''simple docstring'''

from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
		import PIL

if is_torch_available():
		import torch


UpperCAmelCase__ :			Union[str, Any]											=		logging.get_logger(__name__)






class  A  (						SCREAMING_SNAKE_CASE__     ):
	snake_case__   :Any   =  ['pixel_values']





	def __init__(       self      :							List[Any]   ,				__magic_name__      :							bool = True   ,				__magic_name__      :							Dict[str, int] = None   ,				__magic_name__      :							PILImageResampling = PILImageResampling.BILINEAR   ,				__magic_name__      :							bool = True   ,				__magic_name__      :							Union[int, float] = 1 / 255   ,				__magic_name__      :							bool = True   ,				__magic_name__      :							Dict[str, int] = None   ,				__magic_name__      :							bool = True   ,				**__magic_name__      :							Union[str, Any]   ,				):


				"""simple docstring"""

				super().__init__(**__magic_name__				)
				lowerCAmelCase__										=				size if size is not None else {"shortest_edge": 224}
				lowerCAmelCase__										=				get_size_dict(__magic_name__   ,				default_to_square=__magic_name__				)
				lowerCAmelCase__										=				crop_size if crop_size is not None else {"height": 256, "width": 256}
				lowerCAmelCase__										=				get_size_dict(__magic_name__   ,				param_name="crop_size"				)

				lowerCAmelCase__										=				do_resize
				lowerCAmelCase__										=				size
				lowerCAmelCase__										=				resample
				lowerCAmelCase__										=				do_rescale
				lowerCAmelCase__										=				rescale_factor
				lowerCAmelCase__										=				do_center_crop
				lowerCAmelCase__										=				crop_size
				lowerCAmelCase__										=				do_flip_channel_order





	def 				__SCREAMING_SNAKE_CASE	(       self      :							List[str]   ,				__magic_name__      :							np.ndarray   ,				__magic_name__      :							Dict[str, int]   ,				__magic_name__      :							PILImageResampling = PIL.Image.BILINEAR   ,				__magic_name__      :							Optional[Union[str, ChannelDimension]] = None   ,				**__magic_name__      :							str   ,				):


				"""simple docstring"""

				lowerCAmelCase__										=				get_size_dict(__magic_name__   ,				default_to_square=__magic_name__				)
				if "shortest_edge" not in size:
							raise ValueError(f"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""				)
				lowerCAmelCase__										=				get_resize_output_image_size(__magic_name__   ,				size=size["shortest_edge"]   ,				default_to_square=__magic_name__				)
				return resize(__magic_name__   ,				size=__magic_name__   ,				resample=__magic_name__   ,				data_format=__magic_name__   ,				**__magic_name__				)





	def 				__SCREAMING_SNAKE_CASE	(       self      :							List[str]   ,				__magic_name__      :							np.ndarray   ,				__magic_name__      :							Dict[str, int]   ,				__magic_name__      :							Optional[Union[str, ChannelDimension]] = None   ,				**__magic_name__      :							int   ,				):


				"""simple docstring"""

				lowerCAmelCase__										=				get_size_dict(__magic_name__				)
				if "height" not in size or "width" not in size:
							raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""				)
				return center_crop(__magic_name__   ,				size=(size["height"], size["width"])   ,				data_format=__magic_name__   ,				**__magic_name__				)





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Dict   ,				__magic_name__      :							np.ndarray   ,				__magic_name__      :							Union[int, float]   ,				__magic_name__      :							Optional[Union[str, ChannelDimension]] = None   ,				**__magic_name__      :							Optional[int]   ,				):


				"""simple docstring"""

				return rescale(__magic_name__   ,				scale=__magic_name__   ,				data_format=__magic_name__   ,				**__magic_name__				)





	def 				__SCREAMING_SNAKE_CASE	(       self      :							List[Any]   ,				__magic_name__      :							np.ndarray   ,				__magic_name__      :							Optional[Union[str, ChannelDimension]] = None				):


				"""simple docstring"""

				return flip_channel_order(__magic_name__   ,				data_format=__magic_name__				)





	def 				__SCREAMING_SNAKE_CASE	(       self      :							Dict   ,				__magic_name__      :							ImageInput   ,				__magic_name__      :							bool = None   ,				__magic_name__      :							Dict[str, int] = None   ,				__magic_name__      :							PILImageResampling = None   ,				__magic_name__      :							bool = None   ,				__magic_name__      :							float = None   ,				__magic_name__      :							bool = None   ,				__magic_name__      :							Dict[str, int] = None   ,				__magic_name__      :							bool = None   ,				__magic_name__      :							Optional[Union[str, TensorType]] = None   ,				__magic_name__      :							ChannelDimension = ChannelDimension.FIRST   ,				**__magic_name__      :							str   ,				):


				"""simple docstring"""

				lowerCAmelCase__										=				do_resize if do_resize is not None else self.do_resize
				lowerCAmelCase__										=				resample if resample is not None else self.resample
				lowerCAmelCase__										=				do_rescale if do_rescale is not None else self.do_rescale
				lowerCAmelCase__										=				rescale_factor if rescale_factor is not None else self.rescale_factor
				lowerCAmelCase__										=				do_center_crop if do_center_crop is not None else self.do_center_crop
				lowerCAmelCase__										=				(
				    do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
				)

				lowerCAmelCase__										=				size if size is not None else self.size
				lowerCAmelCase__										=				get_size_dict(__magic_name__   ,				default_to_square=__magic_name__				)
				lowerCAmelCase__										=				crop_size if crop_size is not None else self.crop_size
				lowerCAmelCase__										=				get_size_dict(__magic_name__   ,				param_name="crop_size"				)

				lowerCAmelCase__										=				make_list_of_images(__magic_name__				)

				if not valid_images(__magic_name__				):
							raise ValueError(
							    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
							    "torch.Tensor, tf.Tensor or jax.ndarray."				)

				if do_resize and size is None:
							raise ValueError("Size must be specified if do_resize is True."				)

				if do_rescale and rescale_factor is None:
							raise ValueError("Rescale factor must be specified if do_rescale is True."				)

				if do_center_crop and crop_size is None:
							raise ValueError("Crop size must be specified if do_center_crop is True."				)

				# All transformations expect numpy arrays.
				lowerCAmelCase__										=				[to_numpy_array(__magic_name__				) for image in images]

				if do_resize:
							lowerCAmelCase__										=				[self.resize(image=__magic_name__   ,				size=__magic_name__   ,				resample=__magic_name__				) for image in images]

				if do_center_crop:
							lowerCAmelCase__										=				[self.center_crop(image=__magic_name__   ,				size=__magic_name__				) for image in images]

				if do_rescale:
							lowerCAmelCase__										=				[self.rescale(image=__magic_name__   ,				scale=__magic_name__				) for image in images]

				# the pretrained checkpoints assume images are BGR, not RGB
				if do_flip_channel_order:
							lowerCAmelCase__										=				[self.flip_channel_order(image=__magic_name__				) for image in images]

				lowerCAmelCase__										=				[to_channel_dimension_format(__magic_name__   ,				__magic_name__				) for image in images]

				lowerCAmelCase__										=				{"pixel_values": images}
				return BatchFeature(data=__magic_name__   ,				tensor_type=__magic_name__				)







	def 				__SCREAMING_SNAKE_CASE	(       self      :							Dict   ,				__magic_name__      :							Dict   ,				__magic_name__      :							List[Tuple] = None				):


				"""simple docstring"""

				lowerCAmelCase__										=				outputs.logits

				# Resize logits and compute semantic segmentation maps
				if target_sizes is not None:
							if len(__magic_name__				) != len(__magic_name__				):
										raise ValueError(
										    "Make sure that you pass in as many target sizes as the batch dimension of the logits"				)

							if is_torch_tensor(__magic_name__				):
										lowerCAmelCase__										=				target_sizes.numpy()

							lowerCAmelCase__										=				[]

							for idx in range(len(__magic_name__				)				):
										lowerCAmelCase__										=				torch.nn.functional.interpolate(
										    logits[idx].unsqueeze(dim=0				)   ,				size=target_sizes[idx]   ,				mode="bilinear"   ,				align_corners=__magic_name__				)
										lowerCAmelCase__										=				resized_logits[0].argmax(dim=0				)
										semantic_segmentation.append(__magic_name__				)
				else:
							lowerCAmelCase__										=				logits.argmax(dim=1				)
							lowerCAmelCase__										=				[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]				)]

				return semantic_segmentation









import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
__UpperCamelCase     		=							logging.getLogger()

__UpperCamelCase     		=							logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			def 					snake_case__					(				self   ,    snake_case  ):

						'''simple docstring'''


						os.makedirs(snake_case   ,    exist_ok=snake_case  )
						UpperCamelCase__       =							{"source": "What is love ?", "target": "life"}
						UpperCamelCase__       =							{"train": 12, "val": 2, "test": 2}
						for split in ["train", "test", "val"]:
									for field in ["source", "target"]:
												UpperCamelCase__       =							"\n".join([contents[field]] * n_lines[split]  )
												with open(os.path.join(snake_case   ,    F'''{split}.{field}'''  )   ,    "w"  ) as f:
															f.write(snake_case  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case = "pytorch"  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.get_auto_remove_tmp_dir()
						UpperCamelCase__       =							os.path.join(snake_case   ,    "output"  )
						UpperCamelCase__       =							os.path.join(snake_case   ,    "data"  )
						self._create_dummy_data(data_dir=snake_case  )

						UpperCamelCase__       =							F'''
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            '''.split()

						if gpus > 0:
									testargs.append(F'''--gpus={gpus}'''  )
									if is_apex_available():
												testargs.append("--fp16"  )
						else:
									testargs.append("--gpus=0"  )
									testargs.append("--distributed_backend=ddp_cpu"  )
									testargs.append("--num_processes=2"  )

						UpperCamelCase__       =							[sys.executable, str(Path(finetune_rag.__file__  ).resolve()  )] + testargs
						execute_subprocess_async(snake_case   ,    env=self.get_env()  )

						UpperCamelCase__       =							os.path.join(snake_case   ,    "metrics.json"  )
						with open(snake_case  ) as f:
									UpperCamelCase__       =							json.load(snake_case  )
						return result




			@require_torch_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )




			@require_torch_multi_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=2  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )




			@require_torch_gpu
			@require_ray
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1   ,    distributed_retriever="ray"  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )







			@require_torch_multi_gpu
			@require_ray
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self._run_finetune(gpus=1   ,    distributed_retriever="ray"  )
						self.assertGreaterEqual(result["test"][0]["test_avg_em"]   ,    0.2  )


