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.

Whether this model supports additional related context.

forward(x_context: Tensor | TableTensor | EnsembleTable, y_context: Tensor | TableTensor | EnsembleTable, x_query: Tensor | TableTensor | EnsembleTable, related_context_tables: RelatedTables | None = None, related_query_tables: RelatedTables | None = None, *, recipe: Recipe | None = None, num_estimators: int | None = None, callbacks: Sequence[Callback] | None = None, generator: Generator | None = None, **kwargs: Any) → TableTensor#

The in-context learning forward pass.

Parameters:

  • x_context (Tensor | TableTensor | EnsembleTable) – The feature tensor of in-context examples with shape [..., R_context, D] with R_context rows and D columns.

  • y_context (Tensor | TableTensor | EnsembleTable) – The targets of in-context examples with shape [..., R_context, 1].

  • x_query (Tensor | TableTensor | EnsembleTable) – 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).

  • 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(x: Tensor | TableTensor | EnsembleTable, y: Tensor | TableTensor | EnsembleTable, related_tables: RelatedTables | None = None, *, recipe: Recipe | None = None, num_estimators: int | None = None, callbacks: Sequence[Callback] | None = None, generator: Generator | None = None, **kwargs: Any) → None#

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 | EnsembleTable) – The feature tensor of in-context examples with shape [..., R, D] with R rows and C columns.

  • y (Tensor | TableTensor | EnsembleTable) – 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).

  • 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 | EnsembleTable, 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 | EnsembleTable) – 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