'''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


lowerCAmelCase_			:			str    	=   {
    "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"},
}

lowerCAmelCase_			:			List[str]    	=   [
    "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 __a			(						)						->				str:
      '''simple docstring'''


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

      lowercase_     			=   parser.parse_args()

      # Sanitize args
      if args.model_name in DEFAULT_MODELS:
            lowercase_     			=   DEFAULT_MODELS[args.model_name]
            lowercase_     			=   defaults["model"]
            if args.tokenizer_name is None:
                  lowercase_     			=   defaults.get("tokenizer"			,		args.model_name					)
            if args.is_causal is None:
                  lowercase_     			=   defaults["is_causal"]
            if args.model_revision is None:
                  lowercase_     			=   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:
            lowercase_     			=   args.model_name
      if args.model_revision is None:
            lowercase_     			=   "main"

      return args


def __a			(						)						->				Optional[Any]:
      '''simple docstring'''


      transformers.utils.logging.set_verbosity_error()
      lowercase_     			=   parse_args()

      if args.torch_dtype is None:
            lowercase_     			=   AutoConfig.from_pretrained(args.model_name					)
            lowercase_     			=   getattr(__lowerCamelCase			,		"torch_dtype"			,		torch.floataa					)
      else:
            lowercase_     			=   getattr(__lowerCamelCase			,		args.torch_dtype					)
      lowercase_     			=   AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
      lowercase_     			=   {
          "torch_dtype": torch_dtype,
          "revision": args.model_revision,
      }
      if args.disk_offload:
            lowercase_     			=   "tmp_offload"
            lowercase_     			=   True

      lowercase_     			=   start_measure()
      lowercase_     			=   model_cls.from_pretrained(args.model_name			,		device_map="auto"			,		**__lowerCamelCase					)
      lowercase_     			=   end_measure(__lowerCamelCase					)
      log_measures(__lowerCamelCase			,		"Model loading"					)

      lowercase_     			=   compute_module_sizes(__lowerCamelCase					)
      lowercase_     			=   {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]
      lowercase_     			=   "\n".join([f'- {device}: {size // 2**20}MiB' for device, size in device_size.items()]					)
      print(f'\nTheoretical use:\n{message}'					)

      lowercase_     			=   AutoTokenizer.from_pretrained(args.tokenizer_name					)

      lowercase_     			=   start_measure()
      lowercase_     			=   []
      lowercase_     			=   []
      lowercase_     			=   []
      for prompt in PROMPTS:
            lowercase_     			=   tokenizer(__lowerCamelCase			,		return_tensors="pt"					).to(0					)
            lowercase_     			=   inputs["input_ids"][0].tolist()
            lowercase_     			=   time.time()
            lowercase_     			=   model.generate(inputs["input_ids"]					)
            lowercase_     			=   time.time()
            lowercase_     			=   outputs[0].tolist()
            lowercase_     			=   len(__lowerCamelCase					) if outputs[: len(__lowerCamelCase					)] != tokens else len(__lowerCamelCase					) - len(__lowerCamelCase					)
            lowercase_     			=   after_generate - before_generate

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

      lowercase_     			=   end_measure(__lowerCamelCase					)
      log_measures(__lowerCamelCase			,		"Model generation"					)

      lowercase_     			=   [gen / tok for gen, tok in zip(__lowerCamelCase			,		__lowerCamelCase					)]
      lowercase_     			=   sum(__lowerCamelCase					) / len(__lowerCamelCase					)
      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'					)
      lowercase_     			=   sum(generation_times_per_token[1:]					) / (len(__lowerCamelCase					) - 1)
      print(f'Average time of generation per token (excluding the first): {avg_gen:.2f}s'					)


if __name__ == "__main__":
       main()







'''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 ConvNextImageProcessor





class 				lowercase    (    unittest.TestCase			):


   def __init__( self					:						str ,		__lowerCAmelCase					:						Union[str, Any] ,		__lowerCAmelCase					:						Optional[int]=7 ,		__lowerCAmelCase					:						Optional[Any]=3 ,		__lowerCAmelCase					:						Optional[Any]=18 ,		__lowerCAmelCase					:						List[str]=30 ,		__lowerCAmelCase					:						List[str]=400 ,		__lowerCAmelCase					:						Tuple=True ,		__lowerCAmelCase					:						Optional[int]=None ,		__lowerCAmelCase					:						List[str]=0.875 ,		__lowerCAmelCase					:						int=True ,		__lowerCAmelCase					:						Optional[int]=[0.5, 0.5, 0.5] ,		__lowerCAmelCase					:						List[str]=[0.5, 0.5, 0.5] ,		)   ->       Optional[Any]:
         lowercase_     			=   size if size is not None else {"shortest_edge": 20}
         lowercase_     			=   parent
         lowercase_     			=   batch_size
         lowercase_     			=   num_channels
         lowercase_     			=   image_size
         lowercase_     			=   min_resolution
         lowercase_     			=   max_resolution
         lowercase_     			=   do_resize
         lowercase_     			=   size
         lowercase_     			=   crop_pct
         lowercase_     			=   do_normalize
         lowercase_     			=   image_mean
         lowercase_     			=   image_std


   def __UpperCAmelCase							( self					:						List[str])   ->       int:
         return {
             "image_mean": self.image_mean,
             "image_std": self.image_std,
             "do_normalize": self.do_normalize,
             "do_resize": self.do_resize,
             "size": self.size,
             "crop_pct": self.crop_pct,
         }





@require_torch
@require_vision
class 				lowercase    (    __lowerCamelCase			,					unittest.TestCase			):
   lowerCamelCase_            =ConvNextImageProcessor if is_vision_available() else None


   def __UpperCAmelCase							( self					:						Optional[Any])   ->       Tuple:
         lowercase_     			=   ConvNextImageProcessingTester(self)


   @property
   def __UpperCAmelCase							( self					:						Optional[int])   ->       int:
         return self.image_processor_tester.prepare_image_processor_dict()


   def __UpperCAmelCase							( self					:						Optional[int])   ->       Any:
         lowercase_     			=   self.image_processing_class(**self.image_processor_dict)
         self.assertTrue(hasattr(__lowerCAmelCase ,		"do_resize"))
         self.assertTrue(hasattr(__lowerCAmelCase ,		"size"))
         self.assertTrue(hasattr(__lowerCAmelCase ,		"crop_pct"))
         self.assertTrue(hasattr(__lowerCAmelCase ,		"do_normalize"))
         self.assertTrue(hasattr(__lowerCAmelCase ,		"image_mean"))
         self.assertTrue(hasattr(__lowerCAmelCase ,		"image_std"))


   def __UpperCAmelCase							( self					:						List[str])   ->       Any:
         lowercase_     			=   self.image_processing_class.from_dict(self.image_processor_dict)
         self.assertEqual(image_processor.size ,		{"shortest_edge": 20})

         lowercase_     			=   self.image_processing_class.from_dict(self.image_processor_dict ,		size=42)
         self.assertEqual(image_processor.size ,		{"shortest_edge": 42})


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


   def __UpperCAmelCase							( self					:						List[Any])   ->       Any:
         # Initialize image_processing
         lowercase_     			=   self.image_processing_class(**self.image_processor_dict)
         # create random PIL images
         lowercase_     			=   prepare_image_inputs(self.image_processor_tester ,		equal_resolution=__lowerCAmelCase)
         for image in image_inputs:
               self.assertIsInstance(__lowerCAmelCase ,		Image.Image)

         # Test not batched input
         lowercase_     			=   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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)

         # Test batched
         lowercase_     			=   image_processing(__lowerCAmelCase ,		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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)


   def __UpperCAmelCase							( self					:						Union[str, Any])   ->       str:
         # Initialize image_processing
         lowercase_     			=   self.image_processing_class(**self.image_processor_dict)
         # create random numpy tensors
         lowercase_     			=   prepare_image_inputs(self.image_processor_tester ,		equal_resolution=__lowerCAmelCase ,		numpify=__lowerCAmelCase)
         for image in image_inputs:
               self.assertIsInstance(__lowerCAmelCase ,		np.ndarray)

         # Test not batched input
         lowercase_     			=   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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)

         # Test batched
         lowercase_     			=   image_processing(__lowerCAmelCase ,		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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)


   def __UpperCAmelCase							( self					:						List[str])   ->       List[str]:
         # Initialize image_processing
         lowercase_     			=   self.image_processing_class(**self.image_processor_dict)
         # create random PyTorch tensors
         lowercase_     			=   prepare_image_inputs(self.image_processor_tester ,		equal_resolution=__lowerCAmelCase ,		torchify=__lowerCAmelCase)
         for image in image_inputs:
               self.assertIsInstance(__lowerCAmelCase ,		torch.Tensor)

         # Test not batched input
         lowercase_     			=   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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)

         # Test batched
         lowercase_     			=   image_processing(__lowerCAmelCase ,		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.size["shortest_edge"],
                 self.image_processor_tester.size["shortest_edge"],
             ) ,		)






