Merge pull request #1 from Open-Athena/proof-of-concept

Proof of concept

Author: gonzalobenegas

.gitignore

@@ -1,3 +1,5 @@
+.cursorrules
+uv.lock
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]

.pre-commit-config.yaml

@@ -0,0 +1,14 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.11
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.17.1
+    hooks:
+      - id: mypy
+        exclude: ^(tests/|examples/)
+        entry: uv run mypy

.python-version

@@ -0,0 +1 @@
+3.13

README.md

@@ -1,2 +1,64 @@
-# biofoundation
-Biological foundation models
+![BioFoundation Logo](docs/source/_static/logo.png)
+
+# BioFoundation
+
+An ecosystem for biological foundation models.
+
+## Installation
+
+Create and activate a virtual environment:
+
+```bash
+uv venv
+source .venv/bin/activate
+```
+
+Install the package:
+
+```bash
+uv pip install -e .
+```
+
+## Additional Dependencies
+
+```bash
+# PlantCad1
+uv pip install -e .[mamba]
+```
+
+## Usage
+
+Run the example script:
+
+```bash
+source .venv/bin/activate
+python examples/plantcad_evolutionary_constraint.py
+```
+
+## Development Setup
+
+To set up the development environment with linting, formatting, type checking, and testing:
+
+```bash
+# Install development dependencies
+uv pip install --group dev
+
+# Install both main package and dev tools
+uv pip install -e . --group dev
+```
+
+## Development Tools
+
+### Running Quality Checks
+
+```bash
+# Run all pre-commit hooks (linting, formatting, type checking)
+pre-commit run --all-files
+```
+
+### Running Tests
+
+```bash
+# Run all tests
+pytest
+```

biofoundation/__init__.py

@@ -0,0 +1,12 @@
+"""An ecosystem for biological foundation models.
+
+This package provides tools for working with biological foundation models,
+including data transformation, model wrappers, and inference utilities.
+"""
+
+try:
+    from importlib.metadata import version
+
+    __version__ = version("biofoundation")
+except ImportError:
+    __version__ = "unknown"

biofoundation/data.py

@@ -0,0 +1,42 @@
+from transformers import PreTrainedTokenizerBase
+from typing import Any
+
+
+def transform_reflogprob_mlm(
+    example: dict[str, Any],
+    tokenizer: PreTrainedTokenizerBase,
+    pos: int,
+    seq_col: str = "seq",
+) -> dict[str, Any]:
+    """Transform a sequence example for reference log probability MLM inference.
+
+    This function prepares a sequence example for masked language modeling by:
+    1. Tokenizing the input sequence
+    2. Masking a specific position in the sequence
+    3. Recording the reference token at that position
+
+    Args:
+        example: Dictionary containing the sequence data. Must have a key matching
+            `seq_col` that contains the input sequence.
+        tokenizer: HuggingFace tokenizer for converting text to token IDs.
+        pos: Position in the sequence to mask (0-indexed).
+        seq_col: Key in the example dictionary that contains the sequence.
+            Defaults to "seq".
+
+    Returns:
+        Dictionary with three keys:
+        - input_ids_BL: Token IDs with the specified position masked
+        - pos_B: The masked position
+        - ref_B: The reference token ID that was at the masked position
+
+    Example:
+        >>> example = {"seq": "ATCG"}
+        >>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        >>> result = transform_reflogprob_mlm(example, tokenizer, 1)
+        >>> print(result)
+        {'input_ids_BL': tensor([...]), 'pos_B': 1, 'ref_B': 3}
+    """
+    input_ids = tokenizer(example[seq_col], return_tensors="pt")["input_ids"][0]
+    ref = input_ids[pos].item()
+    input_ids[pos] = tokenizer.mask_token_id
+    return dict(input_ids_BL=input_ids, pos_B=pos, ref_B=ref)

biofoundation/inference.py

@@ -0,0 +1,35 @@
+import datasets
+import tempfile
+import torch.nn as nn
+from transformers import Trainer, TrainingArguments
+from typing import Any
+
+
+def run_inference(
+    model: nn.Module,
+    dataset: datasets.Dataset,
+    **kwargs: Any,
+) -> Any:
+    """Run inference on a dataset using a trained model.
+
+    Args:
+        model: A trained PyTorch model that can be used with the HuggingFace Trainer.
+        dataset: HuggingFace dataset to run inference on. The dataset should be
+            compatible with the model's expected input format.
+        **kwargs: Additional keyword arguments to pass to TrainingArguments.
+            Common options include:
+            - per_device_eval_batch_size: Batch size for evaluation
+            - dataloader_num_workers: Number of workers for data loading
+            - torch_compile: Whether to use torch.compile for faster inference
+            - bf16_full_eval: Whether to use bf16 for evaluation
+
+    Returns:
+        The model's predictions on the dataset. The exact format depends on the
+        model and dataset, but typically includes probabilities or embeddings.
+    """
+    training_args = TrainingArguments(
+        output_dir=tempfile.TemporaryDirectory().name,
+        **kwargs,
+    )
+    trainer = Trainer(model=model, args=training_args)
+    return trainer.predict(test_dataset=dataset).predictions

biofoundation/model.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from jaxtyping import Float, Int
+
+
+class RefLogProbMLM(nn.Module):
+    """Reference Log Probability Masked Language Model wrapper.
+
+    This class wraps a pre-trained language model to compute reference log probabilities
+    for masked language modeling tasks. It takes a sequence with a masked position
+    and returns the log probability of the reference token at that position.
+
+    The model expects input tensors with specific shapes:
+    - input_ids_BL: Batch of token ID sequences [batch_size, sequence_length]
+    - pos_B: Positions to evaluate for each sequence in the batch
+    - ref_B: Reference token IDs to compute log probabilities for
+
+    Attributes:
+        model: The underlying language model that provides logits
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(
+        self,
+        input_ids_BL: Int[Tensor, "B L"],
+        pos_B: Int[Tensor, " B"],
+        ref_B: Int[Tensor, " B"],
+    ) -> Float[Tensor, " B"]:
+        """Forward pass to compute reference log probabilities.
+
+        This method:
+        1. Runs the input through the language model to get logits
+        2. Extracts logits at the specified positions
+        3. Computes log probabilities using softmax
+        4. Returns the log probability of the reference tokens
+
+        Args:
+            input_ids_BL: Batch of token ID sequences with shape [B, L] where
+                B is batch size and L is sequence length
+            pos_B: Positions to evaluate for each sequence in the batch with
+                shape [B] where each element is the position index
+            ref_B: Reference token IDs to compute log probabilities for with
+                shape [B] where each element is the token ID
+
+        Returns:
+            Log probabilities of the reference tokens at the specified positions
+            with shape [B] where B is the batch size
+        """
+        B, L = input_ids_BL.shape
+        batch_indices_B = torch.arange(B)
+        logits_BLV = self.model(input_ids_BL).logits
+        logits_BV = logits_BLV[batch_indices_B, pos_B]
+        logprobs_BV = F.log_softmax(logits_BV, dim=-1)
+        res_B = logprobs_BV[batch_indices_B, ref_B]
+        return res_B

docs/source/_static/logo.png

@@ -0,0 +1,42 @@
+from biofoundation.data import transform_reflogprob_mlm
+from biofoundation.model import RefLogProbMLM
+from biofoundation.inference import run_inference
+from datasets import load_dataset
+from functools import partial
+import numpy as np
+from sklearn.metrics import roc_auc_score
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+
+model_name = "kuleshov-group/PlantCaduceus_l20"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)
+
+dataset = load_dataset(
+    "plantcad/evolutionary-constraint-example",
+    "10k",
+    split="validation",
+)
+label = np.array(dataset["label"])
+
+dataset = dataset.map(
+    partial(
+        transform_reflogprob_mlm,
+        tokenizer=tokenizer,
+        pos=255,
+        seq_col="sequences",
+    ),
+    remove_columns=dataset.column_names,
+    num_proc=8,
+)
+
+pred = run_inference(
+    RefLogProbMLM(model),
+    dataset,
+    per_device_eval_batch_size=256,
+    torch_compile=False,
+    bf16_full_eval=True,
+    dataloader_num_workers=8,
+)
+
+print(f"{roc_auc_score(label, pred)=}")