File size: 3,072 Bytes
3155e2a | 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 | 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)
|