



import numpy as np
from scipy.spatial.distance import cdist
from sklearn.metrics import fa_score

import datasets


__lowerCamelCase     :							Any  							=				"\\n    @inproceedings{kakwani2020indicnlpsuite,\n    title={{IndicNLPSuite: Monolingual Corpora, Evaluation Benchmarks and Pre-trained Multilingual Language Models for Indian Languages}},\n    author={Divyanshu Kakwani and Anoop Kunchukuttan and Satish Golla and Gokul N.C. and Avik Bhattacharyya and Mitesh M. Khapra and Pratyush Kumar},\n    year={2020},\n    booktitle={Findings of EMNLP},\n}\n"

__lowerCamelCase     :							Union[str, Any]  							=				"\\n    IndicGLUE is a natural language understanding benchmark for Indian languages. It contains a wide\n    variety of tasks and covers 11 major Indian languages - as, bn, gu, hi, kn, ml, mr, or, pa, ta, te.\n"

__lowerCamelCase     :							str  							=				"\nCompute IndicGLUE evaluation metric associated to each IndicGLUE dataset.\nArgs:\n    predictions: list of predictions to score (as int64),\n        except for 'cvit-mkb-clsr' where each prediction is a vector (of float32).\n    references: list of ground truth labels corresponding to the predictions (as int64),\n        except for 'cvit-mkb-clsr' where each reference is a vector (of float32).\nReturns: depending on the IndicGLUE subset, one or several of:\n    \"accuracy\": Accuracy\n    \"f1\": F1 score\n    \"precision\": Precision@10\nExamples:\n\n    >>> indic_glue_metric = datasets.load_metric('indic_glue', 'wnli')  # 'wnli' or any of [\"copa\", \"sna\", \"csqa\", \"wstp\", \"inltkh\", \"bbca\", \"iitp-mr\", \"iitp-pr\", \"actsa-sc\", \"md\"]\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'accuracy': 1.0}\n\n    >>> indic_glue_metric = datasets.load_metric('indic_glue', 'wiki-ner')\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'accuracy': 1.0, 'f1': 1.0}\n\n    >>> indic_glue_metric = datasets.load_metric('indic_glue', 'cvit-mkb-clsr')\n    >>> references = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]\n    >>> predictions = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'precision@10': 1.0}\n\n"

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Dict:
						return float((preds == labels).mean()    )

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Tuple:
						UpperCAmelCase         =     simple_accuracy(lowerCamelCase_    ,   lowerCamelCase_    )
						UpperCAmelCase         =     float(fa_score(y_true=lowerCamelCase_    ,   y_pred=lowerCamelCase_    )    )
						return {
						    "accuracy": acc,
						    "f1": fa,
						}

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Optional[int]:
						UpperCAmelCase         =     np.array(lowerCamelCase_    )
						UpperCAmelCase         =     np.array(lowerCamelCase_    )
						UpperCAmelCase         =     en_sentvecs.shape[0]

						# mean centering
						UpperCAmelCase         =     en_sentvecs - np.mean(lowerCamelCase_    ,   axis=0    )
						UpperCAmelCase         =     in_sentvecs - np.mean(lowerCamelCase_    ,   axis=0    )

						UpperCAmelCase         =     cdist(lowerCamelCase_    ,   lowerCamelCase_    ,   "cosine"    )
						UpperCAmelCase         =     np.array(range(lowerCamelCase_    )    )
						UpperCAmelCase         =     sim.argsort(axis=1    )[:, :10]
						UpperCAmelCase         =     np.any(preds == actual[:, None]    ,   axis=1    )
						return float(matches.mean()    )





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






		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]      )				->					Optional[Any]:
								'''simple docstring'''







								if self.config_name not in [
								    "wnli",
								    "copa",
								    "sna",
								    "csqa",
								    "wstp",
								    "inltkh",
								    "bbca",
								    "cvit-mkb-clsr",
								    "iitp-mr",
								    "iitp-pr",
								    "actsa-sc",
								    "md",
								    "wiki-ner",
								]:
														raise KeyError(
														    "You should supply a configuration name selected in "
														    "[\"wnli\", \"copa\", \"sna\", \"csqa\", \"wstp\", \"inltkh\", \"bbca\", "
														    "\"cvit-mkb-clsr\", \"iitp-mr\", \"iitp-pr\", \"actsa-sc\", \"md\", "
														    "\"wiki-ner\"]"      )
								return datasets.MetricInfo(
								    description=_DESCRIPTION				,		citation=_CITATION				,		inputs_description=_KWARGS_DESCRIPTION				,		features=datasets.Features(
								        {
								            "predictions": datasets.Value("int64"      )
								            if self.config_name != "cvit-mkb-clsr"
								            else datasets.Sequence(datasets.Value("float32"      )      ),
								            "references": datasets.Value("int64"      )
								            if self.config_name != "cvit-mkb-clsr"
								            else datasets.Sequence(datasets.Value("float32"      )      ),
								        }      )				,		codebase_urls=[]				,		reference_urls=[]				,		format="numpy" if self.config_name != "cvit-mkb-clsr" else None				,		)







		def SCREAMING_SNAKE_CASE_			(						self      :     int				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Any      )				->					List[str]:
								'''simple docstring'''







								if self.config_name == "cvit-mkb-clsr":
														return {"precision@10": precision_at_aa(UpperCamelCase__				,		UpperCamelCase__      )}
								elif self.config_name in ["wiki-ner"]:
														return acc_and_fa(UpperCamelCase__				,		UpperCamelCase__      )
								elif self.config_name in [
								    "wnli",
								    "copa",
								    "sna",
								    "csqa",
								    "wstp",
								    "inltkh",
								    "bbca",
								    "iitp-mr",
								    "iitp-pr",
								    "actsa-sc",
								    "md",
								]:
														return {"accuracy": simple_accuracy(UpperCamelCase__				,		UpperCamelCase__      )}
								else:
														raise KeyError(
														    "You should supply a configuration name selected in "
														    "[\"wnli\", \"copa\", \"sna\", \"csqa\", \"wstp\", \"inltkh\", \"bbca\", "
														    "\"cvit-mkb-clsr\", \"iitp-mr\", \"iitp-pr\", \"actsa-sc\", \"md\", "
														    "\"wiki-ner\"]"      )








import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
		import jax
		import jax.numpy as jnp
		from flax.jax_utils import replicate
		from flax.training.common_utils import shard





@slow
@require_flax
class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Optional[Any]:
								'''simple docstring'''







								super().tearDown()
								gc.collect()






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Dict:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxControlNetModel.from_pretrained(
								    "lllyasviel/sd-controlnet-canny"				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    "runwayml/stable-diffusion-v1-5"				,		controlnet=UpperCamelCase__				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase         =     controlnet_params

								UpperCAmelCase         =     "bird"
								UpperCAmelCase         =     jax.device_count()
								UpperCAmelCase         =     pipe.prepare_text_inputs([prompts] * num_samples      )

								UpperCAmelCase         =     load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"      )
								UpperCAmelCase         =     pipe.prepare_image_inputs([canny_image] * num_samples      )

								UpperCAmelCase         =     jax.random.PRNGKey(0      )
								UpperCAmelCase         =     jax.random.split(UpperCamelCase__				,		jax.device_count()      )

								UpperCAmelCase         =     replicate(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )

								UpperCAmelCase         =     pipe(
								    prompt_ids=UpperCamelCase__				,		image=UpperCamelCase__				,		params=UpperCamelCase__				,		prng_seed=UpperCamelCase__				,		num_inference_steps=50				,		jit=UpperCamelCase__				,		).images
								assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

								UpperCAmelCase         =     images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
								UpperCAmelCase         =     images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCAmelCase         =     jnp.asarray(jax.device_get(image_slice.flatten()      )      )
								UpperCAmelCase         =     jnp.array(
								    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]      )
								print(F'output_slice: {output_slice}'      )
								assert jnp.abs(output_slice - expected_slice      ).max() < 1e-2







		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxControlNetModel.from_pretrained(
								    "lllyasviel/sd-controlnet-openpose"				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    "runwayml/stable-diffusion-v1-5"				,		controlnet=UpperCamelCase__				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase         =     controlnet_params

								UpperCAmelCase         =     "Chef in the kitchen"
								UpperCAmelCase         =     jax.device_count()
								UpperCAmelCase         =     pipe.prepare_text_inputs([prompts] * num_samples      )

								UpperCAmelCase         =     load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png"      )
								UpperCAmelCase         =     pipe.prepare_image_inputs([pose_image] * num_samples      )

								UpperCAmelCase         =     jax.random.PRNGKey(0      )
								UpperCAmelCase         =     jax.random.split(UpperCamelCase__				,		jax.device_count()      )

								UpperCAmelCase         =     replicate(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )

								UpperCAmelCase         =     pipe(
								    prompt_ids=UpperCamelCase__				,		image=UpperCamelCase__				,		params=UpperCamelCase__				,		prng_seed=UpperCamelCase__				,		num_inference_steps=50				,		jit=UpperCamelCase__				,		).images
								assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

								UpperCAmelCase         =     images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
								UpperCAmelCase         =     images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCAmelCase         =     jnp.asarray(jax.device_get(image_slice.flatten()      )      )
								UpperCAmelCase         =     jnp.array(
								    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]      )
								print(F'output_slice: {output_slice}'      )
								assert jnp.abs(output_slice - expected_slice      ).max() < 1e-2



