'''simple docstring'''


import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

 if is_vision_available():
  from transformers import MaskaFormerImageProcessor

if is_vision_available():
 from PIL import Image






class 					__SCREAMING_SNAKE_CASE						:







  '''simple docstring'''






  def __init__(			self				,		snake_case_				,		snake_case_=2				,		snake_case_=True				,		snake_case_=False				,		snake_case_=1_0				,		snake_case_=3				,		snake_case_=3_2 * 8				,		snake_case_=3_2 * 8				,		snake_case_=4				,		snake_case_=6_4				,		):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				parent
       UpperCAmelCase_							:       Dict       =				batch_size
       UpperCAmelCase_							:       Optional[Any]       =				is_training
       UpperCAmelCase_							:       Optional[Any]       =				use_auxiliary_loss
       UpperCAmelCase_							:       int       =				num_queries
       UpperCAmelCase_							:       Tuple       =				num_channels
       UpperCAmelCase_							:       Optional[int]       =				min_size
       UpperCAmelCase_							:       Optional[int]       =				max_size
       UpperCAmelCase_							:       int       =				num_labels
       UpperCAmelCase_							:       List[Any]       =				hidden_dim
       UpperCAmelCase_							:       Any       =				hidden_dim


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]						).to(
           snake_case_						)

       UpperCAmelCase_							:       Optional[Any]       =				torch.ones([self.batch_size, self.min_size, self.max_size]				,		device=snake_case_						)

       UpperCAmelCase_							:       Optional[Any]       =				(
           torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]				,		device=snake_case_						) > 0.5
       ).float()
       UpperCAmelCase_							:       List[str]       =				(torch.rand((self.batch_size, self.num_labels)				,		device=snake_case_						) > 0.5).long()

       UpperCAmelCase_							:       str       =				self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       str       =				MaskaFormerConfig(
           hidden_size=self.hidden_dim				,		)
       UpperCAmelCase_							:       List[str]       =				self.num_queries
       UpperCAmelCase_							:       Any       =				self.num_labels

       UpperCAmelCase_							:       Union[str, Any]       =				[1, 1, 1, 1]
       UpperCAmelCase_							:       List[str]       =				self.num_channels

       UpperCAmelCase_							:       Union[str, Any]       =				6_4
       UpperCAmelCase_							:       Tuple       =				1_2_8
       UpperCAmelCase_							:       Tuple       =				self.hidden_dim
       UpperCAmelCase_							:       str       =				self.hidden_dim
       UpperCAmelCase_							:       List[str]       =				self.hidden_dim
       return config


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_							:       Dict       =				self.prepare_config_and_inputs()
       UpperCAmelCase_							:       Union[str, Any]       =				{'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
       return config, inputs_dict


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[Any]       =				output.encoder_hidden_states
       UpperCAmelCase_							:       Optional[int]       =				output.pixel_decoder_hidden_states
       UpperCAmelCase_							:       Dict       =				output.transformer_decoder_hidden_states

       self.parent.assertTrue(len(snake_case_						)				,		len(config.backbone_config.depths						)						)
       self.parent.assertTrue(len(snake_case_						)				,		len(config.backbone_config.depths						)						)
       self.parent.assertTrue(len(snake_case_						)				,		config.decoder_layers						)


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_=False						):




       '''simple docstring'''

       with torch.no_grad():
            UpperCAmelCase_							:       List[Any]       =				MaskaFormerModel(config=snake_case_						)
            model.to(snake_case_						)
            model.eval()

            UpperCAmelCase_							:       List[str]       =				model(pixel_values=snake_case_				,		pixel_mask=snake_case_						)
            UpperCAmelCase_							:       Any       =				model(snake_case_				,		output_hidden_states=snake_case_						)

       self.parent.assertEqual(
           output.transformer_decoder_last_hidden_state.shape				,		(self.batch_size, self.num_queries, self.hidden_dim)				,		)
       # let's ensure the other two hidden state exists
       self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None						)
       self.parent.assertTrue(output.encoder_last_hidden_state is not None						)

       if output_hidden_states:
            self.check_output_hidden_state(snake_case_				,		snake_case_						)






  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				MaskaFormerForUniversalSegmentation(config=snake_case_						)
       model.to(snake_case_						)
       model.eval()

       def comm_check_on_output(snake_case_						):
            # let's still check that all the required stuff is there
            self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None						)
            self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None						)
            self.parent.assertTrue(result.encoder_last_hidden_state is not None						)
            # okay, now we need to check the logits shape
            # due to the encoder compression, masks have a //4 spatial size
            self.parent.assertEqual(
                result.masks_queries_logits.shape				,		(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)				,		)
            # + 1 for null class
            self.parent.assertEqual(
                result.class_queries_logits.shape				,		(self.batch_size, self.num_queries, self.num_labels + 1)						)

       with torch.no_grad():
            UpperCAmelCase_							:       Union[str, Any]       =				model(pixel_values=snake_case_				,		pixel_mask=snake_case_						)
            UpperCAmelCase_							:       Any       =				model(snake_case_						)

            comm_check_on_output(snake_case_						)

            UpperCAmelCase_							:       Union[str, Any]       =				model(
                pixel_values=snake_case_				,		pixel_mask=snake_case_				,		mask_labels=snake_case_				,		class_labels=snake_case_						)

       comm_check_on_output(snake_case_						)

       self.parent.assertTrue(result.loss is not None						)
       self.parent.assertEqual(result.loss.shape				,		torch.Size([1]						)						)







@require_torch
class 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_						,     lowerCamelCase_						,     unittest.TestCase	):







  '''simple docstring'''






  lowerCamelCase_		:List[str]						   =		(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
  lowerCamelCase_		:str						   =		{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

  lowerCamelCase_		:Optional[int]						   =		False
  lowerCamelCase_		:Any						   =		False
  lowerCamelCase_		:Dict						   =		False
  lowerCamelCase_		:Optional[Any]						   =		False


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Tuple       =				MaskaFormerModelTester(self						)
       UpperCAmelCase_							:       str       =				ConfigTester(self				,		config_class=snake_case_				,		has_text_modality=snake_case_						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       self.config_tester.run_common_tests()


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_							:       Optional[int]       =				self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(snake_case_				,		**snake_case_				,		output_hidden_states=snake_case_						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Union[str, Any]       =				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*snake_case_						)


  @unittest.skip(reason='Mask2Former does not use inputs_embeds'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  @unittest.skip(reason='Mask2Former is not a generative model'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  @unittest.skip(reason='Mask2Former does not use token embeddings'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  @require_torch_multi_gpu
  @unittest.skip(
      reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'						)
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       pass


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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

       for model_class in self.all_model_classes:
            UpperCAmelCase_							:       Dict       =				model_class(snake_case_						)
            UpperCAmelCase_							:       int       =				inspect.signature(model.forward						)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            UpperCAmelCase_							:       List[str]       =				[*signature.parameters.keys()]

            UpperCAmelCase_							:       int       =				['pixel_values']
            self.assertListEqual(arg_names[:1]				,		snake_case_						)


  @slow
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
            UpperCAmelCase_							:       int       =				MaskaFormerModel.from_pretrained(snake_case_						)
            self.assertIsNotNone(snake_case_						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[Any]       =				(self.model_tester.min_size,) * 2
       UpperCAmelCase_							:       Tuple       =				{
           'pixel_values': torch.randn((2, 3, *size)				,		device=snake_case_						),
           'mask_labels': torch.randn((2, 1_0, *size)				,		device=snake_case_						),
           'class_labels': torch.zeros(2				,		1_0				,		device=snake_case_						).long(),
       }
       UpperCAmelCase_							:       Dict       =				self.model_tester.get_config()

       UpperCAmelCase_							:       List[Any]       =				MaskaFormerForUniversalSegmentation(snake_case_						).to(snake_case_						)
       UpperCAmelCase_							:       Tuple       =				model(**snake_case_						)
       self.assertTrue(outputs.loss is not None						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_							:       List[Any]       =				self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(snake_case_				,		**snake_case_				,		output_hidden_states=snake_case_						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_       ,      UpperCAmelCase_							:       Tuple       =				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
            UpperCAmelCase_							:       Union[str, Any]       =				model_class(snake_case_						).to(snake_case_						)
            UpperCAmelCase_							:       Optional[Any]       =				model(**snake_case_				,		output_attentions=snake_case_						)
            self.assertTrue(outputs.attentions is not None						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       if not self.model_tester.is_training:
            return

       UpperCAmelCase_							:       Union[str, Any]       =				self.all_model_classes[1]
       UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_							:       str       =				self.model_tester.prepare_config_and_inputs()

       UpperCAmelCase_							:       Optional[int]       =				model_class(snake_case_						)
       model.to(snake_case_						)
       model.train()

       UpperCAmelCase_							:       int       =				model(snake_case_				,		mask_labels=snake_case_				,		class_labels=snake_case_						).loss
       loss.backward()






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Dict       =				self.all_model_classes[1]
       UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_							:       Union[str, Any]       =				self.model_tester.prepare_config_and_inputs()
       UpperCAmelCase_							:       Any       =				True
       UpperCAmelCase_							:       Union[str, Any]       =				True

       UpperCAmelCase_							:       Tuple       =				model_class(snake_case_						).to(snake_case_						)
       model.train()

       UpperCAmelCase_							:       Optional[Any]       =				model(snake_case_				,		mask_labels=snake_case_				,		class_labels=snake_case_						)

       UpperCAmelCase_							:       int       =				outputs.encoder_hidden_states[0]
       encoder_hidden_states.retain_grad()

       UpperCAmelCase_							:       int       =				outputs.pixel_decoder_hidden_states[0]
       pixel_decoder_hidden_states.retain_grad()

       UpperCAmelCase_							:       str       =				outputs.transformer_decoder_hidden_states[0]
       transformer_decoder_hidden_states.retain_grad()

       UpperCAmelCase_							:       str       =				outputs.attentions[0]
       attentions.retain_grad()

       outputs.loss.backward(retain_graph=snake_case_						)

       self.assertIsNotNone(encoder_hidden_states.grad						)
       self.assertIsNotNone(pixel_decoder_hidden_states.grad						)
       self.assertIsNotNone(transformer_decoder_hidden_states.grad						)
       self.assertIsNotNone(attentions.grad						)







snake_case__					:      int				      =  1e-4


def  _lowerCamelCase   (				):



     """simple docstring"""





     UpperCAmelCase_							:       Optional[Any]       =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
     return image






@require_vision
@slow
class 					__SCREAMING_SNAKE_CASE						(	unittest.TestCase	):







  '''simple docstring'''






  @cached_property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return "facebook/mask2former-swin-small-coco-instance"


  @cached_property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints						) if is_vision_available() else None


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       List[Any]       =				MaskaFormerModel.from_pretrained(self.model_checkpoints						).to(snake_case_						)
       UpperCAmelCase_							:       Any       =				self.default_image_processor
       UpperCAmelCase_							:       int       =				prepare_img()
       UpperCAmelCase_							:       List[str]       =				image_processor(snake_case_				,		return_tensors='pt'						).to(snake_case_						)
       UpperCAmelCase_							:       List[str]       =				inputs['pixel_values'].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0						)
       # check size
       self.assertEqual(snake_case_				,		(1, 3, 3_8_4, 3_8_4)						)

       with torch.no_grad():
            UpperCAmelCase_							:       List[Any]       =				model(**snake_case_						)

       UpperCAmelCase_							:       int       =				torch.tensor(
           [[-0.27_90, -1.07_17, -1.16_68], [-0.51_28, -0.31_28, -0.49_87], [-0.58_32, 0.19_71, -0.01_97]]						).to(snake_case_						)
       self.assertTrue(
           torch.allclose(
               outputs.encoder_last_hidden_state[0, 0, :3, :3]				,		snake_case_				,		atol=snake_case_						)						)

       UpperCAmelCase_							:       Optional[Any]       =				torch.tensor(
           [[0.89_73, 1.18_47, 1.17_76], [1.19_34, 1.50_40, 1.51_28], [1.11_53, 1.44_86, 1.49_51]]						).to(snake_case_						)
       self.assertTrue(
           torch.allclose(
               outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]				,		snake_case_				,		atol=snake_case_						)						)

       UpperCAmelCase_							:       Optional[int]       =				torch.tensor(
           [[2.11_52, 1.70_00, -0.86_03], [1.58_08, 1.80_04, -0.93_53], [1.60_43, 1.74_95, -0.59_99]]						).to(snake_case_						)
       self.assertTrue(
           torch.allclose(
               outputs.transformer_decoder_last_hidden_state[0, :3, :3]				,		snake_case_				,		atol=snake_case_						)						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints						).to(snake_case_						).eval()
       UpperCAmelCase_							:       Union[str, Any]       =				self.default_image_processor
       UpperCAmelCase_							:       Dict       =				prepare_img()
       UpperCAmelCase_							:       Tuple       =				image_processor(snake_case_				,		return_tensors='pt'						).to(snake_case_						)
       UpperCAmelCase_							:       Optional[int]       =				inputs['pixel_values'].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0						)
       # check size
       self.assertEqual(snake_case_				,		(1, 3, 3_8_4, 3_8_4)						)

       with torch.no_grad():
            UpperCAmelCase_							:       int       =				model(**snake_case_						)
       # masks_queries_logits
       UpperCAmelCase_							:       Optional[Any]       =				outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape				,		(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)						)
       UpperCAmelCase_							:       Optional[Any]       =				[
           [-8.78_39, -9.00_56, -8.81_21],
           [-7.41_04, -7.03_13, -6.54_01],
           [-6.61_05, -6.34_27, -6.46_75],
       ]
       UpperCAmelCase_							:       Optional[Any]       =				torch.tensor(snake_case_						).to(snake_case_						)
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]				,		snake_case_				,		atol=snake_case_						)						)
       # class_queries_logits
       UpperCAmelCase_							:       Tuple       =				outputs.class_queries_logits
       self.assertEqual(class_queries_logits.shape				,		(1, model.config.num_queries, model.config.num_labels + 1)						)
       UpperCAmelCase_							:       List[Any]       =				torch.tensor(
           [
               [1.83_24, -8.08_35, -4.19_22],
               [0.84_50, -9.00_50, -3.60_53],
               [0.30_45, -7.72_93, -3.02_75],
           ]						).to(snake_case_						)
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]				,		snake_case_				,		atol=snake_case_						)						)






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints						).to(snake_case_						).eval()
       UpperCAmelCase_							:       Optional[int]       =				self.default_image_processor

       UpperCAmelCase_							:       Union[str, Any]       =				image_processor(
           [np.zeros((3, 8_0_0, 1_3_3_3)						), np.zeros((3, 8_0_0, 1_3_3_3)						)]				,		segmentation_maps=[np.zeros((3_8_4, 3_8_4)						).astype(np.floataa						), np.zeros((3_8_4, 3_8_4)						).astype(np.floataa						)]				,		return_tensors='pt'				,		)

       UpperCAmelCase_							:       Union[str, Any]       =				inputs['pixel_values'].to(snake_case_						)
       UpperCAmelCase_							:       List[str]       =				[el.to(snake_case_						) for el in inputs['mask_labels']]
       UpperCAmelCase_							:       Optional[Any]       =				[el.to(snake_case_						) for el in inputs['class_labels']]

       with torch.no_grad():
            UpperCAmelCase_							:       Any       =				model(**snake_case_						)

       self.assertTrue(outputs.loss is not None						)







'''simple docstring'''


import numpy as np
import torch
from torch.nn import CrossEntropyLoss
from transformers import AutoModelForCausalLM, AutoTokenizer

import datasets
from datasets import logging


snake_case__					:      str				      =  '''\

'''

snake_case__					:      Union[str, Any]				      =  '''
Perplexity (PPL) is one of the most common metrics for evaluating language models.
It is defined as the exponentiated average negative log-likelihood of a sequence.

For more information, see https://huggingface.co/docs/transformers/perplexity
'''

snake_case__					:      str				      =  '''
Args:
    model_id (str): model used for calculating Perplexity
            NOTE: Perplexity can only be calculated for causal language models.
                    This includes models such as gpt2, causal variations of bert,
                    causal versions of t5, and more (the full list can be found
                    in the AutoModelForCausalLM documentation here:
                    https://huggingface.co/docs/transformers/master/en/model_doc/auto#transformers.AutoModelForCausalLM )

    input_texts (list of str): input text, each separate text snippet
        is one list entry.
    batch_size (int): the batch size to run texts through the model. Defaults to 16.
    add_start_token (bool): whether to add the start token to the texts,
        so the perplexity can include the probability of the first word. Defaults to True.
    device (str): device to run on, defaults to \'cuda\' when available
Returns:
    perplexity: dictionary containing the perplexity scores for the texts
        in the input list, as well as the mean perplexity. If one of the input texts is
        longer than the max input length of the model, then it is truncated to the
        max length for the perplexity computation.
Examples:
    Example 1:
        >>> perplexity = datasets.load_metric("perplexity")
        >>> input_texts = ["lorem ipsum", "Happy Birthday!", "Bienvenue"]
        >>> results = perplexity.compute(model_id=\'gpt2\',
        ...                              add_start_token=False,
        ...                              input_texts=input_texts) # doctest:+ELLIPSIS
        >>> print(list(results.keys()))
        [\'perplexities\', \'mean_perplexity\']
        >>> print(round(results["mean_perplexity"], 2))
        78.22
        >>> print(round(results["perplexities"][0], 2))
        11.11

    Example 2:
        >>> perplexity = datasets.load_metric("perplexity")
        >>> input_texts = datasets.load_dataset("wikitext",
        ...                                     "wikitext-2-raw-v1",
        ...                                     split="test")["text"][:50] # doctest:+ELLIPSIS
        [...]
        >>> input_texts = [s for s in input_texts if s!=\'\']
        >>> results = perplexity.compute(model_id=\'gpt2\',
        ...                              input_texts=input_texts) # doctest:+ELLIPSIS
        >>> print(list(results.keys()))
        [\'perplexities\', \'mean_perplexity\']
        >>> print(round(results["mean_perplexity"], 2))
        60.35
        >>> print(round(results["perplexities"][0], 2))
        81.12
'''






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







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return datasets.MetricInfo(
           description=_DESCRIPTION				,		citation=_CITATION				,		inputs_description=_KWARGS_DESCRIPTION				,		features=datasets.Features(
               {
                   'input_texts': datasets.Value('string'						),
               }						)				,		reference_urls=['https://huggingface.co/docs/transformers/perplexity']				,		)






  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_ = 1_6				,		snake_case_ = True				,		snake_case_=None						):




       '''simple docstring'''

       if device is not None:
            assert device in ["gpu", "cpu", "cuda"], "device should be either gpu or cpu."
            if device == "gpu":
                 UpperCAmelCase_							:       int       =				'cuda'
       else:
            UpperCAmelCase_							:       Union[str, Any]       =				'cuda' if torch.cuda.is_available() else 'cpu'

       UpperCAmelCase_							:       Tuple       =				AutoModelForCausalLM.from_pretrained(snake_case_						)
       UpperCAmelCase_							:       Dict       =				model.to(snake_case_						)

       UpperCAmelCase_							:       Dict       =				AutoTokenizer.from_pretrained(snake_case_						)

       # if batch_size > 1 (which generally leads to padding being required), and
       # if there is not an already assigned pad_token, assign an existing
       # special token to also be the padding token
       if tokenizer.pad_token is None and batch_size > 1:
            UpperCAmelCase_							:       Any       =				list(tokenizer.special_tokens_map_extended.values()						)
            # check that the model already has at least one special token defined
            assert (
                len(snake_case_						) > 0
            ), "If batch_size > 1, model must have at least one special token to use for padding. Please use a different model or set batch_size=1."
            # assign one of the special tokens to also be the pad token
            tokenizer.add_special_tokens({'pad_token': existing_special_tokens[0]}						)

       if add_start_token:
            # leave room for <BOS> token to be added:
            assert (
                tokenizer.bos_token is not None
            ), "Input model must already have a BOS token if using add_start_token=True. Please use a different model, or set add_start_token=False"
            UpperCAmelCase_							:       Union[str, Any]       =				model.config.max_length - 1
       else:
            UpperCAmelCase_							:       int       =				model.config.max_length

       UpperCAmelCase_							:       List[str]       =				tokenizer(
           snake_case_				,		add_special_tokens=snake_case_				,		padding=snake_case_				,		truncation=snake_case_				,		max_length=snake_case_				,		return_tensors='pt'				,		return_attention_mask=snake_case_				,		).to(snake_case_						)

       UpperCAmelCase_							:       Tuple       =				encodings['input_ids']
       UpperCAmelCase_							:       Optional[int]       =				encodings['attention_mask']

       # check that each input is long enough:
       if add_start_token:
            assert torch.all(torch.ge(attn_masks.sum(1						)				,		1						)						), "Each input text must be at least one token long."
       else:
            assert torch.all(
                torch.ge(attn_masks.sum(1						)				,		2						)						), "When add_start_token=False, each input text must be at least two tokens long. Run with add_start_token=True if inputting strings of only one token, and remove all empty input strings."

       UpperCAmelCase_							:       Union[str, Any]       =				[]
       UpperCAmelCase_							:       Optional[int]       =				CrossEntropyLoss(reduction='none'						)

       for start_index in logging.tqdm(range(0				,		len(snake_case_						)				,		snake_case_						)						):
            UpperCAmelCase_							:       Union[str, Any]       =				min(start_index + batch_size				,		len(snake_case_						)						)
            UpperCAmelCase_							:       List[Any]       =				encoded_texts[start_index:end_index]
            UpperCAmelCase_							:       Tuple       =				attn_masks[start_index:end_index]

            if add_start_token:
                 UpperCAmelCase_							:       List[str]       =				torch.tensor([[tokenizer.bos_token_id]] * encoded_batch.size(dim=0						)						).to(snake_case_						)
                 UpperCAmelCase_							:       Any       =				torch.cat([bos_tokens_tensor, encoded_batch]				,		dim=1						)
                 UpperCAmelCase_							:       List[Any]       =				torch.cat(
                     [torch.ones(bos_tokens_tensor.size()				,		dtype=torch.intaa						).to(snake_case_						), attn_mask]				,		dim=1						)

            UpperCAmelCase_							:       List[Any]       =				encoded_batch

            with torch.no_grad():
                 UpperCAmelCase_							:       int       =				model(snake_case_				,		attention_mask=snake_case_						).logits

            UpperCAmelCase_							:       Any       =				out_logits[..., :-1, :].contiguous()
            UpperCAmelCase_							:       Dict       =				labels[..., 1:].contiguous()
            UpperCAmelCase_							:       Union[str, Any]       =				attn_mask[..., 1:].contiguous()

            UpperCAmelCase_							:       List[Any]       =				torch.expa(
                (loss_fct(shift_logits.transpose(1				,		2						)				,		snake_case_						) * shift_attention_mask_batch).sum(1						)
                / shift_attention_mask_batch.sum(1						)						)

            ppls += perplexity_batch.tolist()

       return {"perplexities": ppls, "mean_perplexity": np.mean(snake_case_						)}






