





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


lowercase_         =					logging.get_logger(__name__)


class  a_					(    nn.Module	):

    '''simple docstring'''
    UpperCamelCase			   =			42
    UpperCamelCase			   =			jnp.floataa


    def       snake_case_(		self	)		->       Tuple:
        _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			,		A=None			,		A=None			,		A=None			,		A=None			,		A=None			,		A = True			,		A=None			,		A=None			,		A=None			,		)		->       str:
        _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  a_					(    __lowerCamelCase	):

    '''simple docstring'''
    UpperCamelCase			   =			HybridCLIPConfig
    UpperCamelCase			   =			FlaxHybridCLIPModule


    def __init__(		self			,		A			,		A = None			,		A = 0			,		A = jnp.floataa			,		**A			,		)		->       Dict:





        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       snake_case_(		self			,		A			,		A			,		A = None	)		->       FrozenDict:
        _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			,		A			,		A			,		A=None			,		A=None			,		A=None			,		A = None			,		A = None			,		A = False			,		A = None			,		A = None			,		A = None			,		)		->       List[Any]:
        _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       snake_case_(		self			,		A			,		A=None			,		A=None			,		A=None			,		A = None			,		A = None			,		A=False			,		)		->       Union[str, Any]:





        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(A			,		A			,		A			,		A			,		A			,		A	):
            _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       snake_case_(		self			,		A			,		A = None			,		A = None			,		A=False	)		->       List[str]:
        _SCREAMING_SNAKE_CASE     =      {}
        if dropout_rng is not None:
            _SCREAMING_SNAKE_CASE     =      dropout_rng

        def _get_features(A			,		A			,		A	):
            _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       snake_case_(		cls			,		A = None			,		A = None			,		*A			,		**A			,		)		->       FlaxPreTrainedModel:
        _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










'''simple docstring'''

import sacrebleu as scb
from packaging import version

import datasets


lowercase_         =			"""\
@inproceedings{post-2018-call,
    title = \"A Call for Clarity in Reporting {BLEU} Scores\",
    author = \"Post, Matt\",
    booktitle = \"Proceedings of the Third Conference on Machine Translation: Research Papers\",
    month = oct,
    year = \"2018\",
    address = \"Belgium, Brussels\",
    publisher = \"Association for Computational Linguistics\",
    url = \"https://www.aclweb.org/anthology/W18-6319\",
    pages = \"186--191\",
}
"""

lowercase_         =			"""\
SacreBLEU provides hassle-free computation of shareable, comparable, and reproducible BLEU scores.
Inspired by Rico Sennrich's `multi-bleu-detok.perl`, it produces the official WMT scores but works with plain text.
It also knows all the standard test sets and handles downloading, processing, and tokenization for you.

See the [README.md] file at https://github.com/mjpost/sacreBLEU for more information.
"""

lowercase_         =			"""
Produces BLEU scores along with its sufficient statistics
from a source against one or more references.

Args:
    predictions (`list` of `str`): list of translations to score. Each translation should be tokenized into a list of tokens.
    references (`list` of `list` of `str`): A list of lists of references. The contents of the first sub-list are the references for the first prediction, the contents of the second sub-list are for the second prediction, etc. Note that there must be the same number of references for each prediction (i.e. all sub-lists must be of the same length).
    smooth_method (`str`): The smoothing method to use, defaults to `'exp'`. Possible values are:
        - `'none'`: no smoothing
        - `'floor'`: increment zero counts
        - `'add-k'`: increment num/denom by k for n>1
        - `'exp'`: exponential decay
    smooth_value (`float`): The smoothing value. Only valid when `smooth_method='floor'` (in which case `smooth_value` defaults to `0.1`) or `smooth_method='add-k'` (in which case `smooth_value` defaults to `1`).
    tokenize (`str`): Tokenization method to use for BLEU. If not provided, defaults to `'zh'` for Chinese, `'ja-mecab'` for Japanese and `'13a'` (mteval) otherwise. Possible values are:
        - `'none'`: No tokenization.
        - `'zh'`: Chinese tokenization.
        - `'13a'`: mimics the `mteval-v13a` script from Moses.
        - `'intl'`: International tokenization, mimics the `mteval-v14` script from Moses
        - `'char'`: Language-agnostic character-level tokenization.
        - `'ja-mecab'`: Japanese tokenization. Uses the [MeCab tokenizer](https://pypi.org/project/mecab-python3).
    lowercase (`bool`): If `True`, lowercases the input, enabling case-insensitivity. Defaults to `False`.
    force (`bool`): If `True`, insists that your tokenized input is actually detokenized. Defaults to `False`.
    use_effective_order (`bool`): If `True`, stops including n-gram orders for which precision is 0. This should be `True`, if sentence-level BLEU will be computed. Defaults to `False`.

Returns:
    'score': BLEU score,
    'counts': Counts,
    'totals': Totals,
    'precisions': Precisions,
    'bp': Brevity penalty,
    'sys_len': predictions length,
    'ref_len': reference length,

Examples:

    Example 1:
        >>> predictions = [\"hello there general kenobi\", \"foo bar foobar\"]
        >>> references = [[\"hello there general kenobi\", \"hello there !\"], [\"foo bar foobar\", \"foo bar foobar\"]]
        >>> sacrebleu = datasets.load_metric(\"sacrebleu\")
        >>> results = sacrebleu.compute(predictions=predictions, references=references)
        >>> print(list(results.keys()))
        ['score', 'counts', 'totals', 'precisions', 'bp', 'sys_len', 'ref_len']
        >>> print(round(results[\"score\"], 1))
        100.0

    Example 2:
        >>> predictions = [\"hello there general kenobi\",
        ...                 \"on our way to ankh morpork\"]
        >>> references = [[\"hello there general kenobi\", \"hello there !\"],
        ...                 [\"goodbye ankh morpork\", \"ankh morpork\"]]
        >>> sacrebleu = datasets.load_metric(\"sacrebleu\")
        >>> results = sacrebleu.compute(predictions=predictions,
        ...                             references=references)
        >>> print(list(results.keys()))
        ['score', 'counts', 'totals', 'precisions', 'bp', 'sys_len', 'ref_len']
        >>> print(round(results[\"score\"], 1))
        39.8
"""


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION				,   _KWARGS_DESCRIPTION	)
class  a_					(    datasets.Metric	):

    '''simple docstring'''


    def       snake_case_(		self	)		->       Union[str, Any]:
        if version.parse(scb.__version__	) < version.parse("""1.4.12"""	):
            raise ImportWarning(
                """To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn't match this condition.\n"""
                """You can install it with `pip install \"sacrebleu>=1.4.12\"`."""	)
        return datasets.MetricInfo(
            description=_DESCRIPTION			,		citation=_CITATION			,		homepage="""https://github.com/mjpost/sacreBLEU"""			,		inputs_description=_KWARGS_DESCRIPTION			,		features=datasets.Features(
                {
                    """predictions""": datasets.Value("""string"""			,		id="""sequence"""	),
                    """references""": datasets.Sequence(datasets.Value("""string"""			,		id="""sequence"""	)			,		id="""references"""	),
                }	)			,		codebase_urls=["""https://github.com/mjpost/sacreBLEU"""]			,		reference_urls=[
                """https://github.com/mjpost/sacreBLEU""",
                """https://en.wikipedia.org/wiki/BLEU""",
                """https://towardsdatascience.com/evaluating-text-output-in-nlp-bleu-at-your-own-risk-e8609665a213""",
            ]			,		)





    def       snake_case_(		self			,		A			,		A			,		A="exp"			,		A=None			,		A=False			,		A=False			,		A=None			,		A=False			,		)		->       Optional[Any]:
        _SCREAMING_SNAKE_CASE     =      len(references[0]	)
        if any(len(A	) != references_per_prediction for refs in references	):
            raise ValueError("""Sacrebleu requires the same number of references for each prediction"""	)
        _SCREAMING_SNAKE_CASE     =      [[refs[i] for refs in references] for i in range(A	)]
        _SCREAMING_SNAKE_CASE     =      scb.corpus_bleu(
            A			,		A			,		smooth_method=A			,		smooth_value=A			,		force=A			,		lowercase=A			,		use_effective_order=A			,		**({"""tokenize""": tokenize} if tokenize else {})			,		)
        _SCREAMING_SNAKE_CASE     =      {
            """score""": output.score,
            """counts""": output.counts,
            """totals""": output.totals,
            """precisions""": output.precisions,
            """bp""": output.bp,
            """sys_len""": output.sys_len,
            """ref_len""": output.ref_len,
        }
        return output_dict



