File size: 8,936 Bytes
8752091 2ca20a1 8752091 2ca20a1 8752091 2ca20a1 8752091 843a654 8752091 2ca20a1 8752091 2ca20a1 8752091 | 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 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 | """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
from types import MappingProxyType
from ..features.layouts import TopKRow
@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() # (...), equal to 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 # (...), equal to 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
@dataclass(frozen=True, slots=True)
class TapRecord:
"""One sequence's outputs from a tap plan, keyed by tap name."""
id: str
sequence: str
tensors: Mapping[str, Tensor | TopKRow]
retained_positions: tuple[int, ...] | None = None
def __post_init__(self) -> None:
if not isinstance(self.id, str) or not self.id:
raise ValueError("TapRecord.id must be a non-empty string.")
if not isinstance(self.sequence, str) or not self.sequence:
raise ValueError("TapRecord.sequence must be a non-empty string.")
if not isinstance(self.tensors, Mapping) or not self.tensors:
raise TypeError("TapRecord.tensors must be a non-empty mapping of tap name to Tensor.")
if not all(
isinstance(name, str) and isinstance(value, (Tensor, TopKRow))
for name, value in self.tensors.items()
):
raise TypeError("TapRecord.tensors must map tap names to Tensor or TopKRow values.")
object.__setattr__(self, "tensors", MappingProxyType(dict(self.tensors)))
if self.retained_positions is not None:
positions = self.retained_positions
if (type(positions) is not tuple or not positions
or any(type(p) is not int or not 0 <= p < len(self.sequence) for p in positions)
or tuple(sorted(set(positions))) != positions):
raise ValueError(
"TapRecord retained positions must be ordered original-sequence indices."
)
@dataclass(frozen=True, slots=True)
class TapRunReceipt:
"""Completed sink delivery, retaining run metadata but no output tensors."""
record_count: int
metadata: Mapping[str, Any]
class TapResult:
"""Ordered tap records and the metadata needed to reproduce them."""
def __init__(
self,
records: Sequence[TapRecord],
metadata: Mapping[str, Any] | None = None,
) -> None:
self.records: tuple[TapRecord, ...] = tuple(records)
self.metadata = dict(metadata or {})
def __len__(self) -> int:
return len(self.records)
def __iter__(self) -> Iterator[TapRecord]:
return iter(self.records)
def __getitem__(self, index: int) -> TapRecord:
return self.records[index]
__all__ = [
"EmbeddingBatch",
"EmbeddingInput",
"EmbeddingRecord",
"EmbeddingResult",
"LazyTensorReference",
"TapRecord",
"TapResult",
"TapRunReceipt",
"TensorValue",
]
|