

import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu







class  lowercase				(       nn.Module			):




				def __init__(					self       ):
											super().__init__()
											snake_case_							=    nn.Linear(3      ,   4       )
											snake_case_							=    nn.BatchNormad(4       )
											snake_case_							=    nn.Linear(4      ,   5       )




				def     a							(					self      ,   snake_case       ):
											return self.lineara(self.batchnorm(self.lineara(snake_case       )       )       )







class  lowercase				(       lowercase_			):




				def     a							(					self      ,   snake_case      ,   *snake_case      ,   **snake_case       ):
											return (args[0] + 1,) + args[1:], kwargs







class  lowercase				(       lowercase_			):




				def     a							(					self      ,   snake_case      ,   snake_case       ):
											return output + 1







class  lowercase				(       unittest.TestCase			):




				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											snake_case_							=    ModelHook()

											add_hook_to_module(snake_case      ,   snake_case       )
											self.assertEqual(test_model._hf_hook      ,   snake_case       )
											self.assertTrue(hasattr(snake_case      ,   '_old_forward'       )       )

											# Check adding the hook did not change the name or the signature
											self.assertEqual(test_model.forward.__name__      ,   'forward'       )
											self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )      ,   ['x']       )

											remove_hook_from_module(snake_case       )
											self.assertFalse(hasattr(snake_case      ,   '_hf_hook'       )       )
											self.assertFalse(hasattr(snake_case      ,   '_old_forward'       )       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											snake_case_							=    ModelHook()

											add_hook_to_module(snake_case      ,   snake_case       )
											add_hook_to_module(snake_case      ,   snake_case      ,   append=snake_case       )

											self.assertEqual(isinstance(test_model._hf_hook      ,   snake_case       )      ,   snake_case       )
											self.assertEqual(len(test_model._hf_hook.hooks       )      ,   2       )
											self.assertTrue(hasattr(snake_case      ,   '_old_forward'       )       )

											# Check adding the hook did not change the name or the signature
											self.assertEqual(test_model.forward.__name__      ,   'forward'       )
											self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )      ,   ['x']       )

											remove_hook_from_module(snake_case       )
											self.assertFalse(hasattr(snake_case      ,   '_hf_hook'       )       )
											self.assertFalse(hasattr(snake_case      ,   '_old_forward'       )       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    test_model(x + 1       )
											snake_case_							=    test_model(x + 2       )

											snake_case_							=    PreForwardHook()
											add_hook_to_module(snake_case      ,   snake_case       )
											snake_case_							=    test_model(snake_case       )
											self.assertTrue(torch.allclose(snake_case      ,   snake_case      ,   atol=1e-5       )       )

											# Attaching a hook to a model when it already has one replaces, does not chain
											snake_case_							=    PreForwardHook()
											add_hook_to_module(snake_case      ,   snake_case       )
											snake_case_							=    test_model(snake_case       )
											self.assertTrue(torch.allclose(snake_case      ,   snake_case      ,   atol=1e-5       )       )

											# You need to use the sequential hook to chain two or more hooks
											snake_case_							=    SequentialHook(PreForwardHook()      ,   PreForwardHook()       )
											add_hook_to_module(snake_case      ,   snake_case       )

											snake_case_							=    test_model(snake_case       )
											assert torch.allclose(snake_case      ,   snake_case      ,   atol=1e-5       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    test_model(snake_case       )

											snake_case_							=    PostForwardHook()
											add_hook_to_module(snake_case      ,   snake_case       )
											snake_case_							=    test_model(snake_case       )
											self.assertTrue(torch.allclose(snake_case      ,   output + 1      ,   atol=1e-5       )       )

											# Attaching a hook to a model when it already has one replaces, does not chain
											snake_case_							=    PostForwardHook()
											add_hook_to_module(snake_case      ,   snake_case       )
											snake_case_							=    test_model(snake_case       )
											self.assertTrue(torch.allclose(snake_case      ,   output + 1      ,   atol=1e-5       )       )

											# You need to use the sequential hook to chain two or more hooks
											snake_case_							=    SequentialHook(PostForwardHook()      ,   PostForwardHook()       )
											add_hook_to_module(snake_case      ,   snake_case       )

											snake_case_							=    test_model(snake_case       )
											assert torch.allclose(snake_case      ,   output + 2      ,   atol=1e-5       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    test_model(snake_case       )

											snake_case_							=    PostForwardHook()
											add_hook_to_module(snake_case      ,   snake_case       )
											snake_case_							=    test_model(snake_case       )
											self.assertTrue(torch.allclose(snake_case      ,   output + 1       )       )
											self.assertTrue(outputa.requires_grad       )

											snake_case_							=    True
											snake_case_							=    test_model(snake_case       )
											self.assertFalse(outputa.requires_grad       )




				@require_multi_gpu
				def     a							(					self       ):
											snake_case_							=    ModelForTest()
											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# This will move each submodule on different devices
											add_hook_to_module(model.lineara      ,   AlignDevicesHook(execution_device=0       )       )
											add_hook_to_module(model.batchnorm      ,   AlignDevicesHook(execution_device=0       )       )
											add_hook_to_module(model.lineara      ,   AlignDevicesHook(execution_device=1       )       )

											self.assertEqual(model.lineara.weight.device      ,   torch.device(0       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device(0       )       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   torch.device(0       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device(1       )       )

											# We can still make a forward pass. The input does not need to be on any particular device
											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   torch.device(1       )       )

											# We can add a general hook to put back output on same device as input.
											add_hook_to_module(snake_case      ,   AlignDevicesHook(io_same_device=snake_case       )       )
											snake_case_							=    torch.randn(2      ,   3       ).to(0       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   torch.device(0       )       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# This will move each submodule on different devices
											snake_case_							=    {'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

											add_hook_to_module(model.lineara      ,   AlignDevicesHook(**snake_case       )       )
											add_hook_to_module(model.batchnorm      ,   AlignDevicesHook(**snake_case       )       )
											add_hook_to_module(model.lineara      ,   AlignDevicesHook(**snake_case       )       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											snake_case_							=    torch.device(hook_kwargs['execution_device']       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   snake_case       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_module(model.lineara       )
											remove_hook_from_module(model.batchnorm       )
											remove_hook_from_module(model.lineara       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											snake_case_							=    {
											    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
											    'offload': True,
											    'offload_buffers': True,
											}

											add_hook_to_module(model.lineara      ,   AlignDevicesHook(**snake_case       )       )
											add_hook_to_module(model.batchnorm      ,   AlignDevicesHook(**snake_case       )       )
											add_hook_to_module(model.lineara      ,   AlignDevicesHook(**snake_case       )       )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   torch.device('meta'       )       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_module(model.lineara       )
											remove_hook_from_module(model.batchnorm       )
											remove_hook_from_module(model.lineara       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# This will move each submodule on different devices
											snake_case_							=    0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(snake_case      ,   execution_device=snake_case      ,   offload=snake_case       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											snake_case_							=    torch.device(snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   snake_case       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(snake_case       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											attach_align_device_hook(snake_case      ,   execution_device=snake_case      ,   offload=snake_case      ,   offload_buffers=snake_case       )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   torch.device('meta'       )       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(snake_case       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )




				def     a							(					self       ):
											snake_case_							=    ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# This will move each submodule on different devices
											snake_case_							=    0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(
											    snake_case      ,   execution_device=snake_case      ,   offload=snake_case      ,   weights_map=model.state_dict()       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											snake_case_							=    torch.device(snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   snake_case       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(snake_case       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											attach_align_device_hook(
											    snake_case      ,   execution_device=snake_case      ,   offload=snake_case      ,   weights_map=model.state_dict()      ,   offload_buffers=snake_case      ,   )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device      ,   torch.device('meta'       )       )

											snake_case_							=    torch.randn(2      ,   3       )
											snake_case_							=    model(snake_case       )
											self.assertEqual(output.device      ,   snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(snake_case       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device      ,   torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device      ,   torch.device('cpu'       )       )







from typing import Dict, List, Optional

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCAmelCase		:   Any										=       logging.get_logger(__name__)


_UpperCAmelCase		:   Dict										=       {
    """nielsr/canine-s""": 2048,
}

# Unicode defines 1,114,112 total “codepoints”
_UpperCAmelCase		:   Tuple										=       111_4112

# Below: Constants defining canonical codepoints for special, pseudo-characters.
# Copied from https://github.com/google-research/language/blob/master/language/canine/special_codepoints.py
_UpperCAmelCase		:   List[str]										=       0

_UpperCAmelCase		:   Any										=       0xE000
_UpperCAmelCase		:   Dict										=       0xE001
_UpperCAmelCase		:   Optional[int]										=       0xE002
_UpperCAmelCase		:   Tuple										=       0xE003
_UpperCAmelCase		:   Tuple										=       0xE004

# Maps special codepoints to human-readable names.
_UpperCAmelCase		:   Dict[int, str]										=       {
    # Special symbols are represented using codepoints values that are valid,
    # but designated as "Private Use", meaning that they will never be assigned
    # characters by the Unicode Consortium, and are thus safe for use here.
    #
    # NOTE: Do *NOT* add any sort of [UNK_CHAR] here. They are explicitly
    # excluded and should fail with a hard error.
    CLS: "[CLS]",
    SEP: "[SEP]",
    BOS: "[BOS]",
    MASK: "[MASK]",
    PAD: "[PAD]",
    RESERVED: "[RESERVED]",
}

# Maps special codepoint human-readable names to their codepoint values.
_UpperCAmelCase		:   Dict[str, int]										=       {name: codepoint for codepoint, name in SPECIAL_CODEPOINTS.items()}







class  lowercase				(       lowercase_			):
				__SCREAMING_SNAKE_CASE					: str       			=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




				def __init__(					self      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=False      ,   snake_case=2048      ,   **snake_case      ,   ):
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else bos_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else eos_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else sep_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else cls_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else pad_token

											# Mask token behave like a normal word, i.e. include the space before it
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else mask_token

											super().__init__(
											    bos_token=snake_case      ,   eos_token=snake_case      ,   sep_token=snake_case      ,   cls_token=snake_case      ,   pad_token=snake_case      ,   mask_token=snake_case      ,   add_prefix_space=snake_case      ,   model_max_length=snake_case      ,   **snake_case      ,   )

											# Creates a mapping for looking up the IDs of special symbols.
											snake_case_							=    {}
											for codepoint, name in SPECIAL_CODEPOINTS.items():
																		snake_case_							=    codepoint

											# Creates a mapping for looking up the string forms of special symbol IDs.
											snake_case_							=    {
											    codepoint: name for name, codepoint in self._special_codepoints.items()
											}

											snake_case_							=    UNICODE_VOCAB_SIZE
											snake_case_							=    len(self._special_codepoints       )




				@property
				def     a							(					self       ):
											return self._unicode_vocab_size




				def     a							(					self      ,   snake_case       ):
											return list(snake_case       )




				def     a							(					self      ,   snake_case       ):
											try:
																		return ord(snake_case       )
											except TypeError:
																		raise ValueError(F'''invalid token: \'{token}\''''       )




				def     a							(					self      ,   snake_case       ):
											try:
																		if index in SPECIAL_CODEPOINTS:
																									return SPECIAL_CODEPOINTS[index]
																		return chr(snake_case       )
											except TypeError:
																		raise ValueError(F'''invalid id: {index}'''       )




				def     a							(					self      ,   snake_case       ):
											return "".join(snake_case       )




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											snake_case_							=    [self.sep_token_id]
											snake_case_							=    [self.cls_token_id]

											snake_case_							=    cls + token_ids_a + sep
											if token_ids_a is not None:
																		result += token_ids_a + sep
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None      ,   snake_case = False       ):
											if already_has_special_tokens:
																		return super().get_special_tokens_mask(
																		    token_ids_a=snake_case      ,   token_ids_a=snake_case      ,   already_has_special_tokens=snake_case       )

											snake_case_							=    [1] + ([0] * len(snake_case       )) + [1]
											if token_ids_a is not None:
																		result += ([0] * len(snake_case       )) + [1]
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											snake_case_							=    [self.sep_token_id]
											snake_case_							=    [self.cls_token_id]

											snake_case_							=    len(cls + token_ids_a + sep       ) * [0]
											if token_ids_a is not None:
																		result += len(token_ids_a + sep       ) * [1]
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											return ()




