





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


__UpperCAmelCase						=      get_logger(__name__)


def 			__UpperCamelCase				(			lowercase__	:       str		,    lowercase__	:       str		)	->      List[Any]:
    '''simple docstring'''





    if text.startswith(lowercase__		):
        return text[len(lowercase__		) :]
    return text  # or whatever


def 			__UpperCamelCase				(			lowercase__	:       List[str]		)	->      str:
    '''simple docstring'''





    return {remove_prefix(lowercase__		,    """model."""		): v for k, v in sd.items()}


def 			__UpperCamelCase				(			lowercase__	:       List[Dict[str, torch.Tensor]]		)	->      List[Any]:
    '''simple docstring'''





    lowerCAmelCase_					:	Dict						=   {}
    for k in state_dicts[0].keys():
        lowerCAmelCase_					:	str						=   [sd[k] for sd in state_dicts]
        lowerCAmelCase_					:	List[str]						=   sum(lowercase__		) / len(lowercase__		)
        assert isinstance(lowercase__		,    torch.Tensor		)
        lowerCAmelCase_					:	Union[str, Any]						=   new_t
    return new_sd


def 			__UpperCamelCase				(			lowercase__	:       str		,    lowercase__	:       str		,    lowercase__	:       str		)	->      None:
    '''simple docstring'''





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

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

    lowerCAmelCase_					:	str						=   [sanitize(torch.load(lowercase__		,    map_location="""cpu"""		)["""state_dict"""]		) for x in ckpt_files]
    lowerCAmelCase_					:	Dict						=   average_state_dicts(lowercase__		)

    lowerCAmelCase_					:	Dict						=   hf_model.load_state_dict(lowercase__		,    strict=lowercase__		)
    assert not missing, f'missing keys: {missing}'
    hf_model.save_pretrained(lowercase__		)
    try:
        lowerCAmelCase_					:	List[Any]						=   AutoTokenizer.from_pretrained(lowercase__		)
        tok.save_pretrained(lowercase__		)
    except Exception:
        pass
        # dont copy tokenizer if cant


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









def 			__UpperCamelCase				(			lowercase__	:       set		,    lowercase__	:       set		)	->      set:
    '''simple docstring'''





    assert isinstance(lowercase__		,    lowercase__		), f'The input value of [set_a={set_a}] is not a set'
    assert isinstance(lowercase__		,    lowercase__		), f'The input value of [set_b={set_b}] is not a set'

    return {a + b for a in set_a for b in set_b}


if __name__ == "__main__":
  from doctest import testmod

  testmod()


