

'''simple docstring'''
import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()







class 							A__    (       A__     ,   unittest.TestCase  ):
  A__		 =				PriorTransformer
  A__		 =				'hidden_states'





  @property
  def 						A  (		self      :	Any )	->    List[str]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =4
       _SCREAMING_SNAKE_CASE							    =8
       _SCREAMING_SNAKE_CASE							    =7

       _SCREAMING_SNAKE_CASE							    =floats_tensor((batch_size, embedding_dim) ).to(_a )

       _SCREAMING_SNAKE_CASE							    =floats_tensor((batch_size, embedding_dim) ).to(_a )
       _SCREAMING_SNAKE_CASE							    =floats_tensor((batch_size, num_embeddings, embedding_dim) ).to(_a )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }





  def 						A  (		self      :	List[str]							,					_a      :	List[str]=0 )	->    Union[str, Any]:





       '''simple docstring'''






       torch.manual_seed(_a )
       _SCREAMING_SNAKE_CASE							    =4
       _SCREAMING_SNAKE_CASE							    =8
       _SCREAMING_SNAKE_CASE							    =7

       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, embedding_dim) ).to(_a )

       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, embedding_dim) ).to(_a )
       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, num_embeddings, embedding_dim) ).to(_a )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }





  @property
  def 						A  (		self      :	Optional[int] )	->    List[Any]:





       '''simple docstring'''






       return (4, 8)





  @property
  def 						A  (		self      :	Dict )	->    Optional[Any]:





       '''simple docstring'''






       return (4, 8)





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





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    ={
           'num_attention_heads': 2,
           'attention_head_dim': 4,
           'num_layers': 2,
           'embedding_dim': 8,
           'num_embeddings': 7,
           'additional_embeddings': 4,
       }
       _SCREAMING_SNAKE_CASE							    =self.dummy_input
       return init_dict, inputs_dict





  def 						A  (		self      :	Any )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =PriorTransformer.from_pretrained(
           'hf-internal-testing/prior-dummy'							,					output_loading_info=_a )
       self.assertIsNotNone(_a )
       self.assertEqual(len(loading_info['missing_keys'] )							,					0 )

       model.to(_a )
       _SCREAMING_SNAKE_CASE							    =model(**self.dummy_input )[0]

       assert hidden_states is not None, "Make sure output is not None"





  def 						A  (		self      :	int )	->    Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =self.prepare_init_args_and_inputs_for_common()

       _SCREAMING_SNAKE_CASE							    =self.model_class(**_a )
       _SCREAMING_SNAKE_CASE							    =inspect.signature(model.forward )
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       _SCREAMING_SNAKE_CASE							    =[*signature.parameters.keys()]

       _SCREAMING_SNAKE_CASE							    =['hidden_states', 'timestep']
       self.assertListEqual(arg_names[:2]							,					_a )





  def 						A  (		self      :	Optional[int] )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy' )
       _SCREAMING_SNAKE_CASE							    =model.to(_a )

       if hasattr(_a							,					'set_default_attn_processor' ):
            model.set_default_attn_processor()

       _SCREAMING_SNAKE_CASE							    =self.get_dummy_seed_input()

       with torch.no_grad():
            _SCREAMING_SNAKE_CASE							    =model(**_a )[0]

       _SCREAMING_SNAKE_CASE							    =output[0, :5].flatten().cpu()
       print(_a )

       # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
       # the expected output slices are not the same for CPU and GPU.
       _SCREAMING_SNAKE_CASE							    =torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39] )
       self.assertTrue(torch_all_close(_a							,					_a							,					rtol=1e-2 ) )







@slow
class 							A__    (       unittest.TestCase  ):





  def 						A  (		self      :	Union[str, Any]							,					_a      :	Dict=1							,					_a      :	Union[str, Any]=768							,					_a      :	Dict=77							,					_a      :	Optional[Any]=0 )	->    List[Any]:





       '''simple docstring'''






       torch.manual_seed(_a )
       _SCREAMING_SNAKE_CASE							    =batch_size
       _SCREAMING_SNAKE_CASE							    =embedding_dim
       _SCREAMING_SNAKE_CASE							    =num_embeddings

       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, embedding_dim) ).to(_a )

       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, embedding_dim) ).to(_a )
       _SCREAMING_SNAKE_CASE							    =torch.randn((batch_size, num_embeddings, embedding_dim) ).to(_a )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }





  def 						A  (		self      :	str )	->    Optional[int]:





       '''simple docstring'''






       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





  @parameterized.expand(
      [
          # fmt: off
          [13, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
          [37, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
          # fmt: on
      ] )
  def 						A  (		self      :	int							,					_a      :	Optional[Any]							,					_a      :	str )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'							,					subfolder='prior' )
       model.to(_a )
       _SCREAMING_SNAKE_CASE							    =self.get_dummy_seed_input(seed=_a )

       with torch.no_grad():
            _SCREAMING_SNAKE_CASE							    =model(**_a )[0]

       assert list(sample.shape ) == [1, 768]

       _SCREAMING_SNAKE_CASE							    =sample[0, :8].flatten().cpu()
       print(_a )
       _SCREAMING_SNAKE_CASE							    =torch.tensor(_a )

       assert torch_all_close(_a							,					_a							,					atol=1e-3 )



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

import datasets


lowerCamelCase :      List[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 :      Optional[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 :      int					 =     "\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					(						_UpperCamelCase						:     List[str]							, _UpperCamelCase						:     Tuple				)      ->			List[Any]:
     """simple docstring"""




     return float((preds == labels).mean()				)


def        _lowerCAmelCase					(						_UpperCamelCase						:     Union[str, Any]							, _UpperCamelCase						:     Any				)      ->			List[str]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							    =simple_accuracy(_UpperCamelCase							, _UpperCamelCase				)
     _SCREAMING_SNAKE_CASE							    =float(fa_score(y_true=_UpperCamelCase							, y_pred=_UpperCamelCase				)				)
     return {
         "accuracy": acc,
         "f1": fa,
     }


def        _lowerCAmelCase					(						_UpperCamelCase						:     Any							, _UpperCamelCase						:     int				)      ->			str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							    =np.array(_UpperCamelCase				)
     _SCREAMING_SNAKE_CASE							    =np.array(_UpperCamelCase				)
     _SCREAMING_SNAKE_CASE							    =en_sentvecs.shape[0]

     # mean centering
     _SCREAMING_SNAKE_CASE							    =en_sentvecs - np.mean(_UpperCamelCase							, axis=0				)
     _SCREAMING_SNAKE_CASE							    =in_sentvecs - np.mean(_UpperCamelCase							, axis=0				)

     _SCREAMING_SNAKE_CASE							    =cdist(_UpperCamelCase							, _UpperCamelCase							, 'cosine'				)
     _SCREAMING_SNAKE_CASE							    =np.array(range(_UpperCamelCase				)				)
     _SCREAMING_SNAKE_CASE							    =sim.argsort(axis=1				)[:, :10]
     _SCREAMING_SNAKE_CASE							    =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  ):





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





       '''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 						A  (		self      :	List[str]							,					_a      :	Tuple							,					_a      :	Optional[int] )	->    int:





       '''simple docstring'''






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