
import math

import torch
from torch import nn
from torch.nn import init

from .binarizer import MagnitudeBinarizer, ThresholdBinarizer, TopKBinarizer





class   __UpperCamelCase				(    nn.Linear   ):






     """simple docstring"""







     def __init__(							self    :    Tuple      ,		_A    :    int      ,		_A    :    int      ,		_A    :    bool = True      ,		_A    :    str = "constant"      ,		_A    :    float = 0.0      ,		_A    :    str = "topK"      ,		):







       """simple docstring"""




       super(_A      ,		self				).__init__(in_features=_A      ,		out_features=_A      ,		bias=_A				)
       assert pruning_method in ["topK", "threshold", "sigmoied_threshold", "magnitude", "l0"]
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						pruning_method

       if self.pruning_method in ["topK", "threshold", "sigmoied_threshold", "l0"]:
         __SCREAMING_SNAKE_CASE     :					Any			=						mask_scale
         __SCREAMING_SNAKE_CASE     :					int			=						mask_init
         __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						nn.Parameter(torch.empty(self.weight.size()				)				)
         self.init_mask()



     def 				UpperCAmelCase__    (							self    :    Union[str, Any]				):







       """simple docstring"""




       if self.mask_init == "constant":
         init.constant_(self.mask_scores      ,		val=self.mask_scale				)
       elif self.mask_init == "uniform":
         init.uniform_(self.mask_scores      ,		a=-self.mask_scale      ,		b=self.mask_scale				)
       elif self.mask_init == "kaiming":
         init.kaiming_uniform_(self.mask_scores      ,		a=math.sqrt(5				)				)




     def 				UpperCAmelCase__    (							self    :    List[Any]      ,		_A    :    torch.tensor      ,		_A    :    float				):







       """simple docstring"""




       if self.pruning_method == "topK":
         __SCREAMING_SNAKE_CASE     :					Any			=						TopKBinarizer.apply(self.mask_scores      ,		_A				)
       elif self.pruning_method in ["threshold", "sigmoied_threshold"]:
         __SCREAMING_SNAKE_CASE     :					Dict			=						'''sigmoied''' in self.pruning_method
         __SCREAMING_SNAKE_CASE     :					Tuple			=						ThresholdBinarizer.apply(self.mask_scores      ,		_A      ,		_A				)
       elif self.pruning_method == "magnitude":
         __SCREAMING_SNAKE_CASE     :					Any			=						MagnitudeBinarizer.apply(self.weight      ,		_A				)
       elif self.pruning_method == "l0":
         __SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE     :					Dict			=						-0.1, 1.1, 2 / 3
         if self.training:
           __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						torch.zeros_like(self.mask_scores				).uniform_().clamp(0.00_01      ,		0.99_99				)
           __SCREAMING_SNAKE_CASE     :					List[str]			=						torch.sigmoid((u.log() - (1 - u).log() + self.mask_scores) / b				)
         else:
           __SCREAMING_SNAKE_CASE     :					Dict			=						torch.sigmoid(self.mask_scores				)
         __SCREAMING_SNAKE_CASE     :					Tuple			=						s * (r - l) + l
         __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						s_bar.clamp(min=0.0      ,		max=1.0				)
       # Mask weights with computed mask
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						mask * self.weight
       # Compute output (linear layer) with masked weights
       return nn.functional.linear(_A      ,		_A      ,		self.bias				)


import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


lowercase_											=	"""
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
"""





class   __UpperCamelCase				(    unittest.TestCase				,		lowerCAmelCase__   ):






     """simple docstring"""







     def 				UpperCAmelCase__    (							self    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						load_tool('''summarization'''				)
       self.tool.setup()
       __SCREAMING_SNAKE_CASE     :					Any			=						load_tool('''summarization'''      ,		remote=_A				)



     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self.tool(_A				)
       self.assertEqual(
           _A      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)



     def 				UpperCAmelCase__    (							self    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						self.remote_tool(_A				)
       self.assertEqual(
           _A      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)



     def 				UpperCAmelCase__    (							self    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self.tool(text=_A				)
       self.assertEqual(
           _A      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)




     def 				UpperCAmelCase__    (							self    :    Optional[int]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						self.remote_tool(text=_A				)
       self.assertEqual(
           _A      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)
