code: 
  
  
  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()
  
  
  
style_context: 
  
  
  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					)					)
  
  
  