




'''simple docstring'''

import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
		import torch

if is_vision_available():
		from PIL import Image

		from transformers import MobileNetVaImageProcessor





class __magic_name__(						unittest.TestCase      ):






						def __init__(   self						:						int					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						int=7					,   __UpperCamelCase						:						Optional[Any]=3					,   __UpperCamelCase						:						Dict=1_8					,   __UpperCamelCase						:						Tuple=3_0					,   __UpperCamelCase						:						List[Any]=4_0_0					,   __UpperCamelCase						:						List[str]=True					,   __UpperCamelCase						:						Any=None					,   __UpperCamelCase						:						int=True					,   __UpperCamelCase						:						Optional[int]=None					,   ):



									'''simple docstring'''






									snake_case__          =						size if size is not None else {"""shortest_edge""": 2_0}
									snake_case__          =						crop_size if crop_size is not None else {"""height""": 1_8, """width""": 1_8}
									snake_case__          =						parent
									snake_case__          =						batch_size
									snake_case__          =						num_channels
									snake_case__          =						image_size
									snake_case__          =						min_resolution
									snake_case__          =						max_resolution
									snake_case__          =						do_resize
									snake_case__          =						size
									snake_case__          =						do_center_crop
									snake_case__          =						crop_size







						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									return {
									    "do_resize": self.do_resize,
									    "size": self.size,
									    "do_center_crop": self.do_center_crop,
									    "crop_size": self.crop_size,
									}







@require_torch
@require_vision
class __magic_name__(						__lowerCAmelCase    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Optional[Any]    =  MobileNetVaImageProcessor if is_vision_available() else None






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									snake_case__          =						MobileNetVaImageProcessingTester(self       )






						@property
						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									return self.image_processor_tester.prepare_image_processor_dict()






						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						self.image_processing_class(**self.image_processor_dict       )
									self.assertTrue(hasattr(__UpperCamelCase					,   """do_resize"""       )       )
									self.assertTrue(hasattr(__UpperCamelCase					,   """size"""       )       )
									self.assertTrue(hasattr(__UpperCamelCase					,   """do_center_crop"""       )       )
									self.assertTrue(hasattr(__UpperCamelCase					,   """center_crop"""       )       )






						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						self.image_processing_class.from_dict(self.image_processor_dict       )
									self.assertEqual(image_processor.size					,   {"""shortest_edge""": 2_0}       )
									self.assertEqual(image_processor.crop_size					,   {"""height""": 1_8, """width""": 1_8}       )

									snake_case__          =						self.image_processing_class.from_dict(self.image_processor_dict					,   size=4_2					,   crop_size=8_4       )
									self.assertEqual(image_processor.size					,   {"""shortest_edge""": 4_2}       )
									self.assertEqual(image_processor.crop_size					,   {"""height""": 8_4, """width""": 8_4}       )






						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									pass






						def       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__          =						self.image_processing_class(**self.image_processor_dict       )
									# create random PIL images
									snake_case__          =						prepare_image_inputs(self.image_processor_tester					,   equal_resolution=__UpperCamelCase       )
									for image in image_inputs:
												self.assertIsInstance(__UpperCamelCase					,   Image.Image       )

									# Test not batched input
									snake_case__          =						image_processing(image_inputs[0]					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )

									# Test batched
									snake_case__          =						image_processing(__UpperCamelCase					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									snake_case__          =						self.image_processing_class(**self.image_processor_dict       )
									# create random numpy tensors
									snake_case__          =						prepare_image_inputs(self.image_processor_tester					,   equal_resolution=__UpperCamelCase					,   numpify=__UpperCamelCase       )
									for image in image_inputs:
												self.assertIsInstance(__UpperCamelCase					,   np.ndarray       )

									# Test not batched input
									snake_case__          =						image_processing(image_inputs[0]					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )

									# Test batched
									snake_case__          =						image_processing(__UpperCamelCase					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )







						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									snake_case__          =						self.image_processing_class(**self.image_processor_dict       )
									# create random PyTorch tensors
									snake_case__          =						prepare_image_inputs(self.image_processor_tester					,   equal_resolution=__UpperCamelCase					,   torchify=__UpperCamelCase       )
									for image in image_inputs:
												self.assertIsInstance(__UpperCamelCase					,   torch.Tensor       )

									# Test not batched input
									snake_case__          =						image_processing(image_inputs[0]					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )

									# Test batched
									snake_case__          =						image_processing(__UpperCamelCase					,   return_tensors="""pt"""       ).pixel_values
									self.assertEqual(
									    encoded_images.shape					,   (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    )					,   )



"""simple docstring"""





import argparse
import time

import torch
import transformers
from measures_util import end_measure, log_measures, start_measure
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer

from accelerate.utils import compute_module_sizes


snake_case__		:  Optional[Any]     					=					{
    '''gpt-j-6b''': {'''is_causal''': True, '''model''': '''sgugger/sharded-gpt-j-6B''', '''tokenizer''': '''EleutherAI/gpt-j-6B'''},
    '''gpt-neox''': {'''is_causal''': True, '''model''': '''EleutherAI/gpt-neox-20b'''},
    '''opt''': {'''is_causal''': True, '''model''': '''facebook/opt-30b'''},
    '''T0pp''': {'''is_causal''': False, '''model''': '''bigscience/T0pp''', '''model_revision''': '''sharded'''},
}

snake_case__		:  Any     					=					[
    '''Hello, my name is''',
    '''Are unicorns real? Unicorns are''',
    '''For the first time in several years,''',
    '''My name is Julien and I am''',
    '''The goal of life is''',
    '''Whenever I\'m sad, I like to''',
]


def   _snake_case    (    ):
      lowerCAmelCase       :						int							= argparse.ArgumentParser(description='''Run and time generations on a big model using Accelerate.'''      )
      parser.add_argument('''model_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the model to try.'''      )
      parser.add_argument(
          '''--tokenizer_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the tokenizer (if different from the model.'''      )
      parser.add_argument('''--is_causal'''			, type=_snake_case			, default=_snake_case			, help='''Whether or not the model is causal.'''      )
      parser.add_argument(
          '''--model_revision'''			, type=_snake_case			, default=_snake_case			, help='''The revision to use for the model checkpoint.'''      )
      parser.add_argument('''--torch_dtype'''			, type=_snake_case			, default=_snake_case			, help='''The dtype for the model.'''      )
      parser.add_argument('''--disk_offload'''			, action='''store_true'''      )

      lowerCAmelCase       :						Dict							= parser.parse_args()

      # Sanitize args
      if args.model_name in DEFAULT_MODELS:
            lowerCAmelCase       :						Optional[int]							= DEFAULT_MODELS[args.model_name]
            lowerCAmelCase       :						Union[str, Any]							= defaults['''model''']
            if args.tokenizer_name is None:
                  lowerCAmelCase       :						Union[str, Any]							= defaults.get('''tokenizer'''			, args.model_name      )
            if args.is_causal is None:
                  lowerCAmelCase       :						List[Any]							= defaults['''is_causal''']
            if args.model_revision is None:
                  lowerCAmelCase       :						Tuple							= defaults.get('''model_revision'''			, '''main'''      )

      if args.is_causal is None:
            raise ValueError('''Could not infer the default for `--is_causal`, pass either True or False for it.'''      )
      if args.tokenizer_name is None:
            lowerCAmelCase       :						Union[str, Any]							= args.model_name
      if args.model_revision is None:
            lowerCAmelCase       :						List[Any]							= '''main'''

      return args


def   _snake_case    (    ):
      transformers.utils.logging.set_verbosity_error()
      lowerCAmelCase       :						Union[str, Any]							= parse_args()

      if args.torch_dtype is None:
            lowerCAmelCase       :						List[str]							= AutoConfig.from_pretrained(args.model_name      )
            lowerCAmelCase       :						Optional[Any]							= getattr(_snake_case			, '''torch_dtype'''			, torch.floataa      )
      else:
            lowerCAmelCase       :						Optional[int]							= getattr(_snake_case			, args.torch_dtype      )
      lowerCAmelCase       :						Tuple							= AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
      lowerCAmelCase       :						str							= {
          '''torch_dtype''': torch_dtype,
          '''revision''': args.model_revision,
      }
      if args.disk_offload:
            lowerCAmelCase       :						Optional[int]							= '''tmp_offload'''
            lowerCAmelCase       :						Optional[int]							= True

      lowerCAmelCase       :						str							= start_measure()
      lowerCAmelCase       :						List[str]							= model_cls.from_pretrained(args.model_name			, device_map='''auto'''			, **_snake_case      )
      lowerCAmelCase       :						Tuple							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model loading'''      )

      lowerCAmelCase       :						str							= compute_module_sizes(_snake_case      )
      lowerCAmelCase       :						Tuple							= {v: 0 for v in model.hf_device_map.values()}
      for module, device in model.hf_device_map.items():
            device_size[device] += module_sizes[module]
      lowerCAmelCase       :						int							= '''\n'''.join([f'''- {device}: {size // 2**20}MiB''' for device, size in device_size.items()]      )
      print(f'''\nTheoretical use:\n{message}'''      )

      lowerCAmelCase       :						int							= AutoTokenizer.from_pretrained(args.tokenizer_name      )

      lowerCAmelCase       :						List[str]							= start_measure()
      lowerCAmelCase       :						Union[str, Any]							= []
      lowerCAmelCase       :						int							= []
      lowerCAmelCase       :						Dict							= []
      for prompt in PROMPTS:
            lowerCAmelCase       :						str							= tokenizer(_snake_case			, return_tensors='''pt'''      ).to(0      )
            lowerCAmelCase       :						int							= inputs['''input_ids'''][0].tolist()
            lowerCAmelCase       :						Optional[Any]							= time.time()
            lowerCAmelCase       :						Optional[int]							= model.generate(inputs['''input_ids''']      )
            lowerCAmelCase       :						Any							= time.time()
            lowerCAmelCase       :						str							= outputs[0].tolist()
            lowerCAmelCase       :						int							= len(_snake_case      ) if outputs[: len(_snake_case      )] != tokens else len(_snake_case      ) - len(_snake_case      )
            lowerCAmelCase       :						Tuple							= after_generate - before_generate

            lowerCAmelCase       :						Tuple							= tokenizer.decode(_snake_case			, skip_special_tokens=_snake_case      )
            texts_outs.append(_snake_case      )
            generation_times.append(_snake_case      )
            gen_tokens.append(_snake_case      )
            print(f'''Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n'''      )

      lowerCAmelCase       :						List[Any]							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model generation'''      )

      lowerCAmelCase       :						str							= [gen / tok for gen, tok in zip(_snake_case			, _snake_case      )]
      lowerCAmelCase       :						Any							= sum(_snake_case      ) / len(_snake_case      )
      print(f'''Average time of generation per token: {avg_gen:.2f}s'''      )
      print(f'''First generation (avg time per token): {generation_times_per_token[0]:.2f}s'''      )
      lowerCAmelCase       :						List[Any]							= sum(generation_times_per_token[1:]      ) / (len(_snake_case      ) - 1)
      print(f'''Average time of generation per token (excluding the first): {avg_gen:.2f}s'''      )


if __name__ == "__main__":
 main()


