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

import datasets


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

UpperCAmelCase	      = '''\
    IndicGLUE is a natural language understanding benchmark for Indian languages. It contains a wide
    variety of tasks and covers 11 major Indian languages - as, bn, gu, hi, kn, ml, mr, or, pa, ta, te.
'''

UpperCAmelCase	      = '''
Compute IndicGLUE evaluation metric associated to each IndicGLUE dataset.
Args:
    predictions: list of predictions to score (as int64),
        except for \'cvit-mkb-clsr\' where each prediction is a vector (of float32).
    references: list of ground truth labels corresponding to the predictions (as int64),
        except for \'cvit-mkb-clsr\' where each reference is a vector (of float32).
Returns: depending on the IndicGLUE subset, one or several of:
    "accuracy": Accuracy
    "f1": F1 score
    "precision": Precision@10
Examples:

    >>> 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"]
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}

    >>> indic_glue_metric = datasets.load_metric(\'indic_glue\', \'wiki-ner\')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0, \'f1\': 1.0}

    >>> indic_glue_metric = datasets.load_metric(\'indic_glue\', \'cvit-mkb-clsr\')
    >>> references = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]
    >>> predictions = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]
    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'precision@10\': 1.0}

'''






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	return float((preds == labels).mean()							)






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    simple_accuracy(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
	lowercase       =    float(fa_score(y_true=__SCREAMING_SNAKE_CASE ,							y_pred=__SCREAMING_SNAKE_CASE							)							)
	return {
	    "accuracy": acc,
	    "f1": fa,
	}






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    np.array(__SCREAMING_SNAKE_CASE							)
	lowercase       =    np.array(__SCREAMING_SNAKE_CASE							)
	lowercase       =    en_sentvecs.shape[0]

	# mean centering
	lowercase       =    en_sentvecs - np.mean(__SCREAMING_SNAKE_CASE ,							axis=0							)
	lowercase       =    in_sentvecs - np.mean(__SCREAMING_SNAKE_CASE ,							axis=0							)

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







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




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		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		,    snake_case		,    snake_case       ):
		if self.config_name == "cvit-mkb-clsr":
			return {"precision@10": precision_at_aa(snake_case		,    snake_case       )}
		elif self.config_name in ["wiki-ner"]:
			return acc_and_fa(snake_case		,    snake_case       )
		elif self.config_name in [
		    "wnli",
		    "copa",
		    "sna",
		    "csqa",
		    "wstp",
		    "inltkh",
		    "bbca",
		    "iitp-mr",
		    "iitp-pr",
		    "actsa-sc",
		    "md",
		]:
			return {"accuracy": simple_accuracy(snake_case		,    snake_case       )}
		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 re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCAmelCase	      = logging.get_logger(__name__)






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    r'\w+[.]\d+'
	lowercase       =    re.findall(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
	for pat in pats:
		lowercase       =    key.replace(__SCREAMING_SNAKE_CASE ,							'_'.join(pat.split('.'							)							)							)
	return key






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    pt_tuple_key[:-1] + ('scale',)
	if (
	    any('norm' in str_ for str_ in pt_tuple_key							)
	    and (pt_tuple_key[-1] == "bias")
	    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
	    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
	):
		lowercase       =    pt_tuple_key[:-1] + ('scale',)
		return renamed_pt_tuple_key, pt_tensor
	elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
		lowercase       =    pt_tuple_key[:-1] + ('scale',)
		return renamed_pt_tuple_key, pt_tensor

	# embedding
	if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
		lowercase       =    pt_tuple_key[:-1] + ('embedding',)
		return renamed_pt_tuple_key, pt_tensor

	# conv layer
	lowercase       =    pt_tuple_key[:-1] + ('kernel',)
	if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
		lowercase       =    pt_tensor.transpose(2 ,							3 ,							1 ,							0							)
		return renamed_pt_tuple_key, pt_tensor

	# linear layer
	lowercase       =    pt_tuple_key[:-1] + ('kernel',)
	if pt_tuple_key[-1] == "weight":
		lowercase       =    pt_tensor.T
		return renamed_pt_tuple_key, pt_tensor

	# old PyTorch layer norm weight
	lowercase       =    pt_tuple_key[:-1] + ('weight',)
	if pt_tuple_key[-1] == "gamma":
		return renamed_pt_tuple_key, pt_tensor

	# old PyTorch layer norm bias
	lowercase       =    pt_tuple_key[:-1] + ('bias',)
	if pt_tuple_key[-1] == "beta":
		return renamed_pt_tuple_key, pt_tensor

	return pt_tuple_key, pt_tensor






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE=42							):
	# Step 1: Convert pytorch tensor to numpy
	lowercase       =    {k: v.numpy() for k, v in pt_state_dict.items()}

	# Step 2: Since the model is stateless, get random Flax params
	lowercase       =    flax_model.init_weights(PRNGKey(__SCREAMING_SNAKE_CASE							)							)

	lowercase       =    flatten_dict(__SCREAMING_SNAKE_CASE							)
	lowercase       =    {}

	# Need to change some parameters name to match Flax names
	for pt_key, pt_tensor in pt_state_dict.items():
		lowercase       =    rename_key(__SCREAMING_SNAKE_CASE							)
		lowercase       =    tuple(renamed_pt_key.split('.'							)							)

		# Correctly rename weight parameters
		lowercase  ,					lowercase       =    rename_key_and_reshape_tensor(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

		if flax_key in random_flax_state_dict:
			if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
				raise ValueError(
				    F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
				    F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''							)

        # also add unexpected weight so that warning is thrown
		lowercase       =    jnp.asarray(__SCREAMING_SNAKE_CASE							)

	return unflatten_dict(__SCREAMING_SNAKE_CASE							)

