





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
      import PIL


UpperCAmelCase_					:   Any 							=						logging.get_logger(__name__)





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  ["""pixel_values"""]







  def __init__(	self    ,						_lowercase = True    ,						_lowercase = None    ,						_lowercase = PILImageResampling.BICUBIC    ,						_lowercase = True    ,						_lowercase = 1 / 255    ,						_lowercase = True    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = True    ,						**_lowercase    ,						):
       super().__init__(**_lowercase						)
       lowerCAmelCase_					:       int    =		size if size is not None else {"""height""": 384, """width""": 384}
       lowerCAmelCase_					:       Union[str, Any]    =		get_size_dict(_lowercase    ,						default_to_square=_lowercase						)

       lowerCAmelCase_					:       Dict    =		do_resize
       lowerCAmelCase_					:       Tuple    =		size
       lowerCAmelCase_					:       Any    =		resample
       lowerCAmelCase_					:       Union[str, Any]    =		do_rescale
       lowerCAmelCase_					:       List[str]    =		rescale_factor
       lowerCAmelCase_					:       Union[str, Any]    =		do_normalize
       lowerCAmelCase_					:       List[str]    =		image_mean if image_mean is not None else OPENAI_CLIP_MEAN
       lowerCAmelCase_					:       List[Any]    =		image_std if image_std is not None else OPENAI_CLIP_STD
       lowerCAmelCase_					:       Optional[Any]    =		do_convert_rgb







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase = PILImageResampling.BICUBIC    ,						_lowercase = None    ,						**_lowercase    ,						):
       lowerCAmelCase_					:       Tuple    =		get_size_dict(_lowercase    ,						default_to_square=_lowercase						)
       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()}'						)
       lowerCAmelCase_					:       Tuple    =		(size["""height"""], size["""width"""])
       return resize(_lowercase    ,						size=_lowercase    ,						resample=_lowercase    ,						data_format=_lowercase    ,						**_lowercase						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase = None    ,						**_lowercase    ,						):
       return rescale(_lowercase    ,						scale=_lowercase    ,						data_format=_lowercase    ,						**_lowercase						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase = None    ,						**_lowercase    ,						):
       return normalize(_lowercase    ,						mean=_lowercase    ,						std=_lowercase    ,						data_format=_lowercase    ,						**_lowercase						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = None    ,						_lowercase = ChannelDimension.FIRST    ,						**_lowercase    ,						):
       lowerCAmelCase_					:       List[str]    =		do_resize if do_resize is not None else self.do_resize
       lowerCAmelCase_					:       Any    =		resample if resample is not None else self.resample
       lowerCAmelCase_					:       Dict    =		do_rescale if do_rescale is not None else self.do_rescale
       lowerCAmelCase_					:       Any    =		rescale_factor if rescale_factor is not None else self.rescale_factor
       lowerCAmelCase_					:       Tuple    =		do_normalize if do_normalize is not None else self.do_normalize
       lowerCAmelCase_					:       Optional[int]    =		image_mean if image_mean is not None else self.image_mean
       lowerCAmelCase_					:       List[str]    =		image_std if image_std is not None else self.image_std
       lowerCAmelCase_					:       Dict    =		do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

       lowerCAmelCase_					:       Tuple    =		size if size is not None else self.size
       lowerCAmelCase_					:       Optional[int]    =		get_size_dict(_lowercase    ,						default_to_square=_lowercase						)

       lowerCAmelCase_					:       str    =		make_list_of_images(_lowercase						)

       if not valid_images(_lowercase						):
            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 or resample is None:
            raise ValueError("""Size and resample 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_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."""						)

       # PIL RGBA images are converted to RGB
       if do_convert_rgb:
            lowerCAmelCase_					:       Optional[Any]    =		[convert_to_rgb(_lowercase						) for image in images]

       # All transformations expect numpy arrays.
       lowerCAmelCase_					:       Optional[int]    =		[to_numpy_array(_lowercase						) for image in images]

       if do_resize:
            lowerCAmelCase_					:       Optional[Any]    =		[self.resize(image=_lowercase    ,						size=_lowercase    ,						resample=_lowercase						) for image in images]

       if do_rescale:
            lowerCAmelCase_					:       Any    =		[self.rescale(image=_lowercase    ,						scale=_lowercase						) for image in images]

       if do_normalize:
            lowerCAmelCase_					:       str    =		[self.normalize(image=_lowercase    ,						mean=_lowercase    ,						std=_lowercase						) for image in images]

       lowerCAmelCase_					:       Tuple    =		[to_channel_dimension_format(_lowercase    ,						_lowercase						) for image in images]

       lowerCAmelCase_					:       List[str]    =		BatchFeature(data={"""pixel_values""": images}    ,						tensor_type=_lowercase						)

       return encoded_outputs







import torch
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  """M-CLIP"""







  def __init__(	self    ,						_lowercase=1_024    ,						_lowercase=768    ,						**_lowercase						):
       lowerCAmelCase_					:       Tuple    =		transformerDimSize
       lowerCAmelCase_					:       Optional[Any]    =		imageDimSize
       super().__init__(**_lowercase						)





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  MCLIPConfig







  def __init__(	self    ,						_lowercase    ,						*_lowercase    ,						**_lowercase						):
       super().__init__(_lowercase    ,						*_lowercase    ,						**_lowercase						)
       lowerCAmelCase_					:       str    =		XLMRobertaModel(_lowercase						)
       lowerCAmelCase_					:       str    =		torch.nn.Linear(
           in_features=config.transformerDimensions    ,						out_features=config.numDims						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Optional[int]    =		self.transformer(input_ids=_lowercase    ,						attention_mask=_lowercase						)[0]
       lowerCAmelCase_					:       Any    =		(embs * attention_mask.unsqueeze(2						)).sum(dim=1						) / attention_mask.sum(dim=1						)[:, None]
       return self.LinearTransformation(_lowercase						), embs
