




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,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
  import torch


lowerCamelCase    :		List[Any]			   =       logging.get_logger(__name__)







class   A(						__snake_case       ):



 '''simple docstring'''


 UpperCamelCase						       =						['pixel_values']

 def __init__(    self :  List[Any]						,							A_ :  List[Any] = True						,							A_ :  Optional[Any] = None						,							A_ :  Optional[Any] = PILImageResampling.BILINEAR						,							A_ :  Optional[int] = True						,							A_ :  Union[str, Any] = None						,							A_ :  List[str] = True						,							A_ :  Union[str, Any] = 1 / 255						,							A_ :  int = True						,							A_ :  str = None						,							A_ :  int = None						,							**A_ :  Optional[int]						,							)       ->     None:

       """simple docstring"""






       super().__init__(**A_  )
       lowerCamelCase_							=   size if size is not None else {'shortest_edge': 256}
       lowerCamelCase_							=   get_size_dict(A_						,							default_to_square=A_  )
       lowerCamelCase_							=   crop_size if crop_size is not None else {'height': 224, 'width': 224}
       lowerCamelCase_							=   get_size_dict(A_						,							param_name='crop_size'  )
       lowerCamelCase_							=   do_resize
       lowerCamelCase_							=   size
       lowerCamelCase_							=   resample
       lowerCamelCase_							=   do_center_crop
       lowerCamelCase_							=   crop_size
       lowerCamelCase_							=   do_rescale
       lowerCamelCase_							=   rescale_factor
       lowerCamelCase_							=   do_normalize
       lowerCamelCase_							=   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
       lowerCamelCase_							=   image_std if image_std is not None else IMAGENET_STANDARD_STD

 def     a__			(    self :  Union[str, Any]						,							A_ :  Tuple						,							A_ :  List[Any]						,							A_ :  str = PILImageResampling.BICUBIC						,							A_ :  Tuple = None						,							**A_ :  List[Any]						,							)       ->     np.ndarray:

       """simple docstring"""






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

 def     a__			(    self :  Optional[Any]						,							A_ :  List[str]						,							A_ :  List[Any]						,							A_ :  List[Any] = None						,							**A_ :  str						,							)       ->     np.ndarray:

       """simple docstring"""






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

 def     a__			(    self :  Any						,							A_ :  Union[str, Any]						,							A_ :  Dict						,							A_ :  str = None						,							**A_ :  Optional[Any]  )       ->     np.ndarray:

       """simple docstring"""






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

 def     a__			(    self :  Optional[int]						,							A_ :  Optional[int]						,							A_ :  Union[str, Any]						,							A_ :  Optional[int]						,							A_ :  Any = None						,							**A_ :  int						,							)       ->     np.ndarray:

       """simple docstring"""






       return normalize(A_						,							mean=A_						,							std=A_						,							data_format=A_						,							**A_  )

 def     a__			(    self :  Any						,							A_ :  List[str]						,							A_ :  List[str] = None						,							A_ :  Optional[int] = None						,							A_ :  Tuple = None						,							A_ :  Union[str, Any] = None						,							A_ :  Tuple = None						,							A_ :  Any = None						,							A_ :  Union[str, Any] = None						,							A_ :  List[Any] = None						,							A_ :  List[Any] = None						,							A_ :  Dict = None						,							A_ :  Any = None						,							A_ :  int = ChannelDimension.FIRST						,							**A_ :  int						,							)       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   do_resize if do_resize is not None else self.do_resize
       lowerCamelCase_							=   size if size is not None else self.size
       lowerCamelCase_							=   get_size_dict(A_						,							default_to_square=A_  )
       lowerCamelCase_							=   resample if resample is not None else self.resample
       lowerCamelCase_							=   do_center_crop if do_center_crop is not None else self.do_center_crop
       lowerCamelCase_							=   crop_size if crop_size is not None else self.crop_size
       lowerCamelCase_							=   get_size_dict(A_						,							param_name='crop_size'  )
       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_normalize if do_normalize is not None else self.do_normalize
       lowerCamelCase_							=   image_mean if image_mean is not None else self.image_mean
       lowerCamelCase_							=   image_std if image_std is not None else self.image_std

       lowerCamelCase_							=   make_list_of_images(A_  )

       if not valid_images(A_  ):
             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_center_crop and crop_size is None:
             raise ValueError('Crop size must be specified if do_center_crop is True.'  )

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

       if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('Image mean and std must be specified if do_normalize is True.'  )

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

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

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

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

       if do_normalize:
             lowerCamelCase_							=   [self.normalize(image=A_						,							mean=A_						,							std=A_  ) for image in images]

       lowerCamelCase_							=   [to_channel_dimension_format(A_						,							A_  ) for image in images]

       lowerCamelCase_							=   {'pixel_values': images}
       return BatchFeature(data=A_						,							tensor_type=A_  )





 def     a__			(    self :  List[str]						,							A_ :  Any						,							A_ :  Tuple = None  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   outputs.logits

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

             if is_torch_tensor(A_  ):
                   lowerCamelCase_							=   target_sizes.numpy()

             lowerCamelCase_							=   []

             for idx in range(len(A_  )  ):
                   lowerCamelCase_							=   torch.nn.functional.interpolate(
                       logits[idx].unsqueeze(dim=0  )						,							size=target_sizes[idx]						,							mode='bilinear'						,							align_corners=A_  )
                   lowerCamelCase_							=   resized_logits[0].argmax(dim=0  )
                   semantic_segmentation.append(A_  )
       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)
lowerCamelCase    :		int			   =       logging.getLogger()

lowerCamelCase    :		Tuple			   =       logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class   A(						UpperCamelCase       ):



 '''simple docstring'''


 def     a__			(    self :  List[str]						,							A_ :  Optional[Any]  )       ->     int:

       """simple docstring"""






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

 def     a__			(    self :  Optional[Any]						,							A_ :  int						,							A_ :  str = "pytorch"  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   self.get_auto_remove_tmp_dir()
       lowerCamelCase_							=   os.path.join(A_						,							'output'  )
       lowerCamelCase_							=   os.path.join(A_						,							'data'  )
       self._create_dummy_data(data_dir=A_  )

       lowerCamelCase_							=   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'  )

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

       lowerCamelCase_							=   os.path.join(A_						,							'metrics.json'  )
       with open(A_  ) as f:
             lowerCamelCase_							=   json.load(A_  )
       return result

 @require_torch_gpu
 def     a__			(    self :  List[Any]  )       ->     Any:

       """simple docstring"""






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

 @require_torch_multi_gpu
 def     a__			(    self :  List[Any]  )       ->     Union[str, Any]:

       """simple docstring"""






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

 @require_torch_gpu
 @require_ray
 def     a__			(    self :  Tuple  )       ->     int:

       """simple docstring"""






       lowerCamelCase_							=   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     a__			(    self :  Dict  )       ->     Dict:

       """simple docstring"""






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



