

'''simple docstring'''



import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
   import torch


SCREAMING_SNAKE_CASE_    =   logging.get_logger(__name__)

class 			SCREAMING_SNAKE_CASE					:






      '''simple docstring'''






      def __init__(					self :		Union[str, Any]     ,					snake_case :		List[str]     ,					snake_case :		int=21     ,					snake_case :		List[str]=[]     ,					snake_case :		Tuple=None   ):



          """simple docstring"""





          _snake_case : Union[str, Any]												=		model
          _snake_case : Union[str, Any]												=		trace_batch_nums
          _snake_case : Optional[Any]												=		abort_after_batch_num

          # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
          _snake_case : int												=		collections.deque([]     ,					_lowerCAmelCase   )
          _snake_case : str												=		[]
          _snake_case : Optional[int]												=		0
          _snake_case : int												=		0
          _snake_case : int												=		False
          _snake_case : List[str]												=		"""                 """

          self.analyse_model()

          self.register_forward_hook()




      def 	__UpperCAmelCase			(					self :		List[str]     ,					snake_case :		Optional[int]=None   ):



          """simple docstring"""





          if frame is not None:
              self.expand_frame(_lowerCAmelCase   )
          self.frames.append('\n'.join(self.frame   )   )
          _snake_case : str												=		[]  # start a new frame




      def 	__UpperCAmelCase			(					self :		Dict     ,					snake_case :		List[str]   ):



          """simple docstring"""





          self.frame.append(_lowerCAmelCase   )




      def 	__UpperCAmelCase			(					self :		Optional[Any]   ):



          """simple docstring"""





          print('\n'.join(self.frames   )   )
          _snake_case : List[Any]												=		[]




      def 	__UpperCAmelCase			(					self :		List[Any]   ):



          """simple docstring"""





          _snake_case : Optional[int]												=		[]




      def 	__UpperCAmelCase			(					self :		Any   ):



          """simple docstring"""





          print(F"""\nDetected inf/nan during batch_number={self.batch_number}"""   )
          print(F"""Last {len(self.frames   )} forward frames:"""   )
          print(F"""{'abs min':8} {'abs max':8} metadata"""   )
          print('\n'.join(self.frames   )   )
          print('\n\n'   )
          _snake_case : Any												=		[]




      def 	__UpperCAmelCase			(					self :		Dict   ):



          """simple docstring"""





          _snake_case : str												=		{m: name for name, m in self.model.named_modules()}
          # self.longest_module_name = max(len(v) for v in self.module_names.values())




      def 	__UpperCAmelCase			(					self :		Optional[int]     ,					snake_case :		int     ,					snake_case :		Union[str, Any]   ):



          """simple docstring"""





          if torch.is_tensor(_lowerCAmelCase   ):
              self.expand_frame(get_abs_min_max(_lowerCAmelCase     ,					_lowerCAmelCase   )   )
              if detect_overflow(_lowerCAmelCase     ,					_lowerCAmelCase   ):
                  _snake_case : Optional[int]												=		True
          elif var is None:
              self.expand_frame(F"""{'None':>17} {ctx}"""   )
          else:
              self.expand_frame(F"""{'not a tensor':>17} {ctx}"""   )




      def 	__UpperCAmelCase			(					self :		Tuple   ):



          """simple docstring"""





          self.expand_frame(F"""\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"""   )
          self.expand_frame(F"""{'abs min':8} {'abs max':8} metadata"""   )




      def 	__UpperCAmelCase			(					self :		Optional[int]   ):



          """simple docstring"""





          self.expand_frame(F"""{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"""   )




      def 	__UpperCAmelCase			(					self :		int     ,					snake_case :		Optional[Any]     ,					snake_case :		Dict     ,					snake_case :		Optional[Any]   ):



          """simple docstring"""





          self.expand_frame(F"""{self.prefix} {self.module_names[module]} {module.__class__.__name__}"""   )

          # params
          for name, p in module.named_parameters(recurse=_lowerCAmelCase   ):
              self.analyse_variable(_lowerCAmelCase     ,					_lowerCAmelCase   )

          # inputs
          if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase   ):
              for i, x in enumerate(_lowerCAmelCase   ):
                  self.analyse_variable(_lowerCAmelCase     ,					F"""input[{i}]"""   )
          else:
              self.analyse_variable(_lowerCAmelCase     ,					'input'   )

          # outputs
          if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase   ):
              for i, x in enumerate(_lowerCAmelCase   ):
                  # possibly a tuple of tuples
                  if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase   ):
                      for j, y in enumerate(_lowerCAmelCase   ):
                          self.analyse_variable(_lowerCAmelCase     ,					F"""output[{i}][{j}]"""   )
                  else:
                      self.analyse_variable(_lowerCAmelCase     ,					F"""output[{i}]"""   )
          else:
              self.analyse_variable(_lowerCAmelCase     ,					'output'   )

          self.save_frame()




      def 	__UpperCAmelCase			(					self :		int   ):



          """simple docstring"""





          self.model.apply(self._register_forward_hook   )




      def 	__UpperCAmelCase			(					self :		Any     ,					snake_case :		str   ):



          """simple docstring"""





          module.register_forward_hook(self.forward_hook   )




      def 	__UpperCAmelCase			(					self :		Optional[Any]     ,					snake_case :		Tuple     ,					snake_case :		Union[str, Any]     ,					snake_case :		Union[str, Any]   ):



          """simple docstring"""





          _snake_case : Union[str, Any]												=		False

          _snake_case : List[str]												=		True if self.batch_number in self.trace_batch_nums else False
          if trace_mode:
              self.reset_saved_frames()

          if self.total_calls == 0:
              self.batch_start_frame()
          self.total_calls += 1

          # count batch numbers - the very first forward hook of the batch will be called when the
          # batch completes - i.e. it gets called very last - we know this batch has finished
          if module == self.model:
              self.batch_number += 1
              _snake_case : Dict												=		True

          self.create_frame(_lowerCAmelCase     ,					_lowerCAmelCase     ,					_lowerCAmelCase   )

          # if last_frame_of_batch:
          #     self.batch_end_frame()

          if trace_mode:
              self.trace_frames()

          if last_frame_of_batch:
              self.batch_start_frame()

          if self.detected_overflow and not trace_mode:
              self.dump_saved_frames()

              # now we can abort, as it's pointless to continue running
              raise ValueError(
                  'DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. '
                  'Please scroll up above this traceback to see the activation values prior to this event.'   )

          # abort after certain batch if requested to do so
          if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
              raise ValueError(
                  F"""DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"""
                  F""" `abort_after_batch_num={self.abort_after_batch_num}` arg"""   )



def   lowerCamelCase__       (      a__			,      a__)		->    Optional[Any]:



    """simple docstring"""







    _snake_case : Dict												=		var.abs()
    return F"""{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"""

def   lowerCamelCase__       (      a__			,      a__)		->    int:



    """simple docstring"""







    _snake_case : str												=		False
    if torch.isnan(__lowerCamelCase).any().item():
        _snake_case : List[Any]												=		True
        print(F"""{ctx} has nans""")
    if torch.isinf(__lowerCamelCase).any().item():
        _snake_case : List[str]												=		True
        print(F"""{ctx} has infs""")

    # if needed to monitor large elements can enable the following
    if 0:  # and detected:
        _snake_case : Dict												=		var[torch.ge(var.abs()			,      1_0_0)]
        if naaa.numel() > 0:
            print(F"""{ctx}:  n100={naaa.numel()}""")
        _snake_case : Tuple												=		var[torch.ge(var.abs()			,      1_0_0_0)]
        if naaaa.numel() > 0:
            print(F"""{ctx}: n1000={naaaa.numel()}""")
        _snake_case : Optional[int]												=		var[torch.ge(var.abs()			,      1_0_0_0_0)]
        if naaaaa.numel() > 0:
            print(F"""{ctx}: n10000={naaaaa.numel()}""")

    if 0:
        print(F"""min={var.min():9.2e} max={var.max():9.2e}""")

    if 0:
        print(F"""min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})""")

    return detected




class 			SCREAMING_SNAKE_CASE					(							__lowerCamelCase					):






      '''simple docstring'''






      SCREAMING_SNAKE_CASE__ :				Tuple    =				'''underflow_overflow'''
      SCREAMING_SNAKE_CASE__ :				int    =				'''tpu_metrics_debug'''






'''simple docstring'''



import argparse
import logging

import numpy as np
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
SCREAMING_SNAKE_CASE_    =   logging.getLogger(__name__)

def   lowerCamelCase__       (      a__)		->    List[Any]:



    """simple docstring"""







    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def   lowerCamelCase__       (      )		->    Tuple:



    """simple docstring"""







    _snake_case : Optional[int]												=		argparse.ArgumentParser()
    parser.add_argument(
        '--model_name_or_path'			,      default=a__			,      type=a__			,      required=a__			,      )
    parser.add_argument('--prompt'			,      type=a__			,      default='')
    parser.add_argument('--length'			,      type=a__			,      default=2_0)
    parser.add_argument('--stop_token'			,      type=a__			,      default=a__			,      help='Token at which text generation is stopped')
    parser.add_argument(
        '--temperature'			,      type=a__			,      default=1.0			,      help='temperature of 1.0 has no effect, lower tend toward greedy sampling'			,      )
    parser.add_argument(
        '--repetition_penalty'			,      type=a__			,      default=1.0			,      help='primarily useful for CTRL model; in that case, use 1.2')
    parser.add_argument('--k'			,      type=a__			,      default=0)
    parser.add_argument('--penalty_alpha'			,      type=a__			,      default=0.0)
    parser.add_argument('--p'			,      type=a__			,      default=0.9)

    parser.add_argument('--prefix'			,      type=a__			,      default=''			,      help='Text added prior to input.')
    parser.add_argument('--padding_text'			,      type=a__			,      default=''			,      help='Deprecated, the use of `--prefix` is preferred.')
    parser.add_argument('--xlm_language'			,      type=a__			,      default=''			,      help='Optional language when used with the XLM model.')

    parser.add_argument('--seed'			,      type=a__			,      default=4_2			,      help='random seed for initialization')
    parser.add_argument('--no_cuda'			,      action='store_true'			,      help='Avoid using CUDA when available')
    parser.add_argument(
        '--fp16'			,      action='store_true'			,      help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'			,      )
    _snake_case : Optional[int]												=		parser.parse_args()

    _snake_case : Optional[int]												=		torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu')
    _snake_case : int												=		0 if args.no_cuda else torch.cuda.device_count()

    logger.warning(F"""device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}""")

    set_seed(a__)

    # Initialize the model and tokenizer
    _snake_case : Dict												=		AutoTokenizer.from_pretrained(args.model_name_or_path)
    _snake_case : List[Any]												=		AutoModelForCausalLM.from_pretrained(args.model_name_or_path)

    # tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
    # model = OPTForCausalLM.from_pretrained(args.model_name_or_path)
    model.to(args.device)

    if args.fpaa:
        model.half()

    logger.info(a__)
    _snake_case : Any												=		args.prompt if args.prompt else input('Model prompt >>> ')

    _snake_case : List[str]												=		tokenizer(a__			,      return_tensors='pt'			,      add_special_tokens=a__)
    _snake_case : List[str]												=		{key: value.to(args.device) for key, value in inputs.items()}

    _snake_case : Optional[int]												=		model.generate(
        **a__			,      max_length=args.length + len(inputs['input_ids'][0])			,      penalty_alpha=args.penalty_alpha			,      top_k=args.k			,      )

    _snake_case : int												=		[]
    for generated_sequence_idx, generated_sequence in enumerate(a__):
        print(F"""=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===""")
        _snake_case : Any												=		generated_sequence.tolist()

        # Decode text
        _snake_case : str												=		tokenizer.decode(a__			,      clean_up_tokenization_spaces=a__			,      add_special_tokens=a__)

        # Remove all text after the stop token
        _snake_case : str												=		text[: text.find(args.stop_token) if args.stop_token else None]

        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
        _snake_case : Tuple												=		(
            prompt_text + text[len(tokenizer.decode(inputs['input_ids'][0]			,      clean_up_tokenization_spaces=a__)) :]
        )

        generated_sequences.append(a__)
        print(a__)

    return generated_sequences


if __name__ == "__main__":
   main()



