nemo_automodel.components.models.clm

Contrastive Language Model (CLM) decision heads (https://huggingface.co/Contrastive-LM/CLM-v0.1-8B).

Submodules

Package Contents

Classes

Name	Description
`CLMHeadConfig`	Declarative configuration for a CLM state/action head pair.
`CLMHeadPair`	Published CLM state and action projections plus learned temperature.
`EmbeddingTable`	Text -> row lookup into a table of L2-normalised encoder embeddings.

Functions

Name	Description
`load_clm_heads`	Load a CLM head checkpoint (defaults to the released CLM-v0.1-8B heads from the Hugging Face Hub).
`save_clm_heads`	Save heads in the `contrastive-lm` checkpoint format (loadable by `clm-serve` and `load_clm_heads`).

API

class nemo_automodel.components.models.clm.heads.CLMHeadConfig(
    hidden_size: int,
    width: int,
    depth: int,
    projection_dim: int,
    activation: str,
    layer_norm: bool,
    residual: bool
)

Dataclass

Declarative configuration for a CLM state/action head pair.

Parameters:

hidden_size

int

Width of each frozen backbone embedding.

width

int

Hidden width of each projection MLP.

depth

int

Number of linear layers in each projection MLP; at least two.

projection_dim

int

Width of the normalized decision embedding.

activation

str

One of "gelu", "relu", or "silu".

layer_norm

bool

Apply LayerNorm before intermediate hidden layers.

residual

bool

Add residual connections around intermediate hidden layers.

activation

str

depth

int

hidden_size

int

layer_norm

bool

projection_dim

int

residual

bool

width

int

nemo_automodel.components.models.clm.heads.CLMHeadConfig.__post_init__() -> None

nemo_automodel.components.models.clm.heads.CLMHeadConfig.build(
    initial_temperature: float = 0.07
) -> nemo_automodel.components.models.clm.heads.CLMHeadPair

Build the configured head pair.

Parameters:

initial_temperature

floatDefaults to 0.07

Positive InfoNCE softmax temperature.

Returns: CLMHeadPair

Trainable CLM state/action projection heads.

class nemo_automodel.components.models.clm.heads.CLMHeadPair(
    config: nemo_automodel.components.models.clm.heads.CLMHeadConfig,
    initial_temperature: float
)

Bases: Module

Published CLM state and action projections plus learned temperature.

action_head

= _ProjectionHead(config)

logit_scale

scale

Tensor

Return the scalar inverse temperature, capped at the released value 100.

state_head

= _ProjectionHead(config)

nemo_automodel.components.models.clm.heads.CLMHeadPair._validate_embeddings(
    embeddings: torch.Tensor,
    name: str
) -> None

Validate the encoder width without modifying its embeddings.

Parameters:

embeddings

torch.Tensor

Tensor of shape [rows, hidden], with hidden equal to config.hidden_size.

name

str

Input name used in validation errors.

nemo_automodel.components.models.clm.heads.CLMHeadPair.forward(
    state_embeddings: torch.Tensor,
    action_embeddings: torch.Tensor
) -> torch.Tensor

Score every state against every candidate action.

Parameters:

state_embeddings

torch.Tensor

Tensor of shape [batch, hidden].

action_embeddings

torch.Tensor

Tensor of shape [candidates, hidden].

Returns: torch.Tensor

Tensor of shape [batch, candidates] containing scaled cosine logits.

nemo_automodel.components.models.clm.heads.CLMHeadPair.from_upstream_checkpoint(
    path: str | pathlib.Path
) -> nemo_automodel.components.models.clm.heads.CLMHeadPair

classmethod

Load the official Contrastive-LM .pt head checkpoint.

Parameters:

path

str | Path

Local CLM_v0.1-8B.pt-compatible checkpoint path.

Returns: CLMHeadPair

Head pair with the published state, action, and temperature tensors.

nemo_automodel.components.models.clm.heads.CLMHeadPair.project_actions(
    action_embeddings: torch.Tensor
) -> torch.Tensor

Project and normalize action embeddings.

Parameters:

action_embeddings

torch.Tensor

Tensor of shape [candidates, hidden].

Returns: torch.Tensor

Tensor of shape [candidates, projection].

nemo_automodel.components.models.clm.heads.CLMHeadPair.project_states(
    state_embeddings: torch.Tensor
) -> torch.Tensor

Project and normalize state embeddings.

Parameters:

state_embeddings

torch.Tensor

Tensor of shape [batch, hidden].

Returns: torch.Tensor

Tensor of shape [batch, projection].

nemo_automodel.components.models.clm.heads.CLMHeadPair.question_logits(
    state_emb: torch.Tensor,
    option_emb: torch.Tensor,
    option_counts: typing.Sequence[int]
) -> list[torch.Tensor]

Score ragged question options in question-major order.

Parameters:

state_emb

torch.Tensor

Tensor of shape [questions, hidden].

option_emb

torch.Tensor

Tensor of shape [sum_options, hidden].

option_counts

Sequence[int]

Positive option count for each question, summing to sum_options.

Returns: list[torch.Tensor]

Question-major tensors, each of shape [options].

nemo_automodel.components.models.clm.heads.CLMHeadPair.set_logit_scale_trainable(
    enabled: bool
) -> None

Apply the decision-head temperature training policy before optimization.

Parameters:

enabled

bool

Whether the scalar log inverse temperature receives gradients.

class nemo_automodel.components.models.clm.heads.EmbeddingTable(
    texts: typing.Sequence[str],
    embeddings: torch.Tensor
)

Text -> row lookup into a table of L2-normalised encoder embeddings.

index

= {t: i for i, t in (enumerate(texts))}

nemo_automodel.components.models.clm.heads.EmbeddingTable.load(
    path: str | os.PathLike
) -> 'EmbeddingTable'

classmethod

Load a previously published cache using restricted tensor deserialization.

nemo_automodel.components.models.clm.heads.EmbeddingTable.lookup(
    texts: typing.Sequence[str],
    device: torch.device | str | None = None
) -> torch.Tensor

Return a tensor of shape [len(texts), hidden] (a copy, optionally moved to device).

nemo_automodel.components.models.clm.heads.EmbeddingTable.save(
    path: str | os.PathLike
) -> None

Publish the cache atomically, preserving any prior cache on failure.

nemo_automodel.components.models.clm.heads.load_clm_heads(
    path: str | os.PathLike | None = None,
    map_location: str = 'cpu',
    revision: str | None = None
) -> nemo_automodel.components.models.clm.heads.CLMHeadPair

Load a CLM head checkpoint (defaults to the released CLM-v0.1-8B heads from the Hugging Face Hub).

nemo_automodel.components.models.clm.heads.save_clm_heads(
    heads: nemo_automodel.components.models.clm.heads.CLMHeadPair,
    path: str | os.PathLike,
    extra: dict[str, typing.Any] | None = None
) -> None

Save heads in the contrastive-lm checkpoint format (loadable by clm-serve and load_clm_heads).