gemmadecision 0.1.0__py3-none-any.whl

gemmadecision/__init__.py

@@ -0,0 +1,42 @@
+"""Typed decisions without importing torch, vLLM, or loading a model."""
+from .client import AsyncDecisionClient, DecisionClient
+from .types import Choice, Noul, Score, ChoiceAnswer, NoulAnswer, ScoreAnswer, SystemOneResponse, RankingResponse
+from .constants import MODEL_ID, MODEL_REVISION
+
+__version__ = "0.1.0"
+__all__ = ["decide", "rank", "DecisionClient", "AsyncDecisionClient", "DecisionEngine", "GemmaDecisionModel", "Choice", "Noul", "Score",
+           "ChoiceAnswer", "NoulAnswer", "ScoreAnswer", "SystemOneResponse", "RankingResponse"]
+
+
+def __getattr__(name):
+    if name == "DecisionEngine":
+        from .engine import DecisionEngine
+        return DecisionEngine
+    if name == "GemmaDecisionModel":
+        from .integrations.pydantic_ai import GemmaDecisionModel
+        return GemmaDecisionModel
+    raise AttributeError(name)
+
+
+def decide(state, choices, *, question="Choose the best option.") -> str:
+    """Choose locally in one call; first use downloads/loads the pinned model."""
+    from .simple import get_engine
+    from .types import RankRequest
+    if isinstance(choices, (list, tuple)):
+        if any(not isinstance(text, str) for text in choices):
+            raise ValueError("Choices must be strings or a mapping of labels to descriptions")
+        if len(set(choices)) != len(choices):
+            raise ValueError("Choices must be distinct")
+        choices = {text: text for text in choices}
+    choices = RankRequest(state=state, candidates=choices, question=question).candidates
+    if len(set(choices.values())) != len(choices) or any(not text.strip() for text in choices.values()):
+        raise ValueError("Choice descriptions must be distinct and nonempty")
+    return get_engine().decide(state, choices, question=question).choice
+
+
+def rank(state, choices, *, question=""):
+    """Detailed local ranking with lazy model loading, no server required."""
+    from .simple import get_engine
+    from .types import RankRequest
+    choices = RankRequest(state=state, candidates=choices, question=question).candidates
+    return get_engine().rank(state, choices, question=question)

gemmadecision/_clm_schema.py

@@ -0,0 +1,166 @@
+"""Shared question/answer schema for the CLM System One API.
+
+The wire format is the TypeSafe ``POST /v1/systemone`` shape: a ``state``
+(string, object or array) and a map of typed ``questions`` (``noul``,
+``choice``, ``score``), answered with probability distributions.  This module
+holds everything that both the server and the trainer must agree on:
+
+* how a (state, question) pair is turned into the *state text* the state head
+  sees and the *candidate texts* the action head sees (``build_pairs``), and
+* how per-candidate similarities become an Answer of the right type
+  (``answer_from_logits``).
+
+The CLM scores each candidate by the scaled cosine between the projected state
+and the projected candidate, exactly as in the InfoNCE training objective, and
+a softmax over a question's candidates is that question's distribution.
+"""
+from __future__ import annotations
+
+import json
+import math
+from typing import Any
+
+QUESTION_TYPES = ("noul", "choice", "score")
+NOUL_KEYS = ("false", "true")
+
+
+def to_text(x: Any, indent: int = 0) -> str:
+    """Render a state / description that may be a string, object or array as plain text.
+
+    The heads are trained on prose, not on JSON, so an object becomes ``key: value``
+    fields (top-level fields separated by a blank line, nested ones indented) and an
+    array becomes one ``- item`` line per element.  Key order is preserved.
+    """
+    if x is None:
+        return ""
+    if isinstance(x, str):
+        return x
+    if isinstance(x, bool):
+        return "true" if x else "false"
+    if isinstance(x, (int, float)):
+        return str(x)
+    pad = " " * indent
+    if isinstance(x, dict):
+        parts = []
+        for k, v in x.items():
+            if isinstance(v, (dict, list)) and v:
+                parts.append(f"{pad}{k}:\n{to_text(v, indent + 2)}")
+            else:
+                parts.append(f"{pad}{k}: {to_text(v)}")
+        return ("\n\n" if indent == 0 else "\n").join(parts)
+    if isinstance(x, (list, tuple)):
+        parts = []
+        for v in x:
+            if isinstance(v, (dict, list)) and v:
+                parts.append(f"{pad}-\n{to_text(v, indent + 2)}")
+            else:
+                parts.append(f"{pad}- {to_text(v)}")
+        return "\n".join(parts)
+    return json.dumps(x, ensure_ascii=False)
+
+
+def state_text(state: Any, instructions: Any) -> str:
+    """Context first, question last — the layout the heads were trained on."""
+    s, i = to_text(state).strip(), to_text(instructions).strip()
+    return f"{s}\n\n{i}" if s and i else (s or i)
+
+
+def _norm_question(q: dict) -> dict:
+    t = q.get("type")
+    if t not in QUESTION_TYPES:
+        raise ValueError(f"unknown question type {t!r}; expected one of {QUESTION_TYPES}")
+    return q
+
+
+def candidates(q: dict) -> tuple[list[str], list[str]]:
+    """-> (option keys in answer order, candidate text per option)."""
+    q = _norm_question(q)
+    t, crit, ins = q["type"], q.get("criteria"), to_text(q.get("instructions")).strip()
+    if t == "choice":
+        if not isinstance(crit, dict) or not crit:
+            raise ValueError("choice question needs a non-empty 'criteria' object")
+        # The action head embeds the option's own text: its description when one is
+        # given, else the key.  Nothing is prefixed, so a candidate reaches the encoder
+        # exactly as the caller wrote it (the heads are trained on plain answer text).
+        keys = list(crit)
+        texts = [to_text(crit[k]) if crit[k] not in (None, "") else k for k in keys]
+        return keys, texts
+    if t == "score":
+        if not isinstance(crit, list) or len(crit) < 2:
+            raise ValueError("score question needs 'criteria' as an ordered list of >= 2 levels")
+        keys = [str(i) for i in range(len(crit))]
+        return keys, [to_text(c) for c in crit]
+    # noul: optional {"true": ..., "false": ...} descriptions; default to the statement itself
+    crit = crit or {}
+    texts = []
+    for k in NOUL_KEYS:
+        d = crit.get(k) if isinstance(crit, dict) else None
+        if d in (None, ""):
+            d = (f"Yes. This is true: {ins}" if k == "true" else f"No. This is false: {ins}") if ins else k
+        texts.append(f"{k}: {to_text(d)}")
+    return list(NOUL_KEYS), texts
+
+
+def build_pairs(state: Any, questions: dict[str, dict]) -> dict[str, tuple[str, list[str], list[str]]]:
+    """{qid: (state_text, option_keys, candidate_texts)} for every question.
+
+    The state head sees ``context + question``: the state is the context and the
+    question's instructions are the question, appended after a blank line.  That is
+    the layout the heads are trained on, so callers should put the question in
+    ``instructions`` and not repeat it inside the state.
+    """
+    return {qid: (state_text(state, q.get("instructions")), *candidates(q)) for qid, q in questions.items()}
+
+
+def softmax(logits: list[float]) -> list[float]:
+    m = max(logits)
+    e = [math.exp(v - m) for v in logits]
+    z = sum(e)
+    return [v / z for v in e]
+
+
+def confidence(probs: list[float]) -> float:
+    """TypeSafe-style confidence: top probability minus the mean of the rest."""
+    if len(probs) < 2:
+        return 1.0
+    j = max(range(len(probs)), key=probs.__getitem__)
+    rest = [p for i, p in enumerate(probs) if i != j]
+    return max(0.0, min(1.0, probs[j] - sum(rest) / len(rest)))
+
+
+def answer_from_probs(q: dict, keys: list[str], probs: list[float]) -> dict:
+    """Assemble the Answer object for question ``q`` from its option distribution."""
+    t = q["type"]
+    probs = [float(p) for p in probs]
+    dist = {k: p for k, p in zip(keys, probs)}
+    if t == "noul":
+        return {"type": "noul", "noul": dist["true"]}
+    if t == "choice":
+        j = max(range(len(probs)), key=probs.__getitem__)
+        return {"type": "choice", "choice": keys[j], "confidence": confidence(probs), "probabilities": dist}
+    levels = q["criteria"]
+    score = sum(i * p for i, p in enumerate(probs))
+    legend = {str(i): c if isinstance(c, str) else to_text(c) for i, c in enumerate(levels)}
+    return {"type": "score", "score": score, "confidence": confidence(probs), "legend": legend,
+            "probabilities": dist}
+
+
+def answer_from_logits(q: dict, keys: list[str], logits: list[float]) -> dict:
+    return answer_from_probs(q, keys, softmax([float(v) for v in logits]))
+
+
+def label_of(answer: dict) -> str:
+    """Discrete label of an answer (for scoring): choice key, score level, or 'true'/'false'."""
+    t = answer["type"]
+    if t == "choice":
+        return answer["choice"]
+    if t == "noul":
+        return "true" if answer["noul"] >= 0.5 else "false"
+    p = answer["probabilities"]
+    return max(p, key=p.__getitem__)
+
+
+def probabilities_of(answer: dict) -> dict[str, float]:
+    if answer["type"] == "noul":
+        return {"false": 1.0 - answer["noul"], "true": answer["noul"]}
+    return dict(answer["probabilities"])

gemmadecision/backends/__init__.py

@@ -0,0 +1 @@
+"""Optional numerical backends, imported only on explicit model loading."""

gemmadecision/backends/torch.py

@@ -0,0 +1,266 @@
+"""Batched inference for the published full joint GemmaDecision encoder.
+
+Importing this module does not import torch, load weights, or use the network.
+The caller renders the training-format strings and checks the separate state
+and candidate token limits. This backend checks each complete pair's length.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from pathlib import Path
+import threading
+from typing import Any, Sequence
+
+
+MODEL_REPOSITORY = "rajan2k/GemmaDecision-270M"
+MODEL_REVISION = "785d530221c990671f29976902540101bb9c7647"
+
+# These hashes are bound to the published v0.4.0 commit, not downloaded from
+# the same untrusted directory as the model. Loading never executes model code.
+FROZEN_FILES = {
+    "model.safetensors": "d7a3e291bfdfa7cd85b33a8a99ef81a4a7d3192e46c77253f3daf14dfd7d6b95",
+    "joint_head.safetensors": "72ec4e7d1f0908ad684eeae250f0f70128aa4342b46174bf92a3c9851c5dd965",
+    "joint_config.json": "d48f2b5ad5fbea6a3d6723017a6f6a114260a331f0b917b8008de14501091533",
+    "config.json": "c1c64396b2939c76f0fa091aa07815b6e6f5ac1a60bfe3878993f0f575dd3ff3",
+    "tokenizer.json": "7d4046bf0505a327dd5a0abbb427ecd4fc82f99c2ceaa170bc61ecde12809b0c",
+    "tokenizer.model": "1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c",
+    "tokenizer_config.json": "94b03056ec5831e9021c3e3fe9db682778c4f2d081c4188f1fcb8b90541cf1cf",
+    "special_tokens_map.json": "2f7b0adf4fb469770bb1490e3e35df87b1dc578246c5e7e6fc76ecf33213a397",
+    "added_tokens.json": "50b2f405ba56a26d4913fd772089992252d7f942123cc0a034d96424221ba946",
+}
+
+
+def verify_model_files(path: str | Path) -> None:
+    """Verify the files consumed by inference against the pinned release."""
+    root = Path(path)
+    for name, expected in FROZEN_FILES.items():
+        file = root / name
+        if not file.is_file():
+            raise ValueError(f"The model directory is incomplete: missing {name}")
+        digest = hashlib.sha256()
+        with file.open("rb") as stream:
+            for block in iter(lambda: stream.read(1024 * 1024), b""):
+                digest.update(block)
+        if digest.hexdigest() != expected:
+            raise ValueError(
+                f"Model integrity check failed for {name}; expected the published "
+                f"GemmaDecision v0.4.0 revision {MODEL_REVISION}"
+            )
+
+
+def plan_batches(
+    lengths: Sequence[int], max_batch_tokens: int, max_batch_size: int
+) -> list[list[int]]:
+    """Bucket by length, bounding padded tokens and preserving all indices.
+
+    Returned indices are an execution plan; callers restore the original order.
+    The token bound counts padding (batch size times longest input), not merely
+    the sum of the unpadded lengths. No truncation or dropping is permitted.
+    """
+    if max_batch_tokens < 1 or max_batch_size < 1:
+        raise ValueError("Batch token and size limits must be positive")
+    if any(length < 1 for length in lengths):
+        raise ValueError("Every joint input must have at least one token")
+    if any(length > max_batch_tokens for length in lengths):
+        raise ValueError(
+            "A joint input exceeds max_batch_tokens; increase the batch token "
+            "budget to encode this input without truncation"
+        )
+    batches: list[list[int]] = []
+    current: list[int] = []
+    for index in sorted(range(len(lengths)), key=lambda i: (lengths[i], i)):
+        if current and (
+            len(current) == max_batch_size
+            or lengths[index] * (len(current) + 1) > max_batch_tokens
+        ):
+            batches.append(current)
+            current = []
+        current.append(index)
+    if current:
+        batches.append(current)
+    return batches
+
+
+class TorchBackend:
+    """Load one full encoder and score joint inputs using bounded batches.
+
+    CUDA uses BF16 when supported; CPU and MPS use FP32. Scalar-head inputs and
+    parameters always use FP32, matching the frozen reference. ``strict=True``
+    uses one input per forward pass, matching its numerical batch shape. The
+    faster default may produce small batch-shape differences around near ties.
+
+    Instantiate once and reuse it. The lock prevents overlapping forwards from
+    separate caller threads; the serving scheduler should coalesce requests
+    into a single score_pairs call to obtain cross-request batching.
+    """
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        *,
+        device: str = "auto",
+        max_batch_tokens: int = 8192,
+        max_batch_size: int = 32,
+        strict: bool = False,
+        verify: bool = True,
+    ) -> None:
+        self.path = Path(model_path)
+        if not self.path.is_dir():
+            raise ValueError("model_path must be a complete local model directory")
+        if max_batch_tokens < 1 or max_batch_size < 1:
+            raise ValueError("Batch token and size limits must be positive")
+        if verify:
+            verify_model_files(self.path)
+        self.config = json.loads((self.path / "joint_config.json").read_text())
+        if self.config.get("format") != "gemmadecision-full-joint-v4":
+            raise ValueError("This backend requires a full joint GemmaDecision v4 model")
+        self.max_state_tokens = int(self.config.get("max_state_tokens", 2048))
+        self.max_action_tokens = int(self.config.get("max_action_tokens", 768))
+        dimensions = self.config.get("head_config", {"hidden": 640, "width": 512})
+        hidden, width = int(dimensions["hidden"]), int(dimensions["width"])
+        if min(self.max_state_tokens, self.max_action_tokens, hidden, width) < 1:
+            raise ValueError("Token limits and head dimensions must be positive")
+        head_file = self.config.get("head_file", "joint_head.safetensors")
+        if not isinstance(head_file, str) or Path(head_file).name != head_file:
+            raise ValueError("head_file must name a file in the model directory")
+        try:
+            import torch
+            from safetensors.torch import load_file
+            from transformers import AutoModel, AutoTokenizer
+        except ImportError as error:
+            raise ImportError(
+                "Install the local inference dependencies: pip install 'gemmadecision[serve]'"
+            ) from error
+
+        self._torch = torch
+        if device == "auto":
+            device = (
+                "cuda" if torch.cuda.is_available()
+                else "mps" if torch.backends.mps.is_available()
+                else "cpu"
+            )
+        self.device = torch.device(device)
+        if self.device.type not in {"cpu", "mps", "cuda"}:
+            raise ValueError("Supported devices are cpu, mps and cuda")
+        if self.device.type == "cuda" and not torch.cuda.is_available():
+            raise ValueError("CUDA is unavailable on this machine")
+        if self.device.type == "mps" and not torch.backends.mps.is_available():
+            raise ValueError("MPS is unavailable on this machine")
+        dtype = (
+            torch.bfloat16
+            if self.device.type == "cuda" and torch.cuda.is_bf16_supported()
+            else torch.float32
+        )
+        self.dtype = str(dtype).removeprefix("torch.")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.path, local_files_only=True, trust_remote_code=False
+        )
+        # Last-token pooling below assumes right padding; do not inherit an
+        # arbitrary tokenizer padding side from the calling environment.
+        self.tokenizer.padding_side = "right"
+        if self.tokenizer.pad_token_id is None:
+            raise ValueError("The model tokenizer must define a padding token")
+        self.encoder = AutoModel.from_pretrained(
+            self.path,
+            local_files_only=True,
+            trust_remote_code=False,
+            use_safetensors=True,
+            dtype=dtype,
+            attn_implementation="sdpa",
+        ).to(self.device).eval().requires_grad_(False)
+        if self.encoder.config.hidden_size != hidden:
+            raise ValueError("Scalar head width does not match the encoder hidden size")
+        self.context_limit = getattr(self.encoder.config, "max_position_embeddings", None)
+        self.head = torch.nn.Sequential(
+            torch.nn.LayerNorm(hidden),
+            torch.nn.Linear(hidden, width),
+            torch.nn.GELU(),
+            torch.nn.Linear(width, 1),
+        ).to(device=self.device, dtype=torch.float32).eval().requires_grad_(False)
+        weights = load_file(str(self.path / head_file), device="cpu")
+        if any(not torch.isfinite(value).all() for value in weights.values()):
+            raise ValueError("Non-finite scalar head weights")
+        self.head.load_state_dict(weights, strict=True)
+        self.max_batch_tokens = max_batch_tokens
+        self.max_batch_size = 1 if strict else max_batch_size
+        self.strict = strict
+        self.verified = verify
+        self._forward_lock = threading.Lock()
+
+    def count_tokens(self, text: str) -> int:
+        """Count tokens including the same BOS/special tokens used for scoring."""
+        if not isinstance(text, str):
+            raise TypeError("Token-count input must be a string")
+        return len(self.tokenizer(
+            text, add_special_tokens=True, truncation=False, padding=False
+        )["input_ids"])
+
+    def score_pairs(self, texts: Sequence[str]) -> list[float]:
+        """Score full training-format state/candidate strings, in input order.
+
+        The caller must separately enforce max_state_tokens/max_action_tokens
+        on the original fields before combining them. This method cannot recover
+        those field boundaries reliably from arbitrary text. An empty list is
+        allowed for scheduler convenience; empty individual inputs are not.
+        """
+        if isinstance(texts, (str, bytes)) or not isinstance(texts, Sequence):
+            raise TypeError("Joint inputs must be a sequence of strings")
+        if not texts:
+            return []
+        if any(not isinstance(text, str) or not text for text in texts):
+            raise ValueError("Every joint input must be a nonempty string")
+        tokens = self.tokenizer(
+            list(texts), add_special_tokens=True, truncation=False, padding=False
+        )["input_ids"]
+        lengths = [len(row) for row in tokens]
+        if self.context_limit is not None and any(
+            length > self.context_limit for length in lengths
+        ):
+            raise ValueError(f"A joint input exceeds the encoder limit of {self.context_limit} tokens")
+        batches = plan_batches(lengths, self.max_batch_tokens, self.max_batch_size)
+        scores = [0.0] * len(texts)
+        # inference_mode is entered here because it is thread-local: a long-lived
+        # worker may call the backend from different executor threads.
+        with self._forward_lock, self._torch.inference_mode():
+            for indices in batches:
+                batch_scores = self._score_tokens([tokens[i] for i in indices])
+                for index, score in zip(indices, batch_scores, strict=True):
+                    scores[index] = score
+        return scores
+
+    def _score_tokens(self, tokens: Sequence[Sequence[int]]) -> list[float]:
+        torch = self._torch
+        batch = self.tokenizer.pad(
+            {"input_ids": tokens}, padding=True, return_attention_mask=True,
+            return_tensors="pt",
+        )
+        ids = batch["input_ids"].to(self.device)
+        mask = batch["attention_mask"].to(self.device)
+        hidden = self.encoder(
+            input_ids=ids, attention_mask=mask, use_cache=False
+        ).last_hidden_state
+        positions = mask.sum(-1) - 1
+        last = hidden[torch.arange(len(ids), device=self.device), positions].float()
+        vectors = torch.nn.functional.normalize(last, dim=-1)
+        if not torch.isfinite(vectors).all():
+            raise RuntimeError("Non-finite model encoding")
+        scores = self.head(vectors).reshape(-1).cpu().tolist()
+        if any(not math.isfinite(score) for score in scores):
+            raise RuntimeError("Non-finite ranking score")
+        return scores
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        return {
+            "backend": "torch",
+            "device": str(self.device),
+            "dtype": self.dtype,
+            "head_dtype": "float32",
+            "strict": self.strict,
+            "max_batch_tokens": self.max_batch_tokens,
+            "max_batch_size": self.max_batch_size,
+            "model_revision": MODEL_REVISION if self.verified else None,
+            "verified": self.verified,
+        }

gemmadecision/backends/vllm.py

@@ -0,0 +1,144 @@
+"""Optional vLLM pooling backend for the *joint* GemmaDecision encoder.
+
+No dependency, model, CUDA context, or network request is loaded at import.
+This backend is experimental until its deployment-specific GPU parity passes.
+"""
+from __future__ import annotations
+
+from importlib.metadata import version
+import json
+import math
+from pathlib import Path
+from threading import Lock
+
+from ..export_vllm import POOLER_CONFIG, VLLM_VERSION, export_vllm
+
+
+class VLLMBackend:
+    """Batch complete state/candidate pairs in vLLM; run the trained head FP32.
+
+    The decoder stays on CUDA. Its 640-element last-token vectors are returned
+    to a tiny CPU FP32 head. This avoids depending on internal vLLM worker APIs
+    and preserves the original FP32 normalization and scalar-head precision.
+    """
+
+    name = "vllm"
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        *,
+        export_path: str | Path | None = None,
+        gpu_memory_utilization: float = 0.20,
+        max_model_len: int = 3072,
+        max_num_seqs: int = 64,
+        enforce_eager: bool = True,
+    ) -> None:
+        try:
+            installed = version("vllm")
+        except Exception as error:
+            raise ImportError("Install gemmadecision[vllm] on a supported Linux GPU host") from error
+        if installed != VLLM_VERSION:
+            raise RuntimeError(f"This backend requires vllm=={VLLM_VERSION}; found {installed}")
+        if not 0 < gpu_memory_utilization < 1:
+            raise ValueError("gpu_memory_utilization must be between zero and one")
+        if max_num_seqs < 1 or max_model_len < 1:
+            raise ValueError("vLLM sequence limits must be positive")
+        import torch
+        from safetensors.torch import load_file
+        from transformers import AutoTokenizer
+        from vllm import LLM, PoolingParams
+        from vllm.config import PoolerConfig
+
+        if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+            raise RuntimeError("The vLLM backend requires a CUDA GPU with bfloat16 support")
+        self.path = Path(model_path).expanduser().resolve()
+        from .torch import verify_model_files
+        verify_model_files(self.path)
+        config = json.loads((self.path / "joint_config.json").read_text())
+        encoder_config = json.loads((self.path / "config.json").read_text())
+        self.max_state_tokens = int(config["max_state_tokens"])
+        self.max_action_tokens = int(config["max_action_tokens"])
+        self.context_limit = min(int(encoder_config["max_position_embeddings"]), max_model_len)
+        self.device = "cuda"
+        self.dtype = "bfloat16"
+        self.head_device = "cpu"
+        self.hidden = int(config["head_config"]["hidden"])
+        width = int(config["head_config"]["width"])
+        head_file = config["head_file"]
+        if not isinstance(head_file, str) or Path(head_file).name != head_file:
+            raise ValueError("head_file must be a filename in the model directory")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.path, local_files_only=True, trust_remote_code=False)
+        self.head = torch.nn.Sequential(
+            torch.nn.LayerNorm(self.hidden), torch.nn.Linear(self.hidden, width),
+            torch.nn.GELU(), torch.nn.Linear(width, 1),
+        ).to(device="cpu", dtype=torch.float32).eval().requires_grad_(False)
+        weights = load_file(str(self.path / head_file), device="cpu")
+        if any(not torch.isfinite(weight).all() for weight in weights.values()):
+            raise ValueError("Non-finite scalar head weights")
+        self.head.load_state_dict(weights, strict=True)
+        self.encoder_path = export_vllm(self.path, export_path)
+        self._pooling_params = PoolingParams(use_activation=False)
+        self._lock = Lock()
+        self.llm = LLM(
+            model=str(self.encoder_path), tokenizer=str(self.encoder_path),
+            runner="pooling", convert="embed", dtype="bfloat16",
+            trust_remote_code=False, pooler_config=PoolerConfig(**POOLER_CONFIG),
+            max_model_len=self.context_limit, max_num_seqs=max_num_seqs,
+            gpu_memory_utilization=gpu_memory_utilization,
+            enforce_eager=enforce_eager, enable_prefix_caching=False,
+            tensor_parallel_size=1, seed=0,
+        )
+
+    def _tokens(self, text: str) -> list[int]:
+        if not isinstance(text, str):
+            raise TypeError("Input text must be a string")
+        return self.tokenizer(text, add_special_tokens=True, truncation=False,
+                              padding=False)["input_ids"]
+
+    def count_tokens(self, text: str) -> int:
+        return len(self._tokens(text))
+
+    def score_pairs(self, texts: list[str]) -> list[float]:
+        """Score already-rendered joint pairs, preserving input order.
+
+        Explicit token IDs ensure exactly the frozen tokenizer's BOS behavior.
+        No chat template, truncation, probabilities, candidate-only embedding,
+        or independently cached action representation is used.
+        """
+        if not texts:
+            return []
+        tokens = [self._tokens(text) for text in texts]
+        for ids in tokens:
+            if not ids or len(ids) > self.context_limit:
+                raise ValueError(f"Joint input must contain 1–{self.context_limit} tokens")
+        import torch
+        with self._lock, torch.inference_mode():
+            outputs = self.llm.embed(
+                [{"prompt_token_ids": ids} for ids in tokens],
+                pooling_params=self._pooling_params, use_tqdm=False,
+            )
+            if len(outputs) != len(tokens):
+                raise RuntimeError("vLLM returned an unexpected number of embeddings")
+            vectors = torch.tensor([item.outputs.embedding for item in outputs],
+                                   dtype=torch.float32, device="cpu")
+            if vectors.shape != (len(tokens), self.hidden) or not torch.isfinite(vectors).all():
+                raise RuntimeError("vLLM returned invalid last-token hidden states")
+            vectors = torch.nn.functional.normalize(vectors, dim=-1)
+            scores = self.head(vectors).flatten().tolist()
+        if len(scores) != len(texts) or not all(math.isfinite(score) for score in scores):
+            raise RuntimeError("Non-finite or malformed ranking scores")
+        return scores
+
+    score_texts = score_pairs
+
+    @property
+    def metadata(self) -> dict:
+        from .torch import MODEL_REVISION
+        return {"backend": "vllm", "vllm_version": VLLM_VERSION,
+                "device": self.device, "dtype": self.dtype,
+                "head_device": self.head_device, "head_dtype": "float32",
+                "pooling_type": "LAST", "normalization": "float32_after_pooling",
+                "context_limit": self.context_limit, "prefix_caching": False,
+                "model_revision": MODEL_REVISION, "verified": True}