Model API

RelationalTransformer is the high-level entry point for loading checkpoints and computing predictions; the Quickstart shows it in context. RTJModel is the underlying torch.nn.Module for custom training loops and architecture work, and ModelOutput carries every result field.

RelationalTransformer

class relational_transformers.RelationalTransformer(model_name_or_path=None, *, task='classification', backend='torch', device=None, revision=None, compile=False, providers=None)

Load and run an RT-J checkpoint.

Parameters:

  • model_name_or_path (str | Path | None) – Hugging Face repository, local checkpoint directory, or ONNX file. Defaults to DEFAULT_MODEL, or DEFAULT_ONNX_MODEL for the ONNX backend.

  • task (str) – "classification" (default) or "regression".

  • backend (str) – "torch", "triton", "onnx", or "meta".

  • device (str | torch.device | None) – PyTorch device. Auto-selects CUDA, MPS, then CPU.

  • revision (str | None)

  • compile (bool)

forward(inputs, *, output='target_scores', target=None)

Run the model and return the requested ModelOutput view.

Parameters:

  • inputs – A RelationalBatch, a mapping of canonical batch fields, a [cells, 2*d_text] all-text cell array, or a sequence of single-context inputs to collate.

  • output – "target_scores", "token_scores", "target_features", "target_scores_and_text", or "embeddings". The ONNX and Triton backends serve "target_scores" only.

  • target – Target cell position(s), required only for raw cell arrays.

Returns:

A ModelOutput with torch tensors on every backend.

Return type:

ModelOutput

predict(inputs, *, target=None, task_head=None, activation=None, convert_to_numpy=True)

Predict target values for one or more relational contexts.

Parameters:

  • inputs – Anything forward() accepts.

  • target – Target cell position(s) for raw cell arrays.

  • task_head (str | None) – Name of a head fitted with fit_head(). When set, the head runs over frozen target_features instead of the published scalar decoder.

  • activation (str | None) – None selects a default from the problem type: sigmoid for binary, multilabel, and classification tasks, softmax for multiclass heads, and identity for regression. Pass "identity" to obtain raw logits.

  • convert_to_numpy – Return numpy values (a Python float for a single context) instead of a torch tensor.

Returns:

Activated scores with one value per context, or one row per context for multi-label heads.

encode(inputs, *, target=None, output_value='embeddings', convert_to_numpy=True)

Return contextual cell states or the pooled target representation.

Parameters:

  • inputs – Anything forward() accepts.

  • target – Target cell position(s) for raw cell arrays.

  • output_value – "embeddings" for one d_model-wide state per cell, or "target_features" for the summed target state used by task heads.

  • convert_to_numpy – Return numpy arrays instead of torch tensors.

export_onnx(path, example, *, target=None, opset_version=18)

Export target-score inference to ONNX with dynamic batch and cell axes.

Requires the torch backend. The example batch fixes d_text and the architecture; batch size and context length stay dynamic.

save_pretrained(directory)

Write model.safetensors and config.json to a directory.

The directory loads back through the normal constructor and follows the same layout as published Hugging Face checkpoints.

fit_head(examples, **kwargs)

Fit a named TaskHead over frozen features.

Each example is encoded once with the frozen backbone, then a linear head is trained on the resulting [d_model] target features. The fitted head is stored under self.heads[task] and served through predict(..., task_head=task). See relational_transformers.training.fit_head() for keyword options.

get_model_kwargs()

Return dimensions without materializing weights (especially useful on meta).

Return type:

dict

RTJModel

class relational_transformers.RTJModel(num_blocks=12, d_model=512, d_text=384, num_heads=8, d_ff=2048, legacy_attn=False, *, device=None)

The RT-J backbone and published decoder heads.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

  • num_blocks (int)

  • d_model (int)

  • d_text (int)

  • num_heads (int)

  • d_ff (int)

  • legacy_attn (bool)

forward(batch, output='target_scores')

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

ModelOutput

ModelOutput

class relational_transformers.ModelOutput(scores: 'Tensor | None' = None, token_scores: 'Tensor | None' = None, target_text: 'Tensor | None' = None, features: 'Tensor | None' = None, embeddings: 'Tensor | None' = None)
Parameters: