torch_to_nnef.nemo_tract
NVIDIA NeMo export helpers (TractNNEF-focused).
This package splits the original monolithic implementation into submodules while keeping the same public API.
Public re-exports preserved for compatibility: - constants: PARAKEET_V3_SLUG, PARAKEET_110M_SLUG, NEMOTRON_0_6B - functions/classes: iter_export_params_for_generic_nemo_asr_model, export_nemo_asr_model, main
Programmatic (non-CLI) API: - export_nemo_from_model, NemoExportConfig
dataclass
NemoExportConfig(subnet: SubnetSelectionConfig = SubnetSelectionConfig(), compression: CompressionConfig = CompressionConfig(), naming_scheme: VariableNamingScheme = VariableNamingScheme.NATURAL_VERBOSE_CAMEL, data_type: str = 'float32', pretrained_name: str = '')
Configuration for the export_nemo_from_model programmatic API.
All fields have sensible defaults so callers only specify what they actually need to override.
Attributes:
| Name | Type | Description |
|---|---|---|
subnet |
SubnetSelectionConfig
|
Subnet filtering and composition. |
compression |
CompressionConfig
|
Compression options for NNEF artifacts. |
naming_scheme |
VariableNamingScheme
|
NNEF variable naming scheme. |
data_type |
str
|
Export dtype policy: float32 | float16 | mixed. |
pretrained_name |
str
|
Human-readable label embedded in export metadata. |
export_nemo_asr_model(asr_model, inference_target, export_dir: Path, compress_registry: str, compress_method: T.Optional[str] = None, skip_preprocessor: bool = False, split_joint_decoder: bool = False, extra_cfg: T.Optional[T.Dict[str, T.Any]] = None, float_dtype: T.Optional[torch.dtype] = None, remove_unused_inputs: bool = True, dump_checked_io: bool = False, only_subnets: T.Optional[T.Collection[str]] = None, *, omegaconf: InjectedOmegaConfModule = INJECTED, axis_registry=None, **kwargs)
Export a generic NeMo ASR model to NNEF format using TractNNEF.
export_nemo_from_model(*, model, target: InferenceTarget, export_dir: Path, axis_reg: T.Optional[AxisSymbolRegistry] = None, cfg: NemoExportConfig) -> None
Export a NeMo ASR model using a structured configuration.
This is the public programmatic API. Handles dtype preparation
(float16 / mixed) internally so callers don't need to
apply model.half() or WrapPreprocessorCast themselves.
iter_export_params_for_generic_nemo_asr_model(asr_model, inference_target, skip_preprocessor: bool = False, split_joint_decoder: bool = False, remove_unused_inputs: bool = True, float_dtype: T.Optional[torch.dtype] = None, only_subnets: T.Optional[T.Collection[str]] = None, axis_registry=None) -> T.Iterator[ExportParameters]
Iterator over export parameters for a generic NeMo ASR model.