File size: 5,024 Bytes
46bfb9c | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """Pulse Unified DDM — the model (question-conditioned cross-encoder).
CORRECTION (2026-09-24): the model now READS each typed question. It does not use
fixed positional heads. For a CHOICE/BOOLEAN question it scores every declared
choice by encoding `state [SEP] question [SEP] choice`; for a SCORE question it
encodes `state [SEP] question`. So the same model answers arbitrary typed
questions with arbitrary choice sets — this is the general Pulse contract, not a
triage classifier.
"One unified model invocation" = one logical call. Internally it encodes several
short sequences (one per (question,choice) pair, one per (state,plan) pair) and
batches them; it is non-autoregressive and generates no text.
Compute: default encoder is SMALL, to validate the full pipeline first on the
64 GB MLX Mac via torch's MPS backend. The production encoder is chosen by
measured quality/calibration/latency/memory/license in DownshiftEval — NOT copied
from any other project.
Requires torch + transformers (use a supported Python environment; the system
Python 3.14 environment does not currently provide the required torch wheel).
Status: EXPERIMENTAL BASELINE. Instantiating this class creates pretrained
encoder weights plus newly initialized Pulse heads; a trained Banking77
checkpoint exists at ``checkpoints/pulse_banking77.pt`` and must be loaded
explicitly. That checkpoint is not an SCPO or general OutcomeLM model.
"""
from __future__ import annotations
from dataclasses import dataclass
try:
import torch
from torch import nn
from transformers import AutoModel, AutoTokenizer
_HAVE_TORCH = True
except Exception: # pragma: no cover
_HAVE_TORCH = False
nn = object # type: ignore
@dataclass
class PulseConfig:
# SMALL default for first end-to-end pipeline validation (~22M params, 384 hid).
# Swap up (and re-measure) via DownshiftEval; do not hardcode a "winner".
encoder_name: str = "sentence-transformers/all-MiniLM-L6-v2"
hidden_size: int = 384
max_length: int = 256
dropout: float = 0.1
model_version: str = "pulse-1-unified-0.0.0-experimental"
if _HAVE_TORCH:
class PulseUnified(nn.Module):
"""Shared cross-encoder + scalar scorers.
Public scoring API (used by trainer, evaluator, planner):
score_pairs(seq_ids, seq_mask) -> (N,) raw scalar scores
The collator builds the sequences; grouping (below) turns per-choice
scores into per-question distributions.
"""
def __init__(self, cfg: PulseConfig) -> None:
super().__init__()
self.cfg = cfg
self.encoder = AutoModel.from_pretrained(cfg.encoder_name)
h = cfg.hidden_size
self.drop = nn.Dropout(cfg.dropout)
# one scalar scorer reused for choice-scoring AND score-questions;
# a separate abstain scorer; a separate policy scorer.
self.value_head = nn.Linear(h, 1) # choice score / SCORE logit
self.abstain_head = nn.Linear(h, 1) # per (state,question) abstain logit
self.policy_head = nn.Sequential(
nn.Linear(h, h // 2), nn.GELU(), nn.Linear(h // 2, 1)
)
def _pool(self, input_ids, attention_mask):
out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
mask = attention_mask.unsqueeze(-1).float()
pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-6)
return self.drop(pooled)
def score_pairs(self, seq_ids, seq_mask):
"""Raw scalar per input sequence (state+question+choice, or state+question)."""
return self.value_head(self._pool(seq_ids, seq_mask)).squeeze(-1)
def abstain_logits(self, seq_ids, seq_mask):
return self.abstain_head(self._pool(seq_ids, seq_mask)).squeeze(-1)
def policy_scores(self, seq_ids, seq_mask):
"""Raw scalar per (state, plan) sequence; planner ranks by this."""
return self.policy_head(self._pool(seq_ids, seq_mask)).squeeze(-1)
@staticmethod
def grouped_softmax(scores, groups): # type: ignore
"""Turn flat per-choice `scores` into per-question distributions.
`groups` is a list of (start, end) slices, one per CHOICE/BOOLEAN
question, with variable choice counts. Returns list of prob tensors."""
return [torch.softmax(scores[s:e], dim=0) for (s, e) in groups]
def load_tokenizer(cfg: PulseConfig):
return AutoTokenizer.from_pretrained(cfg.encoder_name)
else: # pragma: no cover
class PulseUnified: # type: ignore
def __init__(self, *_: object, **__: object) -> None:
raise RuntimeError(
"torch/transformers not installed. Create a Python 3.12 venv and "
"`pip install -r pulse/requirements.txt` (MPS works on the MLX Mac)."
)
def load_tokenizer(_cfg): # type: ignore
raise RuntimeError("torch/transformers not installed.")
|