





import os
from shutil import copyfile
from typing import TYPE_CHECKING, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging
from ...utils.versions import require_version


if TYPE_CHECKING:
     from transformers.pipelines.conversational import Conversation

require_version('tokenizers>=0.13.3')

if is_sentencepiece_available():
     from .tokenization_llama import LlamaTokenizer
else:
     _UpperCAmelCase				       = None

_UpperCAmelCase				       = logging.get_logger(__name__)
_UpperCAmelCase				       = {'vocab_file': 'tokenizer.model', 'tokenizer_file': 'tokenizer.json'}

_UpperCAmelCase       ,  _UpperCAmelCase				       = '[INST]', '[/INST]'
_UpperCAmelCase       ,  _UpperCAmelCase				       = '<<SYS>>\n', '\n<</SYS>>\n\n'

# fmt: off
_UpperCAmelCase				       = 'You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your\\nanswers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure\\nthat your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not\\ncorrect. If you don\'t know the answer to a question, please don\'t share false information.'



class 							_UpperCamelCase    (		lowerCAmelCase_   ):
       _UpperCamelCase :    Optional[Any]   			=      VOCAB_FILES_NAMES
       _UpperCamelCase :    List[Any]   			=      LlamaTokenizer
       _UpperCamelCase :    str   			=      '''left'''
       _UpperCamelCase :    Tuple   			=      ['''input_ids''', '''attention_mask''']

       def __init__(						self:       int    ,      _SCREAMING_SNAKE_CASE:       Optional[Any]=None    ,      _SCREAMING_SNAKE_CASE:       int=None    ,      _SCREAMING_SNAKE_CASE:       Any=False    ,      _SCREAMING_SNAKE_CASE:       Optional[int]="<unk>"    ,      _SCREAMING_SNAKE_CASE:       Any="<s>"    ,      _SCREAMING_SNAKE_CASE:       Dict="</s>"    ,      _SCREAMING_SNAKE_CASE:       str=True    ,      _SCREAMING_SNAKE_CASE:       Union[str, Any]=False    ,      **_SCREAMING_SNAKE_CASE:       Tuple    ,      )				->       Union[str, Any]:




              """simple docstring"""


              super().__init__(
                  vocab_file=_SCREAMING_SNAKE_CASE    ,      tokenizer_file=_SCREAMING_SNAKE_CASE    ,      clean_up_tokenization_spaces=_SCREAMING_SNAKE_CASE    ,      unk_token=_SCREAMING_SNAKE_CASE    ,      bos_token=_SCREAMING_SNAKE_CASE    ,      eos_token=_SCREAMING_SNAKE_CASE    ,      **_SCREAMING_SNAKE_CASE    ,      )
              UpperCamelCase_      				=     add_bos_token
              UpperCamelCase_      				=     add_eos_token
              self.update_post_processor()

              UpperCamelCase_      				=     vocab_file
              UpperCamelCase_      				=     False if not self.vocab_file else True

       def lowercase       (						self:       str						)				->       List[Any]:




              """simple docstring"""


              UpperCamelCase_      				=     self.bos_token
              UpperCamelCase_      				=     self.bos_token_id

              UpperCamelCase_      				=     self.eos_token
              UpperCamelCase_      				=     self.eos_token_id

              UpperCamelCase_      				=     f'''{(bos+":0 ") * self.add_bos_token}$A:0{(" "+eos+":0") * self.add_eos_token}'''
              UpperCamelCase_      				=     f'''{single}{(" "+bos+":1") * self.add_bos_token} $B:1{(" "+eos+":1") * self.add_eos_token}'''

              UpperCamelCase_      				=     []
              if self.add_bos_token:
                     special_tokens.append((bos, bos_token_id)						)
              if self.add_eos_token:
                     special_tokens.append((eos, eos_token_id)						)
              UpperCamelCase_      				=     processors.TemplateProcessing(
                  single=_SCREAMING_SNAKE_CASE    ,      pair=_SCREAMING_SNAKE_CASE    ,      special_tokens=_SCREAMING_SNAKE_CASE						)

       @property
       def lowercase       (						self:       Union[str, Any]						)				->       Tuple:




              """simple docstring"""


              return self._add_eos_token

       @property
       def lowercase       (						self:       int						)				->       str:




              """simple docstring"""


              return self._add_bos_token

       @add_eos_token.setter
       def lowercase       (						self:       Any    ,      _SCREAMING_SNAKE_CASE:       Optional[Any]						)				->       Tuple:




              """simple docstring"""


              UpperCamelCase_      				=     value
              self.update_post_processor()

       @add_bos_token.setter
       def lowercase       (						self:       List[str]    ,      _SCREAMING_SNAKE_CASE:       Union[str, Any]						)				->       Optional[int]:




              """simple docstring"""


              UpperCamelCase_      				=     value
              self.update_post_processor()

       def lowercase       (						self:       Any    ,      _SCREAMING_SNAKE_CASE:       str    ,      _SCREAMING_SNAKE_CASE:       Optional[str] = None						)				->       Tuple[str]:




              """simple docstring"""


              if not self.can_save_slow_tokenizer:
                     raise ValueError(
                         "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                         "tokenizer."						)

              if not os.path.isdir(_SCREAMING_SNAKE_CASE						):
                     logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''						)
                     return
              UpperCamelCase_      				=     os.path.join(
                  _SCREAMING_SNAKE_CASE    ,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]						)

              if os.path.abspath(self.vocab_file						) != os.path.abspath(_SCREAMING_SNAKE_CASE						):
                     copyfile(self.vocab_file    ,      _SCREAMING_SNAKE_CASE						)

              return (out_vocab_file,)



       def lowercase       (						self:       List[str]    ,      _SCREAMING_SNAKE_CASE:       "Conversation"						)				->       List[Any]:




              """simple docstring"""


              UpperCamelCase_      				=     list(conversation.iter_texts()						)
              if not all([is_user for is_user, msg in dialogue[::2]]						) or not all(
                  [not is_user for is_user, msg in dialogue[1::2]]						):
                     raise ValueError(
                         "The model only supports 'user' and 'assistant' roles, starting with user and alternating (u/a/u/a/u...)"						)

              UpperCamelCase_      				=     []
              if len(conversation.past_user_inputs						) > 0:
                     if not conversation.past_user_inputs[0].startswith(_SCREAMING_SNAKE_CASE						) or E_SYS not in conversation.past_user_inputs[0]:
                            UpperCamelCase_      				=     (
                                B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.past_user_inputs[0]
                            )
              elif not dialogue[0][1].startswith(_SCREAMING_SNAKE_CASE						) or E_SYS not in dialogue[0][1]:
                     UpperCamelCase_      				=     (dialogue[0][0], B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + dialogue[0][1])

              dialog_tokens += sum(
                  [
                      [self.bos_token_id]
                      + self.encode(
                          f'''{B_INST} {(prompt[1]).strip()} {E_INST} {(answer[1]).strip()} '''    ,      add_special_tokens=_SCREAMING_SNAKE_CASE						)
                      + [self.eos_token_id]
                      for prompt, answer in zip(dialogue[::2]    ,      dialogue[1::2]						)
                  ]    ,      []    ,      )
              if not (dialogue[-1][0]):
                     raise ValueError(f'''Last message must be from user, got {dialogue[-1]["role"]}'''						)
              dialog_tokens += [self.bos_token_id] + self.encode(
                  f'''{B_INST} {(dialogue[-1][1]).strip()} {E_INST}'''    ,      add_special_tokens=_SCREAMING_SNAKE_CASE						)
              return dialog_tokens








import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin




def        lowerCAmelCase_       (							UpperCamelCase_=2    ,						UpperCamelCase_=True    ,						UpperCamelCase_=False     )   ->						str:
       UpperCamelCase_      				=     Accelerator(step_scheduler_with_optimizer=UpperCamelCase_    ,						split_batches=UpperCamelCase_     )
       UpperCamelCase_      				=     torch.nn.Linear(2    ,						4     )
       UpperCamelCase_      				=     torch.optim.AdamW(model.parameters()    ,						lr=1.0     )
       UpperCamelCase_      				=     torch.optim.lr_scheduler.OneCycleLR(UpperCamelCase_    ,						max_lr=0.01    ,						steps_per_epoch=2    ,						epochs=1     )
       UpperCamelCase_				,       UpperCamelCase_				,       UpperCamelCase_      				=     accelerator.prepare(UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_     )

       # Optimizer has stepped
       scheduler.step()
       if step_scheduler_with_optimizer or (num_processes == 1):
              assert (
                  scheduler.scheduler.last_epoch == num_processes
              ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'''
       else:
              assert (
                  scheduler.scheduler.last_epoch != num_processes
              ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'''




def        lowerCAmelCase_       (							UpperCamelCase_=2    ,						UpperCamelCase_=True    ,						UpperCamelCase_=False     )   ->						List[Any]:
       UpperCamelCase_      				=     Accelerator(step_scheduler_with_optimizer=UpperCamelCase_    ,						split_batches=UpperCamelCase_     )
       UpperCamelCase_      				=     torch.nn.Linear(2    ,						4     )
       UpperCamelCase_      				=     torch.optim.AdamW(model.parameters()    ,						lr=1.0     )
       UpperCamelCase_      				=     torch.optim.lr_scheduler.LambdaLR(UpperCamelCase_    ,						lr_lambda=lambda UpperCamelCase_     : 1 - n / 10     )
       UpperCamelCase_				,       UpperCamelCase_				,       UpperCamelCase_      				=     accelerator.prepare(UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_     )

       # Optimizer has stepped
       UpperCamelCase_      				=     False
       scheduler.step()
       UpperCamelCase_      				=     1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
       assert (
           scheduler.get_last_lr()[0] == expected_lr
       ), F'''Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''

       # Optimizer has not stepped
       UpperCamelCase_      				=     True
       scheduler.step()
       if not step_scheduler_with_optimizer:
              UpperCamelCase_      				=     1 - 2 / 10
       assert (
           scheduler.get_last_lr()[0] == expected_lr
       ), F'''Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''




def        lowerCAmelCase_       (							UpperCamelCase_ = 2     )   ->						int:
       from transformers import get_linear_schedule_with_warmup

       UpperCamelCase_      				=     [1, 2, 4]
       for num_steps in steps:
              UpperCamelCase_      				=     GradientAccumulationPlugin(num_steps=UpperCamelCase_    ,						adjust_scheduler=num_steps > 1     )
              UpperCamelCase_      				=     Accelerator(gradient_accumulation_plugin=UpperCamelCase_     )
              UpperCamelCase_      				=     torch.nn.Linear(2    ,						4     )
              UpperCamelCase_      				=     torch.optim.AdamW(model.parameters()    ,						lr=10.0     )
              UpperCamelCase_      				=     get_linear_schedule_with_warmup(optimizer=UpperCamelCase_    ,						num_warmup_steps=0    ,						num_training_steps=20     )

              UpperCamelCase_				,       UpperCamelCase_				,       UpperCamelCase_      				=     accelerator.prepare(UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_     )

              for i in range(10 * num_steps     ):
                     with accelerator.accumulate(UpperCamelCase_     ):
                            optimizer.step()
                            scheduler.step()

                     if i == (10 * num_steps - 2):
                            assert (
                                scheduler.get_last_lr()[0] != 0
                            ), F'''Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'''
              assert (
                  scheduler.get_last_lr()[0] == 0
              ), F'''Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'''
              GradientState._reset_state()







@require_cpu
class 							_UpperCamelCase    (		unittest.TestCase   ):

       def lowercase       (						self:       Dict						)				->       Any:




              """simple docstring"""


              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1						)    ,      num_processes=1						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      split_batches=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

       def lowercase       (						self:       Dict						)				->       Dict:




              """simple docstring"""


              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1						)    ,      num_processes=1						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      split_batches=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

       def lowercase       (						self:       Dict						)				->       List[str]:




              """simple docstring"""


              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      step_scheduler_with_optimizer=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

       def lowercase       (						self:       Any						)				->       int:




              """simple docstring"""


              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      step_scheduler_with_optimizer=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

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




              """simple docstring"""


              AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE						)
              debug_launcher(_SCREAMING_SNAKE_CASE						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      split_batches=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

       def lowercase       (						self:       Tuple						)				->       Optional[int]:




              """simple docstring"""


              AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE						)
              debug_launcher(_SCREAMING_SNAKE_CASE						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1    ,      split_batches=_SCREAMING_SNAKE_CASE						)    ,      num_processes=1						)

       def lowercase       (						self:       int						)				->       Tuple:




              """simple docstring"""


              AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      step_scheduler_with_optimizer=_SCREAMING_SNAKE_CASE						)						)

       def lowercase       (						self:       Dict						)				->       Dict:




              """simple docstring"""


              AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      step_scheduler_with_optimizer=_SCREAMING_SNAKE_CASE						)						)



       @require_huggingface_suite
       def lowercase       (						self:       Dict						)				->       Optional[Any]:




              """simple docstring"""


              AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE						)
              debug_launcher(partial(_SCREAMING_SNAKE_CASE    ,      num_processes=1						)						)
              debug_launcher(_SCREAMING_SNAKE_CASE						)

