






from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


UpperCamelCase												=      logging.get_logger(__name__)



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: HybridCLIPConfig
			__snake_case							: jnp.dtype							=				jnp.floataa







			def 				UpperCamelCase (      self:    Tuple	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.config.text_config
									_SCREAMING_SNAKE_CASE       =    self.config.vision_config

									_SCREAMING_SNAKE_CASE       =    self.config.projection_dim
									_SCREAMING_SNAKE_CASE       =    text_config.hidden_size
									_SCREAMING_SNAKE_CASE       =    vision_config.hidden_size

									_SCREAMING_SNAKE_CASE       =    FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
									_SCREAMING_SNAKE_CASE       =    FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__ ,      UpperCAmelCase_	).module_class

									_SCREAMING_SNAKE_CASE       =    text_module(UpperCAmelCase_ ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    vision_module(UpperCAmelCase_ ,      dtype=self.dtype	)

									_SCREAMING_SNAKE_CASE       =    nn.Dense(
									    self.projection_dim ,      dtype=self.dtype ,      kernel_init=jax.nn.initializers.normal(0.02	) ,      use_bias=UpperCAmelCase_ ,      )
									_SCREAMING_SNAKE_CASE       =    nn.Dense(
									    self.projection_dim ,      dtype=self.dtype ,      kernel_init=jax.nn.initializers.normal(0.02	) ,      use_bias=UpperCAmelCase_ ,      )
									_SCREAMING_SNAKE_CASE       =    self.param("""logit_scale""" ,      jax.nn.initializers.ones ,      []	)







			def __call__(      self:    List[Any] ,      UpperCAmelCase_:    Optional[Any]=None ,      UpperCAmelCase_:    int=None ,      UpperCAmelCase_:    Union[str, Any]=None ,      UpperCAmelCase_:    Dict=None ,      UpperCAmelCase_:    Any=None ,      UpperCAmelCase_:    bool = True ,      UpperCAmelCase_:    str=None ,      UpperCAmelCase_:    List[str]=None ,      UpperCAmelCase_:    Union[str, Any]=None ,      ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    return_dict if return_dict is not None else self.config.return_dict

									_SCREAMING_SNAKE_CASE       =    self.vision_model(
									    pixel_values=UpperCAmelCase_ ,      deterministic=UpperCAmelCase_ ,      output_attentions=UpperCAmelCase_ ,      output_hidden_states=UpperCAmelCase_ ,      return_dict=UpperCAmelCase_ ,      )

									_SCREAMING_SNAKE_CASE       =    self.text_model(
									    input_ids=UpperCAmelCase_ ,      attention_mask=UpperCAmelCase_ ,      token_type_ids=UpperCAmelCase_ ,      position_ids=UpperCAmelCase_ ,      deterministic=UpperCAmelCase_ ,      output_attentions=UpperCAmelCase_ ,      output_hidden_states=UpperCAmelCase_ ,      return_dict=UpperCAmelCase_ ,      )

									_SCREAMING_SNAKE_CASE       =    vision_outputs[1]
									_SCREAMING_SNAKE_CASE       =    self.visual_projection(UpperCAmelCase_	)

									_SCREAMING_SNAKE_CASE       =    text_outputs[1]
									_SCREAMING_SNAKE_CASE       =    self.text_projection(UpperCAmelCase_	)

									# normalized features
									_SCREAMING_SNAKE_CASE       =    image_embeds / jnp.linalg.norm(UpperCAmelCase_ ,      axis=-1 ,      keepdims=UpperCAmelCase_	)
									_SCREAMING_SNAKE_CASE       =    text_embeds / jnp.linalg.norm(UpperCAmelCase_ ,      axis=-1 ,      keepdims=UpperCAmelCase_	)

									# cosine similarity as logits
									_SCREAMING_SNAKE_CASE       =    jnp.exp(self.logit_scale	)
									_SCREAMING_SNAKE_CASE       =    jnp.matmul(UpperCAmelCase_ ,      image_embeds.T	) * logit_scale
									_SCREAMING_SNAKE_CASE       =    logits_per_text.T

									if not return_dict:
															return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

									return FlaxCLIPOutput(
									    logits_per_image=UpperCAmelCase_ ,      logits_per_text=UpperCAmelCase_ ,      text_embeds=UpperCAmelCase_ ,      image_embeds=UpperCAmelCase_ ,      text_model_output=UpperCAmelCase_ ,      vision_model_output=UpperCAmelCase_ ,      )



class  __UpperCAmelCase						(_UpperCAmelCase       ):
			__snake_case							: str							=				HybridCLIPConfig
			__snake_case							: Optional[Any]							=				FlaxHybridCLIPModule







			def __init__(      self:    Tuple ,      UpperCAmelCase_:    HybridCLIPConfig ,      UpperCAmelCase_:    Optional[Tuple] = None ,      UpperCAmelCase_:    int = 0 ,      UpperCAmelCase_:    jnp.dtype = jnp.floataa ,      **UpperCAmelCase_:    Dict ,      ):







									'''simple docstring'''




									if input_shape is None:
															_SCREAMING_SNAKE_CASE       =    ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

									_SCREAMING_SNAKE_CASE       =    self.module_class(config=UpperCAmelCase_ ,      dtype=UpperCAmelCase_ ,      **UpperCAmelCase_	)
									super().__init__(UpperCAmelCase_ ,      UpperCAmelCase_ ,      input_shape=UpperCAmelCase_ ,      seed=UpperCAmelCase_ ,      dtype=UpperCAmelCase_	)







			def 				UpperCamelCase (      self:    int ,      UpperCAmelCase_:    jax.random.PRNGKey ,      UpperCAmelCase_:    Tuple ,      UpperCAmelCase_:    FrozenDict = None	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    jnp.zeros(input_shape[0] ,      dtype="""i4"""	)
									_SCREAMING_SNAKE_CASE       =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(UpperCAmelCase_	).shape[-1]	) ,      input_shape[0]	)
									_SCREAMING_SNAKE_CASE       =    jnp.ones_like(UpperCAmelCase_	)
									_SCREAMING_SNAKE_CASE       =    jnp.ones_like(UpperCAmelCase_	)

									_SCREAMING_SNAKE_CASE       =    jax.random.normal(UpperCAmelCase_ ,      input_shape[1]	)

									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    jax.random.split(UpperCAmelCase_	)
									_SCREAMING_SNAKE_CASE       =    {"""params""": params_rng, """dropout""": dropout_rng}

									return self.module.init(UpperCAmelCase_ ,      UpperCAmelCase_ ,      UpperCAmelCase_ ,      UpperCAmelCase_ ,      UpperCAmelCase_ ,      UpperCAmelCase_	)["params"]







			def __call__(      self:    List[str] ,      UpperCAmelCase_:    List[Any] ,      UpperCAmelCase_:    List[str] ,      UpperCAmelCase_:    List[Any]=None ,      UpperCAmelCase_:    Optional[Any]=None ,      UpperCAmelCase_:    Union[str, Any]=None ,      UpperCAmelCase_:    dict = None ,      UpperCAmelCase_:    jax.random.PRNGKey = None ,      UpperCAmelCase_:    bool = False ,      UpperCAmelCase_:    Optional[bool] = None ,      UpperCAmelCase_:    Optional[bool] = None ,      UpperCAmelCase_:    Optional[bool] = None ,      ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    output_attentions if output_attentions is not None else self.config.output_attentions
									_SCREAMING_SNAKE_CASE       =    (
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									_SCREAMING_SNAKE_CASE       =    return_dict if return_dict is not None else self.config.return_dict

									if position_ids is None:
															_SCREAMING_SNAKE_CASE       =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(UpperCAmelCase_	).shape[-1]	) ,      input_ids.shape	)

									if token_type_ids is None:
															_SCREAMING_SNAKE_CASE       =    jnp.zeros_like(UpperCAmelCase_	)

									if attention_mask is None:
															_SCREAMING_SNAKE_CASE       =    jnp.ones_like(UpperCAmelCase_	)

									# Handle any PRNG if needed
									_SCREAMING_SNAKE_CASE       =    {}
									if dropout_rng is not None:
															_SCREAMING_SNAKE_CASE       =    dropout_rng

									return self.module.apply(
									    {"""params""": params or self.params} ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype=jnp.floataa	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      not train ,      UpperCAmelCase_ ,      UpperCAmelCase_ ,      UpperCAmelCase_ ,      rngs=UpperCAmelCase_ ,      )







			def 				UpperCamelCase (      self:    int ,      UpperCAmelCase_:    Any ,      UpperCAmelCase_:    int=None ,      UpperCAmelCase_:    Union[str, Any]=None ,      UpperCAmelCase_:    Optional[Any]=None ,      UpperCAmelCase_:    dict = None ,      UpperCAmelCase_:    jax.random.PRNGKey = None ,      UpperCAmelCase_:    Optional[int]=False ,      ):







									'''simple docstring'''




									if position_ids is None:
															_SCREAMING_SNAKE_CASE       =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(UpperCAmelCase_	).shape[-1]	) ,      input_ids.shape	)

									if token_type_ids is None:
															_SCREAMING_SNAKE_CASE       =    jnp.zeros_like(UpperCAmelCase_	)

									if attention_mask is None:
															_SCREAMING_SNAKE_CASE       =    jnp.ones_like(UpperCAmelCase_	)

									# Handle any PRNG if needed
									_SCREAMING_SNAKE_CASE       =    {}
									if dropout_rng is not None:
															_SCREAMING_SNAKE_CASE       =    dropout_rng

									def _get_features(UpperCAmelCase_:    int ,      UpperCAmelCase_:    Union[str, Any] ,      UpperCAmelCase_:    Any ,      UpperCAmelCase_:    Optional[Any] ,      UpperCAmelCase_:    Dict ,      UpperCAmelCase_:    str	):
															_SCREAMING_SNAKE_CASE       =    module.text_model(
															    input_ids=UpperCAmelCase_ ,      attention_mask=UpperCAmelCase_ ,      position_ids=UpperCAmelCase_ ,      token_type_ids=UpperCAmelCase_ ,      deterministic=UpperCAmelCase_ ,      )
															_SCREAMING_SNAKE_CASE       =    text_outputs[1]
															_SCREAMING_SNAKE_CASE       =    module.text_projection(UpperCAmelCase_	)
															return text_features

									return self.module.apply(
									    {"""params""": params or self.params} ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      jnp.array(UpperCAmelCase_ ,      dtype="""i4"""	) ,      not train ,      method=_get_features ,      rngs=UpperCAmelCase_ ,      )







			def 				UpperCamelCase (      self:    List[Any] ,      UpperCAmelCase_:    Optional[Any] ,      UpperCAmelCase_:    dict = None ,      UpperCAmelCase_:    jax.random.PRNGKey = None ,      UpperCAmelCase_:    Any=False	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    {}
									if dropout_rng is not None:
															_SCREAMING_SNAKE_CASE       =    dropout_rng

									def _get_features(UpperCAmelCase_:    Dict ,      UpperCAmelCase_:    Optional[int] ,      UpperCAmelCase_:    Any	):
															_SCREAMING_SNAKE_CASE       =    module.vision_model(pixel_values=UpperCAmelCase_ ,      deterministic=UpperCAmelCase_	)
															_SCREAMING_SNAKE_CASE       =    vision_outputs[1]  # pooled_output
															_SCREAMING_SNAKE_CASE       =    module.visual_projection(UpperCAmelCase_	)
															return image_features

									return self.module.apply(
									    {"""params""": params or self.params} ,      jnp.array(UpperCAmelCase_ ,      dtype=jnp.floataa	) ,      not train ,      method=_get_features ,      rngs=UpperCAmelCase_ ,      )







			@classmethod
			def 				UpperCamelCase (      cls:    List[str] ,      UpperCAmelCase_:    str = None ,      UpperCAmelCase_:    str = None ,      *UpperCAmelCase_:    Optional[Any] ,      **UpperCAmelCase_:    List[str] ,      ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    {
									    argument[len("""text_"""	) :]: value for argument, value in kwargs.items() if argument.startswith("""text_"""	)
									}

									_SCREAMING_SNAKE_CASE       =    {
									    argument[len("""vision_"""	) :]: value for argument, value in kwargs.items() if argument.startswith("""vision_"""	)
									}

									# remove text, vision kwargs from kwargs
									for key in kwargs_text.keys():
															del kwargs["text_" + key]
									for key in kwargs_vision.keys():
															del kwargs["vision_" + key]

									# Load and initialize the text and vision model
									_SCREAMING_SNAKE_CASE       =    kwargs_text.pop("""model""" ,      UpperCAmelCase_	)
									if text_model is None:
															assert (
															    text_model_name_or_path is not None
															), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
															from transformers import FlaxAutoModel

															if "config" not in kwargs_text:
																					from transformers import AutoConfig

																					_SCREAMING_SNAKE_CASE       =    AutoConfig.from_pretrained(UpperCAmelCase_	)
																					_SCREAMING_SNAKE_CASE       =    text_config

															_SCREAMING_SNAKE_CASE       =    FlaxAutoModel.from_pretrained(UpperCAmelCase_ ,      *UpperCAmelCase_ ,      **UpperCAmelCase_	)

									_SCREAMING_SNAKE_CASE       =    kwargs_vision.pop("""model""" ,      UpperCAmelCase_	)
									if vision_model is None:
															assert (
															    vision_model_name_or_path is not None
															), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
															from transformers import FlaxAutoModel

															if "config" not in kwargs_vision:
																					from transformers import AutoConfig

																					_SCREAMING_SNAKE_CASE       =    AutoConfig.from_pretrained(UpperCAmelCase_	)
																					_SCREAMING_SNAKE_CASE       =    vision_config

															_SCREAMING_SNAKE_CASE       =    FlaxAutoModel.from_pretrained(UpperCAmelCase_ ,      *UpperCAmelCase_ ,      **UpperCAmelCase_	)

									# instantiate config with corresponding kwargs
									_SCREAMING_SNAKE_CASE       =    kwargs.pop("""dtype""" ,      jnp.floataa	)
									_SCREAMING_SNAKE_CASE       =    HybridCLIPConfig.from_text_vision_configs(text_model.config ,      vision_model.config ,      **UpperCAmelCase_	)

									# init model
									_SCREAMING_SNAKE_CASE       =    cls(UpperCAmelCase_ ,      *UpperCAmelCase_ ,      dtype=UpperCAmelCase_ ,      **UpperCAmelCase_	)

									if vision_config.model_type == "clip":
															_SCREAMING_SNAKE_CASE       =    vision_model.params["""vision_model"""]
															_SCREAMING_SNAKE_CASE       =    vision_model.params["""visual_projection"""]["""kernel"""]
									else:
															_SCREAMING_SNAKE_CASE       =    vision_model.params

									_SCREAMING_SNAKE_CASE       =    text_model.params

									return model









import multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader



class  __UpperCAmelCase						(_UpperCAmelCase       ):







			def __init__(      self:    Optional[int] ,      UpperCAmelCase_:    NestedDataStructureLike[PathLike] ,      UpperCAmelCase_:    Optional[NamedSplit] = None ,      UpperCAmelCase_:    Optional[Features] = None ,      UpperCAmelCase_:    str = None ,      UpperCAmelCase_:    bool = False ,      UpperCAmelCase_:    bool = False ,      UpperCAmelCase_:    Optional[int] = None ,      **UpperCAmelCase_:    Union[str, Any] ,      ):







									'''simple docstring'''




									super().__init__(
									    UpperCAmelCase_ ,      split=UpperCAmelCase_ ,      features=UpperCAmelCase_ ,      cache_dir=UpperCAmelCase_ ,      keep_in_memory=UpperCAmelCase_ ,      streaming=UpperCAmelCase_ ,      num_proc=UpperCAmelCase_ ,      **UpperCAmelCase_ ,      )
									_SCREAMING_SNAKE_CASE       =    path_or_paths if isinstance(UpperCAmelCase_ ,      UpperCAmelCase_	) else {self.split: path_or_paths}
									_SCREAMING_SNAKE_CASE       =    Csv(
									    cache_dir=UpperCAmelCase_ ,      data_files=UpperCAmelCase_ ,      features=UpperCAmelCase_ ,      **UpperCAmelCase_ ,      )







			def 				UpperCamelCase (      self:    int	):







									'''simple docstring'''




									if self.streaming:
															_SCREAMING_SNAKE_CASE       =    self.builder.as_streaming_dataset(split=self.split	)
									# Build regular (map-style) dataset
									else:
															_SCREAMING_SNAKE_CASE       =    None
															_SCREAMING_SNAKE_CASE       =    None
															_SCREAMING_SNAKE_CASE       =    None
															_SCREAMING_SNAKE_CASE       =    None

															self.builder.download_and_prepare(
															    download_config=UpperCAmelCase_ ,      download_mode=UpperCAmelCase_ ,      verification_mode=UpperCAmelCase_ ,      base_path=UpperCAmelCase_ ,      num_proc=self.num_proc ,      )
															_SCREAMING_SNAKE_CASE       =    self.builder.as_dataset(
															    split=self.split ,      verification_mode=UpperCAmelCase_ ,      in_memory=self.keep_in_memory	)
									return dataset



class  __UpperCAmelCase						:







			def __init__(      self:    List[str] ,      UpperCAmelCase_:    Dataset ,      UpperCAmelCase_:    Union[PathLike, BinaryIO] ,      UpperCAmelCase_:    Optional[int] = None ,      UpperCAmelCase_:    Optional[int] = None ,      **UpperCAmelCase_:    Any ,      ):







									'''simple docstring'''




									if num_proc is not None and num_proc <= 0:
															raise ValueError(F'num_proc {num_proc} must be an integer > 0.'	)

									_SCREAMING_SNAKE_CASE       =    dataset
									_SCREAMING_SNAKE_CASE       =    path_or_buf
									_SCREAMING_SNAKE_CASE       =    batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
									_SCREAMING_SNAKE_CASE       =    num_proc
									_SCREAMING_SNAKE_CASE       =    """utf-8"""
									_SCREAMING_SNAKE_CASE       =    to_csv_kwargs







			def 				UpperCamelCase (      self:    Optional[int]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.to_csv_kwargs.pop("""path_or_buf""" ,      UpperCAmelCase_	)
									_SCREAMING_SNAKE_CASE       =    self.to_csv_kwargs.pop("""header""" ,      UpperCAmelCase_	)
									_SCREAMING_SNAKE_CASE       =    self.to_csv_kwargs.pop("""index""" ,      UpperCAmelCase_	)

									if isinstance(self.path_or_buf ,      (str, bytes, os.PathLike)	):
															with open(self.path_or_buf ,      """wb+"""	) as buffer:
																					_SCREAMING_SNAKE_CASE       =    self._write(file_obj=UpperCAmelCase_ ,      header=UpperCAmelCase_ ,      index=UpperCAmelCase_ ,      **self.to_csv_kwargs	)
									else:
															_SCREAMING_SNAKE_CASE       =    self._write(file_obj=self.path_or_buf ,      header=UpperCAmelCase_ ,      index=UpperCAmelCase_ ,      **self.to_csv_kwargs	)
									return written







			def 				UpperCamelCase (      self:    str ,      UpperCAmelCase_:    Tuple	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    args

									_SCREAMING_SNAKE_CASE       =    query_table(
									    table=self.dataset.data ,      key=slice(UpperCAmelCase_ ,      offset + self.batch_size	) ,      indices=self.dataset._indices ,      )
									_SCREAMING_SNAKE_CASE       =    batch.to_pandas().to_csv(
									    path_or_buf=UpperCAmelCase_ ,      header=header if (offset == 0) else False ,      index=UpperCAmelCase_ ,      **UpperCAmelCase_	)
									return csv_str.encode(self.encoding	)







			def 				UpperCamelCase (      self:    Optional[Any] ,      UpperCAmelCase_:    BinaryIO ,      UpperCAmelCase_:    Tuple ,      UpperCAmelCase_:    Any ,      **UpperCAmelCase_:    Union[str, Any]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    0

									if self.num_proc is None or self.num_proc == 1:
															for offset in logging.tqdm(
															    range(0 ,      len(self.dataset	) ,      self.batch_size	) ,      unit="""ba""" ,      disable=not logging.is_progress_bar_enabled() ,      desc="""Creating CSV from Arrow format""" ,      ):
																					_SCREAMING_SNAKE_CASE       =    self._batch_csv((offset, header, index, to_csv_kwargs)	)
																					written += file_obj.write(UpperCAmelCase_	)

									else:
															_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    len(self.dataset	), self.batch_size
															with multiprocessing.Pool(self.num_proc	) as pool:
																					for csv_str in logging.tqdm(
																					    pool.imap(
																					        self._batch_csv ,      [(offset, header, index, to_csv_kwargs) for offset in range(0 ,      UpperCAmelCase_ ,      UpperCAmelCase_	)] ,      ) ,      total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size ,      unit="""ba""" ,      disable=not logging.is_progress_bar_enabled() ,      desc="""Creating CSV from Arrow format""" ,      ):
																											written += file_obj.write(UpperCAmelCase_	)

									return written

