BioData

Utilities for working with a PostgreSQL BioData database, protein language model embeddings (pgvector), residue-level probing tasks, and GO/Taxonomy analysis.

What Is In This Repo

---

Quick Start

from CBBIO import connect, Generator, FastaBatcher, EmbeddingWriter, run_embedding_generation

# Connect to the database
client = connect()
client.health_check()

# Generate embeddings from a FASTA file
generator = Generator(model_class="esm2", device="cuda:0")
batcher   = FastaBatcher("proteins.fasta", max_batch_tokens=32_768)
writer    = EmbeddingWriter(format="h5", path="embeddings.h5")
run_embedding_generation(generator, batcher, writer, layer_index=33)

# Nearest-neighbor search
neighbors = client.find_nearest_neighbors(query_emb, embedding_type_id=1, layer_index=33, k=20)

See docs/GettingStarted.md for a full tutorial covering all modules.

---

Installation

From GitHub

pip install "biodata @ git+ssh://[email protected]/cbbio/BioData.git@main"

Optional extras:

pip install "biodata[notebooks] @ git+ssh://[email protected]/cbbio/BioData.git@main"
pip install "biodata[dev] @ git+ssh://[email protected]/cbbio/BioData.git@main"

From a local clone

pip install .           # runtime only
pip install ".[dev]"    # with dev tools

---

Database Setup

1. Start PostgreSQL + pgvector

docker run -d --name pgvectorsql \
  -e POSTGRES_USER=biodata \
  -e POSTGRES_PASSWORD=secret \
  -e POSTGRES_DB=biodata \
  -p 5432:5432 \
  pgvector/pgvector:pg16

2. Create the extension and schema

PGPASSWORD=secret psql -h localhost -U biodata -d biodata \
  -c "CREATE EXTENSION IF NOT EXISTS vector;"

PGPASSWORD=secret psql -h localhost -U biodata -d biodata -f sql/schema.sql

3. Restore a dataset (optional)

PGPASSWORD=secret pg_restore -h localhost -U biodata -d biodata /path/to/BioData.backup

---

Configuration

Copy the untracked template to config.yaml (or pass config_path= to connect()):

cp config.yaml.example config.yaml

database:
  host: localhost
  port: 5432
  user: biodata
  password: ""
  name: biodata

search:
  default_metric: cosine    # l2 | cosine | inner_product
  default_k: 10

Environment variable overrides (highest precedence):

---

Documentation

Browse the published API reference on Read the Docs.

---

Running Tests

poetry run pytest -q

Integration tests (real DB required)

cp tests/assets/config_test.yaml.example config_test.yaml
# edit with your test DB credentials
poetry run pytest -q tests/test_biodata_integration.py

Override config path: BIODATA_TEST_CONFIG=/path/to/config_test.yaml.

Coverage

poetry run pytest --cov=CBBIO --cov-report=term-missing -q

Type checking

poetry run pyright