





'''simple docstring'''


import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


lowerCAmelCase      :  List[str]      							={
    '''vit_b32''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 32,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_b16''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 16,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_l14''': {
        '''embed_dim''': 768,
        '''image_resolution''': 840,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 24,
        '''vision_width''': 1_024,
        '''vision_patch_size''': 14,
        '''transformer_width''': 768,
        '''transformer_heads''': 12,
        '''transformer_layers''': 12,
    },
}
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Union[str, Any]			,__lowerCamelCase      :      str=""			,__lowerCamelCase      :      List[Any]="/"				):
   lowercase_    :Optional[Any]           =				[]

   for k, v in params.items():
      lowercase_    :List[Any]           =				parent_key + sep + k if parent_key else k

      if isinstance(__lowerCamelCase			,collections.MutableMapping				):
         items.extend(flatten_nested_dict(__lowerCamelCase			,__lowerCamelCase			,sep=__lowerCamelCase				).items()				)
      else:
         items.append((new_key, v)				)
   return dict(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str				):
   return jax.tree_util.tree_map(lambda __lowerCamelCase				: x.astype(jnp.floataa				) if x.dtype == jnp.bfloataa else x			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      Dict				):
   lowercase_    :Tuple           =				pt_attn_layer.in_proj_weight.chunk(3			,dim=0				)
   lowercase_    :Tuple           =				pt_attn_layer.in_proj_bias.chunk(3			,dim=0				)

   lowercase_    :Union[str, Any]           =				pt_attn_layer.out_proj.weight
   lowercase_    :List[Any]           =				pt_attn_layer.out_proj.bias

   lowercase_    :Optional[Any]           =				q_proj
   lowercase_    :int           =				q_proj_bias

   lowercase_    :Any           =				k_proj
   lowercase_    :Tuple           =				k_proj_bias

   lowercase_    :List[Any]           =				v_proj
   lowercase_    :Dict           =				v_proj_bias

   lowercase_    :Union[str, Any]           =				out_proj_weights
   lowercase_    :Tuple           =				out_proj_bias
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[Any]			,__lowerCamelCase      :      Tuple				):
   copy_linear(hf_mlp.fca			,pt_mlp.c_fc				)
   copy_linear(hf_mlp.fca			,pt_mlp.c_proj				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[Any]			,__lowerCamelCase      :      Union[str, Any]				):
   lowercase_    :int           =				pt_linear.weight
   lowercase_    :Tuple           =				pt_linear.bias
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Tuple			,__lowerCamelCase      :      Dict				):
   # copy layer norms
   copy_linear(hf_layer.layer_norma			,pt_layer.ln_a				)
   copy_linear(hf_layer.layer_norma			,pt_layer.ln_a				)

   # copy MLP
   copy_mlp(hf_layer.mlp			,pt_layer.mlp				)

   # copy attn
   copy_attn_layer(hf_layer.self_attn			,pt_layer.attn				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      Any				):
   for hf_layer, pt_layer in zip(__lowerCamelCase			,__lowerCamelCase				):
      copy_layer(__lowerCamelCase			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[Any]				):
   # copy  embeds
   lowercase_    :str           =				pt_model.token_embedding.weight
   lowercase_    :Union[str, Any]           =				pt_model.positional_embedding

   # copy layer norm
   copy_linear(hf_encoder.final_layer_norm			,pt_model.ln_final				)

   # copy hidden layers
   copy_layers(hf_encoder.encoder.layers			,pt_model.transformer.resblocks				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[Any]				):
   # copy projection
   lowercase_    :Optional[int]           =				pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Tuple			,__lowerCamelCase      :      Any				):
   # copy projection
   lowercase_    :str           =				pt_model.visual.proj.data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layernorm			,pt_model.visual.ln_pre				)
   copy_linear(hf_model.vision_model.post_layernorm			,pt_model.visual.ln_post				)

   # copy embeds
   lowercase_    :int           =				pt_model.visual.conva.weight.data
   lowercase_    :Any           =				pt_model.visual.class_embedding
   lowercase_    :List[str]           =				pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers			,pt_model.visual.transformer.resblocks				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[str]				):
   lowercase_    :Dict           =				flatten_nested_dict(flax_params["backbone"]["merged_class_token"]				)

   lowercase_    :Union[str, Any]           =				torch.from_numpy(flax_class_token_params["scale"]				)
   lowercase_    :Optional[int]           =				torch.from_numpy(flax_class_token_params["bias"]				)
   lowercase_    :Union[str, Any]           =				nn.Parameter(__lowerCamelCase				)
   lowercase_    :Optional[Any]           =				nn.Parameter(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Any			,__lowerCamelCase      :      Any				):
   lowercase_    :List[Any]           =				hf_model.state_dict()
   lowercase_    :str           =				{}

   # Rename class prediction head flax params to pytorch HF
   lowercase_    :List[Any]           =				flatten_nested_dict(flax_params["class_head"]				)

   for flax_key, v in flax_class_params.items():
      lowercase_    :Tuple           =				flax_key.replace("/"			,"."				)
      lowercase_    :Optional[Any]           =				torch_key.replace(".kernel"			,".weight"				)
      lowercase_    :List[str]           =				torch_key.replace("Dense_0"			,"dense0"				)
      lowercase_    :List[str]           =				"class_head." + torch_key

      if "weight" in torch_key and v.ndim == 2:
         lowercase_    :Tuple           =				v.T

      lowercase_    :Optional[int]           =				nn.Parameter(torch.from_numpy(__lowerCamelCase				)				)

   # Rename box prediction box flax params to pytorch HF
   lowercase_    :Any           =				flatten_nested_dict(flax_params["obj_box_head"]				)

   for flax_key, v in flax_box_params.items():
      lowercase_    :Tuple           =				flax_key.replace("/"			,"."				)
      lowercase_    :Optional[Any]           =				torch_key.replace(".kernel"			,".weight"				)
      lowercase_    :List[Any]           =				torch_key.replace("_"			,""				).lower()
      lowercase_    :Optional[Any]           =				"box_head." + torch_key

      if "weight" in torch_key and v.ndim == 2:
         lowercase_    :Optional[int]           =				v.T

      lowercase_    :List[str]           =				nn.Parameter(torch.from_numpy(__lowerCamelCase				)				)

   # Copy flax params to PyTorch params
   for name, param in new_params.items():
      if name in pt_params.keys():
         pt_params[name].copy_(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      str				):
   for k, v in flax_attn_params.items():
      if k.startswith("transformer"				):
         lowercase_    :List[str]           =				k.replace("transformer.resblocks"			,"text_model.encoder.layers"				)
      else:
         lowercase_    :Dict           =				k.replace("visual.transformer.resblocks"			,"vision_model.encoder.layers"				)

      lowercase_    :int           =				torch_key.replace("attn"			,"self_attn"				)
      lowercase_    :Any           =				torch_key.replace("key"			,"k_proj"				)
      lowercase_    :Tuple           =				torch_key.replace("value"			,"v_proj"				)
      lowercase_    :Any           =				torch_key.replace("query"			,"q_proj"				)
      lowercase_    :Any           =				torch_key.replace("out"			,"out_proj"				)

      if "bias" in torch_key and v.ndim == 2:
         lowercase_    :List[str]           =				v.shape[0] * v.shape[1]
         lowercase_    :Optional[Any]           =				v.reshape(__lowerCamelCase				)

      if "weight" in torch_key and "out" in torch_key:
         lowercase_    :Union[str, Any]           =				(v.shape[0] * v.shape[1], v.shape[2])
         lowercase_    :List[str]           =				v.reshape(__lowerCamelCase				).T

      if "weight" in torch_key and "out" not in torch_key:
         lowercase_    :Any           =				(v.shape[0], v.shape[1] * v.shape[2])
         lowercase_    :Any           =				v.reshape(__lowerCamelCase				).T

      # Copy flax CLIP attn params to HF PyTorch params
      lowercase_    :List[Any]           =				torch.from_numpy(__lowerCamelCase				)
      hf_backbone.state_dict()[torch_key].copy_(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]				):
   lowercase_    :str           =				{}
   lowercase_    :str           =				[]

   for k, v in params.items():
      if "attn." in k:
         lowercase_    :Tuple           =				k[: k.rindex("attn."				) + 5]
         if base in processed_attn_layers:
            continue

         processed_attn_layers.append(__lowerCamelCase				)
         lowercase_    :Dict           =				params[base + "out.weight"].shape[-1]
         lowercase_    :Optional[int]           =				params[base + "out.weight"].reshape(__lowerCamelCase			,__lowerCamelCase				).T
         lowercase_    :Optional[Any]           =				params[base + "out.bias"]
      else:
         lowercase_    :int           =				v
   return new_params
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]			,__lowerCamelCase      :      Dict				):
   lowercase_    :Any           =				CLIP(**__lowerCamelCase				)
   torch_model.eval()
   lowercase_    :Any           =				torch_model.state_dict()

   lowercase_    :Tuple           =				flatten_nested_dict(flax_params["backbone"]["clip"]				)
   lowercase_    :Optional[int]           =				{}

   for flax_key, v in flax_clip_params.items():
      lowercase_    :Dict           =				flax_key.replace("/"			,"."				)
      lowercase_    :Union[str, Any]           =				torch_key.replace("text.token_embedding.embedding"			,"token_embedding.kernel"				)

      if (
          torch_key.startswith("text.transformer"				)
          or torch_key.startswith("text.text_projection"				)
          or torch_key.startswith("text.ln_final"				)
          or torch_key.startswith("text.positional_embedding"				)
      ):
         lowercase_    :str           =				torch_key[5:]

      lowercase_    :Dict           =				torch_key.replace("text_projection.kernel"			,"text_projection"				)
      lowercase_    :List[str]           =				torch_key.replace("visual.proj.kernel"			,"visual.proj"				)
      lowercase_    :Any           =				torch_key.replace(".scale"			,".weight"				)
      lowercase_    :str           =				torch_key.replace(".kernel"			,".weight"				)

      if "conv" in torch_key or "downsample.0.weight" in torch_key:
         lowercase_    :Optional[Any]           =				v.transpose(3			,2			,0			,1				)

      elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
         # Fully connected layers are transposed, embeddings are not
         lowercase_    :str           =				v.T

      lowercase_    :Union[str, Any]           =				v

   lowercase_    :List[str]           =				_convert_attn_layers(__lowerCamelCase				)
   new_torch_params.update(__lowerCamelCase				)
   lowercase_    :List[str]           =				{}

   # Copy flax CLIP backbone params to PyTorch params
   for name, param in new_torch_params.items():
      if name in torch_clip_params.keys():
         lowercase_    :Union[str, Any]           =				torch.from_numpy(new_torch_params[name]				)
         torch_clip_params[name].copy_(__lowerCamelCase				)
      else:
         lowercase_    :Tuple           =				param

   return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      Dict			,__lowerCamelCase      :      List[Any]			,__lowerCamelCase      :      Any			,__lowerCamelCase      :      Optional[int]=None				):
   lowercase_    :Union[str, Any]           =				Repository(__lowerCamelCase			,clone_from=F'google/{pytorch_dump_folder_path}'				)
   repo.git_pull()

   if config_path is not None:
      lowercase_    :Union[str, Any]           =				OwlViTConfig.from_pretrained(__lowerCamelCase				)
   else:
      lowercase_    :str           =				OwlViTConfig()

   lowercase_    :List[str]           =				OwlViTModel(__lowerCamelCase				).eval()
   lowercase_    :str           =				OwlViTForObjectDetection(__lowerCamelCase				).eval()

   copy_text_model_and_projection(__lowerCamelCase			,__lowerCamelCase				)
   copy_vision_model_and_projection(__lowerCamelCase			,__lowerCamelCase				)
   lowercase_    :List[str]           =				pt_backbone.logit_scale
   copy_flax_attn_params(__lowerCamelCase			,__lowerCamelCase				)

   lowercase_    :Dict           =				hf_backbone
   copy_class_merge_token(__lowerCamelCase			,__lowerCamelCase				)
   copy_class_box_heads(__lowerCamelCase			,__lowerCamelCase				)

   # Save HF model
   hf_model.save_pretrained(repo.local_dir				)

   # Initialize image processor
   lowercase_    :Dict           =				OwlViTImageProcessor(
       size=config.vision_config.image_size			,crop_size=config.vision_config.image_size				)
   # Initialize tokenizer
   lowercase_    :Union[str, Any]           =				CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"			,pad_token="!"			,model_max_length=16				)

   # Initialize processor
   lowercase_    :Dict           =				OwlViTProcessor(image_processor=__lowerCamelCase			,tokenizer=__lowerCamelCase				)
   image_processor.save_pretrained(repo.local_dir				)
   processor.save_pretrained(repo.local_dir				)

   repo.git_add()
   repo.git_commit("Upload model and processor"				)
   repo.git_push()


if __name__ == "__main__":
 lowerCAmelCase      :  Union[str, Any]      							=argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--owlvit_version''',
     default=None,
     type=str,
     required=True,
     help='''OWL-ViT model name [clip_b16, clip_b32, clip_l14].''',
 )
 parser.add_argument(
     '''--owlvit_checkpoint''', default=None, type=str, required=True, help='''Path to flax model checkpoint.'''
 )
 parser.add_argument('''--hf_config''', default=None, type=str, required=True, help='''Path to HF model config.''')
 parser.add_argument(
     '''--pytorch_dump_folder_path''', default='''hf_model''', type=str, help='''Path to the output PyTorch model.'''
 )
 lowerCAmelCase      :  Any      							=parser.parse_args()

 # Initialize PyToch clip model
 lowerCAmelCase      :  Union[str, Any]      							=args.owlvit_version
 if model_name == "clip_b16":
  lowerCAmelCase      :  Optional[Any]      							=CONFIGS['''vit_b16''']
 elif model_name == "clip_b32":
  lowerCAmelCase      :  List[str]      							=CONFIGS['''vit_b32''']
 elif model_name == "clip_l14":
  lowerCAmelCase      :  Union[str, Any]      							=CONFIGS['''vit_l14''']

 # Load from checkpoint and convert params to float-32
 lowerCAmelCase      :  Optional[int]      							=checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['''optimizer''']['''target''']
 lowerCAmelCase      :  str      							=jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
 del variables

 # Convert CLIP backbone
 lowerCAmelCase      :  Optional[int]      							=convert_clip_backbone(flax_params, torch_config)

 convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)









'''simple docstring'''


import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]			,__lowerCamelCase      :      Union[str, Any]			,__lowerCamelCase      :      List[Any]			,__lowerCamelCase      :      str=None				):
   lowercase_    :List[Any]					    =				None
   # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
   if str(weight.dtype				) == "torch.bfloat16":
      # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
      lowercase_    :int					    =				weight.view(torch.intaa				)
      lowercase_    :List[Any]					    =				"bfloat16"
   lowercase_    :Optional[Any]					    =				weight.cpu().numpy()
   lowercase_    :Optional[int]					    =				os.path.join(__lowerCamelCase			,F'{weight_name}.dat'				)
   if index is not None:
      if dtype is None:
         lowercase_    :str					    =				str(array.dtype				)
      lowercase_    :int					    =				{"dtype": dtype, "shape": list(array.shape				)}
   if array.ndim == 0:
      lowercase_    :int					    =				array[None]
   lowercase_    :List[Any]					    =				np.memmap(__lowerCamelCase			,dtype=array.dtype			,mode="w+"			,shape=array.shape				)
   lowercase_    :List[Any]					    =				array[:]
   file_array.flush()
   return index
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      Dict				):
   lowercase_    :int					    =				tuple(weight_info["shape"]				)
   if shape == ():
      # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
      lowercase_    :List[str]					    =				(1,)

   lowercase_    :List[str]					    =				weight_info["dtype"]
   if dtype == "bfloat16":
      # NumPy does not support bfloat16 so this was saved as a int16
      lowercase_    :Optional[Any]					    =				"int16"

   lowercase_    :str					    =				np.memmap(__lowerCamelCase			,dtype=__lowerCamelCase			,shape=__lowerCamelCase			,mode="r"				)

   if len(weight_info["shape"]				) == 0:
      lowercase_    :Union[str, Any]					    =				weight[0]
   lowercase_    :str					    =				torch.tensor(__lowerCamelCase				)
   if weight_info["dtype"] == "bfloat16":
      lowercase_    :Any					    =				weight.view(torch.bfloataa				)

   return weight
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Dict			,__lowerCamelCase      :      List[str]				):
   if index is None or len(__lowerCamelCase				) == 0:
      # Nothing to save
      return

   lowercase_    :List[Any]					    =				os.path.join(__lowerCamelCase			,"index.json"				)
   if os.path.isfile(__lowerCamelCase				):
      with open(__lowerCamelCase			,"r"			,encoding="utf-8"				) as f:
         lowercase_    :Any					    =				json.load(__lowerCamelCase				)
   else:
      lowercase_    :Union[str, Any]					    =				{}
   current_index.update(__lowerCamelCase				)

   with open(__lowerCamelCase			,"w"			,encoding="utf-8"				) as f:
      json.dump(__lowerCamelCase			,__lowerCamelCase			,indent=2				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Union[str, os.PathLike]			,__lowerCamelCase      :      Dict[str, torch.Tensor]				):
   os.makedirs(__lowerCamelCase			,exist_ok=__lowerCamelCase				)
   lowercase_    :List[Any]					    =				{}
   for name, parameter in state_dict.items():
      lowercase_    :Tuple					    =				offload_weight(__lowerCamelCase			,__lowerCamelCase			,__lowerCamelCase			,index=__lowerCamelCase				)

   # Update index
   save_offload_index(__lowerCamelCase			,__lowerCamelCase				)
class    a_     (       _lowerCAmelCase							):





     def __init__(				self	:      int			,   lowercase	:      Mapping			,   lowercase	:      str		):





        """simple docstring"""




        lowercase_    :Union[str, Any]					    =				dataset
        lowercase_    :List[str]					    =				prefix





     def __getitem__(				self	:      Any			,   lowercase	:      List[str]		):





        """simple docstring"""




        return self.dataset[F'{self.prefix}{key}']





     def __iter__(				self	:      Optional[int]		):





        """simple docstring"""




        return iter([key for key in self.dataset if key.startswith(self.prefix		)]		)







     def __len__(				self	:      Any		):





        """simple docstring"""




        return len(self.dataset		)







class    a_     (       _lowerCAmelCase							):





     def __init__(				self	:      Tuple			,   lowercase	:      Dict[str, torch.Tensor] = None			,   lowercase	:      Optional[Union[str, os.PathLike]] = None			,   lowercase	:      Mapping = None			,   lowercase	:      Optional[Any]=None			,   ):





        """simple docstring"""




        if state_dict is None and save_folder is None:
           raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights."		)

        lowercase_    :Optional[int]					    =				{} if state_dict is None else state_dict
        lowercase_    :str					    =				save_folder
        if index is None and save_folder is not None:
           with open(os.path.join(lowercase			,   "index.json"		)		) as f:
              lowercase_    :Any					    =				json.load(lowercase		)
        lowercase_    :str					    =				{} if index is None else index
        lowercase_    :List[str]					    =				list(self.state_dict.keys()		)
        self.all_keys.extend([key for key in self.index if key not in self.all_keys]		)
        lowercase_    :List[Any]					    =				device





     def __getitem__(				self	:      Any			,   lowercase	:      str		):





        """simple docstring"""




        if key in self.state_dict:
           return self.state_dict[key]
        lowercase_    :int					    =				self.index[key]
        if weight_info.get("safetensors_file"		) is not None:
           if not is_safetensors_available():
              raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`."		)

           from safetensors import safe_open

           lowercase_    :Dict					    =				"cpu" if self.device is None else self.device
           with safe_open(weight_info["safetensors_file"]			,   framework="pt"			,   device=lowercase		) as f:
              lowercase_    :Tuple					    =				f.get_tensor(weight_info.get("weight_name"			,   lowercase		)		)

           if "dtype" in weight_info:
              return tensor.to(getattr(lowercase			,   weight_info["dtype"]		)		)
           else:
              return tensor

        lowercase_    :str					    =				os.path.join(self.save_folder			,   F'{key}.dat'		)
        return load_offloaded_weight(lowercase			,   lowercase		)





     def __iter__(				self	:      Optional[int]		):





        """simple docstring"""




        return iter(self.all_keys		)







     def __len__(				self	:      Tuple		):





        """simple docstring"""




        return len(self.all_keys		)







def 							UpperCAmelCase_						(		__lowerCamelCase      :      Dict[str, torch.Tensor]			,__lowerCamelCase      :      List[str]				):
   lowercase_    :List[str]					    =				{}
   for module_name in submodule_names:
      # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
      # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
      result.update(
          {
              key: param
              for key, param in state_dict.items()
              if key == module_name or key.startswith(module_name + "."				)
          }				)
   return result


