model_descriptor_factory

Classes

ModelDescriptorFactory

Factory for registering and retrieving ModelDescriptor classes.

Functions

resolve_descriptor_from_pretrained

Resolve the model descriptor by loading the checkpoint config and mapping model_type.
class ModelDescriptorFactory

Bases: object

Factory for registering and retrieving ModelDescriptor classes.

CLASS_MAPPING = {'gpt_oss': <class 'modelopt.torch.puzzletron.anymodel.models.gpt_oss.gpt_oss_model_descriptor.GptOssModelDescriptor'>, 'llama': <class 'modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor.LlamaModelDescriptor'>, 'mistral_small': <class 'modelopt.torch.puzzletron.anymodel.models.mistral_small.mistral_small_model_descriptor.MistralSmallModelDescriptor'>, 'nemotron_h': <class 'modelopt.torch.puzzletron.anymodel.models.nemotron_h.nemotron_h_model_descriptor.NemotronHModelDescriptor'>, 'nemotron_h_v2': <class 'modelopt.torch.puzzletron.anymodel.models.nemotron_h_v2.nemotron_h_v2_model_descriptor.NemotronHV2ModelDescriptor'>, 'qwen2': <class 'modelopt.torch.puzzletron.anymodel.models.qwen2.qwen2_model_descriptor.Qwen2ModelDescriptor'>, 'qwen3': <class 'modelopt.torch.puzzletron.anymodel.models.qwen3.qwen3_model_descriptor.Qwen3ModelDescriptor'>, 'qwen3_vl': <class 'modelopt.torch.puzzletron.anymodel.models.qwen3_vl.qwen3_vl_model_descriptor.Qwen3VLModelDescriptor'>}
classmethod get(value)

Get a registered model descriptor by name or return the descriptor if already resolved.

Parameters:

value (str | ModelDescriptor)

classmethod register(**entries)

Register model descriptor classes.

Raises:

KeyError – if entry key is already in type_dict and points to a different class.

Parameters:

entries (Type)

classmethod register_decorator(name)

Set up a register decorator.

Parameters:

name (str | None) – If specified, the decorated object will be registered with this name.

Returns:

Decorator that registers the callable.

Return type:

Callable

resolve_descriptor_from_pretrained(pretrained, trust_remote_code=False)

Resolve the model descriptor by loading the checkpoint config and mapping model_type.

Parameters:

  • pretrained (str) – Path to a pretrained model checkpoint or HuggingFace model identifier.

  • trust_remote_code (bool) – If True, allows execution of custom code from the model repository. This is a security risk if the model source is untrusted. Only set to True if you trust the source of the model. Defaults to False for security.

Returns:

The resolved ModelDescriptor class for the detected model type.

Raises:

ValueError – If pretrained is not provided or if the model type cannot be auto-detected.