






import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
__magic_name__   :  Optional[Any]        =							logging.get_logger(__name__)
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase=False) ->			Union[str, Any]:
      """simple docstring"""

      UpperCamelCase					       =  []

      # fmt: off
      # stem:
      rename_keys.append(('cls_token', 'vit.embeddings.cls_token'))
      rename_keys.append(('pos_embed', 'vit.embeddings.position_embeddings'))

      rename_keys.append(('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'))
      rename_keys.append(('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'))

      # backbone
      rename_keys.append(('patch_embed.backbone.stem.conv.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight'))
      rename_keys.append(('patch_embed.backbone.stem.norm.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight'))
      rename_keys.append(('patch_embed.backbone.stem.norm.bias', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias'))

      for stage_idx in range(len(config.backbone_config.depths)):
            for layer_idx in range(config.backbone_config.depths[stage_idx]):
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight'))
                  rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias'))

            rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight'))
            rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight'))
            rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias'))

      # transformer encoder
      for i in range(config.num_hidden_layers):
            # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
            rename_keys.append((F'blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight'))
            rename_keys.append((F'blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias'))
            rename_keys.append((F'blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight'))
            rename_keys.append((F'blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias'))
            rename_keys.append((F'blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight'))
            rename_keys.append((F'blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias'))
            rename_keys.append((F'blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight'))
            rename_keys.append((F'blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias'))
            rename_keys.append((F'blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight'))
            rename_keys.append((F'blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias'))

      if base_model:
            # layernorm + pooler
            rename_keys.extend(
                [
                    ('norm.weight', 'layernorm.weight'),
                    ('norm.bias', 'layernorm.bias'),
                    ('pre_logits.fc.weight', 'pooler.dense.weight'),
                    ('pre_logits.fc.bias', 'pooler.dense.bias'),
                ])

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            UpperCamelCase					       =  [(pair[0], pair[1][4:]) if pair[1].startswith('vit') else pair for pair in rename_keys]
      else:
            # layernorm + classification head
            rename_keys.extend(
                [
                    ('norm.weight', 'vit.layernorm.weight'),
                    ('norm.bias', 'vit.layernorm.bias'),
                    ('head.weight', 'classifier.weight'),
                    ('head.bias', 'classifier.bias'),
                ])
      # fmt: on

      return rename_keys
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase=False) ->			Dict:
      """simple docstring"""

      for i in range(config.num_hidden_layers):
            if base_model:
                  UpperCamelCase					       =  ''
            else:
                  UpperCamelCase					       =  'vit.'
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            UpperCamelCase					       =  state_dict.pop(F'blocks.{i}.attn.qkv.weight')
            UpperCamelCase					       =  state_dict.pop(F'blocks.{i}.attn.qkv.bias')
            # next, add query, keys and values (in that order) to the state dict
            UpperCamelCase					       =  in_proj_weight[
                : config.hidden_size, :
            ]
            UpperCamelCase					       =  in_proj_bias[: config.hidden_size]
            UpperCamelCase					       =  in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            UpperCamelCase					       =  in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            UpperCamelCase					       =  in_proj_weight[
                -config.hidden_size :, :
            ]
            UpperCamelCase					       =  in_proj_bias[-config.hidden_size :]
def 		lowercase__  (				_UpperCamelCase) ->			Any:
      """simple docstring"""

      UpperCamelCase					       =  ['head.weight', 'head.bias']
      for k in ignore_keys:
            state_dict.pop(_UpperCamelCase       ,							_UpperCamelCase)
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			str:
      """simple docstring"""

      UpperCamelCase					       =  dct.pop(_UpperCamelCase)
      UpperCamelCase					       =  val
def 		lowercase__  (				) ->			Tuple:
      """simple docstring"""

      UpperCamelCase					       =  'http://images.cocodataset.org/val2017/000000039769.jpg'
      UpperCamelCase					       =  Image.open(requests.get(_UpperCamelCase       ,							stream=_UpperCamelCase).raw)
      return im
@torch.no_grad()
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase=False) ->			Optional[int]:
      """simple docstring"""

      UpperCamelCase					       =  BitConfig(
          global_padding='same'       ,							layer_type='bottleneck'       ,							depths=(3, 4, 9)       ,							out_features=['stage3']       ,							embedding_dynamic_padding=_UpperCamelCase       ,							)
      UpperCamelCase					       =  ViTHybridConfig(backbone_config=_UpperCamelCase       ,							image_size=3_84       ,							num_labels=10_00)
      UpperCamelCase					       =  False

      # load original model from timm
      UpperCamelCase					       =  timm.create_model(_UpperCamelCase       ,							pretrained=_UpperCamelCase)
      timm_model.eval()

      # load state_dict of original model, remove and rename some keys
      UpperCamelCase					       =  timm_model.state_dict()
      if base_model:
            remove_classification_head_(_UpperCamelCase)
      UpperCamelCase					       =  create_rename_keys(_UpperCamelCase       ,							_UpperCamelCase)
      for src, dest in rename_keys:
            rename_key(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)
      read_in_q_k_v(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)

      UpperCamelCase					       =  'huggingface/label-files'
      UpperCamelCase					       =  'imagenet-1k-id2label.json'
      UpperCamelCase					       =  json.load(open(hf_hub_download(_UpperCamelCase       ,							_UpperCamelCase       ,							repo_type='dataset')       ,							'r'))
      UpperCamelCase					       =  {int(_UpperCamelCase): v for k, v in idalabel.items()}
      UpperCamelCase					       =  idalabel
      UpperCamelCase					       =  {v: k for k, v in idalabel.items()}

      # load HuggingFace model
      if vit_name[-5:] == "in21k":
            UpperCamelCase					       =  ViTHybridModel(_UpperCamelCase).eval()
      else:
            UpperCamelCase					       =  ViTHybridForImageClassification(_UpperCamelCase).eval()
      model.load_state_dict(_UpperCamelCase)

      # create image processor
      UpperCamelCase					       =  create_transform(**resolve_data_config({}       ,							model=_UpperCamelCase))
      UpperCamelCase					       =  transform.transforms

      UpperCamelCase					       =  {
          'bilinear': PILImageResampling.BILINEAR,
          'bicubic': PILImageResampling.BICUBIC,
          'nearest': PILImageResampling.NEAREST,
      }

      UpperCamelCase					       =  ViTHybridImageProcessor(
          do_resize=_UpperCamelCase       ,							size={'shortest_edge': timm_transforms[0].size}       ,							resample=pillow_resamplings[timm_transforms[0].interpolation.value]       ,							do_center_crop=_UpperCamelCase       ,							crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]}       ,							do_normalize=_UpperCamelCase       ,							image_mean=timm_transforms[-1].mean.tolist()       ,							image_std=timm_transforms[-1].std.tolist()       ,							)

      UpperCamelCase					       =  prepare_img()
      UpperCamelCase					       =  transform(_UpperCamelCase).unsqueeze(0)
      UpperCamelCase					       =  processor(_UpperCamelCase       ,							return_tensors='pt').pixel_values

      # verify pixel values
      assert torch.allclose(_UpperCamelCase       ,							_UpperCamelCase)

      # verify logits
      with torch.no_grad():
            UpperCamelCase					       =  model(_UpperCamelCase)
            UpperCamelCase					       =  outputs.logits

      print('Predicted class:'       ,							logits.argmax(-1).item())
      if base_model:
            UpperCamelCase					       =  timm_model.forward_features(_UpperCamelCase)
            assert timm_pooled_output.shape == outputs.pooler_output.shape
            assert torch.allclose(_UpperCamelCase       ,							outputs.pooler_output       ,							atol=1e-3)
      else:
            UpperCamelCase					       =  timm_model(_UpperCamelCase)
            assert timm_logits.shape == outputs.logits.shape
            assert torch.allclose(_UpperCamelCase       ,							outputs.logits       ,							atol=1e-3)
      print('Looks ok!')

      if pytorch_dump_folder_path is not None:
            Path(_UpperCamelCase).mkdir(exist_ok=_UpperCamelCase)
            print(F'Saving model {vit_name} to {pytorch_dump_folder_path}')
            model.save_pretrained(_UpperCamelCase)
            print(F'Saving processor to {pytorch_dump_folder_path}')
            processor.save_pretrained(_UpperCamelCase)

      if push_to_hub:
            print(F'Pushing model and processor to the hub {vit_name}')
            model.push_to_hub(F'ybelkada/{vit_name}')
            processor.push_to_hub(F'ybelkada/{vit_name}')


if __name__ == "__main__":
       __magic_name__   :  Tuple        =							argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '''--vit_name''',
           default='''vit_base_r50_s16_384''',
           type=str,
           help='''Name of the hybrid ViT timm model you\'d like to convert.''',
       )
       parser.add_argument(
           '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
       )
       parser.add_argument(
           '''--push_to_hub''', action='''store_true''', help='''Whether to upload the model to the HuggingFace hub.'''
       )

       __magic_name__   :  int        =							parser.parse_args()
       convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)











import torch
from torch import nn





class 					A__     (     nn.Module ):

     '''simple docstring'''






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






           """simple docstring"""




           super().__init__()

           UpperCamelCase					       =  n_token
           UpperCamelCase					       =  d_embed
           UpperCamelCase					       =  d_proj

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

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

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

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

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

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

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

           UpperCamelCase					       =  keep_order







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




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

           return logit







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=False ):






           """simple docstring"""




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

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

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

                       weights.append(_SCREAMING_SNAKE_CASE )
                       biases.append(_SCREAMING_SNAKE_CASE )

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

                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

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

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

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

                             if indices_i.numel() == 0:
                                   continue

                             UpperCamelCase					       =  labels.index_select(0   ,						_SCREAMING_SNAKE_CASE ) - l_idx
                             UpperCamelCase					       =  head_logprob.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  hidden.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                       else:
                             UpperCamelCase					       =  hidden

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

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

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

           return out







     def   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




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

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

                       weights.append(_SCREAMING_SNAKE_CASE )
                       biases.append(_SCREAMING_SNAKE_CASE )

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

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

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

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

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

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

                 return out



