File size: 6,463 Bytes
4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """Public value types for dataset embedding."""
from __future__ import annotations
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Literal, overload
from torch import Tensor
@dataclass(frozen=True, slots=True)
class EmbeddingInput:
"""One named protein sequence supplied to :func:`embed_dataset`."""
id: str
sequence: str
def __post_init__(self) -> None:
if not isinstance(self.id, str) or not self.id:
raise ValueError("EmbeddingInput.id must be a non-empty string.")
if not isinstance(self.sequence, str) or not self.sequence:
raise ValueError("EmbeddingInput.sequence must be a non-empty string.")
@dataclass(frozen=True, slots=True)
class LazyTensorReference:
"""A tensor stored outside memory and loaded only when requested."""
source: str
key: str
dtype: str
shape: tuple[int, ...]
sha256: str
_loader: Callable[[], Tensor] = field(repr=False, compare=False)
def load(self, *, verify: bool = True) -> Tensor:
"""Load X and optionally verify its content digest."""
if not isinstance(verify, bool):
raise TypeError("verify must be a boolean.")
X = self._loader() # self.shape
if not isinstance(X, Tensor):
raise TypeError(f"Stored tensor loader for {self.key!r} must return a Tensor.")
if tuple(X.shape) != self.shape:
raise ValueError(
f"Stored tensor {self.key!r} has shape {tuple(X.shape)}, expected {self.shape}."
)
dtype = str(X.dtype).removeprefix("torch.")
if dtype != self.dtype:
raise ValueError(
f"Stored tensor {self.key!r} has dtype {dtype!r}, expected {self.dtype!r}."
)
if verify:
from .storage import tensor_sha256
digest = tensor_sha256(X)
if digest != self.sha256:
raise ValueError(f"Stored tensor {self.key!r} failed SHA-256 verification.")
return X # self.shape
TensorValue = Tensor | LazyTensorReference
@dataclass(frozen=True, slots=True)
class EmbeddingRecord:
"""One ordered embedding result."""
id: str
sequence: str
tensor: TensorValue
def __post_init__(self) -> None:
if not isinstance(self.id, str) or not self.id:
raise ValueError("EmbeddingRecord.id must be a non-empty string.")
if not isinstance(self.sequence, str) or not self.sequence:
raise ValueError("EmbeddingRecord.sequence must be a non-empty string.")
if not isinstance(self.tensor, (Tensor, LazyTensorReference)):
raise TypeError("EmbeddingRecord.tensor must be a Tensor or LazyTensorReference.")
def load_tensor(self, *, verify: bool = True) -> Tensor:
"""Return X regardless of whether this record is memory-backed or lazy."""
if not isinstance(verify, bool):
raise TypeError("verify must be a boolean.")
if isinstance(self.tensor, LazyTensorReference):
return self.tensor.load(verify=verify) # (...)
return self.tensor # (...)
class EmbeddingResult(Sequence[EmbeddingRecord]):
"""Ordered embedding records and the metadata needed to reproduce them."""
def __init__(
self,
records: Sequence[EmbeddingRecord],
metadata: Mapping[str, Any] | None = None,
) -> None:
self.records: Sequence[EmbeddingRecord] = (
records if getattr(records, "_fastplms_immutable_sequence", False) else tuple(records)
)
self.metadata = dict(metadata or {})
def __len__(self) -> int:
return len(self.records)
def __iter__(self) -> Iterator[EmbeddingRecord]:
return iter(self.records)
@overload
def __getitem__(self, index: int, /) -> EmbeddingRecord: ...
@overload
def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ...
def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]:
return self.records[index]
def as_dict(
self,
*,
key: Literal["id", "sequence"] = "id",
duplicates: Literal["error", "first", "last"] = "error",
materialize: bool = True,
) -> dict[str, TensorValue]:
"""Convert records to a mapping under an explicit duplicate policy."""
if key not in {"id", "sequence"}:
raise ValueError("key must be 'id' or 'sequence'.")
if duplicates not in {"error", "first", "last"}:
raise ValueError("duplicates must be 'error', 'first', or 'last'.")
if not isinstance(materialize, bool):
raise TypeError("materialize must be a boolean.")
output: dict[str, TensorValue] = {}
for record in self.records:
record_key = getattr(record, key)
if record_key in output:
if duplicates == "error":
raise ValueError(
f"Duplicate {key} {record_key!r}; choose duplicates='first' "
"or duplicates='last' explicitly."
)
if duplicates == "first":
continue
output[record_key] = record.load_tensor() if materialize else record.tensor
return output
def materialize(self, *, verify: bool = True) -> EmbeddingResult:
"""Return an equivalent result with every X loaded into CPU memory."""
if not isinstance(verify, bool):
raise TypeError("verify must be a boolean.")
return EmbeddingResult(
[
EmbeddingRecord(
id=record.id,
sequence=record.sequence,
tensor=record.load_tensor(verify=verify),
)
for record in self.records
],
self.metadata,
)
@dataclass(frozen=True, slots=True)
class EmbeddingBatch:
"""Internal model-to-runner contract.
``X`` has shape ``(b, l, d)`` and ``residue_mask`` has shape ``(b, l)``.
``attentions`` may contain layer/head attention matrices for ``parti``.
"""
X: Tensor
residue_mask: Tensor
attentions: Tensor | tuple[Tensor, ...] | None = None
__all__ = [
"EmbeddingBatch",
"EmbeddingInput",
"EmbeddingRecord",
"EmbeddingResult",
"LazyTensorReference",
"TensorValue",
]
|