text: """
   Copyright (c) 2019-2020 Intel Corporation
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
        http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
  """
  
  from typing import List
  
  import torch
  
  from nncf.algo_selector import COMPRESSION_ALGORITHMS
  from nncf.compression_method_api import CompressionAlgorithmController, CompressionLevel
  from nncf.nncf_network import NNCFNetwork
  from nncf.sparsity.base_algo import BaseSparsityAlgoBuilder, BaseSparsityAlgoController, SparseModuleInfo
  from nncf.sparsity.layers import BinaryMask
  from nncf.sparsity.magnitude.functions import WEIGHT_IMPORTANCE_FUNCTIONS, calc_magnitude_binary_mask
  from nncf.sparsity.schedulers import SPARSITY_SCHEDULERS
  
  
  @COMPRESSION_ALGORITHMS.register('magnitude_sparsity')
  class MagnitudeSparsityBuilder(BaseSparsityAlgoBuilder):
      def create_weight_sparsifying_operation(self, module):
          return BinaryMask(module.weight.size())
  
      def build_controller(self, target_model: NNCFNetwork) -> CompressionAlgorithmController:
          params = self.config.get("params", {})
          return MagnitudeSparsityController(target_model, self._sparsified_module_info,
                                             self.config,
                                             params.get('weight_importance', 'normed_abs'))
  
  
  class MagnitudeSparsityController(BaseSparsityAlgoController):
      def __init__(self, target_model: NNCFNetwork,
                   sparsified_module_info: List[SparseModuleInfo],
                   config, weight_importance: str):
          super().__init__(target_model, sparsified_module_info)
          self.config = config
          params = self.config.get("params", {})
          self.weight_importance = WEIGHT_IMPORTANCE_FUNCTIONS.get(weight_importance)
          sparsity_level_mode = params.get("sparsity_level_setting_mode", "global")
          self._scheduler = None
          if sparsity_level_mode == 'global':
              scheduler_cls = SPARSITY_SCHEDULERS.get(params.get("schedule", "polynomial"))
              self._scheduler = scheduler_cls(self, params)
  
      def statistics(self):
          stats = super().statistics()
          stats['sparsity_threshold'] =\
               self._select_threshold(self.sparsity_rate_for_sparsified_modules, self.sparsified_module_info)
          return stats
  
      def freeze(self):
          for layer in self.sparsified_module_info:
              layer.operand.frozen = True
  
      def set_sparsity_level(self, sparsity_level, target_sparsified_module_info: SparseModuleInfo = None):
          if sparsity_level >= 1 or sparsity_level < 0:
              raise AttributeError(
                  'Sparsity level should be within interval [0,1), actual value to set is: {}'.format(sparsity_level))
          if target_sparsified_module_info is None:
              target_sparsified_module_info_list = self.sparsified_module_info # List[SparseModuleInfo]
          else:
              target_sparsified_module_info_list = [target_sparsified_module_info]
          threshold = self._select_threshold(sparsity_level, target_sparsified_module_info_list)
          self._set_masks_for_threshold(threshold, target_sparsified_module_info_list)
          self.run_batchnorm_adaptation(self.config)
  
      def _select_threshold(self, sparsity_level, target_sparsified_module_info_list):
          all_weights = self._collect_all_weights(target_sparsified_module_info_list)
          if not all_weights:
              return 0.0
          all_weights_tensor, _ = torch.cat(all_weights).sort()
          threshold = all_weights_tensor[int((all_weights_tensor.size(0) - 1) * sparsity_level)].item()
          return threshold
  
      def _set_masks_for_threshold(self, threshold_val, target_sparsified_module_info_list):
          for layer in target_sparsified_module_info_list:
              if not layer.operand.frozen:
                  layer.operand.binary_mask = calc_magnitude_binary_mask(layer.module.weight,
                                                                         self.weight_importance,
                                                                         threshold_val)
  
  
      def _collect_all_weights(self, target_sparsified_module_info_list: List[SparseModuleInfo]):
          all_weights = []
          for minfo in target_sparsified_module_info_list:
              all_weights.append(self.weight_importance(minfo.module.weight).view(-1))
          return all_weights
  
      def create_weight_sparsifying_operation(self, module):
          return BinaryMask(module.weight.size())
  
      def compression_level(self) -> CompressionLevel:
          if self.scheduler is not None:
              return self.scheduler.compression_level()
          return CompressionLevel.NONE
  