Function get_model_arch

lmdeploy/archs.py:151–177

Get a model's architecture and configuration. Args: model_path(str): the model path

(model_path: str, trust_remote_code: bool = False)

def get_model_arch(model_path: str, trust_remote_code: bool = False):
    """Get a model's architecture and configuration.

    Args:
        model_path(str): the model path
    """
    try:
        cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
    except Exception as e:  # noqa
        from transformers import PretrainedConfig
        cfg = PretrainedConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)

    _cfg = cfg.to_dict()
    if _cfg.get('architectures', None):
        arch = _cfg['architectures'][0]
        if _cfg.get('auto_map'):
            for _, v in _cfg['auto_map'].items():
                if 'InternLMXComposer2ForCausalLM' in v:
                    arch = 'InternLMXComposer2ForCausalLM'
    elif _cfg.get('auto_map', None) and 'AutoModelForCausalLM' in _cfg['auto_map']:
        arch = _cfg['auto_map']['AutoModelForCausalLM'].split('.')[-1]
    elif _cfg.get('language_config', None) and _cfg['language_config'].get(
            'auto_map', None) and 'AutoModelForCausalLM' in _cfg['language_config']['auto_map']:
        arch = _cfg['language_config']['auto_map']['AutoModelForCausalLM'].split('.')[-1]
    else:
        raise RuntimeError(f'Could not find model architecture from config: {_cfg}')
    return arch, cfg


def search_nested_config(config, key):