



"""simple docstring"""
import torch
from torch import nn


class    SCREAMING_SNAKE_CASE__ (							nn.Module   ):

    def __init__(				self			:				Dict ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				List[str]=1 ,  lowerCAmelCase_			:				int=False):







      """simple docstring"""







      super().__init__()

      lowercase_								= n_token
      lowercase_								= d_embed
      lowercase_								= d_proj

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

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

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

      lowercase_								= nn.ModuleList()
      lowercase_								= 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(lowerCAmelCase_ ,  lowerCAmelCase_)))
          else:
            self.out_projs.append(lowerCAmelCase_)

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

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

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

      lowercase_								= keep_order

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Dict):







      """simple docstring"""







      if proj is None:
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  lowerCAmelCase_ ,  bias=lowerCAmelCase_)
      else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  proj.t().contiguous())
        lowercase_								= nn.functional.linear(lowerCAmelCase_ ,  lowerCAmelCase_ ,  bias=lowerCAmelCase_)
        # else:
        #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
        #     if bias is not None:
        #         logit = logit + bias

      return logit

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Optional[Any]=None ,  lowerCAmelCase_			:				List[Any]=False):







      """simple docstring"""







      if labels is not None:
        # Shift so that tokens < n predict n
        lowercase_								= hidden[..., :-1, :].contiguous()
        lowercase_								= labels[..., 1:].contiguous()
        lowercase_								= hidden.view(-1 ,  hidden.size(-1))
        lowercase_								= 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:
        lowercase_								= hidden.view(-1 ,  hidden.size(-1))

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

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

          weights.append(lowerCAmelCase_)
          biases.append(lowerCAmelCase_)

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

        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
        lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)

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

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

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

            if indices_i.numel() == 0:
              continue

            lowercase_								= labels.index_select(0 ,  lowerCAmelCase_) - l_idx
            lowercase_								= head_logprob.index_select(0 ,  lowerCAmelCase_)
            lowercase_								= hidden.index_select(0 ,  lowerCAmelCase_)
          else:
            lowercase_								= hidden

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

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

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

      return out




    def 					_UpperCAmelCase	(				self			:				List[Any] ,  lowerCAmelCase_			:				str):







      """simple docstring"""







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

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

          weights.append(lowerCAmelCase_)
          biases.append(lowerCAmelCase_)

        lowercase_					,    lowercase_					,    lowercase_								= weights[0], biases[0], self.out_projs[0]
        lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)

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

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

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

            lowercase_								= self._compute_logit(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
            lowercase_								= nn.functional.log_softmax(lowerCAmelCase_ ,  dim=1)

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

        return out






"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


UpperCAmelCase						:		Any			= {
    "configuration_blip": [
        "BLIP_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "BlipConfig",
        "BlipTextConfig",
        "BlipVisionConfig",
    ],
    "processing_blip": ["BlipProcessor"],
}

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase						:		Any			= ["BlipImageProcessor"]


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase						:		Optional[int]			= [
           "BLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
           "BlipModel",
           "BlipPreTrainedModel",
           "BlipForConditionalGeneration",
           "BlipForQuestionAnswering",
           "BlipVisionModel",
           "BlipTextModel",
           "BlipForImageTextRetrieval",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase						:		str			= [
           "TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFBlipModel",
           "TFBlipPreTrainedModel",
           "TFBlipForConditionalGeneration",
           "TFBlipForQuestionAnswering",
           "TFBlipVisionModel",
           "TFBlipTextModel",
           "TFBlipForImageTextRetrieval",
       ]

if TYPE_CHECKING:
       from .configuration_blip import BLIP_PRETRAINED_CONFIG_ARCHIVE_MAP, BlipConfig, BlipTextConfig, BlipVisionConfig
       from .processing_blip import BlipProcessor

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .image_processing_blip import BlipImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_blip import (
                  BLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
                  BlipForConditionalGeneration,
                  BlipForImageTextRetrieval,
                  BlipForQuestionAnswering,
                  BlipModel,
                  BlipPreTrainedModel,
                  BlipTextModel,
                  BlipVisionModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_blip import (
                  TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFBlipForConditionalGeneration,
                  TFBlipForImageTextRetrieval,
                  TFBlipForQuestionAnswering,
                  TFBlipModel,
                  TFBlipPreTrainedModel,
                  TFBlipTextModel,
                  TFBlipVisionModel,
              )

else:
       import sys

       UpperCAmelCase						:		Any			= _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

