`bd96149`

feat(export.gate_fallback): mean-gate static adapter-mix for Ollama/GGUF

The GGUF runtime can't evaluate the learned torch gate at inference,
so export averages the gate's softmax output across training prompts
and emits the averaged weights as --adapter-mix coefficients.
Lossless vs today's shipped behavior. The export manifest will record
gate_mode='static_mean' so downstream tooling can tell these builds
apart from hand-picked mixes. Uniform-mode gates (cold-start) use the
corresponding uniform 1/N split directly.

Authored by

espadonne 3 weeks ago

SHA: bd96149bdc74160d746c55c9edb7fd07d05bff6c
Parents: cbcdb09
Tree: 241ea2e

2 changed files

Status	File	+	-
A	`src/dlm/export/gate_fallback.py`	79	0
A	`tests/unit/export/test_gate_fallback.py`	81	0

src/dlm/export/gate_fallback.pyadded

 +"""Static mean-gate fallback for Ollama / llama.cpp export.
++
 +The learned gate (Sprint 34) runs in PyTorch at `dlm prompt` time. The
 +GGUF runtime (Ollama, llama.cpp) can't evaluate a torch module at
 +inference, so when the user runs `dlm export` on a document with
 +`training.gate.enabled: true` we fall back to:
++
 +1. Compute the gate's softmax output on every training prompt.
 +2. Average those probability vectors across the corpus → one fixed
 +   weight per adapter.
 +3. Emit the averaged weights as the Modelfile's `--adapter-mix`
 +   coefficients.
++
 +The exported model is a statically-weighted merge of the named
 +adapters — lossless vs today's shipped behavior, and strictly better
 +than asking the user to guess coefficients. Dynamic per-prompt routing
 +is the `dlm prompt` / `dlm repl` path only.
++
 +The export manifest records ``gate_mode: "static_mean"`` so downstream
 +tooling can tell an exported-with-mean-gate build apart from a
 +hand-picked `--adapter-mix`.
 +"""
++
 +from __future__ import annotations
++
 +from typing import TYPE_CHECKING
++
 +if TYPE_CHECKING:
 +    import torch
++
 +    from dlm.train.gate.module import Gate, GateMetadata
++
++
 +def mean_gate_weights(
 +    gate: Gate,
 +    metadata: GateMetadata,
 +    prompt_embeddings: list[torch.Tensor],
 +) -> list[tuple[str, float]]:
 +    """Average ``gate(embedding)`` across the training prompts.
++
 +    Returns ``[(adapter_name, weight), ...]`` suitable for direct
 +    substitution into ``dlm export --adapter-mix``. Weights sum to
 +    1.0 (gate output is softmax; average of softmax is still on the
 +    simplex) but we don't renormalize defensively — a numeric-drift
 +    renorm would mask bugs.
++
 +    Raises ``ValueError`` if ``prompt_embeddings`` is empty — a
 +    zero-prompt corpus has nothing to average.
 +    """
 +    import torch
++
 +    if not prompt_embeddings:
 +        raise ValueError("mean_gate_weights requires >= 1 prompt embedding")
++
 +    with torch.no_grad():
 +        stacked = torch.stack([e.detach().to(torch.float32).reshape(-1) for e in prompt_embeddings])
 +        if stacked.shape[1] != metadata.input_dim:
 +            raise ValueError(
 +                f"prompt embedding dim {stacked.shape[1]} != gate input_dim "
 +                f"{metadata.input_dim} (base model mismatch?)"
 +            )
 +        probs = gate(stacked)  # (N, n_adapters)
 +        mean = probs.mean(dim=0)
++
 +    return [(name, float(mean[i].item())) for i, name in enumerate(metadata.adapter_names)]
++
++
 +def uniform_adapter_mix(adapter_names: tuple[str, ...]) -> list[tuple[str, float]]:
 +    """Mean-gate fallback for uniform-mode gates (cold-start).
++
 +    Returns ``[(name, 1/N), ...]`` — the export path for a doc that has
 +    a gate declared but where the gate trainer chose the uniform
 +    fallback because the corpus was too small.
 +    """
 +    n = len(adapter_names)
 +    if n == 0:
 +        return []
 +    w = 1.0 / n
 +    return [(name, w) for name in adapter_names]

tests/unit/export/test_gate_fallback.pyadded

 +"""Static mean-gate fallback for Ollama export."""
++
 +from __future__ import annotations
++
 +import pytest
++
 +from dlm.export.gate_fallback import mean_gate_weights, uniform_adapter_mix
 +from dlm.train.gate.module import Gate, GateMetadata
++
++
 +class TestUniformAdapterMix:
 +    def test_three_adapters_third_each(self) -> None:
 +        mix = uniform_adapter_mix(("a", "b", "c"))
 +        assert mix == [("a", 1 / 3), ("b", 1 / 3), ("c", 1 / 3)]
++
 +    def test_empty_tuple(self) -> None:
 +        assert uniform_adapter_mix(()) == []
++
++
 +class TestMeanGateWeights:
 +    def _gate_and_meta(self) -> tuple[Gate, GateMetadata]:
 +        gate = Gate(input_dim=8, hidden_proj_dim=4, n_adapters=2)
 +        meta = GateMetadata(
 +            input_dim=8,
 +            hidden_proj_dim=4,
 +            adapter_names=("a", "b"),
 +            mode="trained",
 +        )
 +        return gate, meta
++
 +    def test_empty_corpus_refused(self) -> None:
 +        gate, meta = self._gate_and_meta()
 +        with pytest.raises(ValueError, match=">= 1 prompt embedding"):
 +            mean_gate_weights(gate, meta, [])
++
 +    def test_weights_shape_and_sum_to_one(self) -> None:
 +        import torch
++
 +        gate, meta = self._gate_and_meta()
 +        embeddings = [torch.randn(8) for _ in range(16)]
 +        mix = mean_gate_weights(gate, meta, embeddings)
 +        assert [name for name, _ in mix] == ["a", "b"]
 +        total = sum(w for _, w in mix)
 +        assert total == pytest.approx(1.0, abs=1e-5)
 +        for _, w in mix:
 +            assert 0.0 <= w <= 1.0
++
 +    def test_dim_mismatch_refused(self) -> None:
 +        import torch
++
 +        gate, meta = self._gate_and_meta()
 +        # Wrong-dim embedding.
 +        with pytest.raises(ValueError, match="input_dim"):
 +            mean_gate_weights(gate, meta, [torch.randn(4)])
++
 +    def test_mean_reflects_per_prompt_skew(self) -> None:
 +        """Ten prompts near cluster A + one prompt near cluster B should
 +        average out to favor A. Sanity check that mean_gate_weights isn't
 +        just emitting uniform."""
 +        import torch
++
 +        gate = Gate(input_dim=8, hidden_proj_dim=8, n_adapters=2)
 +        meta = GateMetadata(
 +            input_dim=8,
 +            hidden_proj_dim=8,
 +            adapter_names=("a", "b"),
 +            mode="trained",
 +        )
 +        # Force the gate weights so it's (almost) deterministic: class-a
 +        # embeddings near +1, class-b near -1.
 +        torch.manual_seed(0)
 +        a_embeddings = [torch.ones(8) + 0.01 * torch.randn(8) for _ in range(10)]
 +        b_embedding = -torch.ones(8)
 +        # We won't train here — untrained gate may or may not favor A.
 +        # The point is only that the mean is a real average (not uniform
 +        # or fixed), which we check by comparing against a single-prompt
 +        # case.
 +        mix_mixed = mean_gate_weights(gate, meta, a_embeddings + [b_embedding])
 +        mix_single_a = mean_gate_weights(gate, meta, [a_embeddings[0]])
 +        # Different input distributions → different averaged outputs.
 +        assert mix_mixed != mix_single_a