Skip to content

refactor: registry package split #2278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
akoumpa wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

refactor: registry package split #2278

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a38a5b4
refactor(registry): split model registry package
akoumpa May 19, 2026
3824d05
update registry
akoumpa May 19, 2026
e003c16
remove register_retrieval() and has_retrieval_model()
akoumpa May 19, 2026
b5bd2e0
trim
akoumpa May 19, 2026
c2e70b7
trim
akoumpa May 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions nemo_automodel/_transformers/model_init.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,22 @@ def _resolve_custom_model_cls_for_config(config):
return ModelRegistry.resolve_custom_model_cls(arch_name, config)


def _ensure_config_registered_from_config_dict(pretrained_model_name_or_path, **kwargs) -> None:
"""Register a matching local config class before delegating to ``AutoConfig``."""
config_lookup_kwargs = kwargs.copy()
config_lookup_kwargs["_from_auto"] = True
config_lookup_kwargs["name_or_path"] = pretrained_model_name_or_path
try:
config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **config_lookup_kwargs)
except Exception:
logger.debug("Could not inspect config metadata for %s", pretrained_model_name_or_path, exc_info=True)
return

model_type = config_dict.get("model_type")
if isinstance(model_type, str):
ModelRegistry.ensure_config_registered(model_type)


def get_hf_config(pretrained_model_name_or_path, attn_implementation, **kwargs):
"""
Get the HF config for the model.
Expand All @@ -233,6 +249,11 @@ def get_hf_config(pretrained_model_name_or_path, attn_implementation, **kwargs):
# with incomplete dicts, losing all other fields. These nested overrides are
# instead handled by _consume_config_overrides which deep-merges them.
nested_kwargs = {k: kwargs.pop(k) for k in list(kwargs) if isinstance(kwargs[k], dict)} # noqa: F841
_ensure_config_registered_from_config_dict(
pretrained_model_name_or_path,
**kwargs,
attn_implementation=attn_implementation,
)
try:
hf_config = AutoConfig.from_pretrained(
pretrained_model_name_or_path,
Expand Down
Loading
Loading