





"""simple docstring"""





# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


_lowerCamelCase       =	logging.get_logger(__name__)


if is_torch_available():
    import torch
    from torch import nn
    from torch.utils.data import Dataset

    from transformers import Trainer





    class 			snake_case		(   __UpperCAmelCase   ):

          def __init__(			self   :Optional[Any]						,							_lowerCamelCase   :int = 1_0_1						):
           __SCREAMING_SNAKE_CASE :			List[Any]												=       length

          def __len__(			self   :Optional[int]						):
           return self.length



          def __getitem__(			self   :Dict						,							_lowerCamelCase   :Optional[Any]						):
           return i







    class 			snake_case		:

          def __call__(			self   :int						,							_lowerCamelCase   :Tuple						):
           return {"input_ids": torch.tensor(_lowerCamelCase						), "labels": torch.tensor(_lowerCamelCase						)}







    class 			snake_case		(   nn.Module   ):

          def __init__(			self   :Optional[Any]						):
           super().__init__()
           # Add some (unused) params otherwise DDP will complain.
           __SCREAMING_SNAKE_CASE :			Optional[int]												=       nn.Linear(1_2_0						,							8_0						)



          def 						SCREAMING_SNAKE_CASE_			(			self   :str						,							_lowerCamelCase   :Optional[int]						,							_lowerCamelCase   :List[Any]=None						):
           if labels is not None:
            return torch.tensor(0.0						,							device=input_ids.device						), input_ids
           else:
            return input_ids




def 		lowerCAmelCase_	(  ):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE :			Optional[int]												=       HfArgumentParser((TrainingArguments,)		)
 sys.argv += ["--output_dir", "./examples"]
 __SCREAMING_SNAKE_CASE :			Optional[int]												=       parser.parse_args_into_dataclasses()[0]

 logger.warning(
     F'''Process rank: {training_args.local_rank}, device: {training_args.device}, '''
     F'''tpu_num_cores: {training_args.tpu_num_cores}''' ,      )

 # Essentially, what we want to verify in the distributed case is
 # that we get all samples back, in the right order.
 # (this is crucial for prediction for instance)
 for dataset_length in [1001, 256, 15]:
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       DummyDataset(lowercase_		)

  def compute_metrics(lowercase_   :							EvalPrediction		) -> Dict:
   __SCREAMING_SNAKE_CASE :			Any												=       list(range(len(lowercase_		)		)		)
   __SCREAMING_SNAKE_CASE :			Dict												=       p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
   return {"success": success}

  __SCREAMING_SNAKE_CASE :			List[Any]												=       Trainer(
      model=DummyModel() ,      args=lowercase_ ,      data_collator=DummyDataCollator() ,      eval_dataset=lowercase_ ,      compute_metrics=lowercase_ ,      )
  __SCREAMING_SNAKE_CASE :			Any												=       trainer.evaluate()
  logger.info(lowercase_		)
  if metrics["eval_success"] is not True:
   logger.error(lowercase_		)
   exit(1		)

  __SCREAMING_SNAKE_CASE :			str												=       trainer.predict(lowercase_		)
  logger.info(p.metrics		)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics		)
   exit(1		)

  __SCREAMING_SNAKE_CASE :			str												=       2

  __SCREAMING_SNAKE_CASE :			Optional[Any]												=       trainer.evaluate()
  logger.info(lowercase_		)
  if metrics["eval_success"] is not True:
   logger.error(lowercase_		)
   exit(1		)

  __SCREAMING_SNAKE_CASE :			Dict												=       trainer.predict(lowercase_		)
  logger.info(p.metrics		)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics		)
   exit(1		)

  __SCREAMING_SNAKE_CASE :			int												=       None

 logger.info('''🔥 All distributed tests successful'''		)




def 		lowerCAmelCase_	(  lowercase_   :							List[Any]		):

 '''simple docstring'''






 main()


if __name__ == "__main__":
    main()








"""simple docstring"""





from ctypes import c_float, sizeof
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union


if TYPE_CHECKING:
    from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore





class 			snake_case		(   __UpperCAmelCase   ):
      lowerCamelCase__						= c_float

      @property
      def 						SCREAMING_SNAKE_CASE_			(			self   :Optional[Any]						):
       return sizeof(self.value						)







def 		lowerCAmelCase_	(  lowercase_   :							int ,      lowercase_   :							int ,      lowercase_   :							int = 0		):

 '''simple docstring'''






 if dimension <= 0:
  __SCREAMING_SNAKE_CASE :			str												=       fixed_dimension

 dimension -= num_token_to_add
 return dimension




def 		lowerCAmelCase_	(  lowercase_   :							int ,      lowercase_   :							ParameterFormat		):

 '''simple docstring'''






 return num_parameters * dtype.size






def 		lowerCAmelCase_	(  lowercase_   :							str		):

 '''simple docstring'''






 from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

 try:
  return AutoProcessor.from_pretrained(lowercase_		)
 except (ValueError, OSError, KeyError):
  __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			int												=       None, None
  try:
   __SCREAMING_SNAKE_CASE :			Optional[Any]												=       AutoTokenizer.from_pretrained(lowercase_		)
  except (OSError, KeyError):
   pass
  try:
   __SCREAMING_SNAKE_CASE :			Optional[Any]												=       AutoFeatureExtractor.from_pretrained(lowercase_		)
  except (OSError, KeyError):
   pass

  if tokenizer is not None and feature_extractor is not None:
   raise ValueError(
       F'''Couldn\'t auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor.'''		)
  elif tokenizer is None and feature_extractor is None:
   return None
  elif tokenizer is not None:
   return tokenizer
  else:
   return feature_extractor

