ICLModel#

class sdm.models.ICLModel(task: Task | str | Iterable[Task | str] | None = None)#

Bases: Module, ABC

Base model for in-context foundation models on structured data.

ICLModel defines the public interface shared among in-context foundation models on structured data. It enriches models by unified pre-processing and post-processing routines, key/value caching, and ensembling.

Parameters:: task (Task | str | Iterable[Task | str] | None) – The tasks to initialize. If None, all tasks supported by this model are initialized.

supported_feature_stypes: ClassVar[frozenset[Stype]]#: Semantic types supported for input columns in this model.

supported_target_stypes: ClassVar[frozenset[Stype]]#: Semantic types supported for target columns in this model.

supported_tasks: ClassVar[frozenset[Task]]#: Prediction tasks supported in this model.

supports_multi_target: ClassVar[bool]#: Whether this model supports multi-target predictions.

supports_context_only_columns: ClassVar[bool]#: Whether this model supports columns available only in the context.

supports_related_tables: ClassVar[bool]#: Whether this model supports additional related context.

The in-context learning forward pass.

Parameters:

x_context (Tensor | TableTensor) – The feature tensor of in-context examples with shape [..., R_context, D] with R_context rows and D columns.
y_context (Tensor | TableTensor) – The targets of in-context examples with shape [..., R_context, 1].
x_query (Tensor | TableTensor) – The feature tensor of query examples with shape [..., R_query, D] with R_query rows and D columns.
related_context_tables (RelatedTables | None) – Related context for in-context examples.
related_query_tables (RelatedTables | None) – Related context for query examples.
recipe (Recipe | None) – The custom recipe for pre- and post-processing.
num_estimators (int | None) – The number of estimators E for ensembling. If None, the leading dimension of higher-rank inputs is used as the estimator dimension, allowing input data to be customized per estimator (e.g., different in-context examples per estimator).
estimator_batch_size (int | None) – Maximum number of consecutive estimators run through the model in one call. 1 (default) runs estimators one by one; None batches as many as possible. Estimators whose preprocessed tables differ in shape or target class set, or that come with related tables, run in separate calls. Device memory grows with the batch size.
callbacks (Sequence[Callback] | None) – Callbacks applied in sequence to this model call.
generator (Generator | None) – Pseudorandom number generator used for sampling during pre-processing and model execution.
kwargs (Any) – Additional keyword arguments passed to the model.

Returns:

The processed prediction after applying recipe.output to the stacked estimator outputs with shape [E, ..., R_query, *].

Return type:

TableTensor

Fit and cache in-context examples.

Repeated calls to predict() can then reuse the same in-context examples while only providing new query examples.

Parameters:

x (Tensor | TableTensor) – The feature tensor of in-context examples with shape [..., R, D] with R rows and C columns.
y (Tensor | TableTensor) – The targets of in-context examples with shape [..., R, 1].
related_tables (RelatedTables | None) – Related context for in-context examples.
recipe (Recipe | None) – The custom recipe for pre- and post-processing.
num_estimators (int | None) – The number of estimators E for ensembling. If None, the leading dimension of higher-rank inputs is used as the estimator dimension, allowing input data to be customized per estimator (e.g., different in-context examples per estimator).
context_only_columns (Sequence[str] | None) – Feature column names available only in context and omitted from prediction inputs.
estimator_batch_size (int | None) – Maximum number of consecutive estimators run through the model in one call. 1 (default) runs estimators one by one; None batches as many as possible. Estimators whose preprocessed tables differ in shape or target class set, or that come with related tables, run in separate calls. Device memory grows with the batch size. Estimators fitted together are predicted together.
callbacks (Sequence[Callback] | None) – Callbacks applied in sequence to this model call.
generator (Generator | None) – Pseudorandom number generator used for sampling during pre-processing and model execution.
kwargs (Any) – Additional keyword arguments passed to the model.

Return type:

None

predict(x: Tensor | TableTensor, related_tables: RelatedTables | None = None, *, callbacks: Sequence[Callback] | None = None) → TableTensor#

Predict unseen query examples.

Note

This method requires a prior call to fit().

Parameters:

x (Tensor | TableTensor) – The feature tensor of query examples with shape [..., R, D] with R rows and D columns.
related_tables (RelatedTables | None) – Related context for query examples.
callbacks (Sequence[Callback] | None) – Callbacks applied in sequence to this model call.

Returns:

The processed prediction after applying recipe.output to the stacked estimator outputs with shape [E, ..., R, *].

Return type:

TableTensor

clear() → None#

Clear cached context state created by fit().

Return type:: None

abstractmethod classmethod default_recipe() → Recipe#

Return the default processing recipe for this model.

Return type:: Recipe