



"""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"""
from typing import List

from tokenizers import NormalizedString, PreTokenizedString, normalizers
class    a		:



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


     '''simple docstring'''




     SCREAMING_SNAKE_CASE_:	int            =vocab
     SCREAMING_SNAKE_CASE_:	Union[str, Any]            =normalizers.BertNormalizer(
         clean_text=lowerCAmelCase    ,						handle_chinese_chars=lowerCAmelCase    ,						strip_accents=lowerCAmelCase    ,						lowercase=lowerCAmelCase    ,						)
     try:
        import rjieba
     except ImportError:
        raise ImportError(
            """You need to install rjieba to use RoFormerTokenizer. """
            """See https://pypi.org/project/rjieba/ for installation."""						)
     SCREAMING_SNAKE_CASE_:	Dict            =rjieba



  def 		lowerCamelCase__    (  self	:   Tuple    ,						lowerCAmelCase	:   int    ,						lowerCAmelCase	:   NormalizedString						)       ->   List[NormalizedString]:


     '''simple docstring'''




     SCREAMING_SNAKE_CASE_:	Optional[int]            =[]

     # this code slice normalized_string is too slow (6s) but test_alignement_methods can pass
     for token, start, end in self.jieba.tokenize(str(lowerCAmelCase						)    ,						hmm=lowerCAmelCase						):
        if token in self.vocab:
           splits.append(normalized_string[start:end]						)
        else:
           SCREAMING_SNAKE_CASE_:	List[Any]            =self.normalizers.normalize_str(lowerCAmelCase						).split()
           for token in token_list:
              if token:
                 SCREAMING_SNAKE_CASE_:	Dict            =start + len(lowerCAmelCase						)
                 splits.append(normalized_string[start:end]						)
                 SCREAMING_SNAKE_CASE_:	Optional[Any]            =end

        # this code test_alignement_methods can't pass but fast (300ms)
        # for token in self.jieba.cut(str(normalized_string), False):
        #     if token in self.vocab:
        #         splits.append(NormalizedString(token))
        #     else:
        #         token_list = self.normalizers.normalize_str(token).split()
        #         for token in token_list:
        #             if token:
        #                 splits.append(NormalizedString(token))

     return splits





  def 		lowerCamelCase__    (  self	:   Union[str, Any]    ,						lowerCAmelCase	:   PreTokenizedString						)       ->   Tuple:


     '''simple docstring'''




     pretok.split(self.jieba_split						)



