






'''simple docstring'''







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). """       ,   __lowercase       ,   )
class _lowercase			(       __lowercase	):

							'''simple docstring'''





							a          = RobertaConfig
							a          = 'roberta'




							def __init__(							self:       Any    ,							UpperCamelCase__:       List[Any]  ):
														super().__init__(__A  )

														lowerCamelCase__      :			Tuple 						= RobertaEmbeddings(__A  )
														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. """       ,   __lowercase       ,   )
class _lowercase			(       __lowercase	):

							'''simple docstring'''





							a          = RobertaConfig
							a          = 'roberta'




							def __init__(							self:       List[Any]    ,							UpperCamelCase__:       Union[str, Any]  ):
														super().__init__(__A  )
														lowerCamelCase__      :			int 						= config.num_labels
														lowerCamelCase__      :			Optional[Any] 						= config.num_hidden_layers

														lowerCamelCase__      :			Any 						= DeeRobertaModel(__A  )
														lowerCamelCase__      :			Dict 						= nn.Dropout(config.hidden_dropout_prob  )
														lowerCamelCase__      :			Optional[Any] 						= nn.Linear(config.hidden_size    ,							self.config.num_labels  )




							@add_start_docstrings_to_model_forward(__A  )
							def 	lowerCamelCase_       (							self:       Optional[int]    ,							UpperCamelCase__:       Dict=None    ,							UpperCamelCase__:       List[str]=None    ,							UpperCamelCase__:       List[Any]=None    ,							UpperCamelCase__:       Optional[int]=None    ,							UpperCamelCase__:       Optional[Any]=None    ,							UpperCamelCase__:       Any=None    ,							UpperCamelCase__:       str=None    ,							UpperCamelCase__:       Tuple=-1    ,							UpperCamelCase__:       Optional[int]=False    ,							):
														lowerCamelCase__      :			List[str] 						= self.num_layers
														try:
																					lowerCamelCase__      :			str 						= self.roberta(
																					    __A    ,							attention_mask=__A    ,							token_type_ids=__A    ,							position_ids=__A    ,							head_mask=__A    ,							inputs_embeds=__A    ,							)

																					lowerCamelCase__      :			Union[str, Any] 						= outputs[1]

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

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

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

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

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









'''simple docstring'''







from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A			:    Optional[Any]					  =logging.get_logger(__name__)

_A			:    Optional[int]					  ={
    '''RWKV/rwkv-4-169m-pile''': '''https://huggingface.co/RWKV/rwkv-4-169m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-430m-pile''': '''https://huggingface.co/RWKV/rwkv-4-430m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-1b5-pile''': '''https://huggingface.co/RWKV/rwkv-4-1b5-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-3b-pile''': '''https://huggingface.co/RWKV/rwkv-4-3b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-7b-pile''': '''https://huggingface.co/RWKV/rwkv-4-7b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-14b-pile''': '''https://huggingface.co/RWKV/rwkv-4-14b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-raven-1b5''': '''https://huggingface.co/RWKV/rwkv-raven-1b5/resolve/main/config.json''',
    '''RWKV/rwkv-raven-3b''': '''https://huggingface.co/RWKV/rwkv-raven-3b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-7b''': '''https://huggingface.co/RWKV/rwkv-raven-7b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-14b''': '''https://huggingface.co/RWKV/rwkv-raven-14b/resolve/main/config.json''',
}

class _lowercase			(       _lowercase	):
							a          = """rwkv"""
							a          = {"""max_position_embeddings""": """context_length"""}




							def __init__(							self:       Tuple    ,							UpperCamelCase__:       Optional[Any]=50_277    ,							UpperCamelCase__:       Union[str, Any]=1_024    ,							UpperCamelCase__:       Tuple=4_096    ,							UpperCamelCase__:       List[Any]=32    ,							UpperCamelCase__:       Dict=None    ,							UpperCamelCase__:       Dict=None    ,							UpperCamelCase__:       int=1e-5    ,							UpperCamelCase__:       Any=0    ,							UpperCamelCase__:       str=0    ,							UpperCamelCase__:       Union[str, Any]=6    ,							UpperCamelCase__:       Optional[int]=False    ,							UpperCamelCase__:       Dict=True    ,							**UpperCamelCase__:       Dict    ,							):
														lowerCamelCase__      :			Dict 						= vocab_size
														lowerCamelCase__      :			Optional[Any] 						= context_length
														lowerCamelCase__      :			Optional[Any] 						= hidden_size
														lowerCamelCase__      :			Any 						= num_hidden_layers
														lowerCamelCase__      :			int 						= attention_hidden_size if attention_hidden_size is not None else hidden_size
														lowerCamelCase__      :			Union[str, Any] 						= intermediate_size if intermediate_size is not None else 4 * hidden_size
														lowerCamelCase__      :			List[str] 						= layer_norm_epsilon
														lowerCamelCase__      :			int 						= rescale_every
														lowerCamelCase__      :			Optional[int] 						= use_cache

														lowerCamelCase__      :			Dict 						= bos_token_id
														lowerCamelCase__      :			Any 						= eos_token_id

														super().__init__(
														    tie_word_embeddings=UpperCamelCase__    ,							bos_token_id=UpperCamelCase__    ,							eos_token_id=UpperCamelCase__    ,							**UpperCamelCase__  )

