Upload Concept Memory base model

README.md

@@ -22,26 +22,35 @@ not a language model pretrained from scratch. Its text encoder is
 
 - `model.pt`: PyTorch state dictionary and training metadata.
 - `config.json`: encoder, task, and training run metadata.
+- `concept_memory/`: the minimal model runtime modules needed for inference.
 
 ## Run a prediction
 
-Use Python 3.11 or newer. Install the model code and the Hub client:
+Use Python 3.11 or newer. Install the runtime dependencies:
 
 ```bash
-pip install "git+https://github.com/ceylonai/rgcm.git" huggingface_hub
+pip install torch transformers huggingface_hub
 ```
 
-Download the checkpoint, construct a question with the criteria you want the
-model to judge, and inspect its prediction and probabilities:
+Download this model repository. It contains only the small `concept_memory`
+runtime package needed to load the checkpoint, rather than the full training
+project. Then construct a question with the criteria you want the model to
+judge and inspect its prediction and probabilities:
 
 ```python
 import torch
-from huggingface_hub import hf_hub_download
+import sys
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+repo_id = "Syigen/rgcm"
+repo_path = snapshot_download(repo_id)
+sys.path.insert(0, repo_path)
+
 from concept_memory import ConceptMemoryModel
 from concept_memory.types import Question
 
-repo_id = "Syigen/rgcm"
-checkpoint = hf_hub_download(repo_id, "model.pt")
+checkpoint = Path(repo_path) / "model.pt"
 payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
 model = ConceptMemoryModel.from_pretrained(
     payload["encoder"],

concept_memory/__init__.py

@@ -0,0 +1,4 @@
+from .model import ConceptMemoryModel
+from .types import Criterion, Question, QuestionType, TaskOutput
+
+__all__ = ["ConceptMemoryModel", "Criterion", "Question", "QuestionType", "TaskOutput"]

concept_memory/encoding.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+
+from torch import Tensor, nn
+
+
+def serialize_state(state: Mapping[str, object] | str) -> str:
+    """Serialize structured state in stable key order for the text encoder."""
+    if isinstance(state, str):
+        return state
+    parts = []
+    for key in sorted(state):
+        value = state[key]
+        if value is None:
+            continue
+        if isinstance(value, (Mapping, list, tuple)):
+            value = json.dumps(value, ensure_ascii=False, sort_keys=True)
+        parts.append(f"[{key}] {value}")
+    return "\n".join(parts)
+
+
+class TransformerTextEncoder(nn.Module):
+    """Shared transformer encoder with attention-mask-aware mean pooling."""
+
+    def __init__(
+        self, model: nn.Module, tokenizer: object, hidden_size: int, max_length: int = 512
+    ):
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+        self.hidden_size = hidden_size
+        self.max_length = max_length
+
+    @classmethod
+    def from_pretrained(cls, name: str, max_length: int = 512) -> "TransformerTextEncoder":
+        from transformers import AutoModel, AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(name)
+        model = AutoModel.from_pretrained(name)
+        return cls(model, tokenizer, model.config.hidden_size, max_length)
+
+    def forward(self, texts: list[str]) -> Tensor:
+        batch = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+        device = next(self.model.parameters()).device
+        batch = {key: value.to(device) for key, value in batch.items()}
+        outputs = self.model(**batch)
+        token_embeddings = outputs.last_hidden_state
+        mask = batch["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
+        pooled = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1)
+        return nn.functional.layer_norm(pooled, (self.hidden_size,))

concept_memory/experts.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+
+class RawExpert(nn.Module):
+    def __init__(self, hidden_size: int, decision_size: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_size * 2, decision_size),
+            nn.GELU(),
+            nn.LayerNorm(decision_size),
+            nn.Linear(decision_size, decision_size),
+        )
+
+    def forward(self, state: Tensor, question: Tensor) -> Tensor:
+        return self.net(torch.cat((state, question), dim=-1))
+
+
+class CandidateScorer(nn.Module):
+    """Shared pointwise scorer; candidate permutation only permutes its outputs."""
+
+    def __init__(self, decision_size: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(decision_size * 4, decision_size),
+            nn.GELU(),
+            nn.Linear(decision_size, 1),
+        )
+
+    def forward(self, decision: Tensor, candidates: Tensor) -> Tensor:
+        if decision.ndim == 1:
+            decision = decision.unsqueeze(0)
+        if candidates.ndim == 2:
+            candidates = candidates.unsqueeze(0).expand(decision.size(0), -1, -1)
+        decision = decision.unsqueeze(1).expand(-1, candidates.size(1), -1)
+        features = torch.cat(
+            (decision, candidates, decision * candidates, torch.abs(decision - candidates)), dim=-1
+        )
+        return self.net(features).squeeze(-1).squeeze(0)
+
+
+class NoulHead(nn.Module):
+    def __init__(self, decision_size: int):
+        super().__init__()
+        self.linear = nn.Linear(decision_size, 1)
+
+    def forward(self, decision: Tensor) -> Tensor:
+        return self.linear(decision).squeeze(-1)

concept_memory/model.py

@@ -0,0 +1,67 @@
+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)

concept_memory/types.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal, TypeAlias
+
+from torch import Tensor
+
+QuestionType: TypeAlias = Literal["choice", "score", "noul"]
+
+
+@dataclass(frozen=True)
+class Criterion:
+    id: str
+    description: str
+
+
+@dataclass(frozen=True)
+class Question:
+    id: str
+    type: QuestionType
+    instructions: str
+    criteria: list[Criterion | str | dict[str, str]] = field(default_factory=list)
+
+    def normalized_criteria(self) -> list[Criterion]:
+        result: list[Criterion] = []
+        for index, item in enumerate(self.criteria):
+            if isinstance(item, Criterion):
+                result.append(item)
+            elif isinstance(item, str):
+                result.append(Criterion(str(index), item))
+            else:
+                result.append(Criterion(item.get("id", str(index)), item["description"]))
+        if self.type in ("choice", "score") and not result:
+            raise ValueError(f"Question {self.id!r} of type {self.type!r} needs criteria")
+        return result
+
+
+@dataclass
+class TaskOutput:
+    question_id: str
+    type: QuestionType
+    logits: object
+    probabilities: Tensor
+    prediction: object
+    expected_score: object | None = None