





"""simple docstring"""
import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


lowerCAmelCase:	Optional[int]							      =get_logger(__name__)




def 			__snake_case						(      __A ,__A    )		->				str:
   if text.startswith(__A    ):
      return text[len(__A    ) :]
   return text  # or whatever




def 			__snake_case						(      __A    )		->				Dict:
   return {remove_prefix(__A ,"""model."""    ): v for k, v in sd.items()}




def 			__snake_case						(      __A    )		->				Optional[Any]:
   lowercase    :			Any  =       {}
   for k in state_dicts[0].keys():
      lowercase    :			List[Any]  =       [sd[k] for sd in state_dicts]
      lowercase    :			List[Any]  =       sum(__A    ) / len(__A    )
      assert isinstance(__A ,torch.Tensor    )
      lowercase    :			Optional[int]  =       new_t
   return new_sd




def 			__snake_case						(      __A ,__A ,__A    )		->				None:
   lowercase    :			Optional[Any]  =       AutoModelForSeqaSeqLM.from_pretrained(__A    )
   if os.path.isfile(__A    ):
      lowercase    :			Tuple  =       [pl_ckpt_path]
   else:
      assert os.path.isdir(__A    )
      lowercase    :			Tuple  =       list(Path(__A    ).glob("""*.ckpt"""    )    )
      assert ckpt_files, F'''could not find any ckpt files inside the {pl_ckpt_path} directory'''

   if len(__A    ) > 1:
      logger.info(F'''averaging the weights of {ckpt_files}'''    )

   lowercase    :			str  =       [sanitize(torch.load(__A ,map_location="""cpu"""    )["""state_dict"""]    ) for x in ckpt_files]
   lowercase    :			Any  =       average_state_dicts(__A    )

   lowercase  ,					lowercase    :			Union[str, Any]  =       hf_model.load_state_dict(__A ,strict=__A    )
   assert not missing, F'''missing keys: {missing}'''
   hf_model.save_pretrained(__A    )
   try:
      lowercase    :			Optional[Any]  =       AutoTokenizer.from_pretrained(__A    )
      tok.save_pretrained(__A    )
   except Exception:
      pass
      # dont copy tokenizer if cant


if __name__ == "__main__":
   fire.Fire(convert_pl_to_hf)









"""simple docstring"""
import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
   import transformer_engine.pytorch as te




def 			__snake_case						(      __A ,__A=True ,__A=True ,__A=True    )		->				Union[str, Any]:
   if not is_fpa_available():
      raise ImportError("""Using `convert_model` requires transformer_engine to be installed."""    )
   for name, module in model.named_children():
      if isinstance(__A ,nn.Linear    ) and to_transformer_engine and _convert_linear:
         # Return early if the linear layer weights are not multiples of 16
         if any(p % 16 != 0 for p in module.weight.shape    ):
            return
         lowercase    :			Optional[int]  =       module.bias is not None
         lowercase    :			Union[str, Any]  =       te.Linear(
             module.in_features ,module.out_features ,bias=__A ,params_dtype=module.weight.dtype    )
         lowercase    :			Union[str, Any]  =       module.weight.data.clone()
         if has_bias:
            lowercase    :			Dict  =       module.bias.data.clone()

         setattr(__A ,__A ,__A    )
      elif isinstance(__A ,nn.LayerNorm    ) and to_transformer_engine and _convert_ln:
         lowercase    :			List[str]  =       te.LayerNorm(module.normalized_shape[0] ,eps=module.eps ,params_dtype=module.weight.dtype    )
         lowercase    :			Optional[Any]  =       module.weight.data.clone()
         lowercase    :			Any  =       module.bias.data.clone()

         setattr(__A ,__A ,__A    )
      elif isinstance(__A ,te.Linear    ) and not to_transformer_engine and _convert_linear:
         lowercase    :			Any  =       module.bias is not None
         lowercase    :			List[Any]  =       nn.Linear(
             module.in_features ,module.out_features ,bias=__A ,params_dtype=module.weight.dtype    )
         lowercase    :			Union[str, Any]  =       module.weight.data.clone()
         if has_bias:
            lowercase    :			Optional[Any]  =       module.bias.data.clone()

         setattr(__A ,__A ,__A    )
      elif isinstance(__A ,te.LayerNorm    ) and not to_transformer_engine and _convert_ln:
         lowercase    :			Any  =       nn.LayerNorm(module.normalized_shape[0] ,eps=module.eps ,params_dtype=module.weight.dtype    )
         lowercase    :			str  =       module.weight.data.clone()
         lowercase    :			List[str]  =       module.bias.data.clone()

         setattr(__A ,__A ,__A    )
      else:
         convert_model(
             __A ,to_transformer_engine=__A ,_convert_linear=_convert_linear ,_convert_ln=_convert_ln ,)







def 			__snake_case						(      __A    )		->				Union[str, Any]:
   if not is_fpa_available():
      raise ImportError("""Using `has_transformer_engine_layers` requires transformer_engine to be installed."""    )
   for m in model.modules():
      if isinstance(__A ,(te.LayerNorm, te.Linear)    ):
         return True
   return False


