alf_tools 0.1.0__tar.gz

alf_tools-0.1.0/PKG-INFO

@@ -0,0 +1,176 @@
+Metadata-Version: 2.4
+Name: alf_tools
+Version: 0.1.0
+Summary: A Python package with tools for performing active learning experiments.
+Author: InstaDeep Ltd
+License: Apache-2.0
+Project-URL: Homepage, https://instadeepai.github.io/alf/
+Project-URL: Repository, https://github.com/instadeepai/alf
+Project-URL: Documentation, https://instadeepai.github.io/alf/
+Project-URL: Paper, https://arxiv.org/abs/2609.31197
+Project-URL: Bug Tracker, https://github.com/instadeepai/alf/issues
+Keywords: active learning,bayesian optimization,machine learning,scientific discovery,pytorch
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.14.1,>=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: alf_core
+Requires-Dist: requests>=2.25.0
+Requires-Dist: huggingface_hub<2.0,>=0.27.0
+Requires-Dist: pyarrow>=22.0.0
+Requires-Dist: jaxtyping>=0.3.5
+Requires-Dist: torch>=2.10.0
+Requires-Dist: gpytorch<2.0,>=1.9.0
+Requires-Dist: scipy>=1.9.0
+Requires-Dist: botorch<1.0,>=0.17.0
+Provides-Extra: transformers
+Requires-Dist: transformers>=4.36.0; extra == "transformers"
+Provides-Extra: chemprop
+Requires-Dist: chemprop<2.3.0,>=2.0.0; extra == "chemprop"
+Provides-Extra: esm2
+Requires-Dist: alf_tools[transformers]; extra == "esm2"
+Provides-Extra: esmfold
+Requires-Dist: alf_tools[transformers]; extra == "esmfold"
+Requires-Dist: accelerate>=0.26.0; extra == "esmfold"
+Provides-Extra: guacamol
+Requires-Dist: rdkit>=2022.9.0; extra == "guacamol"
+Provides-Extra: mlip
+Requires-Dist: mlip>=0.2.0; extra == "mlip"
+Provides-Extra: mlip-cuda13
+Requires-Dist: mlip[cuda13]>=0.2.0; extra == "mlip-cuda13"
+Provides-Extra: matbench
+Requires-Dist: matbench>=0.6; extra == "matbench"
+Requires-Dist: pymatgen>=2023.0; extra == "matbench"
+Provides-Extra: protein
+Requires-Dist: alf_tools[esm2,esmfold]; extra == "protein"
+Provides-Extra: molecule
+Requires-Dist: alf_tools[chemprop,guacamol,mlip]; extra == "molecule"
+Provides-Extra: materials
+Requires-Dist: alf_tools[matbench]; extra == "materials"
+
+# <img src="https://raw.githubusercontent.com/instadeepai/alf/main/docs/imgs/alf_cover_gradient.png" alt="ALF" height="40" align="top"> alf-tools
+
+[![PyPI](https://img.shields.io/pypi/v/alf-tools.svg)](https://pypi.org/project/alf-tools/)
+[![Python Version](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/license-Apache%202.0-green.svg)](https://github.com/instadeepai/alf/blob/main/LICENSE)
+[![Coverage](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/surana01/25ce4b64d5b9cda177203366146f5bf0/raw/alf-tools-coverage.json)](https://github.com/instadeepai/alf/tree/main/tools)
+[![Docs](https://img.shields.io/badge/docs-instadeepai.github.io%2Falf-blue)](https://instadeepai.github.io/alf/)
+
+**Ready-to-use models, datasets, and acquisition functions for ALF (Active Learning Framework).**
+
+`alf-tools` builds on [alf-core](https://github.com/instadeepai/alf/blob/main/core/README.md)
+to give you everything needed to run active-learning experiments out of the box — example
+datasets (GFP, ProteinGym, FLIP, GuacaMol), models (CNN, Gaussian Process, ESM-2, Chemprop),
+acquisition functions, and search strategies. Use it for quick-start and prototyping; reach for
+`alf-core` alone when you want the lightweight framework with no ML-framework dependencies.
+
+<div align="center">
+  <img src="https://raw.githubusercontent.com/instadeepai/alf/main/docs/imgs/alf_main_figure.png" alt="ALF active-learning loop overview" width="70%">
+</div>
+
+## Installation
+
+```bash
+pip install alf-tools
+```
+
+This installs PyTorch and the core dependencies. Some models need additional dependencies —
+see [Optional Extras](#optional-extras) below.
+
+### Optional Extras
+
+Some models require additional dependencies. Append one or more extras to the package name.
+Per-model extras (`esm2`, `esmfold`, `chemprop`, `guacamol`) install exactly one model's
+dependencies; workflow umbrellas (`protein`, `molecule`) group the extras you are likely to use
+together.
+
+```bash
+pip install "alf-tools[esm2]"      # ESM-2 protein language model (ESM2Model)
+pip install "alf-tools[esmfold]"   # ESMFold structure-prediction oracle (ESMFoldModel)
+pip install "alf-tools[chemprop]"  # Chemprop small-molecule MPNN (ChempropModel)
+pip install "alf-tools[guacamol]"  # GuacaMol RDKit-based dataset/scoring
+pip install "alf-tools[protein]"   # umbrella: esm2 + esmfold
+pip install "alf-tools[molecule]"  # umbrella: chemprop + guacamol
+```
+
+## Documentation
+
+- **Full documentation:** [instadeepai.github.io/alf](https://instadeepai.github.io/alf/)
+- **API reference:** [alf-tools API](https://instadeepai.github.io/alf/api/alf_tools/index.html)
+- **Installation guide:** [instadeepai.github.io/alf/installation.html](https://instadeepai.github.io/alf/installation.html)
+- **Core framework:** [alf-core](https://github.com/instadeepai/alf/blob/main/core/README.md)
+- **Tutorials:** [tutorials/](https://github.com/instadeepai/alf/tree/main/tutorials)
+
+## What's included
+
+`alf-tools` bundles example datasets (GFP, ProteinGym, FLIP, GuacaMol), surrogate and oracle
+models (CNN, Gaussian Process, ESM-2, ESMFold, Chemprop, PyRosetta, plus an ensemble wrapper),
+acquisition functions (Greedy, UCB, Expected Improvement, Thompson Sampling, CoreSet, and a
+BoTorch wrapper), and search strategies. See the
+[API reference](https://instadeepai.github.io/alf/api/alf_tools/index.html) for the full
+catalogue and configuration options.
+
+## Quick example
+
+```python
+from alf_core import (
+    BaseDatasetConfig,
+    DatasetSearch,
+    DesignTask,
+    Optimizer,
+    Oracle,
+    Surrogate,
+    TerminalStateLogger,
+)
+from alf_tools.datasets import GFP
+from alf_tools.models import CNNModel
+from alf_tools.optimizer.acquisition_functions import Greedy
+
+# Configure and load the dataset
+config = BaseDatasetConfig(
+    name="gfp",
+    modality="sequence",
+    seed=42,
+    train_ratio=0.1,
+    validation_frac=0.5,
+    test_ratio=0.2,
+    split_type="random",
+    problem_type="regression",
+)
+dataset = GFP(config)
+surrogate = Surrogate(model=CNNModel())
+optimizer = Optimizer(acquisition_fn=Greedy(), search_fn=DatasetSearch())
+oracle = Oracle(scorer=dataset)
+
+# Run active learning
+task = DesignTask(num_acq_rounds=5, acq_batch_size=100)
+state = task.setup(dataset=dataset, surrogate=surrogate)
+task.run(
+    state=state,
+    state_loggers=[TerminalStateLogger()],
+    optimizer=optimizer,
+    oracle=oracle,
+)
+```
+
+## Normalisation
+
+Models support input normalisation and output standardisation through their train configs
+(`BaseTrainConfig`). `GPModel` enables both by default (`minmax` inputs, standardised outputs);
+the other models default to neither. See the
+[API reference](https://instadeepai.github.io/alf/api/alf_tools/index.html) for per-model
+defaults and configuration.
+
+## Creating custom components
+
+All components subclass the `alf_core` base classes (`BaseDataset`, `BaseModel`,
+`AcquisitionFunction`, `BaseSearch`). See the
+[how-to guides](https://instadeepai.github.io/alf/how-to/index.html) for step-by-step instructions.
+
+## License
+
+Apache License 2.0 — see [LICENSE](https://github.com/instadeepai/alf/blob/main/LICENSE).

alf_tools-0.1.0/README.md

@@ -0,0 +1,122 @@
+# <img src="https://raw.githubusercontent.com/instadeepai/alf/main/docs/imgs/alf_cover_gradient.png" alt="ALF" height="40" align="top"> alf-tools
+
+[![PyPI](https://img.shields.io/pypi/v/alf-tools.svg)](https://pypi.org/project/alf-tools/)
+[![Python Version](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/license-Apache%202.0-green.svg)](https://github.com/instadeepai/alf/blob/main/LICENSE)
+[![Coverage](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/surana01/25ce4b64d5b9cda177203366146f5bf0/raw/alf-tools-coverage.json)](https://github.com/instadeepai/alf/tree/main/tools)
+[![Docs](https://img.shields.io/badge/docs-instadeepai.github.io%2Falf-blue)](https://instadeepai.github.io/alf/)
+
+**Ready-to-use models, datasets, and acquisition functions for ALF (Active Learning Framework).**
+
+`alf-tools` builds on [alf-core](https://github.com/instadeepai/alf/blob/main/core/README.md)
+to give you everything needed to run active-learning experiments out of the box — example
+datasets (GFP, ProteinGym, FLIP, GuacaMol), models (CNN, Gaussian Process, ESM-2, Chemprop),
+acquisition functions, and search strategies. Use it for quick-start and prototyping; reach for
+`alf-core` alone when you want the lightweight framework with no ML-framework dependencies.
+
+<div align="center">
+  <img src="https://raw.githubusercontent.com/instadeepai/alf/main/docs/imgs/alf_main_figure.png" alt="ALF active-learning loop overview" width="70%">
+</div>
+
+## Installation
+
+```bash
+pip install alf-tools
+```
+
+This installs PyTorch and the core dependencies. Some models need additional dependencies —
+see [Optional Extras](#optional-extras) below.
+
+### Optional Extras
+
+Some models require additional dependencies. Append one or more extras to the package name.
+Per-model extras (`esm2`, `esmfold`, `chemprop`, `guacamol`) install exactly one model's
+dependencies; workflow umbrellas (`protein`, `molecule`) group the extras you are likely to use
+together.
+
+```bash
+pip install "alf-tools[esm2]"      # ESM-2 protein language model (ESM2Model)
+pip install "alf-tools[esmfold]"   # ESMFold structure-prediction oracle (ESMFoldModel)
+pip install "alf-tools[chemprop]"  # Chemprop small-molecule MPNN (ChempropModel)
+pip install "alf-tools[guacamol]"  # GuacaMol RDKit-based dataset/scoring
+pip install "alf-tools[protein]"   # umbrella: esm2 + esmfold
+pip install "alf-tools[molecule]"  # umbrella: chemprop + guacamol
+```
+
+## Documentation
+
+- **Full documentation:** [instadeepai.github.io/alf](https://instadeepai.github.io/alf/)
+- **API reference:** [alf-tools API](https://instadeepai.github.io/alf/api/alf_tools/index.html)
+- **Installation guide:** [instadeepai.github.io/alf/installation.html](https://instadeepai.github.io/alf/installation.html)
+- **Core framework:** [alf-core](https://github.com/instadeepai/alf/blob/main/core/README.md)
+- **Tutorials:** [tutorials/](https://github.com/instadeepai/alf/tree/main/tutorials)
+
+## What's included
+
+`alf-tools` bundles example datasets (GFP, ProteinGym, FLIP, GuacaMol), surrogate and oracle
+models (CNN, Gaussian Process, ESM-2, ESMFold, Chemprop, PyRosetta, plus an ensemble wrapper),
+acquisition functions (Greedy, UCB, Expected Improvement, Thompson Sampling, CoreSet, and a
+BoTorch wrapper), and search strategies. See the
+[API reference](https://instadeepai.github.io/alf/api/alf_tools/index.html) for the full
+catalogue and configuration options.
+
+## Quick example
+
+```python
+from alf_core import (
+    BaseDatasetConfig,
+    DatasetSearch,
+    DesignTask,
+    Optimizer,
+    Oracle,
+    Surrogate,
+    TerminalStateLogger,
+)
+from alf_tools.datasets import GFP
+from alf_tools.models import CNNModel
+from alf_tools.optimizer.acquisition_functions import Greedy
+
+# Configure and load the dataset
+config = BaseDatasetConfig(
+    name="gfp",
+    modality="sequence",
+    seed=42,
+    train_ratio=0.1,
+    validation_frac=0.5,
+    test_ratio=0.2,
+    split_type="random",
+    problem_type="regression",
+)
+dataset = GFP(config)
+surrogate = Surrogate(model=CNNModel())
+optimizer = Optimizer(acquisition_fn=Greedy(), search_fn=DatasetSearch())
+oracle = Oracle(scorer=dataset)
+
+# Run active learning
+task = DesignTask(num_acq_rounds=5, acq_batch_size=100)
+state = task.setup(dataset=dataset, surrogate=surrogate)
+task.run(
+    state=state,
+    state_loggers=[TerminalStateLogger()],
+    optimizer=optimizer,
+    oracle=oracle,
+)
+```
+
+## Normalisation
+
+Models support input normalisation and output standardisation through their train configs
+(`BaseTrainConfig`). `GPModel` enables both by default (`minmax` inputs, standardised outputs);
+the other models default to neither. See the
+[API reference](https://instadeepai.github.io/alf/api/alf_tools/index.html) for per-model
+defaults and configuration.
+
+## Creating custom components
+
+All components subclass the `alf_core` base classes (`BaseDataset`, `BaseModel`,
+`AcquisitionFunction`, `BaseSearch`). See the
+[how-to guides](https://instadeepai.github.io/alf/how-to/index.html) for step-by-step instructions.
+
+## License
+
+Apache License 2.0 — see [LICENSE](https://github.com/instadeepai/alf/blob/main/LICENSE).

alf_tools-0.1.0/alf_tools/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2026 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.metadata
+
+__version__ = importlib.metadata.version("alf-tools")

alf_tools-0.1.0/alf_tools/datasets/__init__.py

@@ -0,0 +1,33 @@
+# Copyright 2026 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from alf_tools.datasets.flip import FLIP, FLIPConfig
+from alf_tools.datasets.gfp import GFP
+from alf_tools.datasets.proteingym import ProteinGym
+
+try:
+    from alf_tools.datasets.guacamol import (
+        GuacaMol,
+        GuacaMolConfig,
+        GuacaMolFileInfo,
+        download_guacamol,
+        get_task_scorer,
+    )
+except ImportError:
+    pass
+
+try:
+    from alf_tools.datasets.matbench import Matbench, MatbenchConfig
+except ImportError:
+    pass

alf_tools-0.1.0/alf_tools/datasets/flip.py

@@ -0,0 +1,312 @@
+# Copyright 2026 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import io
+import logging
+import zipfile
+from math import floor
+from pathlib import Path
+from typing import Literal, Self
+
+import numpy as np
+import pandas as pd
+import requests
+from alf_core import BaseDataset, BaseDatasetConfig, Candidate, LabelledCandidates, ProblemType
+from pydantic import model_validator
+
+logger = logging.getLogger("alf-tools")
+
+DATAPATH = Path(__file__).parent / "data"
+FLIP_GITHUB_BASE = "https://github.com/J-SNACKKB/FLIP/raw/main/splits"
+
+FLIP_DATASETS = Literal[
+    "aav",
+    "gb1",
+    "meltome",
+    "scl",
+    "sav",
+]
+
+# Active (green) splits per dataset — orange/red splits and FASTA-only datasets excluded
+FLIP_SPLITS: dict[str, list[str]] = {
+    "aav": ["des_mut", "mut_des", "one_vs_many", "two_vs_many", "seven_vs_many", "low_vs_high"],
+    "gb1": ["one_vs_rest", "two_vs_rest", "three_vs_rest", "low_vs_high"],
+    "meltome": ["mixed", "human", "human_cell"],
+    "scl": ["mixed_soft", "mixed_hard", "human_soft", "human_hard", "balanced", "mixed_vs_human_2"],
+    "sav": ["mixed", "human", "only_savs"],
+}
+
+# A few FLIP split names don't match their CSV file stem exactly
+_SPLIT_FILE_MAP: dict[tuple[str, str], str] = {
+    ("meltome", "mixed"): "mixed_split",
+}
+
+_EXPECTED_COLUMNS = {"sequence", "target", "set", "validation"}
+
+
+class FLIPConfig(BaseDatasetConfig):
+    """Configuration for FLIP benchmark datasets.
+
+    `train_ratio` controls what fraction of the FLIP train sequences form the
+    initial labelled training set. `validation_frac` is applied to that portion
+    to carve out a validation set (same as `BaseDataset`). The remaining FLIP
+    train sequences become the candidate pool (capped at `max_candidate_pool`).
+    `test_ratio` controls what fraction of the FLIP test sequences form the
+    evaluation set. `split_type` defaults to `"random"` and need not be set.
+
+    Attributes:
+        flip_dataset: Name of the FLIP dataset (e.g. `"gb1"`, `"aav"`).
+        flip_split: Name of the split within the dataset (e.g. `"one_vs_rest"`).
+            Must be an active (green-status) split. See `FLIP_SPLITS` for valid options.
+
+    Example::
+
+        config = FLIPConfig(
+            name="gb1_experiment",
+            modality="sequence",
+            seed=42,
+            flip_dataset="gb1",
+            flip_split="one_vs_rest",
+            train_ratio=0.1,
+            validation_frac=0.1,
+            test_ratio=0.8,
+            max_candidate_pool=5000,
+        )
+    """
+
+    flip_dataset: FLIP_DATASETS
+    flip_split: str
+    problem_type: ProblemType = ProblemType.REGRESSION
+
+    @model_validator(mode="after")
+    def validate_config(self) -> Self:
+        """Override base class validator.
+
+        train_ratio and test_ratio apply to separate FLIP pools (FLIP train and FLIP test),
+        so their sum is allowed to exceed 1. Also validates that flip_split is a valid
+        active split for the chosen flip_dataset.
+
+        Returns:
+            The validated configuration instance.
+
+        Raises:
+            ValueError: If `flip_split` is not a valid active split for `flip_dataset`.
+        """
+        if self.flip_split not in FLIP_SPLITS.get(self.flip_dataset, []):
+            raise ValueError(
+                f"'{self.flip_split}' is not a valid active split for dataset "
+                f"'{self.flip_dataset}'. Valid splits: {FLIP_SPLITS[self.flip_dataset]}"
+            )
+        return self
+
+
+class FLIP(BaseDataset):
+    """FLIP benchmark dataset class.
+
+    FLIP provides pre-defined train/test splits for protein fitness landscapes.
+    The CSV columns are:
+        - sequence: amino acid string (may contain special characters)
+        - target: float fitness score
+        - set: "train" or "test"
+        - validation: bool, FLIP-recommended early-stopping flag (stored as a
+          candidate feature but not used for splitting)
+
+    The splits map to alf as follows:
+        - train_ratio × FLIP train                        -> alf train
+        - validation_frac × (train_ratio × FLIP train)   -> alf validation
+        - remaining FLIP train (up to max_candidate_pool) -> alf candidate_pool
+        - test_ratio × FLIP test                          -> alf test
+    """
+
+    config: FLIPConfig  # narrows the inherited BaseDatasetConfig type
+
+    def __init__(self, config: FLIPConfig):
+        """Initialise FLIP dataset."""
+        super().__init__(config)
+        self.setup()
+
+    def __repr__(self) -> str:
+        """Return a string representation identifying dataset and split."""
+        return (
+            f"FLIP(name={self.config.name}, modality={self.modality}, "
+            f"seed={self.config.seed}, "
+            f"flip_dataset={self.config.flip_dataset}, "
+            f"flip_split={self.config.flip_split})"
+        )
+
+    def load_dataset(self) -> LabelledCandidates:
+        """Load FLIP split from local cache or download from GitHub.
+
+        Downloads splits.zip for the configured dataset if not cached locally,
+        then extracts and parses the configured split CSV.
+
+        Returns:
+            LabelledCandidates with sequence data, fitness labels, and features
+            containing 'set' ("train"/"test") and 'validation' (bool).
+
+        Raises:
+            FileNotFoundError: If the split CSV is not found inside splits.zip.
+            requests.HTTPError: If the download from GitHub fails.
+        """
+        df = self._load_split_dataframe()
+
+        candidates = []
+        labels = []
+        for _, row in df.iterrows():
+            candidates.append(
+                Candidate(
+                    data=row["sequence"],
+                    modality=self.modality,
+                    features={
+                        "set": row["set"],
+                        "validation": bool(row["validation"]),
+                    },
+                )
+            )
+            labels.append(row["target"])
+
+        return LabelledCandidates(
+            candidates=candidates,
+            labels=np.array(labels),
+        )
+
+    def _split_dataset(self) -> dict[str, LabelledCandidates]:
+        """Split the raw dataset into train, validation, candidate_pool, and test.
+
+        The FLIP train sequences are shuffled, then partitioned using train_ratio
+        and validation_frac. The remainder becomes the candidate pool (capped at
+        max_candidate_pool). The FLIP test sequences are sampled using test_ratio.
+
+        Returns:
+            Dictionary with keys "train", "validation", "test", "candidate_pool".
+
+        Raises:
+            RuntimeError: If dataset has not been loaded yet.
+        """
+        if self._raw_dataset is None:
+            raise RuntimeError(
+                "Dataset must be loaded before splitting; "
+                "_raw_dataset is None — call dataset.setup() (or load_dataset()) before splitting"
+            )
+
+        # Separate FLIP's pre-defined train and test pools
+        flip_train = LabelledCandidates(candidates=[], labels=np.array([]))
+        flip_test = LabelledCandidates(candidates=[], labels=np.array([]))
+        for candidate, label in self._raw_dataset:
+            if candidate.features and candidate.features["set"] == "train":
+                flip_train.append([candidate], np.array([label]))
+            else:
+                flip_test.append([candidate], np.array([label]))
+
+        logger.debug("FLIP pools — train: %d, test: %d", len(flip_train), len(flip_test))
+
+        # Shuffle FLIP train for reproducible random splitting
+        flip_train = flip_train.shuffle(self.config.seed)
+
+        # Split FLIP train into train, validation, and candidate_pool
+        train_plus_val_size = floor(len(flip_train) * self.split_ratio["train"])
+        validation_size = floor(train_plus_val_size * self.split_ratio["validation_frac"])
+        train_size = train_plus_val_size - validation_size
+
+        train = LabelledCandidates(
+            candidates=flip_train.candidates[:train_size],
+            labels=flip_train.labels[:train_size],
+        )
+        validation = LabelledCandidates(
+            candidates=flip_train.candidates[train_size:train_plus_val_size],
+            labels=flip_train.labels[train_size:train_plus_val_size],
+        )
+
+        pool_candidates = flip_train.candidates[train_plus_val_size:]
+        pool_labels = flip_train.labels[train_plus_val_size:]
+        if self.config.max_candidate_pool is not None:
+            pool_candidates = pool_candidates[: self.config.max_candidate_pool]
+            pool_labels = pool_labels[: self.config.max_candidate_pool]
+        candidate_pool = LabelledCandidates(candidates=pool_candidates, labels=pool_labels)
+
+        # Sample FLIP test using test_ratio
+        test_size = floor(len(flip_test) * self.split_ratio["test"])
+        test = LabelledCandidates(
+            candidates=flip_test.candidates[:test_size],
+            labels=flip_test.labels[:test_size],
+        )
+
+        logger.debug(
+            "Split sizes — train: %d, validation: %d, candidate_pool: %d, test: %d",
+            len(train),
+            len(validation),
+            len(candidate_pool),
+            len(test),
+        )
+
+        self.init_candidate_pool = copy.deepcopy(candidate_pool)
+
+        return {
+            "train": train,
+            "validation": validation,
+            "test": test,
+            "candidate_pool": candidate_pool,
+        }
+
+    def _load_split_dataframe(self) -> pd.DataFrame:
+        """Load the split CSV as a DataFrame, downloading splits.zip if needed.
+
+        Returns:
+            DataFrame with columns: sequence, target, set, validation.
+
+        Raises:
+            FileNotFoundError: If split CSV not found inside the zip archive.
+            ValueError: If the loaded CSV is missing expected columns.
+            requests.HTTPError: If the GitHub download fails.
+        """
+        zip_path = DATAPATH / "FLIP" / self.config.flip_dataset / "splits.zip"
+
+        if not zip_path.exists():
+            zip_path.parent.mkdir(parents=True, exist_ok=True)
+            url = f"{FLIP_GITHUB_BASE}/{self.config.flip_dataset}/splits.zip"
+            logger.info(f"Downloading FLIP splits for '{self.config.flip_dataset}' from {url}")
+            response = requests.get(url, stream=True, timeout=300)
+            response.raise_for_status()
+            with open(zip_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            logger.info("Download complete.")
+        else:
+            logger.debug("Using cached splits.zip at %s", zip_path)
+
+        file_stem = _SPLIT_FILE_MAP.get(
+            (self.config.flip_dataset, self.config.flip_split),
+            self.config.flip_split,
+        )
+        csv_name = f"{file_stem}.csv"
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            available = zf.namelist()
+            match = next((n for n in available if n.endswith(csv_name)), None)
+            if match is None:
+                raise FileNotFoundError(
+                    f"'{csv_name}' not found in splits.zip for dataset "
+                    f"'{self.config.flip_dataset}'. Available files: {available}"
+                )
+            with zf.open(match) as csv_file:
+                df = pd.read_csv(io.BytesIO(csv_file.read()))
+
+        missing = _EXPECTED_COLUMNS - set(df.columns)
+        if missing:
+            raise ValueError(
+                f"FLIP CSV '{csv_name}' is missing expected columns: {missing}. "
+                f"Found: {set(df.columns)}"
+            )
+
+        return df