nemo_automodel.components.models.clm.heads

Contrastive Language Model (CLM) projection heads.

CLM (https://huggingface.co/Contrastive-LM/CLM-v0.1-8B) is a frozen Qwen3-8B encoder (last-token pooled, L2-normalised hidden state) followed by two small MLP projection heads. A (state, candidate) pair is scored by exp(logit_scale) * cos(state_head(e_state), action_head(e_candidate)) and a question’s answer distribution is the softmax of those scores over its options.

Source architecture and checkpoint contract: https://github.com/Contrastive-LM/CLM (Apache-2.0; https://github.com/Contrastive-LM/CLM/blob/main/LICENSE). The NeMo implementation adds typed configuration, validation, and local checkpoint I/O. The modules reproduce the released architecture parameter-for-parameter so that checkpoints round-trip with the contrastive-lm package (clm-serve / clm.Engine): the checkpoint is a torch.save dict with state_head, action_head, logit_scale and cfg.

Module Contents

Classes

Name	Description
`CLMEmbeddingCacheConfig`	Declare the complete frozen-encoder contract used to identify cached features.
`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.
`_ProjectionHead`	-

Functions

Name	Description
`clm_batch`	Gather encoder embeddings for a batch of records.
`collect_clm_texts`	Every unique text the encoder must embed for these records (state texts first, then candidates).
`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`).

Data

CLM_HF_FILE

CLM_HF_REPO

CLM_HF_REVISION

CLM_MAX_LOGIT_SCALE

CLM_MAX_TOKENS

API

class nemo_automodel.components.models.clm.heads.CLMEmbeddingCacheConfig(
    encoder_name: str,
    encoder_revision: str,
    tokenizer_name: str,
    tokenizer_revision: str,
    max_length: int = 2048,
    pooling: str = 'last',
    l2_normalize: bool = True,
    is_causal: bool = True,
    dtype: str = 'bfloat16',
    attention_implementation: str = 'sdpa',
    transformers_version: str = ''
)

Dataclass

Declare the complete frozen-encoder contract used to identify cached features.

Hub revisions must be immutable commit IDs. Different tokenizer, truncation, pooling, normalization, precision, attention, or Transformers settings produce different cache names even when the input texts are unchanged.

attention_implementation

str = 'sdpa'

dtype

str = 'bfloat16'

encoder_name

str

encoder_revision

str

is_causal

bool = True

l2_normalize

bool = True

max_length

int = 2048

pooling

str = 'last'

tokenizer_name

str

tokenizer_revision

str

transformers_version

str = ''

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

nemo_automodel.components.models.clm.heads.CLMEmbeddingCacheConfig.build(
    cache_dir: pathlib.Path,
    texts: typing.Sequence[str]
) -> pathlib.Path

Resolve the cache path without creating directories or fetching model files.

Parameters:

cache_dir

Path

Directory that will own the feature cache.

texts

Sequence[str]

Ordered input strings, aligned with the embedding table rows.

Returns: Path

Cache path identified by all embedding-producing settings and exact texts.

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.

class nemo_automodel.components.models.clm.heads._ProjectionHead(
    config: nemo_automodel.components.models.clm.heads.CLMHeadConfig
)

Bases: Module

act

= activations[config.activation]()

hidden

inp

= nn.Linear(config.hidden_size, config.width)

norms

out

= nn.Linear(config.width, config.projection_dim)

residual

= config.residual

nemo_automodel.components.models.clm.heads._ProjectionHead.forward(
    embeddings: torch.Tensor
) -> torch.Tensor

Project frozen encoder embeddings.

Parameters:

embeddings

torch.Tensor

Tensor of shape […, hidden].

Returns: torch.Tensor

Tensor of shape […, projection].

nemo_automodel.components.models.clm.heads.clm_batch(
    records: typing.Sequence[dict[str, typing.Any]],
    table: nemo_automodel.components.models.clm.heads.EmbeddingTable,
    device: torch.device | str | None = None
) -> tuple[torch.Tensor, torch.Tensor, list[int], list[dict[str, list[str]]]]

Gather encoder embeddings for a batch of records.

Returns: torch.Tensor

(state_emb, option_emb, option_counts, option_orders): state_emb has shape [questions, hidden] (one row

nemo_automodel.components.models.clm.heads.collect_clm_texts(
    records: typing.Sequence[dict[str, typing.Any]]
) -> list[str]

Every unique text the encoder must embed for these records (state texts first, then candidates).

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).

nemo_automodel.components.models.clm.heads.CLM_HF_FILE = 'CLM_v0.1-8B.pt'

nemo_automodel.components.models.clm.heads.CLM_HF_REPO = 'Contrastive-LM/CLM-v0.1-8B'

nemo_automodel.components.models.clm.heads.CLM_HF_REVISION = 'e939398d4556fcd9400c76fa8c5a513202f42b0a'

nemo_automodel.components.models.clm.heads.CLM_MAX_LOGIT_SCALE = math.log(100.0)

nemo_automodel.components.models.clm.heads.CLM_MAX_TOKENS = 2048