
"""simple docstring"""




import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


A__      :	Dict          =    1e-4

if is_torch_available():
    import torch

    from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
    from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder







@require_torch
class   _lowercase						:


     '''simple docstring'''





     def __init__(      self ,			__UpperCamelCase ,			__UpperCamelCase=16 ,			__UpperCamelCase=13 ,			__UpperCamelCase=7 ,			__UpperCamelCase=14 ,			__UpperCamelCase=10 ,			__UpperCamelCase=19 ,			__UpperCamelCase=5 ,			__UpperCamelCase=4 ,			__UpperCamelCase=True ,			__UpperCamelCase=16 ,			__UpperCamelCase=2 ,			__UpperCamelCase=4 ,			__UpperCamelCase=4 ,			__UpperCamelCase="gelu" ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=[1, 2, 3, 4, 5] ,			__UpperCamelCase=25 ,			__UpperCamelCase=5 ,			)->					int:
          UpperCAmelCase__							: Optional[int] 							=						d_model
          UpperCAmelCase__							: str 							=						parent
          UpperCAmelCase__							: Union[str, Any] 							=						batch_size
          UpperCAmelCase__							: List[Any] 							=						prediction_length
          UpperCAmelCase__							: List[Any] 							=						context_length
          UpperCAmelCase__							: Optional[Any] 							=						cardinality
          UpperCAmelCase__							: Any 							=						num_time_features
          UpperCAmelCase__							: Dict 							=						lags_sequence
          UpperCAmelCase__							: Optional[int] 							=						embedding_dimension
          UpperCAmelCase__							: Dict 							=						is_training
          UpperCAmelCase__							: List[Any] 							=						hidden_size
          UpperCAmelCase__							: List[Any] 							=						num_hidden_layers
          UpperCAmelCase__							: List[str] 							=						num_attention_heads
          UpperCAmelCase__							: str 							=						intermediate_size
          UpperCAmelCase__							: Tuple 							=						hidden_act
          UpperCAmelCase__							: List[str] 							=						hidden_dropout_prob
          UpperCAmelCase__							: Dict 							=						attention_probs_dropout_prob

          UpperCAmelCase__							: Any 							=						context_length
          UpperCAmelCase__							: Dict 							=						prediction_length + label_length
          UpperCAmelCase__							: Optional[Any] 							=						label_length

          UpperCAmelCase__							: List[str] 							=						moving_average
          UpperCAmelCase__							: Union[str, Any] 							=						autocorrelation_factor



     def  lowerCAmelCase__       (      self  )->					List[str]:
          return AutoformerConfig(
              d_model=self.d_model ,			encoder_layers=self.num_hidden_layers ,			decoder_layers=self.num_hidden_layers ,			encoder_attention_heads=self.num_attention_heads ,			decoder_attention_heads=self.num_attention_heads ,			encoder_ffn_dim=self.intermediate_size ,			decoder_ffn_dim=self.intermediate_size ,			dropout=self.hidden_dropout_prob ,			attention_dropout=self.attention_probs_dropout_prob ,			prediction_length=self.prediction_length ,			context_length=self.context_length ,			label_length=self.label_length ,			lags_sequence=self.lags_sequence ,			num_time_features=self.num_time_features ,			num_static_categorical_features=1 ,			cardinality=[self.cardinality] ,			embedding_dimension=[self.embedding_dimension] ,			moving_average=self.moving_average ,			)



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[str]:
          UpperCAmelCase__							: str 							=						config.context_length + max(config.lags_sequence  )

          UpperCAmelCase__							: List[Any] 							=						ids_tensor([self.batch_size, 1] ,			config.cardinality[0]  )
          UpperCAmelCase__							: Optional[Any] 							=						floats_tensor([self.batch_size, _past_length, config.num_time_features]  )
          UpperCAmelCase__							: Optional[Any] 							=						floats_tensor([self.batch_size, _past_length]  )
          UpperCAmelCase__							: Tuple 							=						floats_tensor([self.batch_size, _past_length]  ) > 0.5

          # decoder inputs
          UpperCAmelCase__							: Optional[int] 							=						floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]  )
          UpperCAmelCase__							: List[Any] 							=						floats_tensor([self.batch_size, config.prediction_length]  )

          UpperCAmelCase__							: int 							=						{
              "past_values": past_values,
              "static_categorical_features": static_categorical_features,
              "past_time_features": past_time_features,
              "past_observed_mask": past_observed_mask,
              "future_time_features": future_time_features,
              "future_values": future_values,
          }
          return inputs_dict



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							: Tuple 							=						self.get_config()
          UpperCAmelCase__							: int 							=						self.prepare_autoformer_inputs_dict(__UpperCamelCase  )
          return config, inputs_dict



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							,	UpperCAmelCase__							: Any 							=						self.prepare_config_and_inputs()
          return config, inputs_dict



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase  )->					str:
          UpperCAmelCase__							: str 							=						AutoformerModel(config=__UpperCamelCase  ).to(__UpperCamelCase  ).eval()
          UpperCAmelCase__							: Dict 							=						model(**__UpperCamelCase  )

          UpperCAmelCase__							: int 							=						outputs.encoder_last_hidden_state
          UpperCAmelCase__							: Union[str, Any] 							=						outputs.last_hidden_state

          with tempfile.TemporaryDirectory() as tmpdirname:
               UpperCAmelCase__							: int 							=						model.get_encoder()
               encoder.save_pretrained(__UpperCamelCase  )
               UpperCAmelCase__							: Any 							=						AutoformerEncoder.from_pretrained(__UpperCamelCase  ).to(__UpperCamelCase  )

          UpperCAmelCase__							,	UpperCAmelCase__							,	UpperCAmelCase__							,	UpperCAmelCase__							,	UpperCAmelCase__							: int 							=						model.create_network_inputs(**__UpperCamelCase  )
          UpperCAmelCase__							,	UpperCAmelCase__							: Any 							=						model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]  )

          UpperCAmelCase__							: Optional[Any] 							=						torch.cat(
              (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...]) ,			dim=-1 ,			)
          UpperCAmelCase__							: Optional[int] 							=						encoder(inputs_embeds=__UpperCamelCase  )[0]
          self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3  )

          UpperCAmelCase__							: Any 							=						(
              torch.mean(transformer_inputs[:, : config.context_length, ...] ,			dim=1  )
              .unsqueeze(1  )
              .repeat(1 ,			config.prediction_length ,			1  )
          )
          UpperCAmelCase__							: Tuple 							=						torch.zeros(
              [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]] ,			device=enc_input.device ,			)

          UpperCAmelCase__							: Optional[int] 							=						torch.cat(
              (
                  torch.cat((seasonal_input[:, -config.label_length :, ...], zeros) ,			dim=1  ),
                  feature[:, config.context_length - config.label_length :, ...],
              ) ,			dim=-1 ,			)
          UpperCAmelCase__							: int 							=						torch.cat(
              (
                  torch.cat((trend_input[:, -config.label_length :, ...], mean) ,			dim=1  ),
                  feature[:, config.context_length - config.label_length :, ...],
              ) ,			dim=-1 ,			)

          with tempfile.TemporaryDirectory() as tmpdirname:
               UpperCAmelCase__							: Optional[int] 							=						model.get_decoder()
               decoder.save_pretrained(__UpperCamelCase  )
               UpperCAmelCase__							: Union[str, Any] 							=						AutoformerDecoder.from_pretrained(__UpperCamelCase  ).to(__UpperCamelCase  )

          UpperCAmelCase__							: Dict 							=						decoder(
              trend=__UpperCamelCase ,			inputs_embeds=__UpperCamelCase ,			encoder_hidden_states=__UpperCamelCase ,			)[0]

          self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1E-3  )







@require_torch
class   _lowercase						(       lowerCAmelCase_	,   lowerCAmelCase_	,   unittest.TestCase				):


     '''simple docstring'''





     _A								=							(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
     _A								=							(AutoformerForPrediction,) if is_torch_available() else ()
     _A								=							{'feature-extraction': AutoformerModel} if is_torch_available() else {}
     _A								=							False
     _A								=							False
     _A								=							False
     _A								=							False
     _A								=							False
     _A								=							False



     def  lowerCAmelCase__       (      self  )->					Dict:
          UpperCAmelCase__							: Optional[Any] 							=						AutoformerModelTester(self  )
          UpperCAmelCase__							: Optional[int] 							=						ConfigTester(self ,			config_class=__UpperCamelCase ,			has_text_modality=__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					int:
          self.config_tester.run_common_tests()



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							,	UpperCAmelCase__							: List[str] 							=						self.model_tester.prepare_config_and_inputs()
          for model_class in self.all_model_classes:
               UpperCAmelCase__							: Any 							=						model_class(__UpperCamelCase  )

               with tempfile.TemporaryDirectory() as tmpdirname:
                    model.save_pretrained(__UpperCamelCase  )
                    UpperCAmelCase__							,	UpperCAmelCase__							: Tuple 							=						model_class.from_pretrained(__UpperCamelCase ,			output_loading_info=__UpperCamelCase  )
               self.assertEqual(info["missing_keys"] ,			[]  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							: List[Any] 							=						self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.check_encoder_decoder_model_standalone(*__UpperCamelCase  )



     @unittest.skip(reason="Model has no tokens embeddings"  )
     def  lowerCAmelCase__       (      self  )->					List[Any]:
          pass



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          UpperCAmelCase__							: Optional[int] 							=						inspect.signature(getattr(__UpperCamelCase ,			"forward"  )  )
          # The main input is the name of the argument after `self`
          UpperCAmelCase__							: List[Any] 							=						list(model_signature.parameters.keys()  )[1]
          self.assertEqual(AutoformerModel.main_input_name ,			__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							,	UpperCAmelCase__							: Optional[Any] 							=						self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               UpperCAmelCase__							: Tuple 							=						model_class(__UpperCamelCase  )
               UpperCAmelCase__							: Union[str, Any] 							=						inspect.signature(model.forward  )
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               UpperCAmelCase__							: Tuple 							=						[*signature.parameters.keys()]

               UpperCAmelCase__							: Dict 							=						[
                   "past_values",
                   "past_time_features",
                   "past_observed_mask",
                   "static_categorical_features",
                   "static_real_features",
                   "future_values",
                   "future_time_features",
               ]

               if model.__class__.__name__ in ["AutoformerForPrediction"]:
                    expected_arg_names.append("future_observed_mask"  )

               expected_arg_names.extend(
                   [
                       "decoder_attention_mask",
                       "head_mask",
                       "decoder_head_mask",
                       "cross_attn_head_mask",
                       "encoder_outputs",
                       "past_key_values",
                       "output_hidden_states",
                       "output_attentions",
                       "use_cache",
                       "return_dict",
                   ]  )

               self.assertListEqual(arg_names[: len(__UpperCamelCase  )] ,			__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							,	UpperCAmelCase__							: List[Any] 							=						self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__							: Optional[int] 							=						True

          UpperCAmelCase__							: Dict 							=						getattr(self.model_tester ,			"seq_length" ,			__UpperCamelCase  )
          UpperCAmelCase__							: List[Any] 							=						getattr(self.model_tester ,			"decoder_seq_length" ,			__UpperCamelCase  )
          UpperCAmelCase__							: str 							=						getattr(self.model_tester ,			"encoder_seq_length" ,			__UpperCamelCase  )
          UpperCAmelCase__							: Union[str, Any] 							=						getattr(self.model_tester ,			"d_model" ,			__UpperCamelCase  )
          UpperCAmelCase__							: int 							=						getattr(self.model_tester ,			"num_attention_heads" ,			__UpperCamelCase  )
          UpperCAmelCase__							: Union[str, Any] 							=						d_model // num_attention_heads

          for model_class in self.all_model_classes:
               UpperCAmelCase__							: Dict 							=						True
               UpperCAmelCase__							: int 							=						False
               UpperCAmelCase__							: List[Any] 							=						True
               UpperCAmelCase__							: Optional[Any] 							=						model_class(__UpperCamelCase  )
               model.to(__UpperCamelCase  )
               model.eval()
               with torch.no_grad():
                    UpperCAmelCase__							: Optional[int] 							=						model(**self._prepare_for_class(__UpperCamelCase ,			__UpperCamelCase  )  )
               UpperCAmelCase__							: Tuple 							=						outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
               self.assertEqual(len(__UpperCamelCase  ) ,			self.model_tester.num_hidden_layers  )

               # check that output_attentions also work using config
               del inputs_dict["output_attentions"]
               UpperCAmelCase__							: Dict 							=						True
               UpperCAmelCase__							: Any 							=						model_class(__UpperCamelCase  )
               model.to(__UpperCamelCase  )
               model.eval()
               with torch.no_grad():
                    UpperCAmelCase__							: Optional[int] 							=						model(**self._prepare_for_class(__UpperCamelCase ,			__UpperCamelCase  )  )
               UpperCAmelCase__							: str 							=						outputs.encoder_attentions
               self.assertEqual(len(__UpperCamelCase  ) ,			self.model_tester.num_hidden_layers  )

               self.assertListEqual(
                   list(attentions[0].shape[-3:]  ) ,			[self.model_tester.num_attention_heads, encoder_seq_length, dim] ,			)
               UpperCAmelCase__							: Any 							=						len(__UpperCamelCase  )

               UpperCAmelCase__							: int 							=						7

               if "last_hidden_state" in outputs:
                    correct_outlen += 1

               if "trend" in outputs:
                    correct_outlen += 1

               if "past_key_values" in outputs:
                    correct_outlen += 1  # past_key_values have been returned

               if "loss" in outputs:
                    correct_outlen += 1

               if "params" in outputs:
                    correct_outlen += 1

               self.assertEqual(__UpperCamelCase ,			__UpperCamelCase  )

               # decoder attentions
               UpperCAmelCase__							: str 							=						outputs.decoder_attentions
               self.assertIsInstance(__UpperCamelCase ,			(list, tuple)  )
               self.assertEqual(len(__UpperCamelCase  ) ,			self.model_tester.num_hidden_layers  )
               self.assertListEqual(
                   list(decoder_attentions[0].shape[-3:]  ) ,			[self.model_tester.num_attention_heads, decoder_seq_length, dim] ,			)

               # cross attentions
               UpperCAmelCase__							: Union[str, Any] 							=						outputs.cross_attentions
               self.assertIsInstance(__UpperCamelCase ,			(list, tuple)  )
               self.assertEqual(len(__UpperCamelCase  ) ,			self.model_tester.num_hidden_layers  )
               self.assertListEqual(
                   list(cross_attentions[0].shape[-3:]  ) ,			[self.model_tester.num_attention_heads, decoder_seq_length, dim] ,			)

          # Check attention is always last and order is fine
          UpperCAmelCase__							: Optional[int] 							=						True
          UpperCAmelCase__							: Tuple 							=						True
          UpperCAmelCase__							: Tuple 							=						model_class(__UpperCamelCase  )
          model.to(__UpperCamelCase  )
          model.eval()
          with torch.no_grad():
               UpperCAmelCase__							: List[str] 							=						model(**self._prepare_for_class(__UpperCamelCase ,			__UpperCamelCase  )  )

          self.assertEqual(out_len + 2 ,			len(__UpperCamelCase  )  )

          UpperCAmelCase__							: Dict 							=						outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

          self.assertEqual(len(__UpperCamelCase  ) ,			self.model_tester.num_hidden_layers  )
          self.assertListEqual(
              list(self_attentions[0].shape[-3:]  ) ,			[self.model_tester.num_attention_heads, encoder_seq_length, dim] ,			)



     @is_flaky()
     def  lowerCAmelCase__       (      self  )->					Optional[int]:
          super().test_retain_grad_hidden_states_attentions()





def 							a__ (				lowerCAmelCase	:       str="train-batch.pt"		):






     '''simple docstring'''
     UpperCAmelCase__							: int 							=						hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch"       ,    filename=lowerCAmelCase       ,    repo_type="dataset"		)
     UpperCAmelCase__							: Dict 							=						torch.load(lowerCAmelCase       ,    map_location=lowerCAmelCase		)
     return batch







@require_torch
@slow
class   _lowercase						(       unittest.TestCase				):


     '''simple docstring'''





     def  lowerCAmelCase__       (      self  )->					Union[str, Any]:
          UpperCAmelCase__							: int 							=						AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly"  ).to(__UpperCamelCase  )
          UpperCAmelCase__							: int 							=						prepare_batch()

          with torch.no_grad():
               UpperCAmelCase__							: Any 							=						model(
                   past_values=batch["past_values"] ,			past_time_features=batch["past_time_features"] ,			past_observed_mask=batch["past_observed_mask"] ,			static_categorical_features=batch["static_categorical_features"] ,			future_values=batch["future_values"] ,			future_time_features=batch["future_time_features"] ,			)[0]

          UpperCAmelCase__							: Optional[Any] 							=						torch.Size(
              (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)  )
          self.assertEqual(output.shape ,			__UpperCamelCase  )

          UpperCAmelCase__							: Union[str, Any] 							=						torch.tensor(
              [[0.3593, -1.3398, 0.6330], [0.2279, 1.5396, -0.1792], [0.0450, 1.3225, -0.2335]] ,			device=__UpperCamelCase  )
          self.assertTrue(torch.allclose(output[0, :3, :3] ,			__UpperCamelCase ,			atol=__UpperCamelCase  )  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							: List[str] 							=						AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"  ).to(__UpperCamelCase  )
          UpperCAmelCase__							: List[str] 							=						prepare_batch("val-batch.pt"  )
          with torch.no_grad():
               UpperCAmelCase__							: Dict 							=						model(
                   past_values=batch["past_values"] ,			past_time_features=batch["past_time_features"] ,			past_observed_mask=batch["past_observed_mask"] ,			static_categorical_features=batch["static_categorical_features"] ,			).encoder_last_hidden_state
          UpperCAmelCase__							: int 							=						torch.Size((64, model.config.context_length, model.config.d_model)  )
          self.assertEqual(output.shape ,			__UpperCamelCase  )

          UpperCAmelCase__							: Dict 							=						torch.tensor(
              [[-0.0734, -0.9036, 0.8358], [4.7186, 2.4113, 1.9581], [1.7953, 2.3558, 1.2970]] ,			device=__UpperCamelCase  )
          self.assertTrue(torch.allclose(output[0, :3, :3] ,			__UpperCamelCase ,			atol=__UpperCamelCase  )  )



     def  lowerCAmelCase__       (      self  )->					List[str]:
          UpperCAmelCase__							: Optional[int] 							=						AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"  ).to(__UpperCamelCase  )
          UpperCAmelCase__							: Any 							=						prepare_batch("val-batch.pt"  )
          with torch.no_grad():
               UpperCAmelCase__							: Optional[int] 							=						model.generate(
                   static_categorical_features=batch["static_categorical_features"] ,			past_time_features=batch["past_time_features"] ,			past_values=batch["past_values"] ,			future_time_features=batch["future_time_features"] ,			past_observed_mask=batch["past_observed_mask"] ,			)
          UpperCAmelCase__							: int 							=						torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)  )
          self.assertEqual(outputs.sequences.shape ,			__UpperCamelCase  )

          UpperCAmelCase__							: List[str] 							=						torch.tensor([3130.6763, 4056.5293, 7053.0786] ,			device=__UpperCamelCase  )
          UpperCAmelCase__							: Any 							=						outputs.sequences.mean(dim=1  )
          self.assertTrue(torch.allclose(mean_prediction[0, -3:] ,			__UpperCamelCase ,			rtol=1E-1  )  )





"""simple docstring"""




import unicodedata
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np

from transformers.data.data_collator import DataCollatorMixin
from transformers.file_utils import PaddingStrategy
from transformers.tokenization_utils_base import PreTrainedTokenizerBase



def 							a__ (				lowerCAmelCase	:       str       ,    lowerCAmelCase	:       Optional[Any]       ,    lowerCAmelCase	:       Dict       ,    lowerCAmelCase	:       List[Any]		):






     '''simple docstring'''
     if isinstance(lowerCAmelCase       ,    lowerCAmelCase		):
          UpperCAmelCase__							: Optional[int] 							=						np.full((len(lowerCAmelCase		), sequence_length, 2)       ,    lowerCAmelCase		)
     else:
          UpperCAmelCase__							: Optional[Any] 							=						np.full((len(lowerCAmelCase		), sequence_length)       ,    lowerCAmelCase		)

     for i, tensor in enumerate(lowerCAmelCase		):
          if padding_side == "right":
               if isinstance(lowerCAmelCase       ,    lowerCAmelCase		):
                    UpperCAmelCase__							: Dict 							=						tensor[:sequence_length]
               else:
                    UpperCAmelCase__							: Tuple 							=						tensor[:sequence_length]
          else:
               if isinstance(lowerCAmelCase       ,    lowerCAmelCase		):
                    UpperCAmelCase__							: Optional[Any] 							=						tensor[:sequence_length]
               else:
                    UpperCAmelCase__							: int 							=						tensor[:sequence_length]

     return out_tensor.tolist()



def 							a__ (				lowerCAmelCase	:       Optional[int]		):






     '''simple docstring'''
     UpperCAmelCase__							: Tuple 							=						ord(lowerCAmelCase		)
     if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126):
          return True
     UpperCAmelCase__							: Optional[Any] 							=						unicodedata.category(lowerCAmelCase		)
     if cat.startswith("P"		):
          return True
     return False







@dataclass
class   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





     _A								=							42
     _A								=							True
     _A								=							None
     _A								=							None
     _A								=							-100
     _A								=							"pt"



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[str]:
          import torch

          UpperCAmelCase__							: Optional[Any] 							=						"label" if "label" in features[0].keys() else "labels"
          UpperCAmelCase__							: Dict 							=						[feature[label_name] for feature in features] if label_name in features[0].keys() else None
          UpperCAmelCase__							: str 							=						self.tokenizer.pad(
              __UpperCamelCase ,			padding=self.padding ,			max_length=self.max_length ,			pad_to_multiple_of=self.pad_to_multiple_of ,			return_tensors="pt" if labels is None else None ,			)

          if labels is None:
               return batch

          UpperCAmelCase__							: Optional[Any] 							=						torch.tensor(batch["entity_ids"]  ).shape[1]
          UpperCAmelCase__							: int 							=						self.tokenizer.padding_side
          if padding_side == "right":
               UpperCAmelCase__							: int 							=						[
                   list(__UpperCamelCase  ) + [self.label_pad_token_id] * (sequence_length - len(__UpperCamelCase  )) for label in labels
               ]
          else:
               UpperCAmelCase__							: List[Any] 							=						[
                   [self.label_pad_token_id] * (sequence_length - len(__UpperCamelCase  )) + list(__UpperCamelCase  ) for label in labels
               ]

          UpperCAmelCase__							: Optional[Any] 							=						[feature["ner_tags"] for feature in features]
          UpperCAmelCase__							: int 							=						padding_tensor(__UpperCamelCase ,			-1 ,			__UpperCamelCase ,			__UpperCamelCase  )
          UpperCAmelCase__							: List[Any] 							=						[feature["original_entity_spans"] for feature in features]
          UpperCAmelCase__							: int 							=						padding_tensor(__UpperCamelCase ,			(-1, -1) ,			__UpperCamelCase ,			__UpperCamelCase  )
          UpperCAmelCase__							: Optional[int] 							=						{k: torch.tensor(__UpperCamelCase ,			dtype=torch.intaa  ) for k, v in batch.items()}

          return batch



