Training API

The narrative guides for these classes are the Training Overview and the training subpages.

RelationalExample

class relational_transformers.RelationalExample(input, label, target=None)

One labeled context used for training or evaluation.

target is required when input is a raw cell-vector array. Typed RelationalBatch inputs already carry their target mask and leave it unset.

Parameters:

  • input (Any)

  • label (Any)

  • target (int | Sequence[int] | None)

RelationalTrainingArguments

class relational_transformers.RelationalTrainingArguments(output_dir: 'str' = 'relational_model', num_train_epochs: 'int' = 1, per_device_train_batch_size: 'int' = 8, learning_rate: 'float' = 1e-05, weight_decay: 'float' = 0.01, max_grad_norm: 'float' = 1.0, gradient_accumulation_steps: 'int' = 1, seed: 'int' = 42, logging_steps: 'int' = 10, save_strategy: 'str' = 'epoch', training_backend: 'str' = 'torch')
Parameters:

  • output_dir (str)

  • num_train_epochs (int)

  • per_device_train_batch_size (int)

  • learning_rate (float)

  • weight_decay (float)

  • max_grad_norm (float)

  • gradient_accumulation_steps (int)

  • seed (int)

  • logging_steps (int)

  • save_strategy (str)

  • training_backend (str)

RelationalTrainer

class relational_transformers.RelationalTrainer(*, model, args, train_dataset, task=None, problem_type=None)

Small, dependency-free trainer for complete RT-J fine-tuning.

Parameters:

FineTunedHead

class relational_transformers.FineTunedHead(weight, bias, *, task, initial_loss=None, final_loss=None, seconds=None, n_examples=None, classes=(), feat_mu=None, feat_sd=None, column_stats=None, normalization_mode='zero_shot')

A trained task head over the frozen backbone’s target-cell features.

The transformer is never updated; this is the small linear adapter that replaces a released checkpoint’s zero-shot head. predict is plain numpy, so a fitted head serves anywhere. column_stats rides along opaquely (any object with to_dict(), persisted as JSON and returned as a dict on load) so the artifact carries the preprocessing it was fitted under without this package interpreting it.

Parameters:

  • task (str)

  • classes (Sequence[Any])

  • normalization_mode (str)

save(path)

Persist the head plus the preprocessing it was fitted under.

Return type:

str

predict(features)

Score frozen features [N, d_model] -> logits [N, n_outputs].

fit_feature_head

relational_transformers.fit_feature_head(features, labels, task, *, classes=(), group_offsets=None, n_groups=0, epochs=100, learning_rate=0.001, weight_decay=0.0001, class_embeddings=None, text_decoder=None, column_stats=None, normalization_mode='zero_shot')

Fit a task head on frozen features [N, d_model] with AdamW.

task is one of binary, regression, multiclass, or ranking (grouped by group_offsets). Features are standardized per dimension before fitting: the backbone’s target-cell features sit in a very narrow cone (mean pairwise cosine 0.9976 on a 240-issue sample), and without standardization a linear head fits only its bias.

For multiclass, pass class_embeddings (normalized label embeddings, [C, d_text]) together with the checkpoint’s text_decoder (the d_model -> d_text linear, e.g. model.model.dec_dict["text"]) to seed the head in the checkpoint’s own class-embedding basis, so training starts from the zero-shot ordering.

Parameters:

  • task (str)

  • classes (Sequence[Any])

  • n_groups (int)

  • epochs (int)

  • learning_rate (float)

  • weight_decay (float)

  • normalization_mode (str)

Return type:

FineTunedHead

TaskHead

class relational_transformers.TaskHead(d_model, num_labels=1, problem_type='binary')

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

Parameters:

  • d_model (int)

  • num_labels (int)

  • problem_type (str)

forward(features)

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:

features (Tensor)

Return type:

Tensor

fit_head

relational_transformers.training.fit_head(transformer, examples, *, task, num_labels=1, problem_type='binary', epochs=100, learning_rate=0.001, weight_decay=0.0001)

Encode each example once and fit a lightweight task head.

Parameters:

  • examples (Sequence[RelationalExample])

  • task (str)

  • num_labels (int)

  • problem_type (str)

  • epochs (int)

  • learning_rate (float)

  • weight_decay (float)

Return type:

TaskHead