Add TabPFN

.github/workflows/push.yml

@@ -39,7 +39,7 @@ jobs:
         with:
           version: latest
           python-version: ${{ matrix.python-version }}
-      - run: uv sync --extra mgm --extra mcp --extra trl --dev --extra struct-cpu 
+      - run: uv sync --extra mgm --extra mcp --extra trl --dev --extra struct-cpu  --extra tabpfn
       - run: uv run pytest
       - env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

pyproject.toml

@@ -136,6 +136,9 @@ trl = [
     "trl",
     "accelerate",
 ]
+tabpfn = [
+    "tabpfn>=6.0.0",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lobster/callbacks/_inverse_folding_callback.py

@@ -1,8 +1,8 @@
 import lightning
+import logging
 import os
 import torch
 from lobster.model.latent_generator.io import writepdb
-from loguru import logger
 from lobster.model.latent_generator.utils.residue_constants import (
     convert_lobster_aa_tokenization_to_standard_aa,
     restype_order_with_x_inv,
@@ -14,6 +14,8 @@
 from torch.utils.data import DataLoader
 import pooch
 
+logger = logging.getLogger(__name__)
+
 
 class InverseFoldingCallback(lightning.Callback):
     def __init__(

src/lobster/callbacks/_structure_decode.py

@@ -1,9 +1,11 @@
 import lightning
+import logging
 import os
 import torch
 from lobster.model.latent_generator.io import writepdb
 from lobster.model.latent_generator.utils.residue_constants import convert_lobster_aa_tokenization_to_standard_aa
-from loguru import logger
+
+logger = logging.getLogger(__name__)
 
 
 class StructureDecodeCallback(lightning.Callback):

src/lobster/callbacks/_unconditional_generation.py

@@ -1,15 +1,17 @@
 import lightning
+import logging
 import os
 import torch
 from lobster.model.latent_generator.io import writepdb
-from loguru import logger
 from lobster.model.latent_generator.utils.residue_constants import (
     convert_lobster_aa_tokenization_to_standard_aa,
     restype_order_with_x_inv,
 )
 from lobster.model import LobsterPLMFold
 from lobster.metrics import get_folded_structure_metrics
 
+logger = logging.getLogger(__name__)
+
 
 class UnconditionalGenerationCallback(lightning.Callback):
     def __init__(

src/lobster/cmdline/generate.py

@@ -5,7 +5,6 @@
 import hydra
 import torch
 from omegaconf import DictConfig, OmegaConf
-from loguru import logger
 
 from lobster.model.latent_generator.io import writepdb, load_pdb
 from lobster.model.latent_generator.utils.residue_constants import (
@@ -32,6 +31,7 @@
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 @hydra.main(version_base=None, config_path="../hydra_config", config_name="generate")

src/lobster/constants/__init__.py

@@ -38,6 +38,7 @@
     DEFAULT_AF2_PREDICTION_MODELS,
     DEFAULT_AF2_WEIGHTS_DIR,
 )
+from ._tabpfn import MAX_FEATURES_PER_MODEL, MAX_SAMPLES_PER_MODEL
 
 __all__ = [
     "BIOPYTHON_FEATURES",
@@ -82,4 +83,6 @@
     "S3_BUCKET",
     "DEFAULT_AF2_PREDICTION_MODELS",
     "DEFAULT_AF2_WEIGHTS_DIR",
+    "MAX_FEATURES_PER_MODEL",
+    "MAX_SAMPLES_PER_MODEL",
 ]

src/lobster/constants/_tabpfn.py

@@ -0,0 +1,2 @@
+MAX_FEATURES_PER_MODEL = 500
+MAX_SAMPLES_PER_MODEL = 50000

src/lobster/data/__init__.py

@@ -17,6 +17,7 @@
 from ._mmseqs import MMSeqsRunner
 from ._structure_datamodule import PDBDataModule
 
+from ._parquet_datamodule import ParquetLightningDataModule
 from ._ume_datamodule import UMELightningDataModule
 from ._utils import download_from_s3, get_s3_bucket_and_key, load_pickle, upload_to_s3
 
@@ -36,6 +37,7 @@
     "M320MLightningDataModule",
     "LobsterMinHasher",
     "MMSeqsRunner",
+    "ParquetLightningDataModule",
     "PDBDataModule",
     "load_pickle",
     "UMELightningDataModule",