"""Prepare native signals, encode once and score separate candidate answers.""" from dataclasses import dataclass from pathlib import Path import torch from torch import nn from .encoders import SignalEncoder from .head import MLPFusionHead from .inputs import validate_text from .preprocessing import SignalPreprocessor, PreparationConfig, PreparedState from .provenance import RUNTIME_CONTRACT, ARCHITECTURE, verify_assets from .text import TextEncoder from .vendor import signal @dataclass class DecisionBatch: states: list[PreparedState] question: torch.Tensor criteria: torch.Tensor criteria_mask: torch.Tensor descriptions: list[list[str]] def to(self, device): return DecisionBatch(self.states, self.question.to(device), self.criteria.to(device), self.criteria_mask.to(device), self.descriptions) class PayaModel(nn.Module): """Fuse frozen signal embeddings with separate BioLORD text embeddings.""" def __init__(self, signal_encoder, text_encoder, preparation=None, head=None, asset_identity=None): super().__init__() self.encoder = signal_encoder.requires_grad_(False).eval() settings = dict(head or {}) kind = settings.pop("kind", "interaction") if kind != "interaction": raise ValueError("Unknown decision-head kind") self.head = MLPFusionHead(**settings) self.text_encoder = text_encoder self.preprocessor = SignalPreprocessor(preparation) self.asset_identity = asset_identity or {"injected_test_encoders": True} self.register_buffer("temperature", torch.tensor(1.0)) self.calibration = None self.train(False) @classmethod def from_assets(cls, assets, preparation=None, head=None, device="cpu"): sources = verify_assets(assets) settings = PreparationConfig(**(preparation or {})) text = TextEncoder(device=device, model_path=Path(assets["text"])) model = cls(SignalEncoder.load(assets[signal.ASSET_KEY]), text, settings, head, sources) return model.to(device) @property def device(self): return self.temperature.device def identity(self): return {"contract": RUNTIME_CONTRACT, "architecture": ARCHITECTURE, "assets": self.asset_identity, "preparation": self.preprocessor.identity(), "head": self.head.config} def train(self, mode=True): super().train(mode) self.encoder.eval() if mode and hasattr(self, "temperature"): self.temperature.fill_(1.0) self.calibration = None if getattr(self, "text_encoder", None) is not None and self.text_encoder.model is not None: self.text_encoder.model.eval().requires_grad_(False) return self def trainable_parameters(self): return self.head.parameters() def make_batch(self, states, questions, criteria): """Prepare native signals and encode each question and choice separately.""" if not states or len(states) != len(questions) or len(states) != len(criteria): raise ValueError("Supply matching nonempty state/question/criteria batches") for q, c in zip(questions, criteria): validate_text(q, c) prepared = [s if isinstance(s, PreparedState) else self.preprocessor.prepare(s) for s in states] if any(s.audit.get("preparation") != self.preprocessor.identity() for s in prepared): raise ValueError("Prepared signal does not match current preprocessing") strings = [s for q, c in zip(questions, criteria) for s in [q, *c]] vectors = self.text_encoder.encode(strings).detach().cpu().clone() if vectors.shape != (len(strings), 768) or not torch.isfinite(vectors).all(): raise ValueError("BioLORD must return finite, separate 768-dimensional vectors") width, offset = max(map(len, criteria)), 0 qs, cs = [], vectors.new_zeros(len(states), width, 768) mask = torch.zeros(len(states), width, dtype=torch.bool, device=self.device) for i, choices in enumerate(criteria): qs.append(vectors[offset]) cs[i, :len(choices)] = vectors[offset+1:offset+1+len(choices)] mask[i, :len(choices)] = True offset += 1 + len(choices) return DecisionBatch(prepared, torch.stack(qs), cs, mask, [list(c) for c in criteria]).to(self.device) def encode(self, prepared): encoder_input = self.preprocessor.format(prepared) with torch.no_grad(): encoded = self.encoder(encoder_input.to(self.device)) if encoded.shape != (1, 512) or not torch.isfinite(encoded).all(): raise ValueError("Encoder must return one finite signal vector [1,512]") return encoded.float().detach() def forward(self, batch, cache=None): if cache is not None and hasattr(cache, "get_many_or_encode"): signal = cache.get_many_or_encode(self, batch.states)[:, None, :] else: encoded = [cache.get_or_encode(self, state) if cache else self.encode(state) for state in batch.states] signal = torch.stack([e.to(self.device) for e in encoded]) mask = torch.ones(len(batch.states), 1, dtype=torch.bool, device=self.device) logits = self.head(signal, mask, batch.question, batch.criteria, batch.criteria_mask) return {"logits": logits, "criteria_mask": batch.criteria_mask} @torch.no_grad() def encode_batch(self, prepared): samples = [self.preprocessor.format(value).samples for value in prepared] if len({tuple(value.shape) for value in samples}) != 1: raise ValueError("Batch signal encoding requires matching native shapes") if hasattr(self.encoder, "forward_batch"): encoded = self.encoder.forward_batch(torch.stack(samples).to(self.device)) else: encoded = torch.cat([self.encode(value) for value in prepared]) if encoded.shape != (len(prepared), 512) or not torch.isfinite(encoded).all(): raise ValueError("Encoder must return finite signal vectors [B,512]") return encoded.float().detach() @torch.no_grad() def predict(self, state, question, criteria): """Return the chosen answer, its confidence and probabilities keyed by choice.""" self.eval() if not torch.isfinite(self.temperature) or self.temperature <= 0: raise ValueError("Prediction temperature must be finite and positive") batch = self.make_batch([state], [question], [criteria]) logits = self(batch)["logits"][0] probabilities = (logits / self.temperature).softmax(-1).cpu().tolist() chosen = min(range(len(criteria)), key=lambda i: (-probabilities[i], criteria[i])) return {"choice": criteria[chosen], "confidence": probabilities[chosen], "probabilities": dict(zip(criteria, probabilities))}