nemo_automodel.components.models.clef

Cloudflare Clef joint schema head (https://huggingface.co/Cloudflare/clef).

Submodules

Package 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.
`JointSchemaHead`	Cloudflare CLEF joint head over one or more typed questions.

Functions

Name	Description
`load_joint_head`	Build a `JointSchemaHead` from a directory with `joint_head_config.json` / `joint_head.safetensors`.
`save_joint_head`	Save a head in the release layout (`joint_head_config.json` + `joint_head.safetensors`).

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

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