personalitylinmult.inference.onnx_export

Export a trained checkpoint’s model, and its feature extractors, to ONNX.

Three independent exports live here:

• export_to_onnx() – the trained regression model, LinT (single feature) or LinMulT (multiple features, e.g. the AVT champion).

• export_wavlm_to_onnx() – WavLM itself (via transformers, the same model exordium.audio.wavlm.WavlmWrapper wraps), so the full audio -> features -> prediction pipeline can run on ONNX Runtime alone, with no torch/transformers/exordium at inference time.

export_to_onnx handles both model families with two different wrapper classes, since ONNX graphs have a fixed input signature:

• LinT’s forward(x, mask, name) takes two tensors and one constant string, which traces directly (_LinTExportWrapper).

• LinMulT’s forward(inputs, masks, names) takes a list of tensors, one per modality – ONNX/torch.onnx.export do not accept a Python list[Tensor] as a traced input the way LinMulT’s own signature does, so _LinMulTExportWrapper unpacks a fixed-arity *args (built once per checkpoint, sized to that checkpoint’s len(feature_names)) into two equal halves – the first half is inputs, the second is masks, in feature_names order. The consumer (PersonalityModelONNX) already knows this order from the .onnx.json sidecar, so the graph’s input names (x_0, mask_0, x_1, …) need no further metadata to interpret correctly.

Both branches are exported via the legacy TorchScript tracer (dynamo=False): both LinT’s and LinMulT’s forward passes branch on a mask tensor’s content (LinT: if mask is not None and not mask.any(); LinMulT’s _normalize_masks: mask.all()), which torch’s newer torch.export-based exporter (dynamo=True, the default since torch 2.9) cannot trace – it requires a fully static graph. The tracer instead bakes in whichever branch the trace-time input took. For LinT, every current caller passes an all-True mask (FiPredictor._to_batch() never pads), so tracing with an all-True dummy mask is safe in practice, but a caller feeding a genuinely padded (some-False) mask through the ONNX graph would silently get the all-valid-mask behavior regardless of the mask’s real content. For LinMulT, the same caveat applies to partial per-timestep masking – but a whole-modality all-False mask (the “skip this feature entirely” case, e.g. no transcript available) is a different, verified-safe code path: see tests/inference/test_linmult_mask_validation.py, which confirms with a real forward pass (not just by reading the source) that neither attention type this project supports produces NaN for a fully-masked branch – linmult.core.transformer’s TransformerLayer.forward has an explicit guard for exactly this case, independent of attention_type. Document these limitations wherever the ONNX path is offered as an alternative to the PyTorch path, which has no such restriction.

The exported artifact is two files, not one: the .onnx graph itself does not carry Python-level metadata like trait names, so a .onnx.json sidecar travels with it – the same information FiPredictor reads from a checkpoint’s hyperparameters, made explicit here since the ONNX graph cannot hold it. The sidecar schema is always plural (feature_names, input_dims) regardless of family, so PersonalityModelONNX never has to branch on which family produced a given graph – a LinT sidecar simply has length-1 lists.

Note: linmult versions before 2.1.1 cached linmult.core.pe.PositionalEncoding’s table across a TransformerEncoder’s query and key/value calls, which could bake a shared, trace-time sequence length into both calls of an ONNX-exported LinMulT graph – breaking inference for features with genuinely different lengths (WavLM ~50 Hz vs. DINOv2 ~25 fps vs. one row per token for text). Fixed upstream in linmult 2.1.1; this project requires linmult>=2.1.1 accordingly.

Attributes

logger	Module-level logger.
ONNX_OPSET	ONNX opset version to export against; 17 is broadly supported by current
WAVLM_MODEL_ID	HF model id matching WavlmWrapper(model_name="base+"), the variant
WAVLM_LAYER	Which transformer layer's hidden state to export -- matches
DINOV2_MODEL_ID	HF model id matching DINOv2Wrapper(model_name="base"), the variant
MMBERT_MODEL_ID	HF model id matching MmbertWrapper(), the variant this project's
XLM_ROBERTA_MODEL_ID	HF model id matching XlmRobertaWrapper()'s default backbone="mpnet"
EMOTIEFFNET_MODEL_NAME	Variant matching EmotiEffNetWrapper()'s default -- timm
FARL_MODEL_NAME	Variant matching FarlWrapper()'s default -- CLIP ViT-B/16, LAION-Face

Functions

export_to_onnx(→ pathlib.Path)	Export a trained checkpoint's model (LinT or LinMulT) to ONNX.
export_wavlm_to_onnx(→ pathlib.Path)	Export WavLM (one hidden-state layer) to ONNX.
export_dinov2_to_onnx(→ pathlib.Path)	Export DINOv2 (CLS token, L2-normalized) to ONNX.
export_mmbert_to_onnx(→ pathlib.Path)	Export mmBERT (token-level hidden states) to ONNX.
export_xlm_roberta_to_onnx(→ pathlib.Path)	Export XLM-RoBERTa (mpnet backbone, token-level) to ONNX.
export_emotieffnet_to_onnx(→ pathlib.Path)	Export EmotiEffNet (penultimate-layer features) to ONNX.
export_farl_to_onnx(→ pathlib.Path)	Export FaRL (L2-normalized projected embedding) to ONNX.
export_emotion2vec_to_onnx(→ pathlib.Path)	Export emotion2vec+ (frame-level features) to ONNX.

Module Contents

personalitylinmult.inference.onnx_export.logger

Module-level logger.

personalitylinmult.inference.onnx_export.ONNX_OPSET = 17

ONNX opset version to export against; 17 is broadly supported by current onnxruntime releases and covers every op LinT’s forward pass uses.

personalitylinmult.inference.onnx_export.export_to_onnx(module: personalitylinmult.train.module.FiLightningModule, output_path: str | pathlib.Path) → pathlib.Path

Export a trained checkpoint’s model (LinT or LinMulT) to ONNX.

Parameters:

• module (FiLightningModule) – A loaded module (e.g. via FiLightningModule.load_from_checkpoint), in eval mode or not – this function calls .eval() itself.

• output_path (str | Path) – Destination .onnx path. A sidecar .onnx.json is written alongside it with the same stem.

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.WAVLM_MODEL_ID = 'microsoft/wavlm-base-plus'

HF model id matching WavlmWrapper(model_name="base+"), the variant every feature in this project’s fi.h5 was extracted with.

personalitylinmult.inference.onnx_export.WAVLM_LAYER = 9

Which transformer layer’s hidden state to export – matches personalitylinmult.inference.audio_features.WAVLM_LAYER. Indexes into WavlmWrapper’s already-sliced hidden_states[1:] (12 transformer layers, CNN feature-extractor output excluded), so the raw HF model’s hidden_states[WAVLM_LAYER + 1] is the equivalent index – verified numerically equal to WavlmWrapper(...)(waveform)[WAVLM_LAYER] before this export was written.

personalitylinmult.inference.onnx_export.export_wavlm_to_onnx(output_path: str | pathlib.Path, model_id: str = WAVLM_MODEL_ID, layer_index: int = WAVLM_LAYER + 1) → pathlib.Path

Export WavLM (one hidden-state layer) to ONNX.

Downloads the model from HF Hub (transformers) once, at export time – the resulting .onnx graph is self-contained and needs neither transformers nor a network connection to run afterward.

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_id (str) – HF model id to export. Defaults to WAVLM_MODEL_ID (WavLM base+, matching every feature this project’s fi.h5 was built with).

• layer_index (int) – Index into the raw model’s hidden_states tuple (13 entries: CNN output + 12 transformer layers). Defaults to WAVLM_LAYER + 1, the raw-model equivalent of WAVLM_LAYER in WavlmWrapper’s already-sliced list – see WAVLM_LAYER’s docstring for how that offset was confirmed.

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.DINOV2_MODEL_ID = 'facebook/dinov2-base'

HF model id matching DINOv2Wrapper(model_name="base"), the variant this project’s fi.h5 was built with (768-d).

personalitylinmult.inference.onnx_export.export_dinov2_to_onnx(output_path: str | pathlib.Path, model_id: str = DINOV2_MODEL_ID) → pathlib.Path

Export DINOv2 (CLS token, L2-normalized) to ONNX.

Downloads the model from HF Hub (transformers) once, at export time – the resulting .onnx graph is self-contained and needs neither transformers nor a network connection to run afterward. Unlike WavLM (one call per whole clip), a live pipeline calls this graph once per video frame (or batches several frames into one call via the dynamic batch axis) – DINOv2 runs on full, unmodified frames, no face detection or tracking involved (see personalitylinmult.inference.audio_features_onnx’s WavLM-decode pattern for the analogous “preprocessing stays outside the graph” convention this export also follows).

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_id (str) – HF model id to export. Defaults to DINOV2_MODEL_ID (DINOv2 base, matching every feature this project’s fi.h5 was built with).

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.MMBERT_MODEL_ID = 'jhu-clsp/mmBERT-base'

HF model id matching MmbertWrapper(), the variant this project’s fi.h5 (mmbert_gt feature) was built with.

personalitylinmult.inference.onnx_export.XLM_ROBERTA_MODEL_ID = 'sentence-transformers/paraphrase-multilingual-mpnet-base-v2'

HF model id matching XlmRobertaWrapper()’s default backbone="mpnet" – confirmed (not assumed) to be what this project’s fi.h5 (xlm_roberta_gt feature) was built with by reading personalitylinmult/preprocess/fi/textual_xlm_roberta.py, which calls XlmRobertaWrapper(device_id=device_id) with no backbone= override, so the wrapper’s default applies. Not "e5" (intfloat/multilingual-e5-base), the wrapper’s other supported backbone.

personalitylinmult.inference.onnx_export.export_mmbert_to_onnx(output_path: str | pathlib.Path, model_id: str = MMBERT_MODEL_ID) → pathlib.Path

Export mmBERT (token-level hidden states) to ONNX.

Downloads the model from HF Hub (transformers) once, at export time. Tokenization stays outside the graph – see _TokenLevelTextWrapper.

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_id (str) – HF model id to export. Defaults to MMBERT_MODEL_ID.

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.export_xlm_roberta_to_onnx(output_path: str | pathlib.Path, model_id: str = XLM_ROBERTA_MODEL_ID) → pathlib.Path

Export XLM-RoBERTa (mpnet backbone, token-level) to ONNX.

Downloads the model from HF Hub (transformers) once, at export time. Tokenization stays outside the graph – see _TokenLevelTextWrapper.

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_id (str) – HF model id to export. Defaults to XLM_ROBERTA_MODEL_ID (the mpnet backbone, confirmed to be what this project’s xlm_roberta_gt feature actually uses).

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.EMOTIEFFNET_MODEL_NAME = 'enet_b0_8_best_vgaf'

Variant matching EmotiEffNetWrapper()’s default – timm tf_efficientnet_b0.ns_jft_in1k, 224x224 input, 1280-d penultimate features. Matches every feature this project’s fi.h5 was built with.

personalitylinmult.inference.onnx_export.export_emotieffnet_to_onnx(output_path: str | pathlib.Path, model_name: str = EMOTIEFFNET_MODEL_NAME) → pathlib.Path

Export EmotiEffNet (penultimate-layer features) to ONNX.

Instantiates the real exordium.video.deep.emotieffnet .EmotiEffNetWrapper at export time (downloads and correctly remaps its legacy-pickle checkpoint – see that wrapper’s module docstring for why a fresh timm model plus a remapped state_dict is needed rather than loading the pickle directly) and traces its .model – already has classifier = nn.Identity() applied, so the traced graph itself needs no special output selection, unlike WavLM (must select one hidden_states layer) or DINOv2 (must select+normalize the CLS token).

Like DINOv2, a live pipeline calls this graph once per face crop (or batches several into one call via the dynamic batch axis); preprocessing (resize + ImageNet normalize) stays outside the graph.

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_name (str) – EmotiEffNet variant to export. Defaults to EMOTIEFFNET_MODEL_NAME (matching every feature this project’s fi.h5 was built with).

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.FARL_MODEL_NAME = 'ep64'

Variant matching FarlWrapper()’s default – CLIP ViT-B/16, LAION-Face 20M, 64 epochs, 512-d projected embedding. Matches every feature this project’s fi.h5 was built with.

personalitylinmult.inference.onnx_export.export_farl_to_onnx(output_path: str | pathlib.Path, model_name: str = FARL_MODEL_NAME) → pathlib.Path

Export FaRL (L2-normalized projected embedding) to ONNX.

Instantiates the real exordium.video.deep.farl.FarlWrapper at export time – FaRL ships raw OpenAI-CLIP weights with no HF Hub repo to pull an architecture from, so the wrapper’s own _build_vision_tower() + _convert_openai_to_hf(...) construction is reused rather than re-implemented here (the same export-time-only heavy-dependency precedent WavLM’s export set: this only runs once, at export time, never at inference time).

Like DINOv2 and EmotiEffNet, a live pipeline calls this graph once per face crop (or batches several into one call via the dynamic batch axis); preprocessing (resize + CLIP mean/std – not ImageNet’s – normalize) stays outside the graph.

Parameters:

• output_path (str | Path) – Destination .onnx path.

• model_name (str) – FaRL variant to export. Defaults to FARL_MODEL_NAME (matching every feature this project’s fi.h5 was built with).

Returns:

The written .onnx path.

Return type:

Path

personalitylinmult.inference.onnx_export.export_emotion2vec_to_onnx(output_path: str | pathlib.Path) → pathlib.Path

Export emotion2vec+ (frame-level features) to ONNX.

Instantiates the real exordium.audio.emotion2vec .Emotion2vecWrapper at export time – emotion2vec+ is a fully custom, vendored architecture (_Emotion2vecModel, not a thin wrapper over a library model), so its own checkpoint download and key-remapping logic is reused rather than re-implemented here.

ALiBi’s relative position bias (_get_alibi_bias) is recomputed from the input’s time dimension on every call, using torch.arange over a shape-derived count – traced correctly under a dynamic time axis (see test_onnx_export_emotion2vec.py’s explicit different-length check, added specifically because a shape-dependent internal computation like this is a plausible way for the tracer to silently bake in a constant instead of staying dynamic, the same class of risk already documented for LinMulT’s mask-content branching in this module’s docstring).

Parameters:

output_path (str | Path) – Destination .onnx path.

Returns:

The written .onnx path.

Return type:

Path