from __future__ import annotations from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Final, final import numpy as np import numpy.typing as npt from redstack.domain.candidate.archetype import ArchetypeAssignment from redstack.domain.candidate.semantic import SemanticProfile, VectorRef from redstack.domain.ids import AnchorId, ArchetypeId, CandidateId, Similarity, UnitScore from redstack.features.view import clamp_unit from redstack.ports._types import FloatMatrix, FloatVector from redstack.ports.embedding import EmbeddingModelPort from redstack.ports.semantic_index import SemanticVectorStorePort _SIM_LO: Final[float] = -1.0 _SIM_HI: Final[float] = 1.0 def _clamp_similarity(value: float) -> float: return _SIM_LO if value < _SIM_LO else _SIM_HI if value > _SIM_HI else value @final @dataclass(frozen=True) class AnchorSet: """JD anchor vectors split by polarity (rows L2-normalized, ids id-ascending). Authored in ``configs/anchors/jd_anchors.yaml`` → ``anchor_vectors.npy`` (O6), loaded and injected by the pipeline at R0. Ids are sorted ascending so that a plain ``argmax`` resolves ties by ``AnchorId`` ascending (determinism). """ positive_ids: tuple[AnchorId, ...] positive_matrix: FloatMatrix # (P, dim) negative_ids: tuple[AnchorId, ...] negative_matrix: FloatMatrix # (Nn, dim) def __post_init__(self) -> None: if list(self.positive_ids) != sorted(self.positive_ids): raise ValueError("AnchorSet.positive_ids must be ascending") if list(self.negative_ids) != sorted(self.negative_ids): raise ValueError("AnchorSet.negative_ids must be ascending") @final @dataclass(frozen=True) class ArchetypeSpace: """O7 KMeans centroids + target-archetype set (ids id-ascending).""" centroids: FloatMatrix # (K, dim) archetype_ids: tuple[ArchetypeId, ...] target_archetypes: frozenset[ArchetypeId] labels: Mapping[ArchetypeId, str] def __post_init__(self) -> None: if list(self.archetype_ids) != sorted(self.archetype_ids): raise ValueError("ArchetypeSpace.archetype_ids must be ascending") @final @dataclass(frozen=True) class SemanticArrays: """Columnar bulk result aligned to the input matrix row order (for CQV fold).""" positive_fit: npt.NDArray[np.float32] negative_fit: npt.NDArray[np.float32] net_fit: npt.NDArray[np.float32] best_positive_index: npt.NDArray[np.int64] nearest_centroid_index: npt.NDArray[np.int64] centroid_distance: npt.NDArray[np.float32] membership_confidence: npt.NDArray[np.float32] @final class SemanticEngine: """Stateless retrieval engine; all state arrives via ports + injected anchors.""" def __init__( self, *, store: SemanticVectorStorePort, embedder: EmbeddingModelPort, anchors: AnchorSet, archetypes: ArchetypeSpace, vector_store_key: str = "embeddings/candidate_vectors", ) -> None: self._store = store self._embedder = embedder self._anchors = anchors self._archetypes = archetypes self._store_key = vector_store_key # ------------------------------------------------------------ vector source def resolve_vectors( self, ids: Sequence[CandidateId], documents: Mapping[CandidateId, str] ) -> tuple[FloatMatrix, tuple[CandidateId, ...]]: """Lookup-first gather; encode store misses via the fallback embedder. Returns a matrix aligned to ``ids`` order and the tuple of ids that fell back to the encoder (for metrics). A miss with no composed document is a pipeline error (the candidate cannot be scored) and is left to the caller. """ bulk = self._store.get_many(ids) found_set = set(bulk.found) rows: list[FloatVector] = [] found_iter = iter(bulk.vectors) misses: list[CandidateId] = [] miss_docs: list[str] = [] miss_positions: list[int] = [] for position, cid in enumerate(ids): if cid in found_set: rows.append(next(found_iter)) else: misses.append(cid) miss_docs.append(documents[cid]) miss_positions.append(position) rows.append(np.zeros(self._store.dim, dtype=np.float32)) if miss_docs: encoded = self._embedder.encode(miss_docs) for slot, position in enumerate(miss_positions): rows[position] = encoded[slot] matrix = np.vstack(rows).astype(np.float32, copy=False) return matrix, tuple(misses) # ----------------------------------------------------------- vectorized core def compute(self, matrix: FloatMatrix) -> SemanticArrays: """Vectorized similarities + nearest-centroid over an ``(N, dim)`` matrix.""" pos = matrix @ self._anchors.positive_matrix.T # (N, P) neg = matrix @ self._anchors.negative_matrix.T # (N, Nn) positive_fit = np.clip(pos.max(axis=1), _SIM_LO, _SIM_HI).astype(np.float32) best_positive_index = pos.argmax(axis=1).astype(np.int64) # first max == min id negative_fit = np.clip(neg.max(axis=1), _SIM_LO, _SIM_HI).astype(np.float32) net_fit = np.clip( (positive_fit - negative_fit + 1.0) * 0.5, 0.0, 1.0 ).astype(np.float32) nearest_idx, distance, confidence = self._archetype_columns(matrix) return SemanticArrays( positive_fit=positive_fit, negative_fit=negative_fit, net_fit=net_fit, best_positive_index=best_positive_index, nearest_centroid_index=nearest_idx, centroid_distance=distance, membership_confidence=confidence, ) def _archetype_columns( self, matrix: FloatMatrix ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32], npt.NDArray[np.float32]]: centroids = self._archetypes.centroids # (K, dim) # Squared Euclidean distances via the (a-b)^2 expansion (vectorized). sq = ( (matrix * matrix).sum(axis=1, keepdims=True) - 2.0 * (matrix @ centroids.T) + (centroids * centroids).sum(axis=1)[np.newaxis, :] ) sq = np.maximum(sq, 0.0) dists = np.sqrt(sq) # (N, K) nearest = dists.argmin(axis=1).astype(np.int64) # first min == min ArchetypeId k = dists.shape[1] d1 = dists[np.arange(dists.shape[0]), nearest].astype(np.float32) if k == 1: confidence = np.ones(dists.shape[0], dtype=np.float32) else: partitioned = np.partition(dists, 1, axis=1) d2 = partitioned[:, 1].astype(np.float32) denom = d1 + d2 confidence = np.where(denom > 0.0, d2 / denom, 1.0).astype(np.float32) confidence = np.clip(confidence, 0.0, 1.0) return nearest, d1, confidence # ---------------------------------------------------- per-candidate objects def profile_for( self, vector: FloatVector, *, row_index: int ) -> tuple[SemanticProfile, ArchetypeAssignment]: """Materialize the two domain slices for one candidate (survivors/top-K).""" row = vector.reshape(1, -1).astype(np.float32, copy=False) arrays = self.compute(row) anchor_sims = self._anchor_similarity_map(vector) best_idx = int(arrays.best_positive_index[0]) best_anchor: AnchorId | None = ( self._anchors.positive_ids[best_idx] if self._anchors.positive_ids else None ) semantic = SemanticProfile( anchor_similarities=anchor_sims, positive_fit=Similarity(float(arrays.positive_fit[0])), negative_fit=Similarity(float(arrays.negative_fit[0])), net_semantic_fit=UnitScore(clamp_unit(float(arrays.net_fit[0]))), best_positive_anchor=best_anchor, vector_ref=VectorRef( store_key=self._store_key, row_index=row_index, dim=int(vector.shape[0]), ), ) a_idx = int(arrays.nearest_centroid_index[0]) archetype_id = self._archetypes.archetype_ids[a_idx] secondary = self._secondary_archetype(vector, primary_index=a_idx) archetype = ArchetypeAssignment( archetype_id=archetype_id, distance=float(arrays.centroid_distance[0]), membership_confidence=UnitScore( clamp_unit(float(arrays.membership_confidence[0])) ), secondary_archetype=secondary, label=self._archetypes.labels.get(archetype_id), is_target_archetype=archetype_id in self._archetypes.target_archetypes, ) return semantic, archetype def _anchor_similarity_map( self, vector: FloatVector ) -> Mapping[AnchorId, Similarity]: sims: dict[AnchorId, Similarity] = {} pos = vector @ self._anchors.positive_matrix.T for anchor_id, value in zip(self._anchors.positive_ids, pos, strict=True): sims[anchor_id] = Similarity(_clamp_similarity(float(value))) neg = vector @ self._anchors.negative_matrix.T for anchor_id, value in zip(self._anchors.negative_ids, neg, strict=True): sims[anchor_id] = Similarity(_clamp_similarity(float(value))) return sims def batch_profiles_for( self, matrix: FloatMatrix, ) -> tuple[tuple[SemanticProfile, ...], tuple[ArchetypeAssignment, ...]]: """Batch-compute semantic profiles and archetype assignments. Replaces N individual :meth:`profile_for` calls with a single batched ``compute`` pass (one GEMM + archetype distance sweep instead of N individual GEMM + GEMV pairs). Per-candidate anchor sims are still computed via the same :meth:`_anchor_similarity_map` GEMV path as :meth:`profile_for` to maintain bit-identical anchor cosine values. Secondary archetypes are resolved in one ``argsort`` pass via :meth:`_batch_secondary_archetypes`. """ n = matrix.shape[0] if n == 0: return (), () dim = int(matrix.shape[1]) arrays = self.compute(matrix) secondary_archetypes = self._batch_secondary_archetypes( matrix, arrays.nearest_centroid_index ) profiles: list[SemanticProfile] = [] assignments: list[ArchetypeAssignment] = [] for i in range(n): anchor_sims = self._anchor_similarity_map(matrix[i]) best_idx = int(arrays.best_positive_index[i]) best_anchor: AnchorId | None = ( self._anchors.positive_ids[best_idx] if self._anchors.positive_ids else None ) profiles.append( SemanticProfile( anchor_similarities=anchor_sims, positive_fit=Similarity(float(arrays.positive_fit[i])), negative_fit=Similarity(float(arrays.negative_fit[i])), net_semantic_fit=UnitScore(clamp_unit(float(arrays.net_fit[i]))), best_positive_anchor=best_anchor, vector_ref=VectorRef( store_key=self._store_key, row_index=i, dim=dim, ), ) ) a_idx = int(arrays.nearest_centroid_index[i]) archetype_id = self._archetypes.archetype_ids[a_idx] assignments.append( ArchetypeAssignment( archetype_id=archetype_id, distance=float(arrays.centroid_distance[i]), membership_confidence=UnitScore( clamp_unit(float(arrays.membership_confidence[i])) ), secondary_archetype=secondary_archetypes[i], label=self._archetypes.labels.get(archetype_id), is_target_archetype=( archetype_id in self._archetypes.target_archetypes ), ) ) return tuple(profiles), tuple(assignments) def _batch_secondary_archetypes( self, matrix: FloatMatrix, primary_indices: npt.NDArray[np.int64], ) -> tuple[ArchetypeId | None, ...]: """One argsort pass over all N candidates to find secondary archetypes. Uses the same squared-distance expansion as :meth:`_archetype_columns` for consistency. ``secondary_archetype`` is informational only (not used in scoring and not written to ``submission.csv``), so minor float32 differences from the per-candidate subtraction path in :meth:`_secondary_archetype` are acceptable. """ centroids = self._archetypes.centroids n = matrix.shape[0] if centroids.shape[0] < 2: none_list: list[ArchetypeId | None] = [None] * n return tuple(none_list) sq = ( (matrix * matrix).sum(axis=1, keepdims=True) - 2.0 * (matrix @ centroids.T) + (centroids * centroids).sum(axis=1)[np.newaxis, :] ) dists = np.sqrt(np.maximum(sq, 0.0)) sorted_cols: npt.NDArray[np.intp] = np.argsort(dists, axis=1, kind="stable") secondaries: list[ArchetypeId | None] = [] for i in range(n): primary = int(primary_indices[i]) found: ArchetypeId | None = None for idx_np in sorted_cols[i]: idx = int(idx_np) if idx != primary: found = self._archetypes.archetype_ids[idx] break secondaries.append(found) return tuple(secondaries) def _secondary_archetype( self, vector: FloatVector, *, primary_index: int ) -> ArchetypeId | None: centroids = self._archetypes.centroids if centroids.shape[0] < 2: return None diff = centroids - vector[np.newaxis, :] dists = np.sqrt(np.maximum((diff * diff).sum(axis=1), 0.0)) order = np.argsort(dists, kind="stable") for idx in order: if int(idx) != primary_index: return self._archetypes.archetype_ids[int(idx)] return None __all__: tuple[str, ...] = ( "AnchorSet", "ArchetypeSpace", "SemanticArrays", "SemanticEngine", )