



'''simple docstring'''

import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow


if is_torch_available():
				import torch

				from transformers import XLMRobertaModel

@require_sentencepiece
@require_tokenizers
@require_torch
class __snake_case(				unittest.TestCase ):


						'''simple docstring'''



						@slow
						def     __snake_case   ( self				)     ->     int:
							lowerCAmelCase   =     XLMRobertaModel.from_pretrained("""xlm-roberta-base"""				)
							lowerCAmelCase   =     torch.tensor([[0, 581, 1_0269, 83, 9_9942, 136, 6_0742, 23, 70, 8_0583, 1_8276, 2]]				)
							# The dog is cute and lives in the garden house

							lowerCAmelCase   =     torch.Size((1, 12, 768)				)  # batch_size, sequence_length, embedding_vector_dim
							lowerCAmelCase   =     torch.tensor(
							    [[-0.0_1_0_1, 0.1_2_1_8, -0.0_8_0_3, 0.0_8_0_1, 0.1_3_2_7, 0.0_7_7_6, -0.1_2_1_5, 0.2_3_8_3, 0.3_3_3_8, 0.3_1_0_6, 0.0_3_0_0, 0.0_2_5_2]]				)
							#  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')
							#  xlmr.eval()
							#  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
							with torch.no_grad():
								lowerCAmelCase   =     model(A_				)["""last_hidden_state"""].detach()
							self.assertEqual(output.shape     ,	A_				)
							# compare the actual values for a slice of last dim
							self.assertTrue(torch.allclose(output[:, :, -1]     ,	A_     ,	atol=1e-3				)				)




						@slow
						def     __snake_case   ( self				)     ->     Dict:
							lowerCAmelCase   =     XLMRobertaModel.from_pretrained("""xlm-roberta-large"""				)
							lowerCAmelCase   =     torch.tensor([[0, 581, 1_0269, 83, 9_9942, 136, 6_0742, 23, 70, 8_0583, 1_8276, 2]]				)
							# The dog is cute and lives in the garden house

							lowerCAmelCase   =     torch.Size((1, 12, 1024)				)  # batch_size, sequence_length, embedding_vector_dim
							lowerCAmelCase   =     torch.tensor(
							    [[-0.0_6_9_9, -0.0_3_1_8, 0.0_7_0_5, -0.1_2_4_1, 0.0_9_9_9, -0.0_5_2_0, 0.1_0_0_4, -0.1_8_3_8, -0.4_7_0_4, 0.1_4_3_7, 0.0_8_2_1, 0.0_1_2_6]]				)
							#  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.large')
							#  xlmr.eval()
							#  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
							with torch.no_grad():
								lowerCAmelCase   =     model(A_				)["""last_hidden_state"""].detach()
							self.assertEqual(output.shape     ,	A_				)
							# compare the actual values for a slice of last dim
							self.assertTrue(torch.allclose(output[:, :, -1]     ,	A_     ,	atol=1e-3				)				)




'''simple docstring'''

import tensorflow as tf

from ...tf_utils import shape_list

class __snake_case(				tf.keras.layers.Layer ):


						'''simple docstring'''



						def __init__( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_=1     ,	A_=False     ,	**A_				)     ->     Optional[int]:
							super().__init__(**A_				)

							lowerCAmelCase   =     vocab_size
							lowerCAmelCase   =     d_embed
							lowerCAmelCase   =     d_proj

							lowerCAmelCase   =     cutoffs + [vocab_size]
							lowerCAmelCase   =     [0] + self.cutoffs
							lowerCAmelCase   =     div_val

							lowerCAmelCase   =     self.cutoffs[0]
							lowerCAmelCase   =     len(self.cutoffs				) - 1
							lowerCAmelCase   =     self.shortlist_size + self.n_clusters
							lowerCAmelCase   =     keep_order

							lowerCAmelCase   =     []
							lowerCAmelCase   =     []


						def     __snake_case   ( self     ,	A_				)     ->     int:
							if self.n_clusters > 0:
								lowerCAmelCase   =     self.add_weight(
								    shape=(self.n_clusters, self.d_embed)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name="""cluster_weight"""				)
								lowerCAmelCase   =     self.add_weight(
								    shape=(self.n_clusters,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name="""cluster_bias"""				)

							if self.div_val == 1:
								for i in range(len(self.cutoffs				)				):
									if self.d_proj != self.d_embed:
										lowerCAmelCase   =     self.add_weight(
										    shape=(self.d_embed, self.d_proj)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_projs_._{i}'     ,	)
										self.out_projs.append(A_				)
									else:
										self.out_projs.append(A_				)
									lowerCAmelCase   =     self.add_weight(
									    shape=(self.vocab_size, self.d_embed)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._weight'     ,	)
									lowerCAmelCase   =     self.add_weight(
									    shape=(self.vocab_size,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._bias'     ,	)
									self.out_layers.append((weight, bias)				)
							else:
								for i in range(len(self.cutoffs				)				):
									lowerCAmelCase,       lowerCAmelCase   =     self.cutoff_ends[i], self.cutoff_ends[i + 1]
									lowerCAmelCase   =     self.d_embed // (self.div_val**i)

									lowerCAmelCase   =     self.add_weight(
									    shape=(d_emb_i, self.d_proj)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_projs_._{i}'				)
									self.out_projs.append(A_				)
									lowerCAmelCase   =     self.add_weight(
									    shape=(r_idx - l_idx, d_emb_i)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._weight'     ,	)
									lowerCAmelCase   =     self.add_weight(
									    shape=(r_idx - l_idx,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._bias'     ,	)
									self.out_layers.append((weight, bias)				)
							super().build(A_				)


						@staticmethod
						def     __snake_case   ( A_     ,	A_     ,	A_     ,	A_=None				)     ->     List[Any]:
							lowerCAmelCase   =     x
							if proj is not None:
								lowerCAmelCase   =     tf.einsum("""ibd,ed->ibe"""     ,	A_     ,	A_				)
							return tf.einsum("""ibd,nd->ibn"""     ,	A_     ,	A_				) + b


						@staticmethod
						def     __snake_case   ( A_     ,	A_				)     ->     Dict:
							lowerCAmelCase   =     shape_list(A_				)
							lowerCAmelCase   =     tf.range(lp_size[0]     ,	dtype=target.dtype				)
							lowerCAmelCase   =     tf.stack([r, target]     ,	1				)
							return tf.gather_nd(A_     ,	A_				)




						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_=True     ,	A_=False				)     ->     Tuple:
							lowerCAmelCase   =     0
							if self.n_clusters == 0:
								lowerCAmelCase   =     self._logit(A_     ,	self.out_layers[0][0]     ,	self.out_layers[0][1]     ,	self.out_projs[0]				)
								if target is not None:
									lowerCAmelCase   =     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=A_     ,	logits=A_				)
								lowerCAmelCase   =     tf.nn.log_softmax(A_     ,	axis=-1				)
							else:
								lowerCAmelCase   =     shape_list(A_				)
								lowerCAmelCase   =     []
								lowerCAmelCase   =     tf.zeros(hidden_sizes[:2]				)
								for i in range(len(self.cutoffs				)				):
									lowerCAmelCase,       lowerCAmelCase   =     self.cutoff_ends[i], self.cutoff_ends[i + 1]
									if target is not None:
										lowerCAmelCase   =     (target >= l_idx) & (target < r_idx)
										lowerCAmelCase   =     tf.where(A_				)
										lowerCAmelCase   =     tf.boolean_mask(A_     ,	A_				) - l_idx

									if self.div_val == 1:
										lowerCAmelCase   =     self.out_layers[0][0][l_idx:r_idx]
										lowerCAmelCase   =     self.out_layers[0][1][l_idx:r_idx]
									else:
										lowerCAmelCase   =     self.out_layers[i][0]
										lowerCAmelCase   =     self.out_layers[i][1]

									if i == 0:
										lowerCAmelCase   =     tf.concat([cur_W, self.cluster_weight]     ,	0				)
										lowerCAmelCase   =     tf.concat([cur_b, self.cluster_bias]     ,	0				)

										lowerCAmelCase   =     self._logit(A_     ,	A_     ,	A_     ,	self.out_projs[0]				)
										lowerCAmelCase   =     tf.nn.log_softmax(A_				)
										out.append(head_logprob[..., : self.cutoffs[0]]				)
										if target is not None:
											lowerCAmelCase   =     tf.boolean_mask(A_     ,	A_				)
											lowerCAmelCase   =     self._gather_logprob(A_     ,	A_				)
									else:
										lowerCAmelCase   =     self._logit(A_     ,	A_     ,	A_     ,	self.out_projs[i]				)
										lowerCAmelCase   =     tf.nn.log_softmax(A_				)
										lowerCAmelCase   =     self.cutoffs[0] + i - 1  # No probability for the head cluster
										lowerCAmelCase   =     head_logprob[..., cluster_prob_idx, None] + tail_logprob
										out.append(A_				)
										if target is not None:
											lowerCAmelCase   =     tf.boolean_mask(A_     ,	A_				)
											lowerCAmelCase   =     tf.boolean_mask(A_     ,	A_				)
											lowerCAmelCase   =     self._gather_logprob(A_     ,	A_				)
											cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
									if target is not None:
										loss += tf.scatter_nd(A_     ,	-cur_logprob     ,	shape_list(A_				)				)
								lowerCAmelCase   =     tf.concat(A_     ,	axis=-1				)

							if target is not None:
								if return_mean:
									lowerCAmelCase   =     tf.reduce_mean(A_				)
								# Add the training-time loss value to the layer using `self.add_loss()`.
								self.add_loss(A_				)

								# Log the loss as a metric (we could log arbitrary metrics,
								# including different metrics for training and inference.
								self.add_metric(A_     ,	name=self.name     ,	aggregation="""mean""" if return_mean else """"""				)

							return out