biolm.pipeline package

Submodules

biolm.pipeline.async_executor module

Async batch execution utilities with semaphore-based rate limiting.

class biolm.pipeline.async_executor.AsyncBatchExecutor(max_concurrent: int = 10, batch_size: int | None = None, progress_desc: str = 'Processing', show_progress: bool = True)

Bases: object

Execute async tasks in batches with rate limiting.

Features: - Semaphore-based concurrency control - Progress tracking with tqdm - Error handling per item - Dynamic batch sizing

Args:

max_concurrent: Maximum number of concurrent tasks batch_size: Size of each batch (None for no batching) progress_desc: Description for progress bar show_progress: Whether to show progress bar

Example:

default Copy

>>> executor = AsyncBatchExecutor(max_concurrent=10)
>>> results = await executor.execute(items, process_func)

async execute(items: list[T], func: Callable[[T], Coroutine[Any, Any, R]], return_exceptions: bool = False) → list[R]

Execute async function on all items with rate limiting.

Args:

items: List of items to process func: Async function to apply to each item return_exceptions: If True, return exceptions instead of raising

Returns:

List of results (same order as input)

async execute_batched(items: list[T], batch_func: Callable[[list[T]], Coroutine[Any, Any, list[R]]], batch_size: int | None = None, return_exceptions: bool = False) → list[R]

Execute async function on batches of items.

Args:

items: List of items to process batch_func: Async function that processes a batch and returns results batch_size: Size of each batch (uses self.batch_size if None) return_exceptions: If True, return exceptions instead of raising

Returns:

Flattened list of results

class biolm.pipeline.async_executor.CachingExecutor(executor: AsyncBatchExecutor, cache_check_func: Callable[[T], R | None] | None = None, cache_store_func: Callable[[T, R], None] | None = None)

Bases: object

Executor with built-in caching support.

Checks cache before execution and updates cache after completion.

Args:

executor: AsyncBatchExecutor instance cache_check_func: Function to check if result is cached (returns cached result or None) cache_store_func: Function to store result in cache

Example:

default Copy

>>> def cache_check(item):
...     return datastore.get_prediction(item['sequence_id'], 'stability')
>>>
>>> def cache_store(item, result):
...     datastore.add_prediction(item['sequence_id'], 'stability', result)
>>>
>>> executor = CachingExecutor(
...     AsyncBatchExecutor(max_concurrent=10),
...     cache_check, cache_store
... )
>>> results = await executor.execute_with_cache(items, process_func)

async execute_with_cache(items: list[T], func: Callable[[T], Coroutine[Any, Any, R]], return_exceptions: bool = False) → tuple[list[R], int, int]

Execute with caching.

Returns:

Tuple of (results, cached_count, computed_count)

class biolm.pipeline.async_executor.StreamingExecutor(executor: AsyncBatchExecutor, result_callback: Callable[[T, R], Coroutine[Any, Any, None]] | None = None)

Bases: object

Executor that streams results as they complete.

Useful for pipelines where downstream stages can start before all results are ready.

Args:

executor: AsyncBatchExecutor instance result_callback: Optional callback for each completed result

Example:

default Copy

>>> async def on_result(item, result):
...     print(f"Completed: {item}")
...     downstream_queue.put(result)
>>>
>>> executor = StreamingExecutor(
...     AsyncBatchExecutor(max_concurrent=10),
...     result_callback=on_result
... )
>>> await executor.execute_streaming(items, process_func)

async execute_streaming(items: list[T], func: Callable[[T], Coroutine[Any, Any, R]], return_exceptions: bool = False) → list[R]

Execute with streaming results.

Results are passed to callback as they complete.

Returns:

List of results (order may differ from input)

async biolm.pipeline.async_executor.process_dataframe_async(df: DataFrame, process_func: Callable[[Series], Coroutine[Any, Any, Any]], result_column: str, max_concurrent: int = 10, show_progress: bool = True, progress_desc: str = 'Processing') → DataFrame

Process DataFrame rows asynchronously and add results as a new column.

Args:

df: Input DataFrame process_func: Async function to process each row (receives a Series) result_column: Name of column to store results max_concurrent: Maximum concurrent tasks show_progress: Whether to show progress bar progress_desc: Description for progress bar

Returns:

DataFrame with new result column

Example:

default Copy

>>> async def predict_stability(row):
...     return await model.predict(row['sequence'])
>>>
>>> df = await process_dataframe_async(
...     df, predict_stability, 'stability_score'
... )

async biolm.pipeline.async_executor.process_sequences_batched(sequences: list[str], batch_func: Callable[[list[str]], Coroutine[Any, Any, list[Any]]], batch_size: int = 32, max_concurrent: int = 5, show_progress: bool = True, progress_desc: str = 'Processing sequences') → list[Any]

Process sequences in batches.

Args:

sequences: List of sequences batch_func: Async function that processes a batch of sequences batch_size: Size of each batch max_concurrent: Maximum concurrent batches show_progress: Whether to show progress bar progress_desc: Description for progress bar

Returns:

List of results (one per sequence)

Example:

default Copy

>>> async def predict_batch(seqs):
...     return await model.predict_batch(seqs)
>>>
>>> results = await process_sequences_batched(
...     sequences, predict_batch, batch_size=32
... )

biolm.pipeline.base module

Base Pipeline classes for stage management and execution.

class biolm.pipeline.base.BasePipeline(datastore: DuckDBDataStore | str | Path | None = None, run_id: str | None = None, output_dir: str | Path = 'pipeline_outputs', resume: bool = False, verbose: bool = True, input_schema: InputSchema | None = None)

Bases: ABC

Base class for all pipeline types.

Provides: - Stage management and dependency resolution - Async execution with progress tracking - Caching and resumability - Export and visualization

When no datastore is provided, the pipeline automatically creates a DuckDB cache under .biolm/pipelines//. The pipeline_id (and full cache path) is exposed via metadata so users can reconnect to the same cache in later sessions.

Args:

datastore: DataStore instance or path to a DuckDB file. Required. run_id: Unique run identifier (auto-generated if not provided). output_dir: Directory for CSV/Parquet exports (default pipeline_outputs). resume: Whether to resume from a previous run. verbose: Enable verbose output.

add_stage(stage: Stage)

Add a stage to the pipeline.

close()

Close the pipeline’s datastore connection.

Safe to call multiple times. Called automatically by __exit__ and __aexit__. Also called by __del__ for auto-created datastores only (user-provided datastores are not closed on GC so the caller can continue using them after the pipeline is discarded).

export_to_csv(output_path: str | Path | None = None)

Export final results to CSV.

classmethod from_db(db_path: str | Path, definition_id: str | None = None, run_id: str | None = None, verbose: bool = True) → BasePipeline

Reconstruct a pipeline from an existing DuckDB database.

Useful for recovering after a kernel death without re-running already-completed stages.

Args:

db_path: Path to the DuckDB database file. definition_id: Specific definition to load (None = latest). run_id: Run ID for the reconstructed pipeline (None = generate new). verbose: Enable verbose output.

Returns:

Reconstructed BasePipeline subclass instance.

Example:

default Copy

pipeline = DataPipeline.from_db("my_pipeline.duckdb")
pipeline.run(resume=True)

get_final_data() → DataFrame

Get the final output DataFrame.

Materializes from the merged final WorkingSet via DuckDB. In a branched DAG the last-added stage is not necessarily the last executed sink, so we prefer _final_ws (set at end of run()).

The returned DataFrame always contains at minimum:

• sequence_id, sequence, length, hash

• source_label — label from the generation config; None when unset. Always present regardless of whether any labels were set.

• One column per prediction type from upstream prediction stages.

property metadata: PipelineMetadata

Return metadata for reconnecting to this pipeline’s cache later.

query(sql: str, params=None) → DataFrame

Execute arbitrary SQL against the pipeline’s DuckDB datastore.

results() → DataFrame

Return the final output DataFrame. Alias for get_final_data().

The returned DataFrame always contains at minimum:

• sequence_id, sequence, length, hash

• source_label — the label set on the generation config (e.g. DirectGenerationConfig.label). None when no label was set. Always present — do not guard on "source_label" in df.columns.

• One column per prediction type computed by any prediction stages.

run(enable_streaming: bool = True, **kwargs) → dict[str, biolm.pipeline.base.StageResult]

Run the pipeline synchronously.

This is a convenience wrapper around run_async(). Works in both script/notebook environments (detects running event loops).

async run_async(enable_streaming: bool = True, **kwargs) → dict[str, biolm.pipeline.base.StageResult]

Run the pipeline asynchronously.

Args:

enable_streaming: Stream prediction results through per-sequence

filters for better parallelism and lower latency (default True).

Returns:

Dict mapping stage names to StageResults

summary() → DataFrame

Get pipeline summary statistics.

class biolm.pipeline.base.InputSchema(columns: list[str])

Bases: object

Describes the input columns for a pipeline.

When set, these columns are the primary data — sequence is not required. Columns are stored directly on the sequences table via ALTER TABLE ADD COLUMN so that materialize_working_set(), SQL filters, and item_columns all work via direct JOINs.

Hashing uses all columns (sorted alphabetically) joined with \x00 separators so that identical rows produce the same hash regardless of column order.

Args:

columns: List of column names that comprise the primary input.

columns: list[str]

hash_row(row: dict[str, str]) → str

SHA-256 hash of the row values across all input columns (sorted).

class biolm.pipeline.base.PipelineContext(datastore: DuckDBDataStore, run_id: str)

Bases: object

Shared key-value store backed by DuckDB for inter-stage communication.

Stages can read/write arbitrary data through the context. Common use case: stage 1 predicts structures (stored in the structures table), stage 2 reads them for structure-conditioned generation.

Args:

datastore: The pipeline’s DuckDB datastore. run_id: Current pipeline run ID.

get(key: str, default: Any = None) → Any

Retrieve a value from the pipeline context table.

get_structure(sequence_id: int, model_name: str | None = None) → dict | None

Convenience: fetch a structure from the datastore’s structures table.

get_structures_for_ws(ws: WorkingSet, model_name: str | None = None) → DataFrame

Fetch structures for all sequences in a WorkingSet.

set(key: str, value: Any)

Store a value in the pipeline context table.

class biolm.pipeline.base.PipelineMetadata(pipeline_id: str, cache_dir: Path, db_path: Path, run_id: str)

Bases: object

Metadata for a pipeline run — lets users retrieve and reuse cached results.

Attributes:

pipeline_id: Unique identifier for the pipeline’s cache directory. cache_dir: Path to the .biolm/pipelines/ cache directory. db_path: Path to the DuckDB database file inside the cache directory. run_id: The run ID for this execution (there can be multiple runs

sharing the same cache).

Example:

default Copy

pipeline = DataPipeline(sequences=[...])
pipeline.run()
meta = pipeline.metadata
print(meta.pipeline_id)   # "20260302_143022_a1b2c3d4"
print(meta.cache_dir)     # ".biolm/pipelines/20260302_143022_a1b2c3d4"

# Later — reuse the same cache:
pipeline2 = DataPipeline(
    sequences=new_seqs,
    datastore=meta.db_path,   # or str(meta.cache_dir)
    resume=True,
)

cache_dir: Path

db_path: Path

pipeline_id: str

run_id: str

class biolm.pipeline.base.Stage(name: str, cache_key: str | None = None, depends_on: list[str] | None = None, model_name: str | None = None, max_concurrent: int = 10)

Bases: ABC

Abstract base class for pipeline stages.

A stage represents a single processing step in the pipeline. It can filter data, compute predictions, or transform sequences.

Args:

name: Stage name (must be unique within pipeline) cache_key: Unused collision-dedup key (auto-derived by PredictionStage) depends_on: List of stage names this stage depends on model_name: Model name for predictions/structures max_concurrent: Maximum concurrent API calls (for rate limiting)

Class Attributes:

merge_mode: How this stage’s output is merged when it runs in parallel

with other stages. "intersect" (default) = only keep sequences that pass all parallel stages (correct for filters). "union" = keep sequences that appear in any parallel stage output (correct for independent prediction stages that may skip sequences on error but should not drop others).

merge_mode: str = 'intersect'

async process(df: DataFrame, datastore: DuckDBDataStore, **kwargs) → tuple[pandas.core.frame.DataFrame, biolm.pipeline.base.StageResult]

Legacy DataFrame interface — used by streaming mode and GenerationStage.

Subclasses may override this for backward compatibility or for cases where a DataFrame is the natural input (e.g. generation with an empty df).

abstract async process_ws(ws: WorkingSet, datastore: DuckDBDataStore, **kwargs) → tuple[biolm.pipeline.base.WorkingSet, biolm.pipeline.base.StageResult]

Process data using WorkingSet (DuckDB-native).

All stages must implement this method. Stages that need actual sequence data (e.g. ClusteringStage) should call datastore.materialize_working_set(ws) internally — that is the stage’s responsibility, not the pipeline’s.

Args:

ws: Input WorkingSet (set of sequence IDs). datastore: DataStore for reading/writing data. **kwargs: Additional arguments (e.g. run_id).

Returns:

Tuple of (output WorkingSet, StageResult).

to_spec() → dict

Return a serializable dict describing this stage.

Used by BasePipeline.run_async() to persist pipeline definitions to DuckDB (enabling DataPipeline.from_db() recovery after kernel death).

Subclasses must override this. Raises NotImplementedError by default.

class biolm.pipeline.base.StageResult(stage_name: str, input_count: int, output_count: int, filtered_count: int = 0, cached_count: int = 0, computed_count: int = 0, elapsed_time: float = 0.0, metadata: dict[str, typing.Any] = <factory>)

Bases: object

Result from a pipeline stage.

cached_count: int = 0

computed_count: int = 0

elapsed_time: float = 0.0

filtered_count: int = 0

input_count: int

metadata: dict[str, Any]

output_count: int

stage_name: str

class biolm.pipeline.base.WorkingSet(sequence_ids: frozenset[int])

Bases: object

Lightweight set of sequence IDs — replaces DataFrame as inter-stage transport.

Stages operate on DuckDB directly and pass only the set of surviving sequence IDs to the next stage. Materialization to DataFrame happens once at get_final_data() time.

Memory: 1M IDs ≈ 28 MB (frozenset[int]) vs 500 MB+ DataFrame.

difference(other: WorkingSet) → WorkingSet

Return IDs in self but not in other.

classmethod from_ids(ids) → WorkingSet

Create from any iterable of ints.

intersect(other: WorkingSet) → WorkingSet

Return a new WorkingSet containing only IDs present in both sets.

sequence_ids: frozenset[int]

to_list() → list[int]

Return sorted list (useful for DuckDB queries).

union(other: WorkingSet) → WorkingSet

Return a new WorkingSet containing IDs from either set.

biolm.pipeline.clustering module

Sequence clustering and diversity analysis.

Provides tools for grouping similar sequences and measuring sequence space coverage.

Performance Notes: - For large datasets (>10k sequences), use sampling or embedding-based methods - Pairwise distance computations are O(n²) - use max_sample parameter - Embedding-based clustering scales much better than Hamming distance

class biolm.pipeline.clustering.ClusteringResult(cluster_ids: ndarray, centroids: list[str], centroid_indices: ndarray, n_clusters: int, silhouette_score: float | None = None, davies_bouldin_score: float | None = None, cluster_sizes: dict[int, int] | None = None)

Bases: object

Results from sequence clustering.

centroid_indices: ndarray

centroids: list[str]

cluster_ids: ndarray

cluster_sizes: dict[int, int] | None = None

davies_bouldin_score: float | None = None

n_clusters: int

silhouette_score: float | None = None

class biolm.pipeline.clustering.DiversityAnalyzer

Bases: object

Analyze sequence diversity and coverage.

Provides metrics for understanding sequence space exploration.

Performance Notes:

• Shannon entropy: O(n*L) where L is sequence length

• Pairwise distances: O(n²*L) - use max_sample for large n

• All metrics scale linearly except pairwise distances

Example:

default Copy

>>> analyzer = DiversityAnalyzer()
>>> metrics = analyzer.compute_all_metrics(sequences, max_sample=10000)
>>> print(f"Shannon entropy: {metrics['shannon_entropy']:.2f}")

classmethod compute_all_metrics(sequences: list[str], max_sample: int | None = 10000) → dict[str, Union[float, dict]]

Compute all diversity metrics at once.

Args:

sequences: List of protein sequences max_sample: Unused, reserved for future embedding-based pairwise stats

Returns:

Dictionary with all diversity metrics

static motif_diversity(sequences: list[str], k: int = 3) → dict[str, Union[int, float]]

Analyze k-mer (motif) diversity.

Args:

sequences: List of protein sequences k: Length of k-mers to analyze

Returns:

Dictionary with k-mer statistics

static shannon_entropy(sequences: list[str], normalize: bool = True) → float

Calculate Shannon entropy of amino acid distribution.

Measures positional diversity across all sequences. Optimized for large datasets using vectorized operations.

Args:

sequences: List of protein sequences normalize: Normalize by log(20) for 0-1 range

Returns:

Shannon entropy (higher = more diverse)

exception biolm.pipeline.clustering.PerformanceWarning

Bases: UserWarning

Warning for potentially slow operations on large datasets.

class biolm.pipeline.clustering.SequenceClusterer(method: Literal['kmeans', 'dbscan', 'hierarchical'] = 'kmeans', n_clusters: int | None = None, similarity_metric: Literal['embedding'] = 'embedding', eps: float = 0.5, min_samples: int = 5, random_state: int = 42, mini_batch: bool = False, max_sample: int | None = None)

Bases: object

Cluster sequences by similarity.

Supports multiple clustering algorithms and similarity metrics.

Performance Notes:

• For >10k sequences with Hamming distance, consider using max_sample

• Embedding-based clustering scales much better (O(n) with MiniBatch K-means)

• Use mini_batch=True for very large datasets (>50k sequences)

Args:

method: Clustering algorithm (‘kmeans’, ‘dbscan’, ‘hierarchical’) n_clusters: Number of clusters (for kmeans/hierarchical) similarity_metric: How to measure sequence similarity eps: DBSCAN epsilon parameter min_samples: DBSCAN minimum samples per cluster mini_batch: Use MiniBatchKMeans for large datasets (faster, approximate) max_sample: Maximum sequences to use for distance matrix (None = all)

Example:

default Copy

>>> # For large datasets, use sampling or mini-batch
>>> clusterer = SequenceClusterer(
...     method='kmeans',
...     n_clusters=100,
...     mini_batch=True  # Much faster for large N
... )
>>> result = clusterer.cluster(sequences)

cluster(sequences: list[str], embeddings: ndarray | None = None) → ClusteringResult

Cluster sequences and return assignments.

Args:

sequences: List of protein sequences embeddings: Pre-computed embeddings (if using embedding metric)

Returns:

ClusteringResult with cluster assignments and metrics

biolm.pipeline.clustering.analyze_diversity(sequences: list[str], max_sample: int | None = 10000) → dict

Convenience function for analyzing sequence diversity.

Args:

sequences: List of protein sequences max_sample: Reserved for future embedding-based pairwise stats

Returns:

Dictionary of diversity metrics (shannon_entropy, motif diversity, uniqueness)

Example:

default Copy

>>> metrics = analyze_diversity(sequences)
>>> print(f"Entropy: {metrics['shannon_entropy']:.2f}")

biolm.pipeline.clustering.cluster_sequences(sequences: list[str], method: str = 'kmeans', n_clusters: int = 10, embeddings: ndarray | None = None, mini_batch: bool = False, max_sample: int | None = None, **kwargs) → ClusteringResult

Convenience function for clustering sequences using embeddings.

Performance Tips:

• For >50k sequences, use mini_batch=True

• Requires pre-computed embeddings (similarity_metric=’embedding’)

Args:

sequences: List of protein sequences method: Clustering algorithm n_clusters: Number of clusters embeddings: Pre-computed embeddings (required) mini_batch: Use MiniBatchKMeans for faster (approximate) clustering max_sample: Reserved; not yet used **kwargs: Additional arguments for SequenceClusterer

Returns:

ClusteringResult

biolm.pipeline.data module

Data-driven pipeline implementations.

DataPipeline: Load sequences from files/lists and run predictions SingleStepPipeline: Simplified single-step prediction pipeline

class biolm.pipeline.data.ClusteringStage(name: str, method: str = 'kmeans', n_clusters: int | None = None, similarity_metric: str = 'hamming', embedding_model: str | None = None, max_sample: int | None = None, **kwargs)

Bases: Stage

Sequence clustering stage.

Args:

name: Stage name method: Clustering algorithm (‘kmeans’, ‘dbscan’, ‘hierarchical’) n_clusters: Number of clusters (for kmeans/hierarchical) similarity_metric: How to measure similarity (‘hamming’, ‘embedding’) embedding_model: Model to use for embeddings (if similarity_metric=’embedding’)

async process(df: DataFrame, datastore: DuckDBDataStore, **kwargs) → tuple[pandas.core.frame.DataFrame, biolm.pipeline.base.StageResult]

Cluster sequences, add cluster columns, and return the enriched DataFrame.

async process_ws(ws: WorkingSet, datastore: DuckDBDataStore, **kwargs) → tuple[biolm.pipeline.base.WorkingSet, biolm.pipeline.base.StageResult]

Cluster sequences — materializes internally for scikit-learn.

Clustering doesn’t filter rows (all sequences survive), so the returned WorkingSet is the same as the input. Cluster assignments are stored as predictions (prediction_type 'cluster_id') so that materialize_working_set() includes them in get_final_data().

to_spec() → dict

Return a serializable dict for pipeline definition persistence.

class biolm.pipeline.data.CofoldingPredictionStage(name: str, model_name: str, action: str = 'predict', prediction_type: str = 'structure', sequence_chain_id: str = 'A', sequence_entity_type: str = 'protein', static_entities=None, params: dict | None = None, batch_size: int = 1, item_columns: dict[str, str] | None = None, **kwargs)

Bases: Stage

Prediction stage for co-folding models (Boltz2, Chai-1).

Each sequence in the pipeline DataFrame is used as the primary entity (chain) in a multi-molecule folding request. Additional static entities — ligands, cofactors, DNA/RNA, or other protein chains — are injected via static_entities and held constant for every sequence.

The caller is responsible for providing the correct molecule field names and params for the target model. Boltz2 and Chai-1 both expect an item shaped like {'molecules': [{'id': ..., 'type': ..., 'sequence': ...}, ...]}.

Args:

name: Stage name. model_name: BioLM model slug ('boltz2', 'chai1'). action: API action (default 'predict'). prediction_type: Column name written for the confidence score

(default 'structure').

sequence_chain_id: Chain ID / name assigned to the pipeline’s primary

sequence in the molecules list.

sequence_entity_type: Molecule type for the primary sequence

('protein', 'dna', 'rna').

static_entities: List of FoldingEntity

objects appended to every molecules list after the primary chain.

params: Model-specific params dict passed directly to the API

(e.g. {'recycling_steps': 3, 'sampling_steps': 20} for Boltz).

batch_size: Sequences per API call (default 1; co-folding models are

typically limited to batch size 1).

Example:

default Copy

from biolm.pipeline import FoldingEntity

pipeline.add_cofolding_prediction(
    model_name='boltz2',
    static_entities=[
        FoldingEntity(id='L', entity_type='ligand', smiles='c1ccccc1'),
    ],
    params={'recycling_steps': 3, 'sampling_steps': 20},
    depends_on=['filter_top50'],
)

merge_mode: str = 'union'

async process(df: DataFrame, datastore: DuckDBDataStore, **kwargs) → tuple[pandas.core.frame.DataFrame, biolm.pipeline.base.StageResult]

Run co-folding prediction for each sequence in df.

async process_ws(ws: WorkingSet, datastore: DuckDBDataStore, **kwargs) → tuple[biolm.pipeline.base.WorkingSet, biolm.pipeline.base.StageResult]

Run co-folding prediction using WorkingSet — no DataFrame transport.

to_spec() → dict

Serialize to a dict for pipeline definition persistence.

Note: static_entities (FoldingEntity objects) are not serializable and are omitted. A reconstructed pipeline will not have static_entities and must have them re-attached manually after from_db().

class biolm.pipeline.data.DataPipeline(sequences: list[str] | DataFrame | str | Path = None, diff_mode: bool = False, input_columns: list[str] | None = None, **kwargs)

Bases: BasePipeline

Pipeline for processing existing sequences from files or lists.

Load sequences from CSV/FASTA/lists and run predictions/filtering.

Args:

sequences: Input sequences (list of strings, DataFrame, or file path) datastore: DataStore instance or path run_id: Unique run ID output_dir: Output directory resume: Resume from previous run verbose: Enable verbose output diff_mode: If True, merge new sequences with existing cached results.

Only computes predictions for uncached sequences. Use get_merged_results() or query_results() to efficiently access combined data without loading millions of rows into memory (SQL-based queries).

Example:

default Copy

>>> # Standard mode
>>> pipeline = DataPipeline(sequences='sequences.csv')
>>> pipeline.add_prediction('esmfold', extractions='mean_plddt', columns='plddt')
>>> pipeline.add_filter(ThresholdFilter('plddt', min_value=70))
>>> results = pipeline.run()

default Copy

>>> # Diff mode - add new sequences to existing pipeline (SQL-based, efficient)
>>> pipeline = DataPipeline(sequences='new_sequences.csv', diff_mode=True)
>>> pipeline.add_prediction('esmfold', extractions='mean_plddt', columns='plddt')
>>> results = pipeline.run()
>>> # Efficiently query specific data (doesn't load all millions of rows!)
>>> high_quality = pipeline.query_results("s.length > 100 AND p.value > 70")
>>> # Or get merged results with filters
>>> merged = pipeline.get_merged_results(prediction_types=['plddt', 'tm'])

add_clustering(method: str = 'kmeans', n_clusters: int | None = None, similarity_metric: str = 'hamming', embedding_model: str | None = None, stage_name: str | None = None, depends_on: list[str] | None = None, **kwargs)

Add a sequence clustering stage.

Clusters sequences by similarity and adds cluster_id column to DataFrame.

Args:

method: Clustering algorithm (‘kmeans’, ‘dbscan’, ‘hierarchical’) n_clusters: Number of clusters (required for kmeans/hierarchical) similarity_metric: ‘hamming’ or ‘embedding’ embedding_model: Model name if using embedding similarity stage_name: Optional custom stage name depends_on: Optional list of stage names this stage depends on **kwargs: Additional arguments for clustering (eps, min_samples, etc.)

Example:

default Copy

>>> # Cluster by sequence similarity
>>> pipeline.add_clustering(method='kmeans', n_clusters=10)
>>>
>>> # Cluster by embeddings
>>> pipeline.add_prediction('esm2-650m', action='encode', stage_name='embed')
>>> pipeline.add_clustering(
...     method='kmeans',
...     n_clusters=5,
...     similarity_metric='embedding',
...     embedding_model='esm2-650m',
...     depends_on=['embed']
... )

add_cofolding_prediction(model_name: str, action: str = 'predict', stage_name: str | None = None, prediction_type: str = 'structure', sequence_chain_id: str = 'A', sequence_entity_type: str = 'protein', static_entities=None, depends_on: list[str] | None = None, params: dict | None = None, batch_size: int = 1)

Add a co-folding prediction stage (Boltz2, Chai-1).

Each pipeline sequence becomes the primary entity in a multi-molecule folding request. static_entities injects ligands, cofactors, DNA/RNA strands, or additional protein chains that are constant across all designs.

The caller is responsible for providing the correct molecule field names via static_entities and the right params for the model.

Args:

model_name: BioLM model slug ('boltz2', 'chai1'). action: API action (default 'predict'). stage_name: Optional stage name (defaults to model_name). prediction_type: Column name for the confidence score. sequence_chain_id: Chain ID assigned to the primary sequence

(e.g. 'A' for Boltz, molecule name for Chai-1).

sequence_entity_type: Entity type for the primary sequence

('protein', 'dna', 'rna').

static_entities: List of FoldingEntity objects to include

in every request alongside the primary sequence.

depends_on: Upstream stage names. params: Model-specific params (e.g. {'recycling_steps': 3}). batch_size: Sequences per API call (default 1).

Example:

default Copy

from biolm.pipeline import FoldingEntity

pipeline.add_cofolding_prediction(
    model_name='boltz2',
    static_entities=[
        FoldingEntity(id='L', entity_type='ligand', smiles='c1ccccc1'),
    ],
    params={'recycling_steps': 3, 'sampling_steps': 20},
    depends_on=['filter_top50'],
)

add_filter(filter_func: BaseFilter | Callable[[...], Any], stage_name: str | None = None, depends_on: list[str] | None = None, **kwargs)

Add a filter stage.

Args:

filter_func: Filter function or BaseFilter instance stage_name: Custom stage name depends_on: List of stage names this depends on. When None

(default), auto-depends on the last added stage. Pass depends_on=[] for earliest-level execution.

add_prediction(model_name: str, action: str = 'predict', extractions: str | list[Union[str, biolm.pipeline.data.ExtractionSpec]] | None = None, columns: str | dict[str, str] | None = None, params: dict | None = None, stage_name: str | None = None, depends_on: list[str] | None = None, **kwargs)

Add a prediction stage.

Args:

model_name: BioLM model name action: API action (‘predict’, ‘encode’, ‘score’) extractions: API response key(s) to extract. Required for

predict/score actions. Can be a string for a single key or a list of strings / ExtractionSpec objects.

columns: Output column name(s). A string renames a single

extraction; a dict maps response keys to column names (unmapped keys keep their name).

params: Optional API parameters stage_name: Custom stage name (defaults to predict_{first_column}) depends_on: List of stage names this depends on. When None

(default), the stage auto-depends on the last added stage, creating a sequential chain. Pass depends_on=[] to run at the earliest possible level (parallel with other level-0 stages).

Example:

default Copy

pipeline.add_prediction(
    "temberture-regression",
    extractions="prediction",
    columns="tm",
)

add_predictions(models: list[Union[str, dict]], action: str = 'predict', depends_on: list[str] | None = None, **kwargs)

Add multiple prediction stages at the same level (run in parallel).

Args:

models: List of model names or dicts with model configs action: Default action if not specified in dict depends_on: List of stage names all these depend on **kwargs: Default kwargs for all stages

Returns:

self for chaining

Example:

default Copy

>>> pipeline.add_predictions([
...     {'model_name': 'temberture-regression', 'extractions': 'prediction', 'columns': 'tm'},
...     {'model_name': 'biolmsol', 'extractions': 'solubility_score', 'columns': 'solubility'},
... ])

add_structure_prediction(model_name: str, structure_key: str = 'pdb', extractions: str | list[Union[str, biolm.pipeline.data.ExtractionSpec]] | None = None, columns: str | dict[str, str] | None = None, plddt_key: str | None = None, structure_format: str | None = None, stage_name: str | None = None, depends_on: list[str] | None = None, **kwargs)

Add a prediction stage that stores the structure from the response.

Convenience wrapper around add_prediction() with a StructureSpec.

Args:

model_name: BioLM model name (e.g. ‘esmfold’, ‘alphafold2’). structure_key: Response key containing the structure string. extractions: Optional scalar extractions (e.g. ‘mean_plddt’). columns: Output column name(s) for scalar extractions. plddt_key: Optional response key for confidence score. structure_format: ‘pdb’ or ‘cif’ (auto-detected from key if None). stage_name: Custom stage name. depends_on: Stage dependencies.

Example:

default Copy

pipeline.add_structure_prediction(
    "esmfold",
    extractions="mean_plddt", columns="plddt",
    plddt_key="mean_plddt",
)

explore() → dict[str, Any]

Return summary stats for the pipeline’s datastore (all via SQL).

Returns:

Dict with keys: sequences, embeddings, generated, completed_stages, predictions (dict of prediction_type → count).

get_merged_results(prediction_types: list[str] | None = None, sequence_filter: str | None = None) → DataFrame

Get results merged with existing cached data (for diff mode).

This method is SQL-based and efficient - it doesn’t load millions of rows. Instead, it queries only the data you need using DuckDB’s columnar engine.

NOTE (Bug #5): sequence_filter and sql_where in query_results() are interpolated directly into SQL. These parameters are intended for internal/ trusted caller use only — never pass untrusted user input to them.

Args:

prediction_types: List of prediction types to include (None = all) sequence_filter: SQL WHERE clause to filter sequences (e.g., “length > 50”).

TRUSTED CALLERS ONLY — not safe for untrusted user input.

Returns:

DataFrame with requested sequences and predictions

Example:

default Copy

>>> # Get all results (efficient - DuckDB only loads what's needed)
>>> df = pipeline.get_merged_results()

default Copy

>>> # Get only specific predictions (columnar - even faster!)
>>> df = pipeline.get_merged_results(prediction_types=['tm', 'plddt'])

default Copy

>>> # Get sequences matching criteria (predicate pushdown!)
>>> df = pipeline.get_merged_results(sequence_filter="length > 100")

plot(kind: str = 'funnel', **kwargs)

Convenience wrapper around PipelinePlotter.

Args:

kind: One of ‘funnel’, ‘predictions’, ‘distributions’, ‘scatter’,

‘correlation’, ‘diversity’, ‘temperature’.

**kwargs: Forwarded to the underlying plotter method.

scatter requires x_col and y_col. diversity accepts reference_sequence. temperature requires metric_col.

query(sql: str, params=None) → DataFrame

Execute arbitrary SQL against the pipeline’s DuckDB datastore.

Args:

sql: DuckDB SQL query string. params: Optional list of query parameters.

Returns:

DataFrame with results.

Example:

default Copy

>>> pipeline.query("SELECT * FROM sequences WHERE length > 100")

query_results(sql_where: str, columns: list[str] | None = None) → DataFrame

Query results using SQL WHERE clause (for diff mode with large datasets).

This leverages DuckDB’s vectorized engine for maximum performance.

NOTE (Bug #5): sql_where is interpolated directly into SQL. TRUSTED CALLERS ONLY — never pass untrusted user input to this parameter.

Args:

sql_where: SQL WHERE clause using table aliases:

• s.* for sequences table (e.g., “s.length > 100”)

• Column names directly (no p. prefix needed)

TRUSTED CALLERS ONLY — not safe for untrusted user input.

columns: Columns to include (None = all available)

Returns:

DataFrame with matching sequences (only loads what matches!)

Example:

default Copy

>>> # Find long sequences (columnar scan - fast!)
>>> df = pipeline.query_results("s.length > 200")

default Copy

>>> # Complex filter with predictions
>>> df = pipeline.query_results(
...     "s.length > 100",
...     columns=['tm', 'plddt']
... )

stats(stage_name: str | None = None) → DataFrame

Return per-stage counts from stage_completions.

Args:

stage_name: If provided, filter to that stage only.

Returns:

DataFrame with columns: stage_name, status, input_count, output_count, completed_at.

biolm.pipeline.data.Embed(model_name: str, sequences: list[str] | DataFrame | str | Path, layer: int | None = None, key: str | None = None, **kwargs) → DataFrame

Convenience function for generating embeddings.

Args:

model_name: BioLM model name (e.g., 'esm2-8m', 'esm2-650m',

'ablang2'). See biolm.list_models() for the full slugged list — the family name alone ('esm2') is not a valid endpoint.

sequences: Input sequences layer: Optional layer number key: Response dict key containing embeddings. Auto-detected if None:

ESM2 models use “embeddings”, AbLang2 uses “seqcoding”, others default to “embedding”.

**kwargs: Additional arguments

Returns:

DataFrame with ‘sequence’, ‘sequence_id’, and ‘embedding’ columns

Example:

default Copy

>>> df = Embed('esm2-8m', sequences=['MKTAYIAKQRQ', 'MKLAVID'])

class biolm.pipeline.data.EmbeddingSpec(key: str, layer: int | None = None, reduction: str | None = None)

Bases: object

Declarative specification for extracting embeddings from API responses.

Covers common response formats without writing a custom function.

Args:

key: Response dict key containing the embedding data (e.g.

"embedding", "seqcoding", "embeddings").

layer: Which layer to extract when the response contains multiple

layers (list of {layer: int, embedding: [...]} dicts). None stores all layers; an int stores only that layer.

reduction: Reduce per-token 2-D embeddings to a single vector:

"mean", "first", "last", "sum". None stores the full array as-is.

Examples:

default Copy

# ablang2 returns {"seqcoding": [float, ...]}
EmbeddingSpec(key="seqcoding")

# esm2-8m returns {"embeddings": [{embedding: [...], layer: 33}]}
# Store only layer 33:
EmbeddingSpec(key="embeddings", layer=33)

# Per-residue → mean-pool:
EmbeddingSpec(key="embedding", reduction="mean")

key: str

layer: int | None = None

reduction: str | None = None

class biolm.pipeline.data.ExtractionSpec(response_key: str, reduction: str | None = None)

Bases: object

Specification for extracting a value from an API response (with reduction).

Use when you need to apply a reduction (mean, max, min, sum) to an array-valued response key. For simple scalar extractions, pass a plain string to extractions instead.

Args:

response_key: Key in API response dict, e.g. “plddt” reduction: Optional reduction for array values: “mean”, “max”, “min”, “sum”

reduction: str | None = None

response_key: str

class biolm.pipeline.data.FilterStage(name: str, filter_func: BaseFilter | Callable[[...], Any], **kwargs)

Bases: Stage

Generic filtering stage.

Args:

name: Stage name filter_func: Filter function or BaseFilter instance

async process(df: DataFrame, datastore: DuckDBDataStore, **kwargs) → tuple[pandas.core.frame.DataFrame, biolm.pipeline.base.StageResult]

Apply filter to DataFrame and return the filtered result.

async process_ws(ws: WorkingSet, datastore: DuckDBDataStore, **kwargs) → tuple[biolm.pipeline.base.WorkingSet, biolm.pipeline.base.StageResult]

Apply filter using WorkingSet.

Two execution paths — chosen at call time, not as a “fallback”:

1. SQL-native (zero materialization): filters that implement to_sql() return a complete SELECT scoped to the working set. DuckDB executes it directly; no DataFrame is ever created.

2. DataFrame-based: filters that cannot be expressed in SQL (e.g. HammingDistanceFilter, CustomFilter) must materialize a DataFrame. This is the correct path for those filters, not a fallback.

to_spec() → dict

Return a serializable dict for pipeline definition persistence.

class biolm.pipeline.data.MatrixExtractionSpec(prefix: str = 'ddg', values_key: str = 'ddG_matrix.values', row_labels_key: str = 'ddG_matrix.residue_axis', col_labels_key: str = 'ddG_matrix.amino_acid_axis', mutation_key: str | None = None, value_key: str | None = None)

Bases: object

Flattens a per-mutation response into individual prediction rows.

Each mutation becomes a separate prediction row with prediction_type formatted as ‘{prefix}_{label}’ (e.g., ‘ddg_M1A’).

Two modes:

• Matrix mode (SPURS): 2D array + row/col labels.

• List mode (ThermoMPNN): list of dicts with mutation name + value.

Args:

prefix: Prediction type prefix (e.g. “ddg”). values_key: Dot-path to 2D array in response (matrix mode). row_labels_key: Dot-path to row labels (position labels). col_labels_key: Dot-path to column labels (amino acid labels). mutation_key: Dict key for mutation name (list mode). If set, uses list mode. value_key: Dict key for the numeric value (list mode).

col_labels_key: str = 'ddG_matrix.amino_acid_axis'

mutation_key: str | None = None

prefix: str = 'ddg'

row_labels_key: str = 'ddG_matrix.residue_axis'

value_key: str | None = None

values_key: str = 'ddG_matrix.values'

exception biolm.pipeline.data.PipelineAPIAuthError(status_code: int, payload: Any, model_name: str)

Bases: RuntimeError

Raised when the BioLM API returns 401/402 (auth/billing) during a stage.

These errors are not retriable per-item — every batch will hit the same failure — so the pipeline fails fast instead of producing a silent empty result. skip_on_error=True does NOT swallow this; callers always see an unambiguous failure with the upstream error payload.

biolm.pipeline.data.Predict(model_name: str, sequences: list[str] | DataFrame | str | Path, extractions: str | list | None = None, params: dict | None = None, **kwargs) → DataFrame

Convenience function for single-step prediction.

Args:

model_name: BioLM model name sequences: Input sequences extractions: API response key(s) to extract. Required — pass the response

key for your model (e.g. extractions='prediction' for temberture, extractions='mean_plddt' for esmfold). Use a list or ExtractionSpec for multiple extractions.

params: Optional API parameters **kwargs: Additional arguments

Returns:

DataFrame with predictions

Example:

default Copy

>>> df = Predict('temberture-regression', sequences=['MKTAYIAKQRQ'], extractions='prediction')

class biolm.pipeline.data.PredictionStage(name: str, model_name: str, action: str = 'predict', params: dict | None = None, batch_size: int = 32, max_concurrent: int = 5, max_connections: int = 10, item_columns: dict[str, str] | None = None, extractions: str | list[Union[str, biolm.pipeline.data.ExtractionSpec]] | None = None, columns: str | dict[str, str] | None = None, embedding_extractor: EmbeddingSpec | Callable[[dict], Any] | None = None, structure_output: StructureSpec | None = None, structure_input: dict[str, str] | None = None, matrix_extraction: MatrixExtractionSpec | None = None, **kwargs)

Bases: Stage

Generic prediction stage using BioLM API.

Uses merge_mode = "union" so that when multiple prediction stages run in parallel, sequences processed by any stage survive the merge — i.e., independent predictions do not cancel each other out.

Args:

name: Stage name model_name: BioLM model name (e.g., ‘esmfold’, ‘esm2’, ‘temberture-regression’) action: API action (‘predict’, ‘encode’, ‘score’) prediction_type: Type of prediction for caching (e.g., ‘structure’, ‘stability’, ‘embedding’) params: Optional parameters for the API call batch_size: Number of sequences per pipeline batch (default 32).

Each pipeline batch becomes one SDK call, which the SDK may further split by the model’s maxItems schema limit.

max_concurrent: Maximum pipeline batches in flight at once (default 5).

Controls how many batches are dispatched concurrently at the pipeline level. Higher values keep the API more saturated but use more memory (up to max_concurrent * batch_size items plus their responses in memory at once).

max_connections: Maximum concurrent HTTP connections to the API

(default 10). This is the SDK-level semaphore that throttles the actual HTTP requests. Each pipeline batch may be split into multiple sub-requests by the model’s maxItems limit; max_connections caps how many of those sub-requests run simultaneously across all in-flight batches.

structure_output: Store the structure from this model’s response. structure_input: Inject structures from upstream models into API items.

Maps API field name → source model name (e.g. {"pdb": "esmfold"}).

matrix_extraction: Flatten a per-mutation response into individual

prediction rows (e.g. for DMS heatmaps).

merge_mode: str = 'union'

async process(df: DataFrame, datastore: DuckDBDataStore, **kwargs) → tuple[pandas.core.frame.DataFrame, biolm.pipeline.base.StageResult]

Process sequences through prediction model (legacy DataFrame path).

Uses the same bounded-concurrency batching as process_ws — see that method’s docstring for the design rationale.

Performance: single DuckDB anti-join for cache detection, bounded- concurrency API dispatch, single JOIN query to merge predictions back.

async process_streaming(df: DataFrame, datastore: DuckDBDataStore, **kwargs)

Process sequences and yield results as batches complete (streaming).

Yields DataFrames as API batches complete instead of waiting for all results. This allows downstream stages to start processing immediately.

async process_ws(ws: WorkingSet, datastore: DuckDBDataStore, **kwargs) → tuple[biolm.pipeline.base.WorkingSet, biolm.pipeline.base.StageResult]

Process sequences using WorkingSet with bounded-concurrency batching.

Batching strategy (two levels):

This method splits uncached sequences into pipeline-level batches of batch_size items (default 32) and keeps up to max_concurrent batches in flight at once. As each batch completes, its results are written to DuckDB immediately and a new batch is dispatched.

Each pipeline batch is itself an SDK call that may be further split by the model’s maxItems schema limit — the SDK handles that internally via asyncio.gather throttled by a semaphore.

Why bounded concurrency instead of all-at-once or sequential:

• Sequential (old behavior) leaves the API idle between batches. With 300 sequences and 200ms API latency, that’s 10 idle round-trips.

• All-at-once (process_streaming style) has no backpressure — 10k sequences creates 300+ in-flight tasks with all items and response payloads in memory simultaneously.

• Bounded concurrency keeps the API saturated (max_concurrent batches in flight) while bounding memory to at most max_concurrent * batch_size items plus their results. DuckDB writes happen as each batch lands, providing natural backpressure: a slow datastore slows down new batch dispatch.

Steps:

1. Cache check via anti-join (DuckDB)

2. Fetch only uncached (sequence_id, sequence) pairs

3. Dispatch batches with bounded concurrency

4. Write results to DuckDB as each batch completes

5. Return WorkingSet of IDs that have predictions

to_spec() → dict

Return a serializable dict for pipeline definition persistence.

class biolm.pipeline.data.SingleStepPipeline(model_name: str, action: str = 'predict', sequences: list[str] | DataFrame | str | Path = None, params: dict | None = None, extractions=None, columns=None, embedding_extractor=None, **kwargs)

Bases: DataPipeline

Simplified pipeline for single-step predictions.

Convenience class for running a single prediction model on sequences.

Args:

model_name: BioLM model name action: API action (‘predict’, ‘encode’, ‘score’) sequences: Input sequences params: Optional API parameters **kwargs: Additional arguments passed to DataPipeline

Example:

default Copy

>>> pipeline = SingleStepPipeline(
...     model_name='esmfold',
...     sequences=['MKTAYIAKQRQ', 'MKLAVID']
... )
>>> results = pipeline.run()
>>> df = pipeline.get_final_data()

class biolm.pipeline.data.StructureSpec(key: str, format: str | None = None, plddt_key: str | None = None, index: int | None = 0)

Bases: object

Specification for extracting and storing a structure from an API response.

Args:

key: Response dict key containing the structure string (e.g. “pdb”, “cif”, “pdbs”). format: Structure format — “pdb” or “cif”. Auto-detected from key if None. plddt_key: Optional response key for a confidence score to store alongside. index: For list-valued keys (e.g. AF2 “pdbs”), which element to store (default 0).

detect_format() → str

format: str | None = None

index: int | None = 0

key: str

plddt_key: str | None = None

biolm.pipeline.datastore module

Backward-compatibility shim.

Old code imported from biolm.pipeline.datastore; the implementation has moved to biolm.pipeline.datastore_duckdb. Import DataStore from either location — they resolve to the same class.

biolm.pipeline.datastore.DataStore

alias of DuckDBDataStore

biolm.pipeline.datastore_duckdb module

DuckDB-based DataStore with Parquet backend for efficient large-scale data management.

Key features: - Columnar storage (Parquet) - Out-of-core queries (bigger than RAM) - Vectorized anti-join deduplication - 5-50× faster than pandas for aggregations - Diff-mode friendly batch inserts

class biolm.pipeline.datastore_duckdb.DuckDBDataStore(db_path: str | Path | None = None, data_dir: str | Path | None = None)

Bases: object

DuckDB + Parquet based datastore for efficient sequence management.

Optimized for: - Millions of sequences - Complex queries (joins, filters, aggregations) - Out-of-core operations (bigger than RAM) - Diff mode with efficient deduplication

Args:

db_path: Path to DuckDB database file data_dir: Directory for Parquet files and large data

Example:

default Copy

>>> ds = DuckDBDataStore("pipeline.db", "data/")
>>> seq_id = ds.add_sequence("MKLLIV")
>>> ds.add_prediction(seq_id, "tm", "temberture-regression", 65.5)
>>>
>>> # Efficient query - no memory explosion
>>> high_tm = ds.query("SELECT * FROM predictions WHERE value > 60")

add_embedding(sequence_id: int, model_name: str, embedding: ndarray, layer: int | None = None)

Add embedding stored inline in DuckDB as FLOAT[] (no per-file Parquet overhead).

Args:

sequence_id: Sequence ID model_name: Model name embedding: Numpy array layer: Optional layer number

add_embeddings_batch(data: list[dict])

Batch-insert embeddings in a single DuckDB statement.

Each dict must have: sequence_id (int), model_name (str), embedding (np.ndarray). Optional: layer (int or None).

Significantly faster than N individual add_embedding() calls — one INSERT…SELECT vs. N individual row inserts.

add_generation_metadata(sequence_id: int, model_name: str, run_id: str = '', temperature: float | None = None, top_k: int | None = None, top_p: float | None = None, num_return_sequences: int | None = None, do_sample: bool | None = None, repetition_penalty: float | None = None, max_length: int | None = None, sampling_params: dict | None = None) → int

Store generation parameters for a sequence.

Returns:

metadata_id of the inserted row.

add_generation_metadata_batch(rows: list[dict]) → None

Batch-insert generation metadata — one DuckDB round-trip.

Each dict in rows must have at minimum sequence_id and model_name. Optional fields: temperature, top_k, top_p, num_return_sequences, do_sample, repetition_penalty, max_length.

add_prediction(sequence_id: int, prediction_type: str, model_name: str, value: float | None, metadata: dict | None = None)

Add single prediction (convenience wrapper).

add_prediction_by_sequence(sequence: str, prediction_type: str, model_name: str, value: float | None, metadata: dict | None = None) → int

Add a prediction, creating the sequence if it doesn’t exist.

Returns:

prediction_id of the inserted row.

add_predictions_batch(data: list[dict[str, Any]])

Batch add predictions efficiently.

Args:

data: List of dicts with keys: sequence_id, prediction_type,

model_name, value, metadata (optional)

Example:

default Copy

>>> ds.add_predictions_batch([
...     {'sequence_id': 1, 'prediction_type': 'tm',
...      'model_name': 'temberture-regression', 'value': 65.5},
...     {'sequence_id': 2, 'prediction_type': 'tm',
...      'model_name': 'temberture-regression', 'value': 70.2},
... ])

add_sequence(sequence: str) → int

Add single sequence (convenience wrapper).

add_sequences_batch(sequences: list[str] | None = None, deduplicate: bool = True, input_df: DataFrame | None = None, input_columns: list[str] | None = None) → list[int]

Add multiple sequences efficiently using anti-join deduplication.

This is the RECOMMENDED way to add sequences - vectorized and fast!

Two calling conventions:

1. Legacy (sequence-only): add_sequences_batch(["MKLLIV", ...])

2. Multi-column (arbitrary input columns): add_sequences_batch(input_df=df, input_columns=["heavy_chain", "light_chain"]) In this mode, the hash is computed across all input columns, and the column values are stored directly on the sequences table. A sequence column is still written (concatenation of all input columns joined with :) so downstream code has a fallback.

Args:

sequences: List of sequence strings (legacy path). deduplicate: Use anti-join to skip existing sequences. input_df: DataFrame with input columns (multi-column path). input_columns: Column names in input_df to use as primary data.

Returns:

List of sequence_ids (new and existing), preserving input order.

add_structure(sequence_id: int, model_name: str, structure_str: str | None = None, format: str = 'pdb', plddt_mean: float | None = None, plddt: float | None = None) → int

Store a structure gzip-compressed as BLOB (~8-12x smaller than plain TEXT).

Args:

sequence_id: Sequence ID. model_name: Model that produced the structure (e.g. ‘esmfold’). structure_str: Full structure file content as a string. format: ‘pdb’ or ‘cif’ (default ‘pdb’). plddt_mean: Mean pLDDT score (optional). plddt: Alias for plddt_mean.

Returns:

structure_id of the inserted row.

close()

Close the DuckDB connection and release the file lock.

Safe to call multiple times — subsequent calls are no-ops. Called automatically by __exit__ and __del__.

count_matching_sequences(sequences: list[str]) → int

Count how many of the given sequences already exist in the datastore.

Uses a single vectorized hash join instead of N individual lookups. Safe to call from any context — registers DataFrame explicitly.

create_pipeline_run(run_id: str, pipeline_type: str, config: dict, status: str = 'running')

Create or update a pipeline run record (safe for resume runs).

ensure_input_columns(columns: list[str])

Ensure the sequences table has all the given columns.

Uses ALTER TABLE ADD COLUMN for any that don’t already exist. This is idempotent — safe to call on every pipeline run.

Raises:

duckdb.CatalogException: If column creation fails for an unexpected

reason (e.g. type conflict). The “column already exists” case is silently ignored as it is the normal idempotent path.

execute_filter_sql(sequence_ids: list[int], sql_query: str, ws_table_name: str = '_filter_ws') → list[int]

Execute a filter SQL query and return surviving sequence IDs.

The sql_query must be a complete SELECT statement that returns sequence_id values and JOINs against ws_table_name (which contains the input sequence_ids) to scope to the working set.

Pass a unique ws_table_name (use make_filter_ws_name()) when running multiple SQL-native filters concurrently — the default _filter_ws name races otherwise.

Args:

sequence_ids: Input sequence IDs. sql_query: Complete SQL SELECT returning sequence_id values. ws_table_name: Name to register the working-set IDs under. Must

match the ws_table= argument passed to the filter’s to_sql().

Returns:

List of sequence_ids that survive the filter.

Raises:

ValueError: If sql_query is not a single SELECT statement.

export_to_csv(path: str | Path, **kwargs) → None

Export all data to CSV (convenience wrapper around export_to_dataframe).

export_to_dataframe(include_sequences: bool = True, include_predictions: bool = True, include_generation_metadata: bool = False, prediction_types: list[str] | None = None, run_id: str | None = None) → DataFrame

Export data to a flat DataFrame using a single DuckDB SQL query.

Uses conditional aggregation (CASE WHEN pivot) — no per-type queries, no pandas merges, no full table loads.

Args:

include_sequences: Always True; includes sequence_id, sequence, length. include_predictions: Pivot prediction_type values into columns. include_generation_metadata: Join generation_metadata columns. prediction_types: Limit to specific prediction types (None = all).

Returns:

Wide-format DataFrame: one row per sequence, one column per prediction type.

export_to_parquet(table_name: str, output_path: str | Path)

Export table to Parquet file (for sharing/archiving).

Args:

table_name: Table to export (sequences, predictions, etc.) output_path: Path to output Parquet file

Example:

default Copy

>>> ds.export_to_parquet('sequences', 'sequences_backup.parquet')

get_all_sequences() → DataFrame

Return all sequences as a DataFrame with sequence_id, sequence, length columns.

get_column_registry_entry(column_name: str) → dict | None

Return the registry entry for a column, or None if not registered.

get_context(run_id: str, key: str) → Any | None

Retrieve a value from the pipeline context table.

get_embedding(embedding_id: int) → tuple | None

Return (metadata_dict, embedding_array) for a given embedding_id, or None.

get_embeddings_bulk(sequence_ids: list[int], model_name: str | None = None) → dict[int, numpy.ndarray]

Fetch embeddings for multiple sequences in a single JOIN query.

Replaces N individual get_embeddings_by_sequence() calls (O(n) queries → O(1)).

Args:

sequence_ids: List of sequence IDs to fetch. model_name: Optional model filter.

Returns:

Dict mapping sequence_id → numpy embedding array for sequences that have an embedding.

get_embeddings_by_sequence(sequence: str, model_name: str | None = None, load_data: bool = False) → list[dict]

Get embeddings for a sequence.

get_embeddings_concat(sequence_ids: list[int], model_names: list[str]) → dict[int, numpy.ndarray]

Fetch and concatenate embeddings from multiple models per sequence.

For each sequence_id, retrieves the embedding from each model in model_names order and horizontally concatenates them into a single vector. Sequences missing an embedding from any requested model are omitted from the result.

Args:

sequence_ids: Sequence IDs to fetch. model_names: Ordered list of model names whose embeddings will be

concatenated (e.g. ["esm2-8m", "esmc-300m"]).

Returns:

Dict mapping sequence_id → concatenated numpy array.

get_existing_input_columns() → list[str]

Return extra columns on the sequences table (beyond the base schema).

Used for input-schema validation when connecting to an existing DB. Returns an empty list if only the base columns are present.

get_filter_results(run_id: str, stage_name: str) → list[int]

Return sequence_ids that passed a given filter stage in a run.

Args:

run_id: Pipeline run ID. stage_name: Filter stage name.

Returns:

List of sequence_ids that passed the filter, or empty list if no data.

get_generation_metadata(sequence_id: int) → list[dict]

Return generation metadata records for a sequence_id.

get_latest_definition_id() → str | None

Return the definition_id of the most recently created pipeline definition.

get_latest_run_id(definition_id: str | None = None) → str | None

Return the run_id of the most recent pipeline run.

from_db() must reuse the existing run_id so that resume can find already-completed stages. A new run_id means no stages are marked complete, causing everything to re-run.

Args:

definition_id: If provided, restrict to runs that used this definition.

If None, return the most recent run across all definitions.

Returns:

The run_id string, or None if no runs exist.

get_pipeline_metadata(key: str) → Any | None

Retrieve a value from pipeline_metadata by key.

get_pipeline_run(run_id: str) → dict | None

Return pipeline run record as a dict, or None if not found.

get_predictions(sequence_id: int, prediction_type: str | None = None, model_name: str | None = None) → DataFrame

Return predictions for a sequence_id as a DataFrame.

get_predictions_bulk(sequence_ids: list[int], prediction_type: str, model_name: str) → DataFrame

Fetch predictions for multiple sequences in a single JOIN query.

Replaces N individual get_predictions_by_sequence() calls.

Returns:

DataFrame with columns: sequence_id, value, metadata

get_predictions_by_sequence(sequence: str, prediction_type: str | None = None, model_name: str | None = None) → DataFrame

Get predictions for a sequence.

get_sequence(sequence_id: int) → str | None

Return the sequence string for a given sequence_id, or None if not found.

get_sequence_attributes_for_ids(sequence_ids: list[int], attr_names: list[str]) → dict[int, dict[str, str]]

Retrieve per-sequence attributes, returning {seq_id: {attr: value}}.

Args:

sequence_ids: Sequence IDs to look up. attr_names: Attribute names to retrieve.

Returns:

Nested dict: {sequence_id: {attr_name: attr_value}}.

get_sequence_id(sequence: str) → int | None

Get sequence_id for a sequence.

get_sequence_ids_with_prediction(sequence_ids: list[int], prediction_type: str, model_name: str) → list[int]

Return sequence_ids that DO have a given prediction (inverse of uncached check).

Args:

sequence_ids: Candidate sequence IDs. prediction_type: Prediction type key. model_name: Model name.

Returns:

List of sequence_ids that have a cached prediction.

get_sequences_for_ids(sequence_ids: list[int]) → list[tuple[int, str]]

Fetch (sequence_id, sequence) pairs for the given IDs.

Lightweight fetch for building API request items without materializing a full DataFrame.

Args:

sequence_ids: List of sequence IDs to look up.

Returns:

List of (sequence_id, sequence_string) tuples.

get_sequences_for_ids_with_columns(sequence_ids: list[int], columns: list[str]) → dict[int, dict[str, str]]

Fetch column values from the sequences table for given IDs.

This reads columns stored directly on the sequences table (via ensure_input_columns), NOT from sequence_attributes.

Args:

sequence_ids: Sequence IDs to look up. columns: Column names to fetch (must exist on the sequences table).

Returns:

{sequence_id: {col: value, ...}}

get_stats() → dict[str, int]

Return row counts for the main tables.

get_structure(sequence_id: int, model_name: str | None = None) → dict | None

Fetch the most recent structure for a sequence, decompressing on read.

Returns a dict with ‘structure_str’ key (always decompressed string) regardless of whether data was stored compressed (structure_data BLOB) or as legacy plain TEXT.

Args:

sequence_id: Sequence ID. model_name: Optional model filter.

Returns:

Dict with structure record, or None if not found.

get_structure_by_id(structure_id: int) → dict | None

Return a structure record by its structure_id (primary key).

get_structures_bulk(sequence_ids: list[int]) → DataFrame

Fetch structures for multiple sequences, decompressing structure content.

Returns a DataFrame with a ‘structure_str’ column (always plain text) regardless of whether data was stored compressed (structure_data BLOB) or as legacy plain TEXT.

Args:

sequence_ids: List of sequence IDs to look up.

Returns:

DataFrame with one row per structure record.

get_structures_by_sequence(sequence: str, model_name: str | None = None) → list[dict]

Return structure records for a sequence string (decompressed).

get_structures_for_ids(sequence_ids: list[int], model_name: str) → dict[int, dict]

Fetch the most recent structure for each sequence_id in one query.

Returns {sequence_id: record} where record has a ‘structure_str’ key. Replaces N individual get_structure() calls for batch structure injection.

get_uncached_sequence_ids(sequence_ids: list[int], prediction_type: str, model_name: str) → list[int]

Return sequence_ids that do NOT yet have a given prediction (vectorized anti-join).

Replaces N individual has_prediction() calls with a single SQL query.

Args:

sequence_ids: Candidate sequence IDs to check. prediction_type: Prediction type key. model_name: Model name.

Returns:

List of sequence_ids with no cached prediction.

has_prediction(sequence: str, prediction_type: str, model_name: str) → bool

Check if prediction exists for sequence.

is_stage_complete(stage_id: str) → bool

Check if stage is complete.

load_blob(blob_id: str) → str | None

Retrieve a stored blob by its blob_id. Returns None if not found.

load_pipeline_definition(definition_id: str | None = None) → dict | None

Load a pipeline definition by ID, or the latest one if ID is None.

make_filter_ws_name() → str

Mint a unique _filter_ws_ table name for one filter execution.

Required when SQL-native filters can run in parallel at the same DAG level — sharing the canonical _filter_ws registration would let two stages clobber each other’s working set.

mark_stage_complete(run_id: str, stage_name: str, stage_id: str, input_count: int, output_count: int, status: str = 'completed')

Mark stage as complete (or failed/skipped).

uses INSERT … ON CONFLICT DO NOTHING so that a stage that was already marked ‘completed’ does not get its completed_at timestamp overwritten on resume. Only truly new rows are inserted.

materialize_working_set(ws: WorkingSet, include_predictions: bool = True, prediction_types: list[str] | None = None) → pd.DataFrame

Materialize a WorkingSet into a DataFrame via a single DuckDB pivot query.

Args:

ws: WorkingSet containing the sequence IDs to materialize. include_predictions: If True, pivot prediction values into columns. prediction_types: Limit to specific types (None = all available).

Returns:

Wide-format DataFrame: one row per sequence, one column per prediction type. Always includes the following columns regardless of pipeline configuration:

• sequence_id, sequence, length, hash — core sequence data.

• source_label — label set on the generation config (e.g. DirectGenerationConfig.label, SaturationMutagenesisConfig.label). NULL / None when no label was supplied. This column is always present even when no labels have been set, so downstream code should not branch on "source_label" in df.columns — it is always there.

query(sql: str, params: list | None = None) → DataFrame

Execute arbitrary SQL query and return DataFrame.

This is the POWER feature - query directly without loading everything!

Args:

sql: DuckDB SQL query params: Optional query parameters

Returns:

DataFrame with results (only loads what matches query!)

Example:

default Copy

>>> # Find high-quality long sequences
>>> df = ds.query('''
...     SELECT s.sequence, p.value as plddt
...     FROM sequences s
...     JOIN predictions p ON s.sequence_id = p.sequence_id
...     WHERE s.length > 200
...     AND p.prediction_type = 'plddt'
...     AND p.value > 80
... ''')

register_column(column_name: str, model_name: str, action: str, definition_id: str, stage_name: str)

Register an output column in the prediction_column_registry. Idempotent.

save_filter_results(run_id: str, stage_name: str, passed_sequence_ids: list[int])

Record which sequence_ids passed a filter stage (for resume support).

Args:

run_id: Pipeline run ID. stage_name: Filter stage name. passed_sequence_ids: IDs of sequences that passed the filter.

save_pipeline_definition(definition_id: str, pipeline_type: str, input_schema_json: str | None, stages_json: str)

Persist a pipeline definition. Updates stages/schema if definition_id already exists.

set_context(run_id: str, key: str, value: Any)

Store a key-value pair in the pipeline context table.

set_pipeline_metadata(key: str, value: Any)

Upsert a key/value pair in the pipeline_metadata table.

store_blob(content: str) → str

Store a large string value, returning its blob_id (SHA-256[:32]).

Content-addressed: identical content always maps to the same blob_id. Safe to call multiple times with the same content (INSERT OR IGNORE).

store_sequence_attributes(seq_ids: list[int], attr_name: str, attr_values: list[str])

Persist a per-sequence attribute column (e.g. heavy_chain, light_chain).

Args:

seq_ids: Sequence IDs. attr_name: Attribute name (column name from the input DataFrame). attr_values: Corresponding values (one per sequence_id).

update_pipeline_run_status(run_id: str, status: str)

Update pipeline run status.