




from __future__ import absolute_import, division, print_function, unicode_literals

from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy


@add_start_docstrings(
    'The RoBERTa Model transformer with early exiting (DeeRoBERTa). '	,  _UpperCamelCase	,  )
class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



	"""simple docstring"""


	__magic_name__   :       Tuple      	=		RobertaConfig
	__magic_name__   :       Tuple      	=		'roberta'






	def __init__(					self		:						Optional[int]					,     lowerCAmelCase		:						List[Any]	)	->       Dict:







						"""simple docstring"""







						super().__init__(lowerCAmelCase	)

						__UpperCamelCase					:	Tuple             =      RobertaEmbeddings(lowerCAmelCase	)
						self.init_weights()






@add_start_docstrings(
    'RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. '	,  _UpperCamelCase	,  )
class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



	"""simple docstring"""


	__magic_name__   :       int      	=		RobertaConfig
	__magic_name__   :       List[Any]      	=		'roberta'






	def __init__(					self		:						Any					,     lowerCAmelCase		:						Tuple	)	->       Dict:







						"""simple docstring"""







						super().__init__(lowerCAmelCase	)
						__UpperCamelCase					:	Optional[int]             =      config.num_labels
						__UpperCamelCase					:	Tuple             =      config.num_hidden_layers

						__UpperCamelCase					:	Union[str, Any]             =      DeeRobertaModel(lowerCAmelCase	)
						__UpperCamelCase					:	List[str]             =      nn.Dropout(config.hidden_dropout_prob	)
						__UpperCamelCase					:	List[str]             =      nn.Linear(config.hidden_size					,     self.config.num_labels	)






	@add_start_docstrings_to_model_forward(lowerCAmelCase	)
	def 					lowerCamelCase__     (					self		:						Tuple					,     lowerCAmelCase		:						Optional[int]=None					,     lowerCAmelCase		:						str=None					,     lowerCAmelCase		:						Optional[Any]=None					,     lowerCAmelCase		:						Tuple=None					,     lowerCAmelCase		:						int=None					,     lowerCAmelCase		:						Any=None					,     lowerCAmelCase		:						Dict=None					,     lowerCAmelCase		:						Union[str, Any]=-1					,     lowerCAmelCase		:						str=False					,     )	->       List[str]:







						"""simple docstring"""







						__UpperCamelCase					:	List[Any]             =      self.num_layers
						try:
											__UpperCamelCase					:	List[Any]             =      self.roberta(
											    lowerCAmelCase					,     attention_mask=lowerCAmelCase					,     token_type_ids=lowerCAmelCase					,     position_ids=lowerCAmelCase					,     head_mask=lowerCAmelCase					,     inputs_embeds=lowerCAmelCase					,     )

											__UpperCamelCase					:	Optional[Any]             =      outputs[1]

											__UpperCamelCase					:	int             =      self.dropout(lowerCAmelCase	)
											__UpperCamelCase					:	Any             =      self.classifier(lowerCAmelCase	)
											__UpperCamelCase					:	Any             =      (logits,) + outputs[2:]  # add hidden states and attention if they are here
						except HighwayException as e:
											__UpperCamelCase					:	str             =      e.message
											__UpperCamelCase					:	str             =      e.exit_layer
											__UpperCamelCase					:	Any             =      outputs[0]

						if not self.training:
											__UpperCamelCase					:	List[Any]             =      entropy(lowerCAmelCase	)
											__UpperCamelCase					:	str             =      []
											__UpperCamelCase					:	List[str]             =      []
						if labels is not None:
											if self.num_labels == 1:
																#  We are doing regression
																__UpperCamelCase					:	List[Any]             =      MSELoss()
																__UpperCamelCase					:	Optional[Any]             =      loss_fct(logits.view(-1	)					,     labels.view(-1	)	)
											else:
																__UpperCamelCase					:	Dict             =      CrossEntropyLoss()
																__UpperCamelCase					:	Optional[int]             =      loss_fct(logits.view(-1					,     self.num_labels	)					,     labels.view(-1	)	)

											# work with highway exits
											__UpperCamelCase					:	List[Any]             =      []
											for highway_exit in outputs[-1]:
																__UpperCamelCase					:	List[str]             =      highway_exit[0]
																if not self.training:
																					highway_logits_all.append(lowerCAmelCase	)
																					highway_entropy.append(highway_exit[2]	)
																if self.num_labels == 1:
																					#  We are doing regression
																					__UpperCamelCase					:	Tuple             =      MSELoss()
																					__UpperCamelCase					:	Union[str, Any]             =      loss_fct(highway_logits.view(-1	)					,     labels.view(-1	)	)
																else:
																					__UpperCamelCase					:	Dict             =      CrossEntropyLoss()
																					__UpperCamelCase					:	Any             =      loss_fct(highway_logits.view(-1					,     self.num_labels	)					,     labels.view(-1	)	)
																highway_losses.append(lowerCAmelCase	)

											if train_highway:
																__UpperCamelCase					:	str             =      (sum(highway_losses[:-1]	),) + outputs
																# exclude the final highway, of course
											else:
																__UpperCamelCase					:	Any             =      (loss,) + outputs
						if not self.training:
											__UpperCamelCase					:	List[str]             =      outputs + ((original_entropy, highway_entropy), exit_layer)
											if output_layer >= 0:
																__UpperCamelCase					:	Any             =      (
																    (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
																)  # use the highway of the last layer

						return outputs  # (loss), logits, (hidden_states), (attentions), entropy












import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / '''utils'''))

from test_module.custom_configuration import CustomConfig  # noqa E402


a__		     =     get_tests_dir('''fixtures/dummy-config.json''')


class    SCREAMING_SNAKE_CASE_				(       unittest.TestCase   ):



	"""simple docstring"""


	def 					lowerCamelCase__     (					self		:						str	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[Any]             =      0






	def 					lowerCamelCase__     (					self		:						Union[str, Any]	)	->       List[str]:







						"""simple docstring"""







						self.assertIsNotNone(transformers.models.auto.__spec__	)
						self.assertIsNotNone(importlib.util.find_spec("""transformers.models.auto"""	)	)






	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	int             =      AutoConfig.from_pretrained("""bert-base-uncased"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						List[Any]	)	->       Optional[Any]:







						"""simple docstring"""







						__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Any	)	->       Optional[int]:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[int]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       Any:







						"""simple docstring"""







						__UpperCamelCase					:	List[Any]             =      AutoConfig.for_model("""roberta"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						int	)	->       List[Any]:







						"""simple docstring"""







						with tempfile.TemporaryDirectory() as tmp_dir:
											# This model name contains bert and roberta, but roberta ends up being picked.
											__UpperCamelCase					:	Union[str, Any]             =      os.path.join(lowerCAmelCase					,     """fake-roberta"""	)
											os.makedirs(lowerCAmelCase					,     exist_ok=lowerCAmelCase	)
											with open(os.path.join(lowerCAmelCase					,     """config.json"""	)					,     """w"""	) as f:
																f.write(json.dumps({}	)	)
											__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
											self.assertEqual(type(lowerCAmelCase	)					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						int	)	->       str:







						"""simple docstring"""







						try:
											AutoConfig.register("""custom"""					,     lowerCAmelCase	)
											# Wrong model type will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""model"""					,     lowerCAmelCase	)
											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""bert"""					,     lowerCAmelCase	)

											# Now that the config is registered, it can be used as any other config with the auto-API
											__UpperCamelCase					:	Optional[Any]             =      CustomConfig()
											with tempfile.TemporaryDirectory() as tmp_dir:
																config.save_pretrained(lowerCAmelCase	)
																__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
																self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)

						finally:
											if "custom" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["custom"]






	def 					lowerCamelCase__     (					self		:						int	)	->       Optional[Any]:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """bert-base is not a local folder and is not a valid model identifier"""	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""bert-base"""	)






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Any:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""	):
											__UpperCamelCase					:	Optional[Any]             =      AutoConfig.from_pretrained(lowerCAmelCase					,     revision="""aaaaaa"""	)






	def 					lowerCamelCase__     (					self		:						str	)	->       List[str]:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """hf-internal-testing/no-config-test-repo does not appear to have a file named config.json."""					,     ):
											__UpperCamelCase					:	Union[str, Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/no-config-test-repo"""	)






	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       Optional[int]:







						"""simple docstring"""







						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
						# If remote code is disabled, we can't load this config.
						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)

						__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						# Test config can be reloaded.
						with tempfile.TemporaryDirectory() as tmp_dir:
											config.save_pretrained(lowerCAmelCase	)
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained(lowerCAmelCase					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(reloaded_config.__class__.__name__					,     """NewModelConfig"""	)






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







						"""simple docstring"""







						class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



							"""simple docstring"""


							__magic_name__   :       int      	=		'new-model'

						try:
											AutoConfig.register("""new-model"""					,     lowerCAmelCase	)
											# If remote code is not set, the default is to use local
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote code is disabled, we load the local one.
											__UpperCamelCase					:	Any             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote is enabled, we load from the Hub
											__UpperCamelCase					:	List[Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						finally:
											if "new-model" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["new-model"]






