


"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


lowercase  		=			{"""configuration_vilt""": ["""VILT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViltConfig"""]}

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  		=			["""ViltFeatureExtractor"""]
					lowercase  		=			["""ViltImageProcessor"""]
					lowercase  		=			["""ViltProcessor"""]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  		=			[
					    """VILT_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """ViltForImageAndTextRetrieval""",
					    """ViltForImagesAndTextClassification""",
					    """ViltForTokenClassification""",
					    """ViltForMaskedLM""",
					    """ViltForQuestionAnswering""",
					    """ViltLayer""",
					    """ViltModel""",
					    """ViltPreTrainedModel""",
					]


if TYPE_CHECKING:
					from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_vilt import ViltFeatureExtractor
										from .image_processing_vilt import ViltImageProcessor
										from .processing_vilt import ViltProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_vilt import (
										    VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ViltForImageAndTextRetrieval,
										    ViltForImagesAndTextClassification,
										    ViltForMaskedLM,
										    ViltForQuestionAnswering,
										    ViltForTokenClassification,
										    ViltLayer,
										    ViltModel,
										    ViltPreTrainedModel,
										)


else:
					import sys

					lowercase  		=			_LazyModule(__name__, globals()["""__file__"""], _import_structure)






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


_SCREAMING_SNAKE_CASE     =     get_logger(__name__)
def 					__lowercase						(    lowerCamelCase__			,					lowerCamelCase__       )					->Optional[int]:

  '''simple docstring'''




  if text.startswith(lowerCamelCase__       ):
    return text[len(lowerCamelCase__       ) :]
  return text  # or whatever
def 					__lowercase						(    lowerCamelCase__       )					->int:

  '''simple docstring'''




  return {remove_prefix(lowerCamelCase__			,					"model."       ): v for k, v in sd.items()}
def 					__lowercase						(    lowerCamelCase__       )					->Any:

  '''simple docstring'''




  a__  		=  {}
  for k in state_dicts[0].keys():
    a__  		=  [sd[k] for sd in state_dicts]
    a__  		=  sum(lowerCamelCase__       ) / len(lowerCamelCase__       )
    assert isinstance(lowerCamelCase__			,					torch.Tensor       )
    a__  		=  new_t
  return new_sd
def 					__lowercase						(    lowerCamelCase__			,					lowerCamelCase__			,					lowerCamelCase__       )					->None:

  '''simple docstring'''




  a__  		=  AutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__       )
  if os.path.isfile(lowerCamelCase__       ):
    a__  		=  [pl_ckpt_path]
  else:
    assert os.path.isdir(lowerCamelCase__       )
    a__  		=  list(Path(lowerCamelCase__       ).glob("*.ckpt"       )       )
    assert ckpt_files, f"could not find any ckpt files inside the {pl_ckpt_path} directory"

  if len(lowerCamelCase__       ) > 1:
    logger.info(f"averaging the weights of {ckpt_files}"       )

  a__  		=  [sanitize(torch.load(lowerCamelCase__			,					map_location="cpu"       )["state_dict"]       ) for x in ckpt_files]
  a__  		=  average_state_dicts(lowerCamelCase__       )

  a__		,a__  		=  hf_model.load_state_dict(lowerCamelCase__			,					strict=lowerCamelCase__       )
  assert not missing, f"missing keys: {missing}"
  hf_model.save_pretrained(lowerCamelCase__       )
  try:
    a__  		=  AutoTokenizer.from_pretrained(lowerCamelCase__       )
    tok.save_pretrained(lowerCamelCase__       )
  except Exception:
    pass
    # dont copy tokenizer if cant


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

