diffbio 0.1.0__py3-none-any.whl

diffbio/operators/foundation_models/adapters.py

@@ -0,0 +1,114 @@
+"""Shared adapter interfaces for benchmark-facing foundation-model integrations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+import jax.numpy as jnp
+
+from diffbio.operators.foundation_models.contracts import (
+    FoundationArtifactSpec,
+    build_foundation_model_metadata,
+)
+
+_CANONICAL_BENCHMARK_METADATA_KEYS = (
+    "embedding_source",
+    "foundation_source_name",
+)
+
+
+def _validate_benchmark_text(value: str, *, field_name: str) -> None:
+    """Require non-empty adapter metadata text for stable benchmark contracts."""
+    if not value:
+        raise ValueError(f"{field_name} must be non-empty.")
+
+
+class FoundationBenchmarkAdapter(Protocol):
+    """Common benchmark-facing contract for foundation-model adapters."""
+
+    def result_data(self) -> dict[str, Any]:
+        """Return canonical operator metadata for benchmark tagging."""
+        ...
+
+    def benchmark_metadata(self) -> dict[str, Any]:
+        """Return benchmark metadata describing the adapter source."""
+        ...
+
+
+class SequenceFoundationAdapter(FoundationBenchmarkAdapter, Protocol):
+    """Shared contract for sequence foundation-model adapters."""
+
+    def load_dataset_embeddings(
+        self,
+        *,
+        reference_sequence_ids: Sequence[str],
+        one_hot_sequences: Any,
+    ) -> jnp.ndarray:
+        """Return embeddings aligned to a benchmark dataset order."""
+        ...
+
+
+class FoundationBenchmarkAdapterBase:
+    """Base implementation for stable benchmark metadata handling."""
+
+    def __init__(
+        self,
+        *,
+        artifact_spec: FoundationArtifactSpec,
+        source_name: str,
+        embedding_source: str,
+        extra_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        _validate_benchmark_text(source_name, field_name="source_name")
+        _validate_benchmark_text(embedding_source, field_name="embedding_source")
+        normalized_extra_metadata = (
+            {} if extra_metadata is None else dict(sorted(extra_metadata.items()))
+        )
+        conflicting_keys = [
+            key for key in _CANONICAL_BENCHMARK_METADATA_KEYS if key in normalized_extra_metadata
+        ]
+        if conflicting_keys:
+            keys = ", ".join(conflicting_keys)
+            raise ValueError(f"extra_metadata cannot override canonical key(s): {keys}")
+
+        self.artifact_spec = artifact_spec
+        self.source_name = source_name
+        self.embedding_source = embedding_source
+        self.extra_metadata = normalized_extra_metadata
+
+    def result_data(self) -> dict[str, Any]:
+        """Return benchmark-ready foundation-model metadata."""
+        return {"foundation_model": build_foundation_model_metadata(self.artifact_spec)}
+
+    def benchmark_metadata(self) -> dict[str, Any]:
+        """Return benchmark metadata describing the adapter source."""
+        metadata: dict[str, Any] = {
+            "embedding_source": self.embedding_source,
+            "foundation_source_name": self.source_name,
+        }
+        metadata.update(self.extra_metadata)
+        return metadata
+
+
+_FOUNDATION_ADAPTER_REGISTRY: dict[str, type[object]] = {}
+
+
+def register_foundation_adapter(adapter_key: str, adapter_cls: type[object]) -> None:
+    """Register a benchmark-facing foundation-model adapter class."""
+    _validate_benchmark_text(adapter_key, field_name="adapter_key")
+    _FOUNDATION_ADAPTER_REGISTRY[adapter_key] = adapter_cls
+
+
+def get_foundation_adapter_cls(adapter_key: str) -> type[object]:
+    """Return the registered adapter class for a canonical adapter key."""
+    try:
+        return _FOUNDATION_ADAPTER_REGISTRY[adapter_key]
+    except KeyError as exc:
+        raise KeyError(f"No adapter registered for {adapter_key!r}.") from exc
+
+
+def create_foundation_adapter(adapter_key: str, *args: Any, **kwargs: Any) -> object:
+    """Instantiate a registered foundation-model adapter."""
+    adapter_cls = get_foundation_adapter_cls(adapter_key)
+    return adapter_cls(*args, **kwargs)

diffbio/operators/foundation_models/contracts.py

@@ -0,0 +1,245 @@
+"""Shared contracts for DiffBio foundation-model operators.
+
+This module centralizes the common metadata, output schema, and registry used
+by DiffBio foundation-model operators. The contract is intentionally shared
+across sequence, single-cell, and future imported biological foundation models
+so downstream code can rely on one stable interface instead of task-specific
+ad hoc keys.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Any, Mapping
+
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from jaxtyping import Array, PyTree
+
+
+class FoundationModelKind(StrEnum):
+    """Supported high-level foundation-model families."""
+
+    SEQUENCE_TRANSFORMER = "sequence_transformer"
+    SINGLE_CELL_TRANSFORMER = "single_cell_transformer"
+
+
+class AdapterMode(StrEnum):
+    """How DiffBio integrates the underlying foundation model."""
+
+    PRECOMPUTED = "precomputed"
+    FROZEN_ENCODER = "frozen_encoder"
+    NATIVE_TRAINABLE = "native_trainable"
+
+
+class PoolingStrategy(StrEnum):
+    """Canonical pooling strategies for foundation-model outputs."""
+
+    NONE = "none"
+    MEAN = "mean"
+    CLS = "cls"
+
+
+FOUNDATION_TAG_KEYS = (
+    "model_family",
+    "adapter_mode",
+    "artifact_id",
+    "preprocessing_version",
+)
+FOUNDATION_METADATA_KEYS = (*FOUNDATION_TAG_KEYS, "pooling_strategy")
+FOUNDATION_BENCHMARK_COMPARISON_AXES = (
+    "dataset",
+    "task",
+    *FOUNDATION_TAG_KEYS,
+)
+
+
+@dataclass(frozen=True)
+class FoundationEmbeddingOperatorConfig(OperatorConfig):
+    """Shared config fields for foundation-model operators."""
+
+    adapter_mode: AdapterMode = AdapterMode.NATIVE_TRAINABLE
+    artifact_id: str = "diffbio.builtin"
+    preprocessing_version: str = "native_v1"
+
+    def __post_init__(self) -> None:
+        """Validate foundation-model metadata fields."""
+        super().__post_init__()
+        _validate_ascii_text(self.artifact_id, field_name="artifact_id")
+        _validate_ascii_text(
+            self.preprocessing_version,
+            field_name="preprocessing_version",
+        )
+
+
+@dataclass(frozen=True)
+class FoundationArtifactSpec:
+    """Immutable spec describing a foundation-model artifact and interface."""
+
+    model_family: FoundationModelKind
+    artifact_id: str
+    preprocessing_version: str
+    adapter_mode: AdapterMode
+    pooling_strategy: PoolingStrategy
+
+    def __post_init__(self) -> None:
+        """Validate artifact fields for JAX-safe metadata encoding."""
+        _validate_ascii_text(self.artifact_id, field_name="artifact_id")
+        _validate_ascii_text(
+            self.preprocessing_version,
+            field_name="preprocessing_version",
+        )
+
+
+def _validate_ascii_text(value: str, *, field_name: str) -> None:
+    """Require non-empty ASCII metadata for JAX-safe output encoding."""
+    if not value:
+        raise ValueError(f"{field_name} must be non-empty.")
+    try:
+        value.encode("ascii")
+    except UnicodeEncodeError as exc:
+        raise ValueError(f"{field_name} must be ASCII-only.") from exc
+
+
+def encode_foundation_text(value: str) -> Array:
+    """Encode ASCII metadata text as a JAX array for jit-safe outputs."""
+    _validate_ascii_text(value, field_name="metadata_text")
+    encoded = value.encode("ascii")
+    return jnp.asarray(list(encoded), dtype=jnp.uint8)
+
+
+def decode_foundation_text(value: Array) -> str:
+    """Decode a JAX uint8 text array back into an ASCII string."""
+    data = bytes(int(item) for item in value.tolist())
+    return data.decode("ascii")
+
+
+def decode_foundation_model_metadata(
+    metadata: Mapping[str, Array | str],
+) -> dict[str, str]:
+    """Decode canonical foundation-model metadata into benchmark-friendly strings."""
+    decoded: dict[str, str] = {}
+    for key in FOUNDATION_METADATA_KEYS:
+        value = metadata.get(key)
+        if value is None:
+            continue
+        if isinstance(value, str):
+            decoded[key] = value
+            continue
+        decoded[key] = decode_foundation_text(value)
+    return decoded
+
+
+def build_foundation_benchmark_metadata(
+    metadata: Mapping[str, Array | str],
+    *,
+    dataset: str,
+    task: str,
+) -> dict[str, str]:
+    """Build the shared benchmark-facing foundation metadata contract."""
+    _validate_ascii_text(dataset, field_name="dataset")
+    _validate_ascii_text(task, field_name="task")
+    decoded = decode_foundation_model_metadata(metadata)
+    missing = [key for key in FOUNDATION_TAG_KEYS if key not in decoded]
+    if missing:
+        missing_fields = ", ".join(missing)
+        raise ValueError(f"foundation metadata missing required field(s): {missing_fields}")
+
+    return {
+        "dataset": dataset,
+        "task": task,
+        **decoded,
+    }
+
+
+def build_foundation_model_metadata(
+    artifact_spec: FoundationArtifactSpec,
+) -> dict[str, Array]:
+    """Build a jit-safe metadata payload for a foundation-model result."""
+    return {
+        "model_family": encode_foundation_text(artifact_spec.model_family.value),
+        "artifact_id": encode_foundation_text(artifact_spec.artifact_id),
+        "preprocessing_version": encode_foundation_text(artifact_spec.preprocessing_version),
+        "adapter_mode": encode_foundation_text(artifact_spec.adapter_mode.value),
+        "pooling_strategy": encode_foundation_text(artifact_spec.pooling_strategy.value),
+    }
+
+
+class FoundationEmbeddingMixin:
+    """Mixin providing canonical outputs for foundation-model operators."""
+
+    config: FoundationEmbeddingOperatorConfig
+    foundation_model_kind: FoundationModelKind
+
+    def foundation_pooling_strategy(self) -> PoolingStrategy:
+        """Return the pooling strategy used for the global embedding."""
+        return PoolingStrategy.NONE
+
+    def foundation_artifact_spec(self) -> FoundationArtifactSpec:
+        """Build the artifact spec for the current operator."""
+        return FoundationArtifactSpec(
+            model_family=self.foundation_model_kind,
+            artifact_id=self.config.artifact_id,
+            preprocessing_version=self.config.preprocessing_version,
+            adapter_mode=self.config.adapter_mode,
+            pooling_strategy=self.foundation_pooling_strategy(),
+        )
+
+    def foundation_result(
+        self,
+        data: PyTree,
+        embeddings: Array,
+        *,
+        token_embeddings: Array | None = None,
+        extra_outputs: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Build the canonical operator result payload."""
+        transformed_data = {
+            **data,
+            "embeddings": embeddings,
+            "foundation_model": build_foundation_model_metadata(self.foundation_artifact_spec()),
+        }
+
+        if token_embeddings is not None:
+            transformed_data["token_embeddings"] = token_embeddings
+
+        if extra_outputs:
+            for key in ("embeddings", "token_embeddings", "foundation_model"):
+                if key in extra_outputs:
+                    raise ValueError(f"extra_outputs cannot override canonical key {key!r}.")
+            transformed_data.update(extra_outputs)
+
+        return transformed_data
+
+
+_FOUNDATION_MODEL_REGISTRY: dict[FoundationModelKind, type[OperatorModule]] = {}
+
+
+def register_foundation_model(
+    model_family: FoundationModelKind,
+    operator_cls: type[OperatorModule],
+) -> None:
+    """Register an operator class for a foundation-model family."""
+    _FOUNDATION_MODEL_REGISTRY[model_family] = operator_cls
+
+
+def get_foundation_model_cls(
+    model_family: FoundationModelKind,
+) -> type[OperatorModule]:
+    """Return the registered operator class for a foundation-model family."""
+    try:
+        return _FOUNDATION_MODEL_REGISTRY[model_family]
+    except KeyError as exc:
+        raise KeyError(f"No operator registered for {model_family.value!r}.") from exc
+
+
+def create_foundation_model(
+    model_family: FoundationModelKind,
+    *args: Any,
+    **kwargs: Any,
+) -> OperatorModule:
+    """Instantiate a registered foundation-model operator."""
+    operator_cls = get_foundation_model_cls(model_family)
+    return operator_cls(*args, **kwargs)

diffbio/operators/foundation_models/embedding_probe.py

@@ -0,0 +1,83 @@
+"""Lightweight probing operators for foundation-model embeddings."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+from jaxtyping import PyTree
+
+
+@dataclass(frozen=True, kw_only=True)
+class EmbeddingProbeConfig(OperatorConfig):
+    """Configuration for a lightweight embedding probe."""
+
+    input_dim: int
+    n_classes: int
+    hidden_dim: int | None = None
+
+
+class LinearEmbeddingProbe(OperatorModule):
+    """Small classifier for probing embedding quality on downstream tasks."""
+
+    def __init__(
+        self,
+        config: EmbeddingProbeConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        super().__init__(config, rngs=rngs, name=name)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.hidden = None
+        if config.hidden_dim is not None:
+            self.hidden = nnx.Linear(
+                in_features=config.input_dim,
+                out_features=config.hidden_dim,
+                rngs=rngs,
+            )
+            classifier_in_dim = config.hidden_dim
+        else:
+            classifier_in_dim = config.input_dim
+
+        self.classifier = nnx.Linear(
+            in_features=classifier_in_dim,
+            out_features=config.n_classes,
+            rngs=rngs,
+        )
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Predict class probabilities from input embeddings."""
+        del random_params, stats
+
+        embeddings = data["embeddings"]
+        features = embeddings
+        if self.hidden is not None:
+            features = nnx.relu(self.hidden(features))
+
+        logits = self.classifier(features)
+        probabilities = jax.nn.softmax(logits, axis=-1)
+        predicted_labels = jnp.argmax(probabilities, axis=-1)
+
+        transformed_data = {
+            **data,
+            "logits": logits,
+            "probabilities": probabilities,
+            "predicted_labels": predicted_labels,
+        }
+        return transformed_data, state, metadata

diffbio/operators/foundation_models/experimental.py

@@ -0,0 +1,128 @@
+"""Experimental foundation-model boundary policy.
+
+This module is a policy namespace, not a stable implementation surface. It
+keeps speculative foundation-model capabilities explicit until benchmarks,
+provenance, regression guards, and docs promote them into stable support.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Mapping
+
+from diffbio.operators.foundation_models.contracts import _validate_ascii_text
+
+EXPERIMENTAL_FOUNDATION_MODEL_NAMESPACE = "diffbio.operators.foundation_models.experimental"
+FOUNDATION_EXPERIMENTAL_PROMOTION_CRITERIA = (
+    "explicit experimental namespace",
+    "canonical artifact provenance",
+    "downstream benchmark suite",
+    "Calibrax regression guard",
+    "shared audit bundle",
+    "stable documentation update",
+)
+
+
+@dataclass(frozen=True)
+class ExperimentalFoundationCapability:
+    """Fail-closed policy record for speculative foundation-model scope."""
+
+    key: str
+    display_name: str
+    stable_exclusion: str
+    namespace: str = EXPERIMENTAL_FOUNDATION_MODEL_NAMESPACE
+    stable_support: bool = False
+    benchmark_status: str = "benchmark_unverified"
+    required_promotion_criteria: tuple[str, ...] = FOUNDATION_EXPERIMENTAL_PROMOTION_CRITERIA
+
+    def __post_init__(self) -> None:
+        """Validate the policy record cannot silently become stable."""
+        for field_name, value in (
+            ("key", self.key),
+            ("display_name", self.display_name),
+            ("stable_exclusion", self.stable_exclusion),
+            ("namespace", self.namespace),
+            ("benchmark_status", self.benchmark_status),
+        ):
+            _validate_ascii_text(value, field_name=field_name)
+        if self.namespace != EXPERIMENTAL_FOUNDATION_MODEL_NAMESPACE:
+            raise ValueError(
+                "experimental foundation capabilities must use the experimental namespace."
+            )
+        if self.stable_support:
+            raise ValueError("experimental foundation capabilities cannot be stable.")
+        if self.required_promotion_criteria != FOUNDATION_EXPERIMENTAL_PROMOTION_CRITERIA:
+            raise ValueError(
+                "experimental foundation capabilities must use the shared promotion criteria."
+            )
+
+
+_FOUNDATION_EXPERIMENTAL_CAPABILITIES = {
+    "long_context_sequence_models": ExperimentalFoundationCapability(
+        key="long_context_sequence_models",
+        display_name="Long-context sequence models",
+        stable_exclusion=(
+            "No stable long-context sequence model support until promotion criteria are satisfied."
+        ),
+    ),
+    "hyena_style_sequence_models": ExperimentalFoundationCapability(
+        key="hyena_style_sequence_models",
+        display_name="Hyena-style sequence models",
+        stable_exclusion=(
+            "No stable Hyena-style runtime support until promotion criteria are satisfied."
+        ),
+    ),
+    "external_native_trainable_checkpoint_import": ExperimentalFoundationCapability(
+        key="external_native_trainable_checkpoint_import",
+        display_name="External native_trainable checkpoint import",
+        stable_exclusion=(
+            "The stable native_trainable adapter mode only covers DiffBio-native "
+            "operators, not external checkpoint conversion."
+        ),
+    ),
+    "peft_finetuning": ExperimentalFoundationCapability(
+        key="peft_finetuning",
+        display_name="PEFT fine-tuning utilities",
+        stable_exclusion=(
+            "No stable PEFT utilities are shipped until promotion criteria are satisfied."
+        ),
+    ),
+    "lora_adaptation": ExperimentalFoundationCapability(
+        key="lora_adaptation",
+        display_name="LoRA adaptation utilities",
+        stable_exclusion=(
+            "No stable LoRA adaptation utilities are shipped until promotion "
+            "criteria are satisfied."
+        ),
+    ),
+}
+FOUNDATION_EXPERIMENTAL_CAPABILITIES: Mapping[
+    str,
+    ExperimentalFoundationCapability,
+] = MappingProxyType(_FOUNDATION_EXPERIMENTAL_CAPABILITIES)
+
+
+def is_experimental_foundation_capability(key: str) -> bool:
+    """Return whether a capability is explicitly fenced as experimental."""
+    return key in FOUNDATION_EXPERIMENTAL_CAPABILITIES
+
+
+def get_experimental_foundation_capability(
+    key: str,
+) -> ExperimentalFoundationCapability:
+    """Return the experimental policy record for a capability key."""
+    try:
+        return FOUNDATION_EXPERIMENTAL_CAPABILITIES[key]
+    except KeyError as exc:
+        raise KeyError(f"No experimental foundation capability registered for {key!r}.") from exc
+
+
+__all__ = [
+    "EXPERIMENTAL_FOUNDATION_MODEL_NAMESPACE",
+    "FOUNDATION_EXPERIMENTAL_CAPABILITIES",
+    "FOUNDATION_EXPERIMENTAL_PROMOTION_CRITERIA",
+    "ExperimentalFoundationCapability",
+    "get_experimental_foundation_capability",
+    "is_experimental_foundation_capability",
+]