import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
	import jax
	import jax.numpy as jnp

	from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor
class        __snake_case  (unittest.TestCase		):



				def __init__(      self      :		int		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		int=3		,	_UpperCAmelCase      :		Union[str, Any]=32		,	_UpperCAmelCase      :		List[str]=3		,	_UpperCAmelCase      :		Dict=10		,	_UpperCAmelCase      :		Tuple=[10, 20, 30, 40]		,	_UpperCAmelCase      :		str=[1, 1, 2, 1]		,	_UpperCAmelCase      :		Optional[int]=True		,	_UpperCAmelCase      :		str=True		,	_UpperCAmelCase      :		str="relu"		,	_UpperCAmelCase      :		List[str]=3		,	_UpperCAmelCase      :		List[Any]=None		,	)  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      parent
								_lowerCAmelCase    :    Optional[Any]										=      batch_size
								_lowerCAmelCase    :    str										=      image_size
								_lowerCAmelCase    :    int										=      num_channels
								_lowerCAmelCase    :    List[str]										=      embeddings_size
								_lowerCAmelCase    :    Dict										=      hidden_sizes
								_lowerCAmelCase    :    str										=      depths
								_lowerCAmelCase    :    str										=      is_training
								_lowerCAmelCase    :    List[str]										=      use_labels
								_lowerCAmelCase    :    List[str]										=      hidden_act
								_lowerCAmelCase    :    str										=      num_labels
								_lowerCAmelCase    :    Tuple										=      scope
								_lowerCAmelCase    :    Tuple										=      len(_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

								_lowerCAmelCase    :    Optional[int]										=      self.get_config()

								return config, pixel_values



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Any:





								'''simple docstring'''




								return RegNetConfig(
								    num_channels=self.num_channels		,	embeddings_size=self.embeddings_size		,	hidden_sizes=self.hidden_sizes		,	depths=self.depths		,	hidden_act=self.hidden_act		,	num_labels=self.num_labels		,	image_size=self.image_size		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		int		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Any      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      FlaxRegNetModel(config=_UpperCAmelCase      )
								_lowerCAmelCase    :    Optional[Any]										=      model(_UpperCAmelCase      )

								# Output shape (b, c, h, w)
								self.parent.assertEqual(
								    result.last_hidden_state.shape		,	(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		str		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		List[Any]      )  ->    Optional[int]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      self.num_labels
								_lowerCAmelCase    :    Dict										=      FlaxRegNetForImageClassification(config=_UpperCAmelCase      )
								_lowerCAmelCase    :    Dict										=      model(_UpperCAmelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_labels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      self.prepare_config_and_inputs()
								_lowerCAmelCase    :    List[str]										=      config_and_inputs
								_lowerCAmelCase    :    Optional[Any]										=      {"""pixel_values""": pixel_values}
								return config, inputs_dict






@require_flax
class        __snake_case  (_a						,	unittest.TestCase		):
				lowerCAmelCase__				=     (FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

				lowerCAmelCase__				=     False
				lowerCAmelCase__				=     False
				lowerCAmelCase__				=     False



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    None:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      FlaxRegNetModelTester(self      )
								_lowerCAmelCase    :    Dict										=      ConfigTester(self		,	config_class=_UpperCAmelCase		,	has_text_modality=_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    Any:





								'''simple docstring'''




								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()



				def       SCREAMING_SNAKE_CASE       (      self      :		Any      )  ->    Any:





								'''simple docstring'''




								return



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Union[str, Any]      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase      )



				@unittest.skip(reason="""RegNet does not use inputs_embeds"""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Any:





								'''simple docstring'''




								pass



				@unittest.skip(reason="""RegNet does not support input and output embeddings"""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    Dict:





								'''simple docstring'''




								pass



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    str:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												_lowerCAmelCase    :    Any										=      model_class(_UpperCAmelCase      )
												_lowerCAmelCase    :    Any										=      inspect.signature(model.__call__      )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												_lowerCAmelCase    :    Optional[int]										=      [*signature.parameters.keys()]

												_lowerCAmelCase    :    Union[str, Any]										=      ["""pixel_values"""]
												self.assertListEqual(arg_names[:1]		,	_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    List[str]:





								'''simple docstring'''




								def check_hidden_states_output(_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		int      ):
												_lowerCAmelCase    :    Tuple										=      model_class(_UpperCAmelCase      )

												_lowerCAmelCase    :    List[Any]										=      model(**self._prepare_for_class(_UpperCAmelCase		,	_UpperCAmelCase      )      )

												_lowerCAmelCase    :    int										=      outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

												_lowerCAmelCase    :    Union[str, Any]										=      self.model_tester.num_stages
												self.assertEqual(len(_UpperCAmelCase      )		,	expected_num_stages + 1      )

								_lowerCAmelCase    :    Any										=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												_lowerCAmelCase    :    int										=      True
												check_hidden_states_output(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )

												# check that output_hidden_states also work using config
												del inputs_dict["output_hidden_states"]
												_lowerCAmelCase    :    List[Any]										=      True

												check_hidden_states_output(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    Tuple:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__      ):
																_lowerCAmelCase    :    Union[str, Any]										=      self._prepare_for_class(_UpperCAmelCase		,	_UpperCAmelCase      )
																_lowerCAmelCase    :    Optional[int]										=      model_class(_UpperCAmelCase      )

																@jax.jit
																def model_jitted(_UpperCAmelCase      :		Union[str, Any]		,	**_UpperCAmelCase      :		Dict      ):
																				return model(pixel_values=_UpperCAmelCase		,	**_UpperCAmelCase      )

																with self.subTest("""JIT Enabled"""      ):
																				_lowerCAmelCase    :    List[str]										=      model_jitted(**_UpperCAmelCase      ).to_tuple()

																with self.subTest("""JIT Disabled"""      ):
																				with jax.disable_jit():
																								_lowerCAmelCase    :    Optional[int]										=      model_jitted(**_UpperCAmelCase      ).to_tuple()

																self.assertEqual(len(_UpperCAmelCase      )		,	len(_UpperCAmelCase      )      )
																for jitted_output, output in zip(_UpperCAmelCase		,	_UpperCAmelCase      ):
																				self.assertEqual(jitted_output.shape		,	output.shape      )






def 							_UpperCAmelCase  ():
				'''simple docstring'''
				_lowerCAmelCase    :    Any										=      Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""       )
				return image





@require_flax
class        __snake_case  (unittest.TestCase		):



				@cached_property
				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    List[Any]:





								'''simple docstring'''




								return AutoImageProcessor.from_pretrained("""facebook/regnet-y-040"""      ) if is_vision_available() else None



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      FlaxRegNetForImageClassification.from_pretrained("""facebook/regnet-y-040"""      )

								_lowerCAmelCase    :    Any										=      self.default_image_processor
								_lowerCAmelCase    :    Any										=      prepare_img()
								_lowerCAmelCase    :    int										=      image_processor(images=_UpperCAmelCase		,	return_tensors="""np"""      )

								_lowerCAmelCase    :    str										=      model(**_UpperCAmelCase      )

								# verify the logits
								_lowerCAmelCase    :    List[Any]										=      (1, 1000)
								self.assertEqual(outputs.logits.shape		,	_UpperCAmelCase      )

								_lowerCAmelCase    :    List[str]										=      jnp.array([-0.4180, -1.5051, -3.4836]      )

								self.assertTrue(jnp.allclose(outputs.logits[0, :3]		,	_UpperCAmelCase		,	atol=1E-4      )      )







import comet  # From: unbabel-comet
import torch

import datasets


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

_lowerCamelCase      :  Optional[Any]      						=						"\\n@inproceedings{rei-EtAl:2020:WMT,\n   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},\n   title     = {Unbabel's Participation in the WMT20 Metrics Shared Task},\n   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},\n   month          = {November},\n   year           = {2020},\n   address        = {Online},\n   publisher      = {Association for Computational Linguistics},\n   pages     = {909--918},\n}\n@inproceedings{rei-etal-2020-comet,\n   title = \"{COMET}: A Neural Framework for {MT} Evaluation\",\n   author = \"Rei, Ricardo  and\n      Stewart, Craig  and\n      Farinha, Ana C  and\n      Lavie, Alon\",\n   booktitle = \"Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)\",\n   month = nov,\n   year = \"2020\",\n   address = \"Online\",\n   publisher = \"Association for Computational Linguistics\",\n   url = \"https://www.aclweb.org/anthology/2020.emnlp-main.213\",\n   pages = \"2685--2702\",\n}\n"

_lowerCamelCase      :  Union[str, Any]      						=						"\\nCrosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA's or MQM).\nWith the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.\n\nSee the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.\n"

_lowerCamelCase      :  Optional[int]      						=						"\nCOMET score.\n\nArgs:\n\n`sources` (list of str): Source sentences\n`predictions` (list of str): candidate translations\n`references` (list of str): reference translations\n`cuda` (bool): If set to True, runs COMET using GPU\n`show_progress` (bool): Shows progress\n`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.\n\nReturns:\n    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.\n    `scores`: List of scores.\n\nExamples:\n\n    >>> comet_metric = datasets.load_metric('comet')\n    >>> # comet_metric = load_metric('comet', 'wmt20-comet-da')  # you can also choose which model to use\n    >>> source = [\"Dem Feuer konnte Einhalt geboten werden\", \"Schulen und Kindergärten wurden eröffnet.\"]\n    >>> hypothesis = [\"The fire could be stopped\", \"Schools and kindergartens were open\"]\n    >>> reference = [\"They were able to control the fire.\", \"Schools and kindergartens opened\"]\n    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)\n    >>> print([round(v, 2) for v in results[\"scores\"]])\n    [0.19, 0.92]\n"
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,	_KWARGS_DESCRIPTION		)
class        __snake_case  (datasets.Metric		):



				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    Tuple:





								'''simple docstring'''




								return datasets.MetricInfo(
								    description=_DESCRIPTION		,	citation=_CITATION		,	homepage="""https://unbabel.github.io/COMET/html/index.html"""		,	inputs_description=_KWARGS_DESCRIPTION		,	features=datasets.Features(
								        {
								            """sources""": datasets.Value("""string"""		,	id="""sequence"""      ),
								            """predictions""": datasets.Value("""string"""		,	id="""sequence"""      ),
								            """references""": datasets.Value("""string"""		,	id="""sequence"""      ),
								        }      )		,	codebase_urls=["""https://github.com/Unbabel/COMET"""]		,	reference_urls=[
								        """https://github.com/Unbabel/COMET""",
								        """https://www.aclweb.org/anthology/2020.emnlp-main.213/""",
								        """http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6""",
								    ]		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict		,	_UpperCAmelCase      :		List[str]      )  ->    Dict:





								'''simple docstring'''




								if self.config_name == "default":
												_lowerCAmelCase    :    Union[str, Any]										=      comet.load_from_checkpoint(comet.download_model("""wmt20-comet-da"""      )      )
								else:
												_lowerCAmelCase    :    Dict										=      comet.load_from_checkpoint(comet.download_model(self.config_name      )      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Dict=None		,	_UpperCAmelCase      :		Tuple=False      )  ->    int:





								'''simple docstring'''




								if gpus is None:
												_lowerCAmelCase    :    Tuple										=      1 if torch.cuda.is_available() else 0
								_lowerCAmelCase    :    Optional[int]										=      {"""src""": sources, """mt""": predictions, """ref""": references}
								_lowerCAmelCase    :    List[Any]										=      [dict(zip(_UpperCAmelCase		,	_UpperCAmelCase      )      ) for t in zip(*data.values()      )]
								_lowerCAmelCase					,    _lowerCAmelCase    :    Optional[Any]										=      self.scorer.predict(_UpperCAmelCase		,	gpus=_UpperCAmelCase		,	progress_bar=_UpperCAmelCase      )
								return {"mean_score": mean_score, "scores": scores}

