tabforge-ml 2026.8.9__py3-none-any.whl

tabforge/__init__.py

@@ -0,0 +1,50 @@
+"""TabFORGE v1 public API."""
+
+from ._metadata import MODEL_ID, __version__
+from .api import (
+    TabFORGEAnomalyDetector,
+    TabFORGEClassifier,
+    TabFORGEClusterer,
+    TabFORGEEmbedder,
+    TabFORGEGenerator,
+    TabFORGEImputer,
+    TabFORGERegressor,
+    restore_from_checkpoint,
+)
+from .checkpoints import (
+    CHECKPOINT_FORMAT_VERSION,
+    detokeniser_compatibility_reason,
+    is_detokeniser_compatible,
+    official_checkpoint_path,
+)
+from .config import (
+    TabFORGEArchitectureConfig,
+    TabFORGEDiffusionConfig,
+    TabFORGEEmbeddingConfig,
+    TabFORGERuntimeConfig,
+    TabFORGESchedulerConfig,
+    TabFORGETrainingConfig,
+)
+
+__all__ = [
+    "CHECKPOINT_FORMAT_VERSION",
+    "MODEL_ID",
+    "__version__",
+    "TabFORGEArchitectureConfig",
+    "TabFORGEAnomalyDetector",
+    "TabFORGEClassifier",
+    "TabFORGEClusterer",
+    "TabFORGEDiffusionConfig",
+    "TabFORGEEmbedder",
+    "TabFORGEEmbeddingConfig",
+    "TabFORGEGenerator",
+    "TabFORGEImputer",
+    "TabFORGERuntimeConfig",
+    "TabFORGESchedulerConfig",
+    "TabFORGERegressor",
+    "TabFORGETrainingConfig",
+    "official_checkpoint_path",
+    "restore_from_checkpoint",
+    "detokeniser_compatibility_reason",
+    "is_detokeniser_compatible",
+]

tabforge/_metadata.py

@@ -0,0 +1,4 @@
+"""Package and model identifiers shared by artifacts and public metadata."""
+
+__version__ = "2026.8.9"
+MODEL_ID = "tabforge-v1"

tabforge/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Private vendored runtimes."""

tabforge/_vendor/tabpfn/__init__.py

@@ -0,0 +1,23 @@
+from importlib.metadata import version
+
+from tabforge._vendor.tabpfn.classifier import TabPFNClassifier
+from tabforge._vendor.tabpfn.misc.debug_versions import display_debug_info
+from tabforge._vendor.tabpfn.model_loading import (
+    load_fitted_tabpfn_model,
+    save_fitted_tabpfn_model,
+)
+from tabforge._vendor.tabpfn.regressor import TabPFNRegressor
+
+try:
+    __version__ = version(__name__)
+except ImportError:
+    __version__ = "unknown"
+
+__all__ = [
+    "TabPFNClassifier",
+    "TabPFNRegressor",
+    "__version__",
+    "display_debug_info",
+    "load_fitted_tabpfn_model",
+    "save_fitted_tabpfn_model",
+]

tabforge/_vendor/tabpfn/architectures/__init__.py

@@ -0,0 +1,20 @@
+"""Contains a collection of different model architectures.
+
+"Architecture" refers to a PyTorch module, which is then wrapped by e.g.
+TabPFNClassifier or TabPFNRegressor to form the complete model.
+
+Each submodule in this module should contain an architecture. Each may be a directory,
+or just a single file.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from . import base
+
+if TYPE_CHECKING:
+    from tabforge._vendor.tabpfn.architectures.interface import ArchitectureModule
+
+ARCHITECTURES: dict[str, ArchitectureModule] = {"base": base}
+"""Map from architecture names to the corresponding module."""

tabforge/_vendor/tabpfn/architectures/base/__init__.py

@@ -0,0 +1,226 @@
+"""The base architecture.
+
+This is the original model before we switched to the multiple architecture arrangement.
+It is a single architecture which implements many different functionalities. Other
+architectures can import components from here to reuse them, and over time we should
+refactor this architecture to improve reusability.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from tabforge._vendor.tabpfn.architectures.base.config import ModelConfig
+from tabforge._vendor.tabpfn.architectures.base.encoders import (
+    InputNormalizationEncoderStep,
+    LinearInputEncoderStep,
+    MLPInputEncoderStep,
+    MulticlassClassificationTargetEncoder,
+    NanHandlingEncoderStep,
+    RemoveDuplicateFeaturesEncoderStep,
+    RemoveEmptyFeaturesEncoderStep,
+    SeqEncStep,
+    SequentialEncoder,
+    VariableNumFeaturesEncoderStep,
+)
+from tabforge._vendor.tabpfn.architectures.base.transformer import PerFeatureTransformer
+
+if TYPE_CHECKING:
+    from torch import nn
+
+    from tabforge._vendor.tabpfn.architectures.interface import ArchitectureConfig
+
+
+def parse_config(config: dict[str, Any]) -> tuple[ArchitectureConfig, dict[str, Any]]:
+    """Parse the config dict into a ModelConfig, a subclass of ArchitectureConfig.
+
+    This method implements the interface defined in
+    tabpfn.architectures.interface.ArchitectureModule.parse_config().
+
+    Unrecognised keys should be ignored during parsing, and returned in the `unused
+    config items` dict.
+
+    Args:
+        config: Config dict to parse. This function should use Pydantic to
+            verify that it matches the expected schema.
+
+    Returns: a tuple (the parsed config, dict containing unused config items)
+
+    Raises:
+        pydantic.ValidationError: one or more of the values have the wrong type
+    """
+    upgraded_dict = ModelConfig.upgrade_config(config)
+    parsed_config = ModelConfig(**upgraded_dict)
+    return parsed_config, parsed_config.get_unused_config(upgraded_dict)
+
+
+def get_architecture(
+    config: ArchitectureConfig,
+    *,
+    n_out: int,
+    cache_trainset_representation: bool,
+) -> PerFeatureTransformer:
+    """Construct the base architecture following the given config.
+
+    This factory method implements the interface defined in
+    tabpfn.architectures.interface.ArchitectureModule.get_architecture().
+
+    Args:
+        config: The config returned by parse_config(). This method should use a
+            runtime isinstance() check to downcast the config to this architecture's
+            specific config class.
+        n_out: The number of output classes that the model should predict.
+        cache_trainset_representation: If True, the model should be configured to
+            cache the training data during inference to improve speed.
+
+    Returns: the constructed architecture
+    """
+    assert isinstance(config, ModelConfig)
+    return PerFeatureTransformer(
+        config=config,
+        # Things that were explicitly passed inside `build_model()`
+        encoder=get_encoder(
+            num_features=config.features_per_group,
+            embedding_size=config.emsize,
+            remove_empty_features=config.remove_empty_features,
+            remove_duplicate_features=config.remove_duplicate_features,
+            nan_handling_enabled=config.nan_handling_enabled,
+            normalize_on_train_only=config.normalize_on_train_only,
+            normalize_to_ranking=config.normalize_to_ranking,
+            normalize_x=config.normalize_x,
+            remove_outliers=config.remove_outliers,
+            normalize_by_used_features=config.normalize_by_used_features,
+            encoder_use_bias=config.encoder_use_bias,
+            encoder_type=config.encoder_type,
+            encoder_mlp_hidden_dim=config.encoder_mlp_hidden_dim,
+            encoder_mlp_num_layers=config.encoder_mlp_num_layers,
+        ),
+        y_encoder=get_y_encoder(
+            num_inputs=1,
+            embedding_size=config.emsize,
+            nan_handling_y_encoder=config.nan_handling_y_encoder,
+            max_num_classes=config.max_num_classes,
+        ),
+        cache_trainset_representation=cache_trainset_representation,
+        use_encoder_compression_layer=False,
+        n_out=n_out,
+        #
+        # These are things that had default values from config.get() but were not
+        # present in any config.
+        layer_norm_with_elementwise_affine=False,
+    )
+
+
+def get_encoder(  # noqa: PLR0913
+    *,
+    num_features: int,
+    embedding_size: int,
+    remove_empty_features: bool,
+    remove_duplicate_features: bool,
+    nan_handling_enabled: bool,
+    normalize_on_train_only: bool,
+    normalize_to_ranking: bool,
+    normalize_x: bool,
+    remove_outliers: bool,
+    normalize_by_used_features: bool,
+    encoder_use_bias: bool,
+    encoder_type: Literal["linear", "mlp"] = "linear",
+    encoder_mlp_hidden_dim: int | None = None,
+    encoder_mlp_num_layers: int = 2,
+) -> nn.Module:
+    inputs_to_merge = {"main": {"dim": num_features}}
+
+    encoder_steps: list[SeqEncStep] = []
+    if remove_empty_features:
+        encoder_steps += [RemoveEmptyFeaturesEncoderStep()]
+
+    if remove_duplicate_features:
+        encoder_steps += [RemoveDuplicateFeaturesEncoderStep()]
+
+    encoder_steps += [NanHandlingEncoderStep(keep_nans=nan_handling_enabled)]
+
+    if nan_handling_enabled:
+        inputs_to_merge["nan_indicators"] = {"dim": num_features}
+
+        encoder_steps += [
+            VariableNumFeaturesEncoderStep(
+                num_features=num_features,
+                normalize_by_used_features=False,
+                in_keys=["nan_indicators"],
+                out_keys=["nan_indicators"],
+            ),
+        ]
+
+    encoder_steps += [
+        InputNormalizationEncoderStep(
+            normalize_on_train_only=normalize_on_train_only,
+            normalize_to_ranking=normalize_to_ranking,
+            normalize_x=normalize_x,
+            remove_outliers=remove_outliers,
+        ),
+    ]
+
+    encoder_steps += [
+        VariableNumFeaturesEncoderStep(
+            num_features=num_features,
+            normalize_by_used_features=normalize_by_used_features,
+        ),
+    ]
+
+    num_input_features = sum(i["dim"] for i in inputs_to_merge.values())
+    if encoder_type == "mlp":
+        encoder_steps += [
+            MLPInputEncoderStep(
+                num_features=num_input_features,
+                emsize=embedding_size,
+                hidden_dim=encoder_mlp_hidden_dim,
+                activation="gelu",
+                num_layers=encoder_mlp_num_layers,
+                bias=encoder_use_bias,
+                in_keys=tuple(inputs_to_merge),
+                out_keys=("output",),
+            ),
+        ]
+    elif encoder_type == "linear":
+        encoder_steps += [
+            LinearInputEncoderStep(
+                num_features=num_input_features,
+                emsize=embedding_size,
+                bias=encoder_use_bias,
+                in_keys=tuple(inputs_to_merge),
+                out_keys=("output",),
+            ),
+        ]
+    else:
+        raise ValueError(
+            f"Invalid encoder type: {encoder_type} (expected 'linear' or 'mlp')"
+        )
+
+    return SequentialEncoder(*encoder_steps, output_key="output")
+
+
+def get_y_encoder(
+    *,
+    num_inputs: int,
+    embedding_size: int,
+    nan_handling_y_encoder: bool,
+    max_num_classes: int,
+) -> nn.Module:
+    steps: list[SeqEncStep] = []
+    inputs_to_merge = [{"name": "main", "dim": num_inputs}]
+    if nan_handling_y_encoder:
+        steps += [NanHandlingEncoderStep()]
+        inputs_to_merge += [{"name": "nan_indicators", "dim": num_inputs}]
+
+    if max_num_classes >= 2:
+        steps += [MulticlassClassificationTargetEncoder()]
+
+    steps += [
+        LinearInputEncoderStep(
+            num_features=sum([i["dim"] for i in inputs_to_merge]),  # type: ignore
+            emsize=embedding_size,
+            in_keys=tuple(i["name"] for i in inputs_to_merge),  # type: ignore
+            out_keys=("output",),
+        ),
+    ]
+    return SequentialEncoder(*steps, output_key="output")

tabforge/_vendor/tabpfn/architectures/base/attention/__init__.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing_extensions import override
+
+import torch
+from torch import nn
+
+
+class Attention(ABC, nn.Module):
+    """Base class for attention layers."""
+
+    @override
+    @abstractmethod
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_kv: torch.Tensor | None = None,
+        *,
+        cache_kv: bool = False,
+        use_cached_kv: bool = False,
+        reuse_first_head_kv: bool = False,
+        only_cache_first_head_kv: bool = False,
+        save_peak_mem_factor: int | None = None,
+        add_input: bool = False,
+        allow_inplace: bool = False,
+    ) -> torch.Tensor:
+        """Performs the attention layer.
+
+        Args:
+            x: Input sequence of embeddings with shape
+                [batch... x query seq len x embedding dim].
+                If `x_kv` is None, this is used to compute the queries, keys, and
+                values.
+                If `x_kv` is not None, this is used to compute the queries only.
+            x_kv: If not None, an input sequence of embeddings with shape
+                [batch... x kv seq len x embedding dim].
+                It will be used to compute the keys and the values, with `x` used only
+                to compute the queries. This is useful to avoid some sequence positions
+                attending to others.
+            cache_kv: If True, replaces the current key-value cache with the keys and
+                values computed during this forward pass. Otherwise, the KV cache is
+                left unchanged. If True, `use_cached_kv` must be False.
+            use_cached_kv: If True, uses the keys and values cached during a previous
+                forward pass when `cache_kv` was True. If True, `cache_kv` must be
+                False.
+            reuse_first_head_kv: If True, then uses the keys and values projected in the
+                first head for all the heads.
+            only_cache_first_head_kv: When True, only caches the keys and values of the
+                first head.
+            save_peak_mem_factor: Loops over the batch dimension rather than processing
+                it in parallel, to reduce memory usage.
+            add_input: If True, returns (x+the output of attention), i.e. enables a
+                residual connection. If False, just returns the output of attention.
+            allow_inplace: By setting this to True, the caller indicates that 'x' is not
+                used after the call and its buffer can be reused for the output.
+                The operation is not guaranteed to be inplace.
+
+        Returns:
+            Output hidden states of shape [batch x query seq len x embedding dim]
+        """
+        ...