| from __future__ import annotations |
|
|
| from collections.abc import Mapping, Sequence |
|
|
| import torch |
| from torch import Tensor, nn |
|
|
| from .encoding import TransformerTextEncoder, serialize_state |
| from .experts import CandidateScorer, NoulHead, RawExpert |
| from .types import Question, TaskOutput |
|
|
|
|
| class ConceptMemoryModel(nn.Module): |
| """Raw-only typed judgment model; state representation is shared per request.""" |
|
|
| def __init__(self, encoder: TransformerTextEncoder, decision_size: int = 256): |
| super().__init__() |
| self.encoder = encoder |
| self.raw_expert = RawExpert(encoder.hidden_size, decision_size) |
| self.criterion_projection = nn.Linear(encoder.hidden_size, decision_size) |
| self.candidate_scorer = CandidateScorer(decision_size) |
| self.noul_head = NoulHead(decision_size) |
| self.decision_size = decision_size |
|
|
| @classmethod |
| def from_pretrained( |
| cls, name: str = "distilbert/distilbert-base-uncased", decision_size: int = 256, |
| max_length: int = 512, |
| ) -> ConceptMemoryModel: |
| return cls(TransformerTextEncoder.from_pretrained(name, max_length), decision_size) |
|
|
| def encode_state(self, state: Mapping[str, object] | str) -> Tensor: |
| return self.encoder([serialize_state(state)])[0] |
|
|
| def forward_question(self, state_embedding: Tensor, question: Question) -> TaskOutput: |
| question_embedding = self.encoder([question.instructions])[0] |
| decision = self.raw_expert(state_embedding, question_embedding) |
| if question.type in ("choice", "score"): |
| criteria = question.normalized_criteria() |
| descriptions = [criterion.description for criterion in criteria] |
| criteria_embeddings = self.criterion_projection(self.encoder(descriptions)) |
| logits = self.candidate_scorer(decision, criteria_embeddings) |
| probabilities = logits.softmax(dim=-1) |
| prediction = logits.argmax(dim=-1) |
| expected_score = None |
| if question.type == "score": |
| positions = torch.arange(logits.numel(), device=logits.device, dtype=logits.dtype) |
| expected_score = (probabilities * positions).sum() |
| return TaskOutput(question.id, question.type, logits, probabilities, prediction, expected_score) |
| if question.type == "noul": |
| logit = self.noul_head(decision) |
| probability = logit.sigmoid() |
| return TaskOutput(question.id, question.type, logit, probability, probability >= 0.5) |
| raise ValueError(f"Unsupported question type: {question.type!r}") |
|
|
| def forward( |
| self, state: Mapping[str, object] | str, questions: Sequence[Question] |
| ) -> list[TaskOutput]: |
| state_embedding = self.encode_state(state) |
| return [self.forward_question(state_embedding, question) for question in questions] |
|
|
| @torch.inference_mode() |
| def predict( |
| self, state: Mapping[str, object] | str, questions: Sequence[Question] |
| ) -> list[TaskOutput]: |
| self.eval() |
| return self(state, questions) |
|
|
