
import torch
from torch import nn





class       __A(				nn.Module       ):

					def __init__(	self					:     List[str]   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     Dict   ,					__UpperCamelCase					:     Union[str, Any]=1   ,					__UpperCamelCase					:     Any=False   ):
												super().__init__()

												lowerCamelCase_							  =					n_token
												lowerCamelCase_							  =					d_embed
												lowerCamelCase_							  =					d_proj

												lowerCamelCase_							  =					cutoffs + [n_token]
												lowerCamelCase_							  =					[0] + self.cutoffs
												lowerCamelCase_							  =					div_val

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

												if self.n_clusters > 0:
																			lowerCamelCase_							  =					nn.Parameter(torch.zeros(self.n_clusters   ,					self.d_embed   )   )
																			lowerCamelCase_							  =					nn.Parameter(torch.zeros(self.n_clusters   )   )

												lowerCamelCase_							  =					nn.ModuleList()
												lowerCamelCase_							  =					nn.ParameterList()

												if div_val == 1:
																			for i in range(len(self.cutoffs   )   ):
																										if d_proj != d_embed:
																																	self.out_projs.append(nn.Parameter(torch.FloatTensor(__UpperCamelCase   ,					__UpperCamelCase   )   )   )
																										else:
																																	self.out_projs.append(__UpperCamelCase   )

																			self.out_layers.append(nn.Linear(__UpperCamelCase   ,					__UpperCamelCase   )   )
												else:
																			for i in range(len(self.cutoffs   )   ):
																										lowerCamelCase_  ,   lowerCamelCase_							  =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
																										lowerCamelCase_							  =					d_embed // (div_val**i)

																										self.out_projs.append(nn.Parameter(torch.FloatTensor(__UpperCamelCase   ,					__UpperCamelCase   )   )   )

																										self.out_layers.append(nn.Linear(__UpperCamelCase   ,					r_idx - l_idx   )   )

												lowerCamelCase_							  =					keep_order

					def    lowercase__      (	self					:     int   ,					__UpperCamelCase					:     List[str]   ,					__UpperCamelCase					:     str   ,					__UpperCamelCase					:     List[Any]   ,					__UpperCamelCase					:     List[Any]   ):
												if proj is None:
																			lowerCamelCase_							  =					nn.functional.linear(__UpperCamelCase   ,					__UpperCamelCase   ,					bias=__UpperCamelCase   )
												else:
																			# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
																			lowerCamelCase_							  =					nn.functional.linear(__UpperCamelCase   ,					proj.t().contiguous()   )
																			lowerCamelCase_							  =					nn.functional.linear(__UpperCamelCase   ,					__UpperCamelCase   ,					bias=__UpperCamelCase   )
																			# else:
																			#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
																			#     if bias is not None:
																			#         logit = logit + bias

												return logit

					def    lowercase__      (	self					:     Optional[int]   ,					__UpperCamelCase					:     Optional[Any]   ,					__UpperCamelCase					:     List[str]=None   ,					__UpperCamelCase					:     Union[str, Any]=False   ):

												if labels is not None:
																			# Shift so that tokens < n predict n
																			lowerCamelCase_							  =					hidden[..., :-1, :].contiguous()
																			lowerCamelCase_							  =					labels[..., 1:].contiguous()
																			lowerCamelCase_							  =					hidden.view(-1   ,					hidden.size(-1   )   )
																			lowerCamelCase_							  =					labels.view(-1   )
																			if hidden.size(0   ) != labels.size(0   ):
																										raise RuntimeError("""Input and labels should have the same size in the batch dimension."""   )
												else:
																			lowerCamelCase_							  =					hidden.view(-1   ,					hidden.size(-1   )   )

												if self.n_clusters == 0:
																			lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					self.out_layers[0].weight   ,					self.out_layers[0].bias   ,					self.out_projs[0]   )
																			if labels is not None:
																										lowerCamelCase_							  =					labels != -1_0_0
																										lowerCamelCase_							  =					torch.zeros_like(__UpperCamelCase   ,					dtype=hidden.dtype   ,					device=hidden.device   )
																										lowerCamelCase_							  =					(
																										    -nn.functional.log_softmax(__UpperCamelCase   ,					dim=-1   )[mask].gather(1   ,					labels[mask].unsqueeze(1   )   ).squeeze(1   )
																										)
																			else:
																										lowerCamelCase_							  =					nn.functional.log_softmax(__UpperCamelCase   ,					dim=-1   )
												else:
																			# construct weights and biases
																			lowerCamelCase_  ,   lowerCamelCase_							  =					[], []
																			for i in range(len(self.cutoffs   )   ):
																										if self.div_val == 1:
																																	lowerCamelCase_  ,   lowerCamelCase_							  =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
																																	lowerCamelCase_							  =					self.out_layers[0].weight[l_idx:r_idx]
																																	lowerCamelCase_							  =					self.out_layers[0].bias[l_idx:r_idx]
																										else:
																																	lowerCamelCase_							  =					self.out_layers[i].weight
																																	lowerCamelCase_							  =					self.out_layers[i].bias

																										if i == 0:
																																	lowerCamelCase_							  =					torch.cat([weight_i, self.cluster_weight]   ,					dim=0   )
																																	lowerCamelCase_							  =					torch.cat([bias_i, self.cluster_bias]   ,					dim=0   )

																										weights.append(__UpperCamelCase   )
																										biases.append(__UpperCamelCase   )

																			lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_							  =					weights[0], biases[0], self.out_projs[0]

																			lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   )
																			lowerCamelCase_							  =					nn.functional.log_softmax(__UpperCamelCase   ,					dim=1   )

																			if labels is None:
																										lowerCamelCase_							  =					hidden.new_empty((head_logit.size(0   ), self.n_token)   )
																			else:
																										lowerCamelCase_							  =					torch.zeros_like(__UpperCamelCase   ,					dtype=hidden.dtype   ,					device=hidden.device   )

																			lowerCamelCase_							  =					0
																			lowerCamelCase_							  =					[0] + self.cutoffs
																			for i in range(len(__UpperCamelCase   ) - 1   ):
																										lowerCamelCase_  ,   lowerCamelCase_							  =					cutoff_values[i], cutoff_values[i + 1]

																										if labels is not None:
																																	lowerCamelCase_							  =					(labels >= l_idx) & (labels < r_idx)
																																	lowerCamelCase_							  =					mask_i.nonzero().squeeze()

																																	if indices_i.numel() == 0:
																																								continue

																																	lowerCamelCase_							  =					labels.index_select(0   ,					__UpperCamelCase   ) - l_idx
																																	lowerCamelCase_							  =					head_logprob.index_select(0   ,					__UpperCamelCase   )
																																	lowerCamelCase_							  =					hidden.index_select(0   ,					__UpperCamelCase   )
																										else:
																																	lowerCamelCase_							  =					hidden

																										if i == 0:
																																	if labels is not None:
																																								lowerCamelCase_							  =					head_logprob_i.gather(1   ,					target_i[:, None]   ).squeeze(1   )
																																	else:
																																								lowerCamelCase_							  =					head_logprob[:, : self.cutoffs[0]]
																										else:
																																	lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_							  =					weights[i], biases[i], self.out_projs[i]

																																	lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   )
																																	lowerCamelCase_							  =					nn.functional.log_softmax(__UpperCamelCase   ,					dim=1   )
																																	lowerCamelCase_							  =					self.cutoffs[0] + i - 1  # No probability for the head cluster
																																	if labels is not None:
																																								lowerCamelCase_							  =					head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																																								    1   ,					target_i[:, None]   ).squeeze(1   )
																																	else:
																																								lowerCamelCase_							  =					head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																																								lowerCamelCase_							  =					logprob_i

																										if labels is not None:
																																	if (hasattr(self   ,					"""keep_order"""   ) and self.keep_order) or keep_order:
																																								out.index_copy_(0   ,					__UpperCamelCase   ,					-logprob_i   )
																																	else:
																																								out[offset : offset + logprob_i.size(0   )].copy_(-logprob_i   )
																																	offset += logprob_i.size(0   )

												return out



					def    lowercase__      (	self					:     Optional[int]   ,					__UpperCamelCase					:     int   ):
												if self.n_clusters == 0:
																			lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					self.out_layers[0].weight   ,					self.out_layers[0].bias   ,					self.out_projs[0]   )
																			return nn.functional.log_softmax(__UpperCamelCase   ,					dim=-1   )
												else:
																			# construct weights and biases
																			lowerCamelCase_  ,   lowerCamelCase_							  =					[], []
																			for i in range(len(self.cutoffs   )   ):
																										if self.div_val == 1:
																																	lowerCamelCase_  ,   lowerCamelCase_							  =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
																																	lowerCamelCase_							  =					self.out_layers[0].weight[l_idx:r_idx]
																																	lowerCamelCase_							  =					self.out_layers[0].bias[l_idx:r_idx]
																										else:
																																	lowerCamelCase_							  =					self.out_layers[i].weight
																																	lowerCamelCase_							  =					self.out_layers[i].bias

																										if i == 0:
																																	lowerCamelCase_							  =					torch.cat([weight_i, self.cluster_weight]   ,					dim=0   )
																																	lowerCamelCase_							  =					torch.cat([bias_i, self.cluster_bias]   ,					dim=0   )

																										weights.append(__UpperCamelCase   )
																										biases.append(__UpperCamelCase   )

																			lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_							  =					weights[0], biases[0], self.out_projs[0]
																			lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   )

																			lowerCamelCase_							  =					hidden.new_empty((head_logit.size(0   ), self.n_token)   )
																			lowerCamelCase_							  =					nn.functional.log_softmax(__UpperCamelCase   ,					dim=1   )

																			lowerCamelCase_							  =					[0] + self.cutoffs
																			for i in range(len(__UpperCamelCase   ) - 1   ):
																										lowerCamelCase_  ,   lowerCamelCase_							  =					cutoff_values[i], cutoff_values[i + 1]

																										if i == 0:
																																	lowerCamelCase_							  =					head_logprob[:, : self.cutoffs[0]]
																										else:
																																	lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_							  =					weights[i], biases[i], self.out_projs[i]

																																	lowerCamelCase_							  =					self._compute_logit(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   )
																																	lowerCamelCase_							  =					nn.functional.log_softmax(__UpperCamelCase   ,					dim=1   )

																																	lowerCamelCase_							  =					head_logprob[:, -i] + tail_logprob_i
																																	lowerCamelCase_							  =					logprob_i

																			return out







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowercase        =  {
    '''configuration_rag''': ['''RagConfig'''],
    '''retrieval_rag''': ['''RagRetriever'''],
    '''tokenization_rag''': ['''RagTokenizer'''],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowercase        =  [
		    '''RagModel''',
		    '''RagPreTrainedModel''',
		    '''RagSequenceForGeneration''',
		    '''RagTokenForGeneration''',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowercase        =  [
		    '''TFRagModel''',
		    '''TFRagPreTrainedModel''',
		    '''TFRagSequenceForGeneration''',
		    '''TFRagTokenForGeneration''',
		]


if TYPE_CHECKING:
		from .configuration_rag import RagConfig
		from .retrieval_rag import RagRetriever
		from .tokenization_rag import RagTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_rag import (
				    TFRagModel,
				    TFRagPreTrainedModel,
				    TFRagSequenceForGeneration,
				    TFRagTokenForGeneration,
				)

else:
		import sys

		lowercase        =  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





