





"""simple docstring"""






import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE=False							)		->       Any:






       """simple docstring"""




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

       # projection layer + position embeddings
       rename_keys.extend(
           [
               ('module.cls_token', 'vit.embeddings.cls_token'),
               ('module.patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'),
               ('module.patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'),
               ('module.pos_embed', 'vit.embeddings.position_embeddings'),
           ]							)

       if base_model:
              # layernorm + pooler
              rename_keys.extend(
                  [
                      ('module.norm.weight', 'layernorm.weight'),
                      ('module.norm.bias', 'layernorm.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'),
                  ]							)

       return rename_keys

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE=False							)		->       List[str]:






       """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"""module.blocks.{i}.attn.qkv.weight"""							)
              _UpperCAmelCase     =     state_dict.pop(F"""module.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       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     ['head.weight', 'head.bias']
       for k in ignore_keys:
              state_dict.pop(_A,_A							)

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       int:






       """simple docstring"""




       _UpperCAmelCase     =     [
           'module.fc.fc1.weight',
           'module.fc.fc1.bias',
           'module.fc.bn1.weight',
           'module.fc.bn1.bias',
           'module.fc.bn1.running_mean',
           'module.fc.bn1.running_var',
           'module.fc.bn1.num_batches_tracked',
           'module.fc.fc2.weight',
           'module.fc.fc2.bias',
           'module.fc.bn2.weight',
           'module.fc.bn2.bias',
           'module.fc.bn2.running_mean',
           'module.fc.bn2.running_var',
           'module.fc.bn2.num_batches_tracked',
           'module.fc.fc3.weight',
           'module.fc.fc3.bias',
       ]
       for k in ignore_keys:
              state_dict.pop(_A,_A							)

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       int:






       """simple docstring"""




       _UpperCAmelCase     =     dct.pop(_A							)
       _UpperCAmelCase     =     val

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       Any:






       """simple docstring"""




       _UpperCAmelCase     =     ViTMSNConfig()
       _UpperCAmelCase     =     1000

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

       if "s16" in checkpoint_url:
              _UpperCAmelCase     =     384
              _UpperCAmelCase     =     1536
              _UpperCAmelCase     =     6
       elif "l16" in checkpoint_url:
              _UpperCAmelCase     =     1024
              _UpperCAmelCase     =     4096
              _UpperCAmelCase     =     24
              _UpperCAmelCase     =     16
              _UpperCAmelCase     =     0.1
       elif "b4" in checkpoint_url:
              _UpperCAmelCase     =     4
       elif "l7" in checkpoint_url:
              _UpperCAmelCase     =     7
              _UpperCAmelCase     =     1024
              _UpperCAmelCase     =     4096
              _UpperCAmelCase     =     24
              _UpperCAmelCase     =     16
              _UpperCAmelCase     =     0.1

       _UpperCAmelCase     =     ViTMSNModel(_A							)

       _UpperCAmelCase     =     torch.hub.load_state_dict_from_url(_A,map_location='cpu'							)['target_encoder']

       _UpperCAmelCase     =     ViTImageProcessor(size=config.image_size							)

       remove_projection_head(_A							)
       _UpperCAmelCase     =     create_rename_keys(_A,base_model=_A							)

       for src, dest in rename_keys:
              rename_key(_A,_A,_A							)
       read_in_q_k_v(_A,_A,base_model=_A							)

       model.load_state_dict(_A							)
       model.eval()

       _UpperCAmelCase     =     'http://images.cocodataset.org/val2017/000000039769.jpg'

       _UpperCAmelCase     =     Image.open(requests.get(_A,stream=_A							).raw							)
       _UpperCAmelCase     =     ViTImageProcessor(
           size=config.image_size,image_mean=_A,image_std=_A							)
       _UpperCAmelCase     =     image_processor(images=_A,return_tensors='pt'							)

       # forward pass
       torch.manual_seed(2							)
       _UpperCAmelCase     =     model(**_A							)
       _UpperCAmelCase     =     outputs.last_hidden_state

       # The following Colab Notebook was used to generate these outputs:
       # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
       if "s16" in checkpoint_url:
              _UpperCAmelCase     =     torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]							)
       elif "b16" in checkpoint_url:
              _UpperCAmelCase     =     torch.tensor([[14.2889, -18.9045, 11.7281]]							)
       elif "l16" in checkpoint_url:
              _UpperCAmelCase     =     torch.tensor([[41.5028, -22.8681, 45.6475]]							)
       elif "b4" in checkpoint_url:
              _UpperCAmelCase     =     torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]							)
       else:
              _UpperCAmelCase     =     torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]							)

       # verify logits
       assert torch.allclose(last_hidden_state[:, 0, :3],_A,atol=1E-4							)

       print(F"""Saving model to {pytorch_dump_folder_path}"""							)
       model.save_pretrained(_A							)

       print(F"""Saving image processor to {pytorch_dump_folder_path}"""							)
       image_processor.save_pretrained(_A							)


if __name__ == "__main__":
    lowerCAmelCase_			=							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--checkpoint_url''',
        default='''https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar''',
        type=str,
        help='''URL of the checkpoint you\'d like to convert.''',
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
    )

    lowerCAmelCase_			=							parser.parse_args()
    convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)













"""simple docstring"""






from itertools import product

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       list[int]:






       """simple docstring"""




       _UpperCAmelCase     =     sides_number
       _UpperCAmelCase     =     max_face_number * dice_number
       _UpperCAmelCase     =     [0] * (max_total + 1)

       _UpperCAmelCase     =     1
       _UpperCAmelCase     =     range(SCREAMING_SNAKE_CASE,max_face_number + 1							)
       for dice_numbers in product(SCREAMING_SNAKE_CASE,repeat=SCREAMING_SNAKE_CASE							):
              _UpperCAmelCase     =     sum(SCREAMING_SNAKE_CASE							)
              totals_frequencies[total] += 1

       return totals_frequencies

def       __lowerCamelCase    (       )		->       float:






       """simple docstring"""




       _UpperCAmelCase     =     total_frequency_distribution(
           sides_number=4,dice_number=9							)
       _UpperCAmelCase     =     total_frequency_distribution(
           sides_number=6,dice_number=6							)

       _UpperCAmelCase     =     0
       _UpperCAmelCase     =     9
       _UpperCAmelCase     =     4 * 9
       _UpperCAmelCase     =     6
       for peter_total in range(SCREAMING_SNAKE_CASE,max_peter_total + 1							):
              peter_wins_count += peter_totals_frequencies[peter_total] * sum(
                  colin_totals_frequencies[min_colin_total:peter_total]							)

       _UpperCAmelCase     =     (4**9) * (6**6)
       _UpperCAmelCase     =     peter_wins_count / total_games_number

       _UpperCAmelCase     =     round(SCREAMING_SNAKE_CASE,ndigits=7							)

       return rounded_peter_win_probability


if __name__ == "__main__":
    print(F'''{solution() = }''')






