


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


A_														=      {
    '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'},
}

A_														=      [
    '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 							UpperCAmelCase				(     )->		Tuple:







     '''simple docstring'''



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

     SCREAMING_SNAKE_CASE_	      =       parser.parse_args()

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

     return args







def 							UpperCAmelCase				(     )->		List[str]:







     '''simple docstring'''



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

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

     SCREAMING_SNAKE_CASE_	      =       start_measure()
     SCREAMING_SNAKE_CASE_	      =       model_cls.from_pretrained(args.model_name				,device_map='''auto'''				,**lowercase__      )
     SCREAMING_SNAKE_CASE_	      =       end_measure(lowercase__      )
     log_measures(lowercase__				,'''Model loading'''      )

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

     SCREAMING_SNAKE_CASE_	      =       AutoTokenizer.from_pretrained(args.tokenizer_name      )

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

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

     SCREAMING_SNAKE_CASE_	      =       end_measure(lowercase__      )
     log_measures(lowercase__				,'''Model generation'''      )

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


if __name__ == "__main__":
  main()






import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types

class 						snake_case			(       unittest.TestCase	, lowerCAmelCase__	):
			'''simple docstring'''

			def _lowercase (	self     :     List[Any]					)					->					Optional[Any]:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       load_tool('''translation'''					)
								self.tool.setup()
								SCREAMING_SNAKE_CASE_	      =       load_tool('''translation'''							,							remote=lowerCAmelCase_					)

			def _lowercase (	self     :     Tuple					)					->					str:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       self.tool('''Hey, what\'s up?'''							,							src_lang='''English'''							,							tgt_lang='''French'''					)
								self.assertEqual(lowerCAmelCase_							,							'''- Hé, comment ça va?'''					)

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




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       self.remote_tool('''Hey, what\'s up?'''							,							src_lang='''English'''							,							tgt_lang='''French'''					)
								self.assertEqual(lowerCAmelCase_							,							'''- Hé, comment ça va?'''					)

			def _lowercase (	self     :     Optional[int]					)					->					str:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       self.tool(text='''Hey, what\'s up?'''							,							src_lang='''English'''							,							tgt_lang='''French'''					)
								self.assertEqual(lowerCAmelCase_							,							'''- Hé, comment ça va?'''					)

			def _lowercase (	self     :     str					)					->					List[str]:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       self.remote_tool(text='''Hey, what\'s up?'''							,							src_lang='''English'''							,							tgt_lang='''French'''					)
								self.assertEqual(lowerCAmelCase_							,							'''- Hé, comment ça va?'''					)

			def _lowercase (	self     :     Optional[int]					)					->					Dict:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       ['''Hey, what\'s up?''', '''English''', '''Spanish''']
								SCREAMING_SNAKE_CASE_	      =       self.tool(*lowerCAmelCase_					)

								# There is a single output
								if len(self.tool.outputs					) == 1:
													SCREAMING_SNAKE_CASE_	      =       [outputs]

								self.assertListEqual(output_types(lowerCAmelCase_					)							,							self.tool.outputs					)

			def _lowercase (	self     :     Optional[int]					)					->					List[str]:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       ['''Hey, what\'s up?''', '''English''', '''Spanish''']
								SCREAMING_SNAKE_CASE_	      =       self.tool(*lowerCAmelCase_					)

								if not isinstance(lowerCAmelCase_							,							lowerCAmelCase_					):
													SCREAMING_SNAKE_CASE_	      =       [outputs]

								self.assertEqual(len(lowerCAmelCase_					)							,							len(self.tool.outputs					)					)

								for output, output_type in zip(lowerCAmelCase_							,							self.tool.outputs					):
													SCREAMING_SNAKE_CASE_	      =       AGENT_TYPE_MAPPING[output_type]
													self.assertTrue(isinstance(lowerCAmelCase_							,							lowerCAmelCase_					)					)




			def _lowercase (	self     :     Tuple					)					->					int:




								"""simple docstring"""






								SCREAMING_SNAKE_CASE_	      =       ['''Hey, what\'s up?''', '''English''', '''Spanish''']

								SCREAMING_SNAKE_CASE_	      =       []

								for _input, input_type in zip(lowerCAmelCase_							,							self.tool.inputs					):
													if isinstance(lowerCAmelCase_							,							lowerCAmelCase_					):
																		_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input					) for _input_type in input_type]					)
													else:
																		_inputs.append(AGENT_TYPE_MAPPING[input_type](_input					)					)

        # Should not raise an error
								SCREAMING_SNAKE_CASE_	      =       self.tool(*lowerCAmelCase_					)

								if not isinstance(lowerCAmelCase_							,							lowerCAmelCase_					):
													SCREAMING_SNAKE_CASE_	      =       [outputs]

								self.assertEqual(len(lowerCAmelCase_					)							,							len(self.tool.outputs					)					)


