nemo_automodel.components.models.clef.joint_schema_head

Clef joint schema head and record encoding.

Adapted from Cloudflare, Inc.’s Apache-2.0 joint_schema_model.py: https://huggingface.co/Cloudflare/clef/blob/ed3eed331870db2eff4b0db01237128ede8a00ce/joint_schema_model.py The publisher’s license is https://huggingface.co/Cloudflare/clef/blob/ed3eed331870db2eff4b0db01237128ede8a00ce/LICENSE. Parameter names remain unchanged. NeMo adaptations add typed configuration, tensor validation, a callable lexical-row table for sharded backbones, and fp32 head training.

Clef turns a state and a schema of typed questions into one logit per allowed option per question, in a single forward pass of a Qwen3.5-architecture backbone followed by this small transformer head.

Module Contents

Classes

Name	Description
`ClefHeadConfig`	Declarative configuration matching joint_head_config.json.
`EncodedQuestion`	Token spans of one question and its options inside an encoded record (end-exclusive).
`EncodedRecord`	A tokenised record: input ids, per-question spans, and optional media tensors.
`EvidenceRoutingLayer`	Route state evidence into per-option queries.
`JointSchemaHead`	Cloudflare CLEF joint head over one or more typed questions.
`LexicalTable`	Frozen cache of the output-embedding rows the head needs, indexed by token id.

Functions

Name	Description
`_encode_media`	-
`_tokens`	-
`build_lexical_table`	Cache the frozen output-embedding rows consumed by the CLEF readout.
`clef_option_orders`	`{question_id: option ids}` in the order the Clef head returns logits.
`collate_records`	Right-pad encoded records into a batch.
`encode_record`	Tokenise a SystemOne record into Clef’s prompt layout and record question / option token spans.
`load_joint_head`	Build a `JointSchemaHead` from a directory with `joint_head_config.json` / `joint_head.safetensors`.
`option_token_ids`	Sorted unique token ids inside option spans (the only output-embedding rows the head reads).
`question_options`	`(option id, description)` pairs in Clef order: noul `[true, false]`, choice sorted, score by level.
`render`	Render a state / instruction value as text (strings verbatim, everything else as compact sorted JSON).
`save_joint_head`	Save a head in the release layout (`joint_head_config.json` + `joint_head.safetensors`).

Data

API

class nemo_automodel.components.models.clef.joint_schema_head.ClefHeadConfig(
    hidden_size: int,
    width: int,
    routing_layers: int,
    layers: int,
    heads: int,
    feedforward: int,
    dropout: float = 0.0
)

Dataclass

Declarative configuration matching joint_head_config.json.

Parameters:

hidden_size

int

Backbone hidden-state and output-embedding width.

width

int

Internal decision-head width.

routing_layers

int

Number of option-to-evidence cross-attention blocks.

layers

int

Number of joint field decoder blocks.

heads

int

Attention head count.

feedforward

int

Feed-forward width in attention blocks.

dropout

floatDefaults to 0.0

Dropout probability used while post-training the head.

dropout

float = 0.0

feedforward

int

heads

int

hidden_size

int

layers

int

routing_layers

int

width

int

nemo_automodel.components.models.clef.joint_schema_head.ClefHeadConfig.__post_init__() -> None

nemo_automodel.components.models.clef.joint_schema_head.ClefHeadConfig.build() -> nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead

Build a checkpoint-compatible joint schema head.

Returns: JointSchemaHead

Newly initialized CLEF decision head.

class nemo_automodel.components.models.clef.joint_schema_head.EncodedQuestion(
    question_id: str,
    question_type: int,
    question_span: tuple[int, int],
    option_spans: tuple[tuple[int, int], ...],
    option_ids: tuple[str, ...]
)

Dataclass

Token spans of one question and its options inside an encoded record (end-exclusive).

option_ids

tuple[str, ...]

option_spans

tuple[tuple[int, int], ...]

question_id

str

question_span

tuple[int, int]

question_type

int

class nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord(
    input_ids: tuple[int, ...],
    questions: tuple[nemo_automodel.components.models.clef.joint_schema_head.EncodedQuestion, ...],
    record_id: str,
    media: dict[str, typing.Any] | None = None
)

Dataclass

A tokenised record: input ids, per-question spans, and optional media tensors.

input_ids

tuple[int, ...]

media

dict[str, Any] | None

questions

tuple[EncodedQuestion, ...]

record_id

str

class nemo_automodel.components.models.clef.joint_schema_head.EvidenceRoutingLayer(
    config: nemo_automodel.components.models.clef.joint_schema_head.ClefHeadConfig
)

Bases: Module

Route state evidence into per-option queries.

attention

attention_dropout

= nn.Dropout(config.dropout)

feedforward

feedforward_norm

= nn.LayerNorm(config.width)

memory_norm

= nn.LayerNorm(config.width)

query_norm

= nn.LayerNorm(config.width)

nemo_automodel.components.models.clef.joint_schema_head.EvidenceRoutingLayer.forward(
    queries: torch.Tensor,
    memory: torch.Tensor
) -> torch.Tensor

Cross-attend option queries to the encoded record.

Parameters:

queries

torch.Tensor

Tensor of shape [batch, options, width].

memory

torch.Tensor

Tensor of shape [batch, sequence, width].

Returns: torch.Tensor

Tensor of shape [batch, options, width].

class nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead(
    config: nemo_automodel.components.models.clef.joint_schema_head.ClefHeadConfig
)

Bases: Module

Cloudflare CLEF joint head over one or more typed questions.

evidence_layers

field_norm

= nn.LayerNorm(config.width)

global_projection

hidden_norm

= nn.LayerNorm(config.hidden_size)

joint_logit_scale

= nn.Parameter(torch.zeros(()))

layers

memory_projection

option_context_projection

option_lexical_projection

option_norm

= nn.LayerNorm(config.width)

option_question_projection

option_summary_norm

= nn.LayerNorm(config.width)

prior_logit_scale

= nn.Parameter(torch.zeros(()))

question_projection

residual_gate

= nn.Parameter(torch.zeros(()))

residual_scorer

type_embedding

= nn.Embedding(3, config.width)

nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead._mean_span(
    values: torch.Tensor,
    span: tuple[int, int]
) -> torch.Tensor

staticmethod

Mean-pool a non-empty token span.

Parameters:

values

torch.Tensor

Tensor of shape [sequence, hidden].

span

tuple[int, int]

Half-open (start, end) token interval.

Returns: torch.Tensor

Tensor of shape [hidden].

nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    records: list[nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord],
    output_embedding_weight: torch.Tensor | typing.Callable[[torch.Tensor], torch.Tensor]
) -> list[list[torch.Tensor]]

Produce one variable-length option-logit tensor per question.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

input_ids

torch.Tensor

Tensor of shape [batch, sequence].

attention_mask

torch.Tensor

Tensor of shape [batch, sequence], with one for real tokens.

records

list[EncodedRecord]

Batch-aligned token-span metadata.

output_embedding_weight

torch.Tensor | Callable[[torch.Tensor], torch.Tensor]

Tensor of shape [vocab, hidden], or a callable mapping long token IDs [tokens] to rows [tokens, hidden].

Returns: list[list[torch.Tensor]]

Nested batch/question list. Each Tensor has shape [options].

nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead.score_records(
    backbone: torch.nn.Module,
    records: list[nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord],
    lexical: nemo_automodel.components.models.clef.joint_schema_head.LexicalTable,
    pad_token_id: int,
    device: torch.device
) -> list[list[torch.Tensor]]

Run CLEF’s backbone-to-head composition through the public hidden-state interface.

Parameters:

backbone

nn.Module

Language model returning hidden states of shape [batch, sequence, hidden].

records

list[EncodedRecord]

Publisher-ordered token spans. Optional media contains pixel_values of shape [patches, features], pixel_values_videos of shape [video_patches, features], and image/video_grid_thw of shape [items, 3] (temporal, height, width).

lexical

LexicalTable

Frozen token-row lookup returning tensors of shape [tokens, hidden].

pad_token_id

int

Token used for right padding.

device

torch.device

Execution device for collated input tensors.

Returns: list[list[torch.Tensor]]

Batch/question lists of fp32 logit tensors of shape [options].

class nemo_automodel.components.models.clef.joint_schema_head.LexicalTable(
    token_ids: typing.Sequence[int],
    rows: torch.Tensor
)

Bases: Module

Frozen cache of the output-embedding rows the head needs, indexed by token id.

The released head reads lm_head.weight[token_ids] for option tokens. When the LM head is sharded we gather the full weight once to select these rows (they never change, the LM head is frozen), then retain only the selected rows in this table.

nemo_automodel.components.models.clef.joint_schema_head.LexicalTable.forward(
    token_ids: torch.Tensor
) -> torch.Tensor

Look up LM-head rows.

Parameters:

token_ids

torch.Tensor

Long tensor of shape [tokens].

Returns: torch.Tensor

Tensor of shape [tokens, hidden] on token_ids.device.

nemo_automodel.components.models.clef.joint_schema_head._encode_media(
    processor: typing.Any,
    record: dict[str, typing.Any]
) -> tuple[list[int], dict[str, typing.Any] | None]

nemo_automodel.components.models.clef.joint_schema_head._tokens(
    tokenizer: typing.Any,
    text: str
) -> list[int]

nemo_automodel.components.models.clef.joint_schema_head.build_lexical_table(
    model: torch.nn.Module,
    records: typing.Sequence[nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord]
) -> 'LexicalTable'

Cache the frozen output-embedding rows consumed by the CLEF readout.

All model-parallel ranks must participate when the output weights are sharded.

Parameters:

model

torch.nn.Module

Backbone exposing output embeddings of shape [vocabulary, hidden].

records

Sequence[EncodedRecord]

Encoded records whose option spans determine the cached token rows.

Returns: 'LexicalTable'

Frozen table of shape [unique_option_tokens, hidden] on the embedding device.

Raises:

ValueError: If the output embeddings are trainable or no option tokens exist.

nemo_automodel.components.models.clef.joint_schema_head.clef_option_orders(
    encoded: nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord
) -> dict[str, tuple[str, ...]]

{question_id: option ids} in the order the Clef head returns logits.

nemo_automodel.components.models.clef.joint_schema_head.collate_records(
    records: list[nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord],
    pad_token_id: int,
    device: torch.device
) -> dict[str, typing.Any]

Right-pad encoded records into a batch.

Parameters:

records

list[EncodedRecord]

Batch-aligned token spans. Optional media contains pixel_values of shape [patches, patch_features], where patch_features flattens each pixel patch; image_grid_thw of shape [images, 3] (temporal, height, width patch counts); pixel_values_videos and video_grid_thw with corresponding video layouts; and mm_token_type_ids, a list of media-token types starting at token_offset.

pad_token_id

int

Token used to fill right-padded positions.

device

torch.device

Device for input IDs, attention masks, and media token-type tensors.

Returns: dict[str, Any]

Mapping with input_ids and attention_mask, long Tensors of shape [batch, sequence]

nemo_automodel.components.models.clef.joint_schema_head.encode_record(
    tokenizer: typing.Any,
    record: dict[str, typing.Any],
    max_length: int = 16384,
    max_state_tokens: int | None = None,
    processor: typing.Any | None = None
) -> nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord

Tokenise a SystemOne record into Clef’s prompt layout and record question / option token spans.

The state is truncated (never the schema) so the prompt fits in max_length tokens.

nemo_automodel.components.models.clef.joint_schema_head.load_joint_head(
    path: str | pathlib.Path,
    dtype: torch.dtype = torch.float32,
    dropout: float = 0.0,
    revision: str | None = None
) -> nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead

Build a JointSchemaHead from a directory with joint_head_config.json / joint_head.safetensors.

path may also be a Hugging Face repo id (e.g. "Cloudflare/clef"); only the head files are downloaded.

nemo_automodel.components.models.clef.joint_schema_head.option_token_ids(
    encoded_records: typing.Sequence[nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord]
) -> list[int]

Sorted unique token ids inside option spans (the only output-embedding rows the head reads).

nemo_automodel.components.models.clef.joint_schema_head.question_options(
    question: dict[str, typing.Any]
) -> list[tuple[str, typing.Any]]

(option id, description) pairs in Clef order: noul [true, false], choice sorted, score by level.

nemo_automodel.components.models.clef.joint_schema_head.render(
    value: typing.Any
) -> str

Render a state / instruction value as text (strings verbatim, everything else as compact sorted JSON).

nemo_automodel.components.models.clef.joint_schema_head.save_joint_head(
    head: nemo_automodel.components.models.clef.joint_schema_head.JointSchemaHead,
    path: str | pathlib.Path,
    dtype: torch.dtype = torch.bfloat16
) -> None

Save a head in the release layout (joint_head_config.json + joint_head.safetensors).

nemo_automodel.components.models.clef.joint_schema_head.IMAGE_PLACEHOLDER = '<|vision_start|><|image_pad|><|vision_end|>'

nemo_automodel.components.models.clef.joint_schema_head.MEDIA_BATCH_KEYS = ('pixel_values', 'image_grid_thw', 'pixel_values_videos', 'video_grid_thw')

nemo_automodel.components.models.clef.joint_schema_head.MEDIA_TOKEN_KEYS = ('mm_token_type_ids',)

nemo_automodel.components.models.clef.joint_schema_head.QUESTION_TYPES = {'noul': 0, 'choice': 1, 'score': 2}

nemo_automodel.components.models.clef.joint_schema_head.SYSTEM_PROMPT = "Read the complete state and schema. Decide every field jointly. Each answer mus...

nemo_automodel.components.models.clef.joint_schema_head.VIDEO_PLACEHOLDER = '<|vision_start|><|video_pad|><|vision_end|>'