




'''simple docstring'''







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


_lowerCamelCase				:		str				      =    logging.get_logger(__name__)





class 	lowerCamelCase_				(			nn.Module				):




					'''simple docstring'''






					lowercase__       =					42
					lowercase__       =					jnp.floataa






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			self.config.text_config
							_snake_case      			=			self.config.vision_config

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

							_snake_case      			=			FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							_snake_case      			=			FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__,    __UpperCamelCase					).module_class

							_snake_case      			=			text_module(__UpperCamelCase,    dtype=self.dtype					)
							_snake_case      			=			vision_module(__UpperCamelCase,    dtype=self.dtype					)

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






					def __call__(       self,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase = True,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    ):
							_snake_case      			=			return_dict if return_dict is not None else self.config.return_dict

							_snake_case      			=			self.vision_model(
							    pixel_values=__UpperCamelCase,    deterministic=__UpperCamelCase,    output_attentions=__UpperCamelCase,    output_hidden_states=__UpperCamelCase,    return_dict=__UpperCamelCase,    )

							_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,    )

							_snake_case      			=			vision_outputs[1]
							_snake_case      			=			self.visual_projection(__UpperCamelCase					)

							_snake_case      			=			text_outputs[1]
							_snake_case      			=			self.text_projection(__UpperCamelCase					)

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

							# cosine similarity as logits
							_snake_case      			=			jnp.exp(self.logit_scale					)
							_snake_case      			=			jnp.matmul(__UpperCamelCase,    image_embeds.T					) * logit_scale
							_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 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					lowercase__       =					HybridCLIPConfig
					lowercase__       =					FlaxHybridCLIPModule






					def __init__(       self,    __UpperCamelCase,    __UpperCamelCase = None,    __UpperCamelCase = 0,    __UpperCamelCase = jnp.floataa,    **__UpperCamelCase,    ):
							if input_shape is None:
									_snake_case      			=			((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

							_snake_case      			=			self.module_class(config=__UpperCamelCase,    dtype=__UpperCamelCase,    **__UpperCamelCase					)
							super().__init__(__UpperCamelCase,    __UpperCamelCase,    input_shape=__UpperCamelCase,    seed=__UpperCamelCase,    dtype=__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase = None					):
							# init input tensor
							_snake_case      			=			jnp.zeros(input_shape[0],    dtype="i4"					)
							_snake_case      			=			jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase					).shape[-1]					),    input_shape[0]					)
							_snake_case      			=			jnp.ones_like(__UpperCamelCase					)
							_snake_case      			=			jnp.ones_like(__UpperCamelCase					)

							_snake_case      			=			jax.random.normal(__UpperCamelCase,    input_shape[1]					)

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

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






					def __call__(       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase = None,    __UpperCamelCase = None,    __UpperCamelCase = False,    __UpperCamelCase = None,    __UpperCamelCase = None,    __UpperCamelCase = None,    ):
							_snake_case      			=			output_attentions if output_attentions is not None else self.config.output_attentions
							_snake_case      			=			(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							_snake_case      			=			return_dict if return_dict is not None else self.config.return_dict

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

							if token_type_ids is None:
									_snake_case      			=			jnp.zeros_like(__UpperCamelCase					)

							if attention_mask is None:
									_snake_case      			=			jnp.ones_like(__UpperCamelCase					)

							# Handle any PRNG if needed
							_snake_case      			=			{}
							if dropout_rng is not None:
									_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 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase = None,    __UpperCamelCase = None,    __UpperCamelCase=False,    ):
							if position_ids is None:
									_snake_case      			=			jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase					).shape[-1]					),    input_ids.shape					)

							if token_type_ids is None:
									_snake_case      			=			jnp.zeros_like(__UpperCamelCase					)

							if attention_mask is None:
									_snake_case      			=			jnp.ones_like(__UpperCamelCase					)

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

							def _get_features(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase					):
									_snake_case      			=			module.text_model(
									    input_ids=__UpperCamelCase,    attention_mask=__UpperCamelCase,    position_ids=__UpperCamelCase,    token_type_ids=__UpperCamelCase,    deterministic=__UpperCamelCase,    )
									_snake_case      			=			text_outputs[1]
									_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 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase = None,    __UpperCamelCase = None,    __UpperCamelCase=False					):
							_snake_case      			=			{}
							if dropout_rng is not None:
									_snake_case      			=			dropout_rng

							def _get_features(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase					):
									_snake_case      			=			module.vision_model(pixel_values=__UpperCamelCase,    deterministic=__UpperCamelCase					)
									_snake_case      			=			vision_outputs[1]  # pooled_output
									_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 							__SCREAMING_SNAKE_CASE      (       cls,    __UpperCamelCase = None,    __UpperCamelCase = None,    *__UpperCamelCase,    **__UpperCamelCase,    ):
							_snake_case      			=			{
							    argument[len("text_"					) :]: value for argument, value in kwargs.items() if argument.startswith("text_"					)
							}

							_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
							_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

											_snake_case      			=			AutoConfig.from_pretrained(__UpperCamelCase					)
											_snake_case      			=			text_config

									_snake_case      			=			FlaxAutoModel.from_pretrained(__UpperCamelCase,    *__UpperCamelCase,    **__UpperCamelCase					)

							_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

											_snake_case      			=			AutoConfig.from_pretrained(__UpperCamelCase					)
											_snake_case      			=			vision_config

									_snake_case      			=			FlaxAutoModel.from_pretrained(__UpperCamelCase,    *__UpperCamelCase,    **__UpperCamelCase					)

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

							# init model
							_snake_case      			=			cls(__UpperCamelCase,    *__UpperCamelCase,    dtype=__UpperCamelCase,    **__UpperCamelCase					)

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

							_snake_case      			=			text_model.params

							return model











'''simple docstring'''







import subprocess
from argparse import ArgumentParser
from typing import List, Union

from huggingface_hub.hf_api import HfFolder, create_repo, whoami
from requests.exceptions import HTTPError

from . import BaseTransformersCLICommand





class 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					@staticmethod
					def 							__SCREAMING_SNAKE_CASE      (       __UpperCamelCase					):
							_snake_case      			=			parser.add_parser("login",    help="Log in using the same credentials as on huggingface.co"					)
							login_parser.set_defaults(func=lambda __UpperCamelCase					: LoginCommand(__UpperCamelCase					)					)
							_snake_case      			=			parser.add_parser("whoami",    help="Find out which huggingface.co account you are logged in as."					)
							whoami_parser.set_defaults(func=lambda __UpperCamelCase					: WhoamiCommand(__UpperCamelCase					)					)
							_snake_case      			=			parser.add_parser("logout",    help="Log out"					)
							logout_parser.set_defaults(func=lambda __UpperCamelCase					: LogoutCommand(__UpperCamelCase					)					)

							# new system: git-based repo system
							_snake_case      			=			parser.add_parser(
							    "repo",    help="Deprecated: use `huggingface-cli` instead. Commands to interact with your huggingface.co repos.",    )
							_snake_case      			=			repo_parser.add_subparsers(
							    help="Deprecated: use `huggingface-cli` instead. huggingface.co repos related commands"					)
							_snake_case      			=			repo_subparsers.add_parser(
							    "create",    help="Deprecated: use `huggingface-cli` instead. Create a new repo on huggingface.co"					)
							repo_create_parser.add_argument(
							    "name",    type=__UpperCamelCase,    help="Name for your model's repo. Will be namespaced under your username to build the model id.",    )
							repo_create_parser.add_argument("--organization",    type=__UpperCamelCase,    help="Optional: organization namespace."					)
							repo_create_parser.add_argument("-y",    "--yes",    action="store_true",    help="Optional: answer Yes to the prompt"					)
							repo_create_parser.set_defaults(func=lambda __UpperCamelCase					: RepoCreateCommand(__UpperCamelCase					)					)





class 	lowerCamelCase_				:




					'''simple docstring'''






					lowercase__       =					"""\u001b[1m"""
					lowercase__       =					"""\u001b[31m"""
					lowercase__       =					"""\u001b[90m"""
					lowercase__       =					"""\u001b[0m"""






					@classmethod
					def 							__SCREAMING_SNAKE_CASE      (       cls,    __UpperCamelCase					):
							return f'''{cls._bold}{s}{cls._reset}'''






					@classmethod
					def 							__SCREAMING_SNAKE_CASE      (       cls,    __UpperCamelCase					):
							return f'''{cls._bold}{cls._red}{s}{cls._reset}'''






					@classmethod
					def 							__SCREAMING_SNAKE_CASE      (       cls,    __UpperCamelCase					):
							return f'''{cls._gray}{s}{cls._reset}'''






def 						_lowercase						(      lowercase_		,				lowercase_    ):



		'''simple docstring'''







		_snake_case      			=			[max(len(str(lowercase_    )    ) for x in col    ) for col in zip(*lowercase_		,				lowercase_    )]
		_snake_case      			=			("{{:{}}} " * len(lowercase_    )).format(*lowercase_    )
		_snake_case      			=			[]
		lines.append(row_format.format(*lowercase_    )    )
		lines.append(row_format.format(*["-" * w for w in col_widths]    )    )
		for row in rows:
				lines.append(row_format.format(*lowercase_    )    )
		return "\n".join(lowercase_    )





class 	lowerCamelCase_				:




					'''simple docstring'''






					def __init__(       self,    __UpperCamelCase					):
							_snake_case      			=			args





class 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							print(
							    ANSI.red(
							        "ERROR! `huggingface-cli login` uses an outdated login mechanism "
							        "that is not compatible with the Hugging Face Hub backend anymore. "
							        "Please use `huggingface-cli login instead."					)					)





class 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							print(
							    ANSI.red(
							        "WARNING! `transformers-cli whoami` is deprecated and will be removed in v5. Please use "
							        "`huggingface-cli whoami` instead."					)					)
							_snake_case      			=			HfFolder.get_token()
							if token is None:
									print("Not logged in"					)
									exit()
							try:
									_snake_case			,       _snake_case      			=			whoami(__UpperCamelCase					)
									print(__UpperCamelCase					)
									if orgs:
											print(ANSI.bold("orgs: "					),    ",".join(__UpperCamelCase					)					)
							except HTTPError as e:
									print(__UpperCamelCase					)
									print(ANSI.red(e.response.text					)					)
									exit(1					)





class 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							print(
							    ANSI.red(
							        "ERROR! `transformers-cli logout` uses an outdated logout mechanism "
							        "that is not compatible with the Hugging Face Hub backend anymore. "
							        "Please use `huggingface-cli logout instead."					)					)





class 	lowerCamelCase_				(			UpperCAmelCase_				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							print(
							    ANSI.red(
							        "WARNING! Managing repositories through transformers-cli is deprecated. "
							        "Please use `huggingface-cli` instead."					)					)
							_snake_case      			=			HfFolder.get_token()
							if token is None:
									print("Not logged in"					)
									exit(1					)
							try:
									_snake_case      			=			subprocess.check_output(["git", "--version"]					).decode("utf-8"					)
									print(ANSI.gray(stdout.strip()					)					)
							except FileNotFoundError:
									print("Looks like you do not have git installed, please install."					)

							try:
									_snake_case      			=			subprocess.check_output(["git-lfs", "--version"]					).decode("utf-8"					)
									print(ANSI.gray(stdout.strip()					)					)
							except FileNotFoundError:
									print(
									    ANSI.red(
									        "Looks like you do not have git-lfs installed, please install."
									        " You can install from https://git-lfs.github.com/."
									        " Then run `git lfs install` (you only have to do this once)."					)					)
							print(""					)

							_snake_case			,       _snake_case      			=			whoami(__UpperCamelCase					)
							_snake_case      			=			self.args.organization if self.args.organization is not None else user
							_snake_case      			=			f'''{namespace}/{self.args.name}'''
							print(f'''You are about to create {ANSI.bold(__UpperCamelCase					)}'''					)

							if not self.args.yes:
									_snake_case      			=			input("Proceed? [Y/n] "					).lower()
									if not (choice == "" or choice == "y" or choice == "yes"):
											print("Abort"					)
											exit()
							try:
									_snake_case      			=			create_repo(__UpperCamelCase,    name=self.args.name,    organization=self.args.organization					)
							except HTTPError as e:
									print(__UpperCamelCase					)
									print(ANSI.red(e.response.text					)					)
									exit(1					)
							print("\nYour repo now lives at:"					)
							print(f'''  {ANSI.bold(__UpperCamelCase					)}'''					)
							print("\nYou can clone it locally with the command below, and commit/push as usual."					)
							print(f'''\n  git clone {url}'''					)
							print(""					)





