



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


_UpperCAmelCase       =     {
    """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"""},
}

_UpperCAmelCase       =     [
    """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        __magic_name__			(							):
   SCREAMING_SNAKE_CASE_:	List[Any]            =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_:	Tuple            =parser.parse_args()

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

   return args





def        __magic_name__			(							):
   transformers.utils.logging.set_verbosity_error()
   SCREAMING_SNAKE_CASE_:	List[str]            =parse_args()

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

   SCREAMING_SNAKE_CASE_:	int            =start_measure()
   SCREAMING_SNAKE_CASE_:	List[str]            =model_cls.from_pretrained(args.model_name							,     device_map="""auto"""							,     **lowercase	)
   SCREAMING_SNAKE_CASE_:	str            =end_measure(lowercase	)
   log_measures(lowercase							,     """Model loading"""	)

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

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

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

      SCREAMING_SNAKE_CASE_:	List[Any]            =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_:	Dict            =end_measure(lowercase	)
   log_measures(lowercase							,     """Model generation"""	)

   SCREAMING_SNAKE_CASE_:	Union[str, Any]            =[gen / tok for gen, tok in zip(lowercase							,     lowercase	)]
   SCREAMING_SNAKE_CASE_:	Any            =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_:	str            =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()








"""simple docstring"""
import copy

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging


_UpperCAmelCase       =     logging.get_logger(__name__)
class    a		(					UpperCAmelCase__    ):
  UpperCamelCase		:			int         =   'hybrid-clip'
  UpperCamelCase		:			Optional[int]         =   True



  def __init__(  self	:   List[Any]    ,						lowerCAmelCase	:   Any=512    ,						**lowerCAmelCase	:   Optional[Any]						)       ->   List[str]:


     '''simple docstring'''




     super().__init__(**lowerCAmelCase						)

     if "text_config" not in kwargs:
        raise ValueError("""`text_config` can not be `None`."""						)

     if "vision_config" not in kwargs:
        raise ValueError("""`vision_config` can not be `None`."""						)

     SCREAMING_SNAKE_CASE_:	Optional[int]            =kwargs.pop("""text_config"""						)
     SCREAMING_SNAKE_CASE_:	Optional[Any]            =kwargs.pop("""vision_config"""						)

     SCREAMING_SNAKE_CASE_:	int            =text_config.pop("""model_type"""						)
     SCREAMING_SNAKE_CASE_:	str            =vision_config.pop("""model_type"""						)

     from transformers import AutoConfig

     SCREAMING_SNAKE_CASE_:	Optional[int]            =AutoConfig.for_model(lowerCAmelCase    ,						**lowerCAmelCase						)

     if vision_model_type == "clip":
        SCREAMING_SNAKE_CASE_:	Tuple            =AutoConfig.for_model(lowerCAmelCase    ,						**lowerCAmelCase						).vision_config
     elif vision_model_type == "clip_vision_model":
        from transformers import CLIPVisionConfig

        SCREAMING_SNAKE_CASE_:	List[str]            =CLIPVisionConfig(**lowerCAmelCase						)
     else:
        SCREAMING_SNAKE_CASE_:	Optional[Any]            =AutoConfig.for_model(lowerCAmelCase    ,						**lowerCAmelCase						)

     SCREAMING_SNAKE_CASE_:	Any            =projection_dim
     SCREAMING_SNAKE_CASE_:	int            =1.0



  @classmethod
  def 		lowerCamelCase__    (  cls	:   List[Any]    ,						lowerCAmelCase	:   PretrainedConfig    ,						lowerCAmelCase	:   PretrainedConfig    ,						**lowerCAmelCase	:   str						)       ->   Optional[Any]:


     '''simple docstring'''




     return cls(text_config=text_config.to_dict()    ,						vision_config=vision_config.to_dict()    ,						**lowerCAmelCase						)





  def 		lowerCamelCase__    (  self	:   List[str]						)       ->   Tuple:


     '''simple docstring'''




     SCREAMING_SNAKE_CASE_:	int            =copy.deepcopy(self.__dict__						)
     SCREAMING_SNAKE_CASE_:	List[str]            =self.text_config.to_dict()
     SCREAMING_SNAKE_CASE_:	Optional[int]            =self.vision_config.to_dict()
     SCREAMING_SNAKE_CASE_:	Union[str, Any]            =self.__class__.model_type
     return output



