
"""simple docstring"""

import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

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


SCREAMING_SNAKE_CASE_    					=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_    					=						'''▁'''

SCREAMING_SNAKE_CASE_    					=						{
    '''vocab_file''': '''vocab.json''',
    '''spm_file''': '''sentencepiece.bpe.model''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

SCREAMING_SNAKE_CASE_    					=						{
    '''vocab_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json''',
    },
    '''spm_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model''',
    },
    '''tokenizer_config_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json''',
    },
}

SCREAMING_SNAKE_CASE_    					=						{
    '''facebook/m2m100_418M''': 1024,
}

# fmt: off
SCREAMING_SNAKE_CASE_    					=						{
    '''m2m100''': ['''af''', '''am''', '''ar''', '''ast''', '''az''', '''ba''', '''be''', '''bg''', '''bn''', '''br''', '''bs''', '''ca''', '''ceb''', '''cs''', '''cy''', '''da''', '''de''', '''el''', '''en''', '''es''', '''et''', '''fa''', '''ff''', '''fi''', '''fr''', '''fy''', '''ga''', '''gd''', '''gl''', '''gu''', '''ha''', '''he''', '''hi''', '''hr''', '''ht''', '''hu''', '''hy''', '''id''', '''ig''', '''ilo''', '''is''', '''it''', '''ja''', '''jv''', '''ka''', '''kk''', '''km''', '''kn''', '''ko''', '''lb''', '''lg''', '''ln''', '''lo''', '''lt''', '''lv''', '''mg''', '''mk''', '''ml''', '''mn''', '''mr''', '''ms''', '''my''', '''ne''', '''nl''', '''no''', '''ns''', '''oc''', '''or''', '''pa''', '''pl''', '''ps''', '''pt''', '''ro''', '''ru''', '''sd''', '''si''', '''sk''', '''sl''', '''so''', '''sq''', '''sr''', '''ss''', '''su''', '''sv''', '''sw''', '''ta''', '''th''', '''tl''', '''tn''', '''tr''', '''uk''', '''ur''', '''uz''', '''vi''', '''wo''', '''xh''', '''yi''', '''yo''', '''zh''', '''zu'''],
    '''wmt21''': ['''en''', '''ha''', '''is''', '''ja''', '''cs''', '''ru''', '''zh''', '''de''']
}
class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							Optional[int]							      =	VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE   :							Optional[Any]							      =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE   :							List[Any]							      =	PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE   :							Any							      =	["input_ids", "attention_mask"]

      __SCREAMING_SNAKE_CASE   :							List[int]							      =	[]
      __SCREAMING_SNAKE_CASE   :							List[int]							      =	[]




      def __init__(							self   ,       lowercase_   ,       lowercase_   ,       lowercase_=None   ,       lowercase_=None   ,       lowercase_="<s>"   ,       lowercase_="</s>"   ,       lowercase_="</s>"   ,       lowercase_="<pad>"   ,       lowercase_="<unk>"   ,       lowercase_="m2m100"   ,       lowercase_ = None   ,       lowercase_=8   ,       **lowercase_   ,       )						->					None:
          UpperCAmelCase									=					{} if sp_model_kwargs is None else sp_model_kwargs

          UpperCAmelCase									=					language_codes
          UpperCAmelCase									=					FAIRSEQ_LANGUAGE_CODES[language_codes]
          UpperCAmelCase									=					{lang_code: F"__{lang_code}__" for lang_code in fairseq_language_code}

          UpperCAmelCase									=					kwargs.get('additional_special_tokens'   ,       []							)
          kwargs["additional_special_tokens"] += [
              self.get_lang_token(lowercase_							)
              for lang_code in fairseq_language_code
              if self.get_lang_token(lowercase_							) not in kwargs["additional_special_tokens"]
          ]

          super().__init__(
              src_lang=lowercase_   ,       tgt_lang=lowercase_   ,       bos_token=lowercase_   ,       eos_token=lowercase_   ,       sep_token=lowercase_   ,       unk_token=lowercase_   ,       pad_token=lowercase_   ,       language_codes=lowercase_   ,       sp_model_kwargs=self.sp_model_kwargs   ,       num_madeup_words=lowercase_   ,       **lowercase_   ,       )

          UpperCAmelCase									=					vocab_file
          UpperCAmelCase									=					load_json(lowercase_							)
          UpperCAmelCase									=					{v: k for k, v in self.encoder.items()}
          UpperCAmelCase									=					spm_file
          UpperCAmelCase									=					load_spm(lowercase_   ,       self.sp_model_kwargs							)

          UpperCAmelCase									=					len(self.encoder							)

          UpperCAmelCase									=					{
              self.get_lang_token(lowercase_							): self.encoder_size + i for i, lang_code in enumerate(lowercase_							)
          }
          UpperCAmelCase									=					{lang_code: self.encoder_size + i for i, lang_code in enumerate(lowercase_							)}
          UpperCAmelCase									=					{v: k for k, v in self.lang_token_to_id.items()}

          UpperCAmelCase									=					src_lang if src_lang is not None else 'en'
          UpperCAmelCase									=					tgt_lang
          UpperCAmelCase									=					self.get_lang_id(self._src_lang							)
          self.set_src_lang_special_tokens(self._src_lang							)

          UpperCAmelCase									=					num_madeup_words




      @property
      def     a_      (							self							)						->					int:
          return len(self.encoder							) + len(self.lang_token_to_id							)




      @property
      def     a_      (							self							)						->					str:
          return self._src_lang




      @src_lang.setter
      def     a_      (							self   ,       lowercase_							)						->					None:
          UpperCAmelCase									=					new_src_lang
          self.set_src_lang_special_tokens(self._src_lang							)




      def     a_      (							self   ,       lowercase_							)						->					List[str]:
          return self.sp_model.encode(lowercase_   ,       out_type=lowercase_							)




      def     a_      (							self   ,       lowercase_							)						->					Union[str, Any]:
          if token in self.lang_token_to_id:
              return self.lang_token_to_id[token]
          return self.encoder.get(lowercase_   ,       self.encoder[self.unk_token]							)




      def     a_      (							self   ,       lowercase_							)						->					str:
          if index in self.id_to_lang_token:
              return self.id_to_lang_token[index]
          return self.decoder.get(lowercase_   ,       self.unk_token							)




      def     a_      (							self   ,       lowercase_							)						->					List[Any]:
          UpperCAmelCase									=					[]
          UpperCAmelCase									=					''
          for token in tokens:
              # make sure that special tokens are not decoded using sentencepiece model
              if token in self.all_special_tokens:
                  out_string += self.sp_model.decode(lowercase_							) + token
                  UpperCAmelCase									=					[]
              else:
                  current_sub_tokens.append(lowercase_							)
          out_string += self.sp_model.decode(lowercase_							)
          return out_string.strip()




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = None   ,       lowercase_ = False							)						->					List[int]:

          if already_has_special_tokens:
              return super().get_special_tokens_mask(
                  token_ids_a=lowercase_   ,       token_ids_a=lowercase_   ,       already_has_special_tokens=lowercase_							)

          UpperCAmelCase									=					[1] * len(self.prefix_tokens							)
          UpperCAmelCase									=					[1] * len(self.suffix_tokens							)
          if token_ids_a is None:
              return prefix_ones + ([0] * len(lowercase_							)) + suffix_ones
          return prefix_ones + ([0] * len(lowercase_							)) + ([0] * len(lowercase_							)) + suffix_ones




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = None							)						->					List[int]:
          if token_ids_a is None:
              return self.prefix_tokens + token_ids_a + self.suffix_tokens
          # We don't expect to process pairs, but leave the pair logic for API consistency
          return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens




      def     a_      (							self							)						->					Dict:
          UpperCAmelCase									=					{self.convert_ids_to_tokens(lowercase_							): i for i in range(self.vocab_size							)}
          vocab.update(self.added_tokens_encoder							)
          return vocab




      def __getstate__(							self							)						->					Dict:
          UpperCAmelCase									=					self.__dict__.copy()
          UpperCAmelCase									=					None
          return state




      def __setstate__(							self   ,       lowercase_							)						->					None:
          UpperCAmelCase									=					d

          # for backward compatibility
          if not hasattr(self   ,       'sp_model_kwargs'							):
              UpperCAmelCase									=					{}

          UpperCAmelCase									=					load_spm(self.spm_file   ,       self.sp_model_kwargs							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = None							)						->					Tuple[str]:
          UpperCAmelCase									=					Path(lowercase_							)
          if not save_dir.is_dir():
              raise OSError(F"{save_directory} should be a directory"							)
          UpperCAmelCase									=					save_dir / (
              (filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['vocab_file']
          )
          UpperCAmelCase									=					save_dir / (
              (filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['spm_file']
          )

          save_json(self.encoder   ,       lowercase_							)

          if os.path.abspath(self.spm_file							) != os.path.abspath(lowercase_							) and os.path.isfile(self.spm_file							):
              copyfile(self.spm_file   ,       lowercase_							)
          elif not os.path.isfile(self.spm_file							):
              with open(lowercase_   ,       'wb'							) as fi:
                  UpperCAmelCase									=					self.sp_model.serialized_model_proto()
                  fi.write(lowercase_							)

          return (str(lowercase_							), str(lowercase_							))




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = "en"   ,       lowercase_ = None   ,       lowercase_ = "ro"   ,       **lowercase_   ,       )						->					BatchEncoding:
          UpperCAmelCase									=					src_lang
          UpperCAmelCase									=					tgt_lang
          self.set_src_lang_special_tokens(self.src_lang							)
          return super().prepare_seqaseq_batch(lowercase_   ,       lowercase_   ,       **lowercase_							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       **lowercase_							)						->					List[str]:
          if src_lang is None or tgt_lang is None:
              raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model'							)
          UpperCAmelCase									=					src_lang
          UpperCAmelCase									=					self(lowercase_   ,       add_special_tokens=lowercase_   ,       **lowercase_							)
          UpperCAmelCase									=					self.get_lang_id(lowercase_							)
          UpperCAmelCase									=					tgt_lang_id
          return inputs




      def     a_      (							self							)						->					str:
          self.set_src_lang_special_tokens(self.src_lang							)




      def     a_      (							self							)						->					Union[str, Any]:
          self.set_tgt_lang_special_tokens(self.tgt_lang							)




      def     a_      (							self   ,       lowercase_							)						->					None:
          UpperCAmelCase									=					self.get_lang_token(lowercase_							)
          UpperCAmelCase									=					self.lang_token_to_id[lang_token]
          UpperCAmelCase									=					[self.cur_lang_id]
          UpperCAmelCase									=					[self.eos_token_id]




      def     a_      (							self   ,       lowercase_							)						->					None:
          UpperCAmelCase									=					self.get_lang_token(lowercase_							)
          UpperCAmelCase									=					self.lang_token_to_id[lang_token]
          UpperCAmelCase									=					[self.cur_lang_id]
          UpperCAmelCase									=					[self.eos_token_id]




      def     a_      (							self   ,       lowercase_							)						->					str:
          return self.lang_code_to_token[lang]




      def     a_      (							self   ,       lowercase_							)						->					int:
          UpperCAmelCase									=					self.get_lang_token(lowercase_							)
          return self.lang_token_to_id[lang_token]




def     lowercase__						(					lowerCAmelCase				:					str ,				lowerCAmelCase				:					Dict[str, Any] ) ->      sentencepiece.SentencePieceProcessor:






    """simple docstring"""




    UpperCAmelCase									=					sentencepiece.SentencePieceProcessor(**lowerCAmelCase )
    spm.Load(str(lowerCAmelCase ) )
    return spm




def     lowercase__						(					lowerCAmelCase				:					str ) ->      Union[Dict, List]:






    """simple docstring"""




    with open(lowerCAmelCase ,				'r' ) as f:
        return json.load(lowerCAmelCase )







def     lowercase__						(					lowerCAmelCase				:					List[str] ,				lowerCAmelCase				:					str ) ->      None:






    """simple docstring"""




    with open(lowerCAmelCase ,				'w' ) as f:
        json.dump(lowerCAmelCase ,				lowerCAmelCase ,				indent=2 )





"""simple docstring"""

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      _UpperCAmelCase							(		nn.Module							):




      def __init__(							self							)						->					Any:
          super().__init__()
          UpperCAmelCase									=					nn.Linear(3   ,       4							)
          UpperCAmelCase									=					nn.BatchNormad(4							)
          UpperCAmelCase									=					nn.Linear(4   ,       5							)




      def     a_      (							self   ,       lowercase_							)						->					Any:
          return self.lineara(self.batchnorm(self.lineara(lowercase_							)							)							)

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):




      def     a_      (							self   ,       lowercase_   ,       *lowercase_   ,       **lowercase_							)						->					List[str]:
          return (args[0] + 1,) + args[1:], kwargs

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):




      def     a_      (							self   ,       lowercase_   ,       lowercase_							)						->					Union[str, Any]:
          return output + 1

class      _UpperCAmelCase							(		unittest.TestCase							):




      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					ModelHook()

          add_hook_to_module(lowercase_   ,       lowercase_							)
          self.assertEqual(test_model._hf_hook   ,       lowercase_							)
          self.assertTrue(hasattr(lowercase_   ,       '_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(lowercase_							)
          self.assertFalse(hasattr(lowercase_   ,       '_hf_hook'							)							)
          self.assertFalse(hasattr(lowercase_   ,       '_old_forward'							)							)




      def     a_      (							self							)						->					Any:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					ModelHook()

          add_hook_to_module(lowercase_   ,       lowercase_							)
          add_hook_to_module(lowercase_   ,       lowercase_   ,       append=lowercase_							)

          self.assertEqual(isinstance(test_model._hf_hook   ,       lowercase_							)   ,       lowercase_							)
          self.assertEqual(len(test_model._hf_hook.hooks							)   ,       2							)
          self.assertTrue(hasattr(lowercase_   ,       '_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(lowercase_							)
          self.assertFalse(hasattr(lowercase_   ,       '_hf_hook'							)							)
          self.assertFalse(hasattr(lowercase_   ,       '_old_forward'							)							)




      def     a_      (							self							)						->					Any:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(x + 1							)
          UpperCAmelCase									=					test_model(x + 2							)

          UpperCAmelCase									=					PreForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       lowercase_   ,       atol=1E-5							)							)

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

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

          UpperCAmelCase									=					test_model(lowercase_							)
          assert torch.allclose(lowercase_   ,       lowercase_   ,       atol=1E-5							)




      def     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(lowercase_							)

          UpperCAmelCase									=					PostForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       output + 1   ,       atol=1E-5							)							)

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

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

          UpperCAmelCase									=					test_model(lowercase_							)
          assert torch.allclose(lowercase_   ,       output + 2   ,       atol=1E-5							)




      def     a_      (							self							)						->					Tuple:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(lowercase_							)

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

          UpperCAmelCase									=					True
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertFalse(outputa.requires_grad							)




      @require_multi_gpu
      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					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
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          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(lowercase_   ,       AlignDevicesHook(io_same_device=lowercase_							)							)
          UpperCAmelCase									=					torch.randn(2   ,       3							).to(0							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       torch.device(0							)							)




      def     a_      (							self							)						->					Dict:
          UpperCAmelCase									=					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
          UpperCAmelCase									=					{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

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

          # 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
          UpperCAmelCase									=					torch.device(hook_kwargs['execution_device']							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # 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
          UpperCAmelCase									=					{
              'execution_device': 0 if torch.cuda.is_available() else 'cpu',
              'offload': True,
              'offload_buffers': True,
          }

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

          # 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'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # 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							)						->					Optional[int]:
          UpperCAmelCase									=					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
          UpperCAmelCase									=					0 if torch.cuda.is_available() else 'cpu'
          attach_align_device_hook(lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_							)

          # 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
          UpperCAmelCase									=					torch.device(lowercase_							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          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(lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       offload_buffers=lowercase_							)

          # 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'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          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							)						->					Dict:
          UpperCAmelCase									=					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
          UpperCAmelCase									=					0 if torch.cuda.is_available() else 'cpu'
          attach_align_device_hook(
              lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       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
          UpperCAmelCase									=					torch.device(lowercase_							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          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(
              lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       weights_map=model.state_dict()   ,       offload_buffers=lowercase_   ,       )

          # 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'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          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'							)							)



