plumbify 0.2.0__py3-none-any.whl

plumbify/__init__.py

@@ -0,0 +1,20 @@
+"""Plumbify: a fast decision path (System 1) for open language models, served natively by vLLM.
+
+    plumbify train --base Qwen/Qwen3.5-9B --rows train.jsonl --dev dev.jsonl --out plumbed-qwen3.5-9b
+    vllm serve plumbed-qwen3.5-9b
+
+In Python, ``System1`` runs a plumb with transformers (training, evaluation, the reference implementation).
+"""
+
+__version__ = "0.2.0"
+
+
+def __getattr__(name):  # System1 pulls in torch and transformers: load it on first use
+    if name == "System1":
+        from .system1 import System1
+
+        return System1
+    raise AttributeError(name)
+
+
+__all__ = ["System1", "__version__"]

plumbify/artifact.py

@@ -0,0 +1,102 @@
+"""The plumb: the trained decision branch of a model, stored next to (never inside) the base weights.
+
+    plumb.json                    spec: base model, taps, head shape, calibration
+    head.safetensors              the decision head
+    suffix_adapter.json/.safetensors   the suffix-only LoRA (absent when trained with --no_lora)
+
+``plumbify train-head`` writes this layout; ``plumbify package`` (plumbify/plumbed.py) combines it with the base model
+into a plumbed model directory. Paths can be local directories or Hugging Face Hub repos.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import os
+from dataclasses import dataclass, field
+
+FORMAT = "plumb/2"
+RENDER_VERSION = 2  # plumbify.core.render: bump when the decision text layout changes
+
+
+@dataclass
+class Conformal:
+    """Split-conformal threshold fitted on held-out rows: Set = {k : p_k >= 1 - qhat}, P(gold in Set) >= 1 - alpha."""
+
+    alpha: float
+    qhat: float
+    n: int
+
+
+@dataclass
+class Calibration:
+    temperature: float = 1.0
+    conformal: Conformal | None = None
+
+
+@dataclass
+class PlumbSpec:
+    base_model: str
+    d: int
+    d_proj: int = 512
+    has_adapter: bool = True
+    render_version: int = RENDER_VERSION
+    head_type: str = "decision"
+    head_config: dict = field(default_factory=dict)
+    taps: list = field(default_factory=list)  # layers the head reads; -1 is the final normed output
+    calibration: Calibration = field(default_factory=Calibration)
+    name: str = ""
+    format: str = FORMAT
+    extra: dict = field(default_factory=dict)
+
+    def to_json(self) -> str:
+        return json.dumps(dataclasses.asdict(self), indent=1)
+
+    @staticmethod
+    def from_dict(d: dict) -> PlumbSpec:
+        cal = d.get("calibration") or {}
+        conf = cal.get("conformal")
+        d = {
+            **d,
+            "calibration": Calibration(
+                cal.get("temperature", 1.0), Conformal(**conf) if conf else None
+            ),
+        }
+        known = {f.name for f in dataclasses.fields(PlumbSpec)}
+        return PlumbSpec(**{k: v for k, v in d.items() if k in known})
+
+
+def fetch(path: str, name: str) -> str:
+    """Local path of a file in a plumb directory or a Hub repo."""
+    if os.path.isdir(path):
+        return os.path.join(path, name)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(path, name)
+
+
+def read_spec(path: str) -> PlumbSpec:
+    try:
+        spec_file = fetch(path, "plumb.json")
+    except Exception as e:  # Hub errors come in several types; the message says which
+        raise FileNotFoundError(f"{path}: cannot read plumb.json ({e})") from e
+    if not os.path.exists(spec_file):
+        raise FileNotFoundError(
+            f"{path}: no plumb.json; is this a plumb or plumbed model directory?"
+        )
+    spec = PlumbSpec.from_dict(json.load(open(spec_file)))
+    if spec.head_type != "decision":
+        raise ValueError(f"{path}: unsupported plumb head type {spec.head_type!r}")
+    return spec
+
+
+def save(out: str, spec: PlumbSpec, head) -> None:
+    """Write the spec and the head (with its calibrated temperature)."""
+    from safetensors.torch import save_file
+
+    os.makedirs(out, exist_ok=True)
+    sd = {k: v.detach().contiguous().cpu() for k, v in head.state_dict().items()}
+    sd["temperature"] = sd["temperature"].new_tensor(spec.calibration.temperature)
+    save_file(sd, os.path.join(out, "head.safetensors"))
+    with open(os.path.join(out, "plumb.json"), "w") as f:
+        f.write(dataclasses.replace(spec, format=FORMAT).to_json())

plumbify/branch.py

@@ -0,0 +1,370 @@
+"""System 1 inside System 2's generation: a decision reads the live KV cache, then the cache is rolled back.
+
+    Qwen generating ... <tool_call>{"name": "plumb_decide", "arguments": {question, options}}</tool_call>
+        -> checkpoint the cache (clone the small recurrent / conv states; remember the attention length)
+        -> append the decision suffix (close the turn, "Decision: ... Options: ...", assistant-open), adapter ON
+        -> the head reads the suffix at the tapped layers -> probabilities
+        -> restore the cache exactly, append <tool_response>{answer}</tool_response>, keep generating
+
+The tool call is only the trigger; nothing external is called. The decision costs its ~200 suffix tokens, never a
+re-read of the conversation. That is valid because the plumb's adapter is suffix-only: the context in the live cache
+was computed by the base weights, which is exactly what the adapter was trained on top of. Without a cache (a fresh
+request) the same decision is ``System1.decide`` over the rendered text.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from .calibration import prediction_set
+from .core.answers import answer
+from .core.decision_head import SuffixBatch
+from .core.render import LEADS
+from .core.row import Row
+from .core.tool import PLUMB_TOOL, row_from_tool_args, tool_schema
+
+_A, _D, _T = "\x00ASSISTANT\x00", "\x00DECISION\x00", "\x00TOOL\x00"
+_XML_FN = re.compile(r"<function=([^>\s]+)>(.*?)</function>", re.S)
+_XML_PARAM = re.compile(r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>", re.S)
+
+
+_GEMMA_CALL = re.compile(r"^call:([^\s{]+)\s*(\{.*\})$", re.S)
+_GEMMA_KEY = re.compile(r"([{,]\s*)([A-Za-z_][\w\-]*)\s*:")
+_GEMMA_BARE = re.compile(r"(:\s*)([A-Za-z_][\w\-]*)(\s*[,}\]])")
+_GEMMA_STR = '<|"|>'
+
+
+@dataclass(frozen=True)
+class Markup:
+    """The model's own tool-call and thinking markup, read from its chat template."""
+
+    call_open: str
+    call_close: str
+    think_open: str | None
+    think_close: str | None
+
+    # Gemma 4; Hermes / Qwen
+    CALLS = (
+        ("<|tool_call>", "<tool_call|>"),
+        ("<tool_call>", "</tool_call>"),
+    )
+    # Gemma 4; Qwen, DeepSeek, GLM
+    THINKS = (
+        ("<|channel>thought", "<channel|>"),
+        ("<think>", "</think>"),
+    )
+
+    @staticmethod
+    def of(tok) -> Markup:
+        tmpl = getattr(tok, "chat_template", None) or ""
+        if isinstance(tmpl, dict):
+            tmpl = " ".join(tmpl.values())
+        call = next((c for c in Markup.CALLS if c[0] in tmpl and c[1] in tmpl), Markup.CALLS[1])
+        think = next((t for t in Markup.THINKS if t[0] in tmpl), (None, None))
+        return Markup(*call, *think)
+
+    def calls(self, text: str) -> list[str]:
+        """Bodies of the complete tool calls in ``text``."""
+        pat = re.escape(self.call_open) + "(.*?)" + re.escape(self.call_close)
+        return re.findall(pat, text, re.S)
+
+    def opens_thinking(self, prompt: str) -> bool:
+        """The prompt ends inside an open thinking block (generation starts as reasoning)."""
+        return bool(self.think_open) and prompt.rstrip().endswith(self.think_open.rstrip())
+
+    def final(self, text: str) -> str:
+        """The reply after the last tool call and the last thinking block, without later call markup."""
+        text = text.split(self.call_close)[-1]
+        if self.think_close:
+            text = text.split(self.think_close)[-1]
+        return text.split(self.call_open)[0]
+
+
+def _gemma_args(s: str) -> dict | None:
+    """Gemma 4's argument syntax -> dict: bare keys, strings between ``<|"|>`` marks, JSON-like nesting."""
+    parts = s.split(_GEMMA_STR)
+    if len(parts) % 2 == 0:
+        return None
+    code = []
+    for i, p in enumerate(parts):
+        if i % 2:
+            code.append(json.dumps(p))
+        else:
+            p = _GEMMA_KEY.sub(r'\1"\2":', p)
+            p = _GEMMA_BARE.sub(
+                lambda m: m[0] if m[2] in ("true", "false", "null") else f'{m[1]}"{m[2]}"{m[3]}', p
+            )
+            code.append(p)
+    try:
+        v = json.loads("".join(code))
+    except ValueError:
+        return None
+    return v if isinstance(v, dict) else None
+
+
+def parse_tool_call(body: str) -> tuple[str, dict] | None:
+    """One tool-call body (between the model's call markers) -> (name, arguments). Formats open models emit:
+    Hermes JSON ``{"name": ..., "arguments": {...}}`` (Qwen3, most templates), XML
+    ``<function=NAME><parameter=KEY>VALUE</parameter>...</function>`` (Qwen3.5, Qwen3-Coder) and Gemma 4's
+    ``call:NAME{key:<|"|>value<|"|>,...}``. XML values that parse as JSON (lists, numbers) are decoded; others stay
+    strings."""
+    body = body.strip()
+    g = _GEMMA_CALL.match(body)
+    if g:
+        args = _gemma_args(g[2])
+        return (g[1], args) if args is not None else None
+    if body.startswith("{"):
+        try:
+            c = json.loads(body)
+        except ValueError:
+            return None
+        args = c.get("arguments") or {}
+        if isinstance(args, str):
+            try:
+                args = json.loads(args)
+            except ValueError:
+                return None
+        return c.get("name"), args
+    m = _XML_FN.search(body)
+    if not m:
+        return None
+    args = {}
+    for k, v in _XML_PARAM.findall(m.group(2)):
+        try:
+            args[k] = json.loads(v)
+        except ValueError:
+            args[k] = v.strip()
+    return m.group(1), args
+
+
+class CacheCheckpoint:
+    """Everything needed to put a transformers cache back as it was: the attention length (keys/values grow by
+    concatenation, so a slice undoes them) and clones of linear-attention / Mamba states (updated in place)."""
+
+    def __init__(self, cache):
+        self.length = cache.get_seq_length()
+        self.states = []
+        for layer in cache.layers:
+            for attr in ("conv_states", "recurrent_states"):
+                d = getattr(layer, attr, None)
+                if isinstance(d, dict):
+                    self.states += [(d, i, t.clone()) for i, t in d.items() if t is not None]
+
+    def restore(self, cache) -> None:
+        for d, i, t in self.states:
+            if d[i].shape == t.shape:
+                d[i].copy_(t)  # keep the tensor's address (static for cudagraphs)
+            else:
+                d[i] = t
+        for layer in cache.layers:
+            keys = getattr(layer, "keys", None)
+            if isinstance(keys, torch.Tensor) and keys.dim() == 4 and keys.shape[-2] > self.length:
+                layer.keys = keys[..., : self.length, :]
+                layer.values = layer.values[..., : self.length, :]
+
+
+@dataclass
+class Frames:
+    """Chat-template text that follows an assistant turn in progress, derived from the tokenizer's own template."""
+
+    to_decision: str  # close the assistant turn, open a user turn
+    after_decision: str  # close the user turn, open the assistant (non-thinking, as in training)
+    to_tool: str  # close the assistant turn, open the tool response
+    after_tool: str  # close the tool response, open the assistant
+
+    @staticmethod
+    def of(tok, template_kwargs: dict | None = None, markup: Markup | None = None) -> Frames:
+        kw = {"enable_thinking": False, **(template_kwargs or {})}
+        markup = markup or Markup.of(tok)
+        u = {"role": "user", "content": "hi"}
+        a = {"role": "assistant", "content": _A}
+        t1 = tok.apply_chat_template(
+            [u, a, {"role": "user", "content": _D}],
+            tokenize=False,
+            add_generation_prompt=True,
+            **kw,
+        )
+        # the tool response follows a real tool call: some templates (Gemma 4) render tool messages only after one,
+        # inside the same model turn
+        call = {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": "call_0",
+                    "type": "function",
+                    "function": {"name": PLUMB_TOOL, "arguments": {"question": "q"}},
+                }
+            ],
+        }
+        tool = {"role": "tool", "tool_call_id": "call_0", "name": PLUMB_TOOL, "content": _T}
+        t2 = tok.apply_chat_template(
+            [u, call, tool], tokenize=False, add_generation_prompt=True, **kw
+        )
+        close = t2.rfind(markup.call_close, 0, t2.index(_T)) if _T in t2 else -1
+        if close >= 0:
+            start = close + len(markup.call_close)
+        else:  # a template that doesn't render tool calls: the tool turn follows plain assistant text
+            t2 = tok.apply_chat_template(
+                [u, a, {"role": "tool", "content": _T}],
+                tokenize=False,
+                add_generation_prompt=True,
+                **kw,
+            )
+            start = t2.index(_A) + len(_A)
+        return Frames(
+            t1[t1.index(_A) + len(_A) : t1.index(_D)],
+            t1[t1.index(_D) + len(_D) :],
+            t2[start : t2.index(_T)],
+            t2[t2.index(_T) + len(_T) :],
+        )
+
+
+def decision_suffix(tok, frames: Frames, row: Row) -> tuple[list[int], list[tuple[int, int]], int]:
+    """Token ids of the decision suffix after a live assistant turn, with option spans and DECIDE relative to it.
+    Same pieces, same order and same instruction as ``render_decision`` (training)."""
+    lead, instruction = LEADS[row.qtype]
+    ids = tok(
+        frames.to_decision + f"Decision: {row.question}\n{lead}\n", add_special_tokens=False
+    ).input_ids
+    spans = []
+    for o in row.options:
+        s = len(ids)
+        ids += tok(
+            f"- {o.name}: {o.desc}\n" if o.desc else f"- {o.name}\n", add_special_tokens=False
+        ).input_ids
+        spans.append((s, len(ids)))
+    ids += tok(instruction + frames.after_decision, add_special_tokens=False).input_ids
+    return ids, spans, len(ids) - 1
+
+
+@torch.no_grad()
+def decide_on_cache(s1, cache, row: Row, frames: Frames) -> list[float]:
+    """Probabilities over ``row.options`` from the live cache; the cache is left exactly as it was."""
+    ids, spans, decide = decision_suffix(s1.tok, frames, row)
+    ck = CacheCheckpoint(cache)
+    dev = next(s1.trunk.parameters()).device
+    x = torch.tensor([ids], device=dev)
+    try:
+        s1._set_mask(torch.ones(1, len(ids), device=dev))  # every token here is suffix: adapter on
+        feats = s1.reader.run(
+            x,
+            None,
+            [0],
+            [len(ids)],
+            past_key_values=cache,
+            use_cache=True,
+            cache_position=torch.arange(ck.length, ck.length + len(ids), device=dev),
+        )
+    finally:
+        s1._set_mask(None)
+        ck.restore(cache)
+    item = {"feats": feats[0], "opt_spans": spans, "decide": decide}
+    logits = s1.head(SuffixBatch.collate([item], device=next(s1.head.parameters()).device))
+    return torch.softmax(logits[0, : len(row.options)].float(), -1).tolist()
+
+
+@dataclass
+class Turn:
+    text: str
+    decisions: list[dict] = field(default_factory=list)
+    tokens: int = 0
+
+
+class Assistant:
+    """Generation with the untouched base model; ``plumb_decide`` tool calls are answered by System 1 on the live
+    cache. ``sample(logits) -> token id`` defaults to greedy (a scripted sampler drives the tests)."""
+
+    def __init__(
+        self, s1, conformal_qhat: float | None = None, template_kwargs: dict | None = None
+    ):
+        self.s1, self.tok, self.lm = s1, s1.tok, s1.lm
+        self.qhat = conformal_qhat
+        self.template_kwargs = template_kwargs or {}
+        self.markup = Markup.of(self.tok)
+        self.frames = Frames.of(self.tok, s1.template_kwargs, self.markup)
+
+    def _feed(self, ids: list[int], cache):
+        dev = next(self.lm.parameters()).device
+        start = cache.get_seq_length() if cache is not None else 0
+        out = self.lm(
+            input_ids=torch.tensor([ids], device=dev),
+            past_key_values=cache,
+            use_cache=True,
+            cache_position=torch.arange(start, start + len(ids), device=dev),
+            logits_to_keep=1,
+        )
+        return out.logits[0, -1], out.past_key_values
+
+    @torch.no_grad()
+    def chat(
+        self,
+        messages: list[dict],
+        max_new_tokens: int = 512,
+        tools: list | None = None,
+        sample: Callable[[torch.Tensor], int] | None = None,
+        max_decisions: int = 8,
+    ) -> Turn:
+        sample = sample or (lambda z: int(z.argmax()))
+        gen_eos = getattr(getattr(self.lm, "generation_config", None), "eos_token_id", None)
+        eos = (
+            {
+                self.tok.convert_tokens_to_ids(t)
+                for t in ("<|im_end|>", "<|endoftext|>")
+                if hasattr(self.tok, "convert_tokens_to_ids")
+            }
+            | {getattr(self.tok, "eos_token_id", None)}
+            | set(gen_eos if isinstance(gen_eos, list) else [gen_eos])
+        )
+        prompt = self.tok.apply_chat_template(
+            messages,
+            tools=[tool_schema(), *(tools or [])],
+            tokenize=False,
+            add_generation_prompt=True,
+            **self.template_kwargs,
+        )
+        logits, cache = self._feed(self.tok(prompt, add_special_tokens=False).input_ids, None)
+        turn, gen, handled, sampled = Turn(""), [], 0, 0
+        while sampled < max_new_tokens:  # injected tool responses don't count against the budget
+            t = sample(logits)
+            if t in eos:
+                break
+            gen.append(t)
+            sampled += 1
+            logits, cache = self._feed([t], cache)
+            text = self.tok.decode(gen, skip_special_tokens=False)
+            calls = self.markup.calls(text)
+            if len(calls) > handled and len(turn.decisions) < max_decisions:
+                handled = len(calls)
+                parsed = parse_tool_call(calls[-1])
+                if parsed is None:
+                    continue
+                name, args = parsed
+                if name != PLUMB_TOOL:
+                    break  # a client tool: hand the turn back
+                try:
+                    row = row_from_tool_args(args)
+                except (KeyError, TypeError, ValueError):
+                    continue  # malformed arguments: let generation carry on
+                probs = decide_on_cache(self.s1, cache, row, self.frames)
+                ans = answer(row, probs)
+                if self.qhat is not None:
+                    ans["set"] = [row.options[k].name for k in prediction_set(probs, self.qhat)]
+                turn.decisions.append(
+                    {
+                        "arguments": args,
+                        "result": ans,
+                        "context_tokens": cache.get_seq_length(),
+                    }
+                )
+                resp = self.frames.to_tool + json.dumps(ans) + self.frames.after_tool
+                resp_ids = self.tok(resp, add_special_tokens=False).input_ids
+                gen += resp_ids
+                logits, cache = self._feed(resp_ids, cache)
+        turn.text, turn.tokens = self.tok.decode(gen, skip_special_tokens=True), len(gen)
+        return turn

plumbify/calibration.py

@@ -0,0 +1,109 @@
+"""Calibration: one temperature, a split-conformal threshold, and the System 1 -> System 2 escalation it implies.
+
+Everything works on prediction rows ``{"probs": [...], "gold": int}`` as written by ``plumbify eval-head`` (probs at T = 1).
+
+- ``fit_temperature``: NLL-optimal T; moves confidence, never the winning option.
+- ``fit_conformal``:   split conformal with the LAC score s = 1 - p_gold. ``prediction_set`` then contains the gold
+  option with probability >= 1 - alpha on exchangeable data, distribution-free.
+- ``escalation_curve``: for each confidence threshold, the share of traffic System 1 keeps and its accuracy there.
+  A request whose conformal set is not a singleton is exactly one System 1 should hand to System 2.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+
+def rescale(probs: Sequence[float], temperature: float) -> list[float]:
+    z = [math.log(max(p, 1e-12)) / temperature for p in probs]
+    m = max(z)
+    e = [math.exp(x - m) for x in z]
+    s = sum(e)
+    return [x / s for x in e]
+
+
+def fit_temperature(rows: list[dict], lo: float = 0.25, hi: float = 8.0, iters: int = 60) -> float:
+    """Golden-section search of the mean NLL over log T (the NLL is unimodal in log T)."""
+    rows = [r for r in rows if r.get("gold") is not None]
+
+    def nll(log_t: float) -> float:
+        t = math.exp(log_t)
+        return -sum(math.log(max(rescale(r["probs"], t)[r["gold"]], 1e-12)) for r in rows) / max(
+            len(rows), 1
+        )
+
+    a, b = math.log(lo), math.log(hi)
+    g = (math.sqrt(5) - 1) / 2
+    c, d = b - g * (b - a), a + g * (b - a)
+    fc, fd = nll(c), nll(d)
+    for _ in range(iters):
+        if fc < fd:
+            b, d, fd = d, c, fc
+            c = b - g * (b - a)
+            fc = nll(c)
+        else:
+            a, c, fc = c, d, fd
+            d = a + g * (b - a)
+            fd = nll(d)
+    return math.exp((a + b) / 2)
+
+
+def fit_conformal(rows: list[dict], alpha: float = 0.1, temperature: float = 1.0) -> dict:
+    """qhat = the ceil((n+1)(1-alpha))/n empirical quantile of s = 1 - p_gold on held-out rows."""
+    scores = sorted(
+        1.0 - rescale(r["probs"], temperature)[r["gold"]] for r in rows if r.get("gold") is not None
+    )
+    n = len(scores)
+    if n == 0:
+        raise ValueError("conformal calibration needs labelled rows")
+    k = math.ceil((n + 1) * (1 - alpha))
+    qhat = 1.0 if k > n else scores[k - 1]
+    return {"alpha": alpha, "qhat": qhat, "n": n}
+
+
+def prediction_set(probs: Sequence[float], qhat: float) -> list[int]:
+    """Options whose probability clears 1 - qhat; never empty (falls back to the argmax)."""
+    s = [k for k, p in enumerate(probs) if p >= 1.0 - qhat]
+    return s or [max(range(len(probs)), key=probs.__getitem__)]
+
+
+def coverage(rows: list[dict], qhat: float, temperature: float = 1.0) -> dict:
+    """Empirical coverage and mean set size of the conformal sets on another labelled set."""
+    hit = size = singles = 0
+    rows = [r for r in rows if r.get("gold") is not None]
+    for r in rows:
+        s = prediction_set(rescale(r["probs"], temperature), qhat)
+        hit += r["gold"] in s
+        size += len(s)
+        singles += len(s) == 1
+    n = max(len(rows), 1)
+    return {
+        "coverage": hit / n,
+        "mean_set_size": size / n,
+        "singleton_rate": singles / n,
+        "n": len(rows),
+    }
+
+
+def escalation_curve(
+    rows: list[dict],
+    temperature: float = 1.0,
+    thresholds: Sequence[float] = (0.0, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99),
+) -> list[dict]:
+    """For each p_max threshold: share kept by System 1, System 1 accuracy on what it keeps."""
+    rows = [r for r in rows if r.get("gold") is not None]
+    out = []
+    for t in thresholds:
+        kept = [r for r in rows if max(rescale(r["probs"], temperature)) >= t]
+        acc = sum(
+            max(range(len(r["probs"])), key=r["probs"].__getitem__) == r["gold"] for r in kept
+        )
+        out.append(
+            {
+                "threshold": t,
+                "system1_share": len(kept) / max(len(rows), 1),
+                "system1_accuracy": acc / max(len(kept), 1),
+            }
+        )
+    return out

plumbify/cli.py

@@ -0,0 +1,61 @@
+"""``plumbify <command>``: each command is a module with ``main()``."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+
+COMMANDS = {
+    "train": (
+        "plumbify.training.train",
+        "train a plumb on an open model and write a plumbed model directory for vLLM",
+    ),
+    "package": (None, "turn a trained plumb into a plumbed model directory"),
+    "train-head": (
+        "plumbify.training.train_head",
+        "train and calibrate a plumb only (the first step of train)",
+    ),
+    "eval-head": ("plumbify.training.eval_head", "evaluate a plumb, or the base model zero-shot"),
+}
+
+
+def _usage() -> str:
+    w = max(map(len, COMMANDS))
+    return "usage: plumbify <command> [args]\n\n" + "\n".join(
+        f"  {k:<{w}}  {v[1]}" for k, v in COMMANDS.items()
+    )
+
+
+def _package(rest: list[str]) -> int:
+    from .plumbed import package
+
+    ap = argparse.ArgumentParser("plumbify package")
+    ap.add_argument("plumb", help="trained plumb directory (plumb.json, head, suffix adapter)")
+    ap.add_argument("out", help="plumbed model directory to write")
+    ap.add_argument("--base", default=None, help="override the base model recorded in plumb.json")
+    ap.add_argument(
+        "--copy", action="store_true", help="copy the base weights instead of linking them"
+    )
+    a = ap.parse_args(rest)
+    print(package(a.plumb, a.out, a.base, a.copy))
+    return 0
+
+
+def main(argv=None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if not argv or argv[0] in ("-h", "--help"):
+        print(_usage())
+        return 0
+    cmd, rest = argv[0], argv[1:]
+    if cmd not in COMMANDS:
+        print(f"unknown command {cmd!r}\n\n{_usage()}")
+        return 2
+    if cmd == "package":
+        return _package(rest)
+    sys.argv = [f"plumbify {cmd}", *rest]  # command modules parse sys.argv
+    return importlib.import_module(COMMANDS[cmd][0]).main() or 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

plumbify/core/__init__.py

@@ -0,0 +1,6 @@
+"""The pieces shared by training and serving: the decision row, its rendering, the suffix-only LoRA, the decision
+head and the layer taps."""
+
+from .row import QTYPES, Option, Row
+
+__all__ = ["QTYPES", "Option", "Row"]