from typing import OrderedDict


def get_partial_state_dict(model_state_dict, modules, 
                            init_from_decoder_asr=False,
                            init_from_decoder_mt=False):
    """Create state_dict with specified modules matching input model modules.

    Note that get_partial_lm_state_dict is used if a LM specified.

    Args:
        model_state_dict (OrderedDict): trained model state_dict
        modules (list): specified module list for transfer

    Return:
        new_state_dict (OrderedDict): the updated state_dict

    """
    new_state_dict = OrderedDict()

    for key, value in model_state_dict.items():
        if init_from_decoder_asr:
            key = key.replace("_asr", "")
        elif init_from_decoder_mt:
            key = key.replace("decoder_mt", "decoder")
        if any(key.startswith(m) for m in modules):
            new_state_dict[key] = value

    return new_state_dict