Encoder path, quantized router index, and a catalog SPI - #111
Merged
Conversation
Replaces hand-authored exemplars with prompts drawn from public benchmarks: MBPP, GSM8K, gretel text-to-sql, ARC-Challenge and BIG-Bench Hard, dolly-15k, XSum, OPUS-100, Hermes function-calling and OASST1. 2391 prompts over 10 tasks, 80/20 train/eval, with per-source provenance and licence in sources.json. Hand-authored prompts describe how one person phrases a task. They are internally consistent and separate cleanly, which makes a classifier trained on them look better than it is; real prompts carry the phrasing variety the classifier actually has to survive. The corpus is not shipped in the jar. Only the input side of each benchmark is taken and reference answers are discarded. Summarization and translation have no benchmark of user-style requests, so those rows are documents or sentence pairs wrapped in an instruction, recorded as wrapped: true. TaskExemplars becomes a reader for this format rather than a loader for a bundled resource, so callers can classify against their own tasks by supplying a corpus in the same shape.
…lection Embedding a few thousand prompts takes long enough that doing it per process would dominate the cost of routing, and the result is identical every time, so it belongs in the artifact rather than the application's boot path. Vectors are only comparable within the model that produced them, so the builder writes a manifest recording the embedding model id and dimension beside the index. An index searched with the wrong model returns confident nonsense rather than an error, which is the failure the manifest exists to prevent. Flat scan rather than an approximate index: a few thousand vectors search exactly in well under a millisecond, and an approximate index would make classification non-deterministic for no gain at this size. The embedder is a one-method interface so models-router takes no dependency on an embedding backend; the build wires in whichever ModelJar the index is pinned to.
A query takes the label of its single nearest indexed prompt. Nearest-neighbour rather than a per-task centroid because a task is not one region of embedding space: 'write a Dockerfile' and 'fix this segfault' are both code and sit far apart, and their midpoint is neither. A query farther than the threshold classifies as null, and ModelRouter then routes on cost, latency and reliability with quality averaged across tasks. Guessing a task for an unfamiliar request would silently apply the wrong quality column. TaskIndex reports plain cosine rather than the collection's score. SimilarityFunction.COSINE scores as (1 + cosine) / 2 over [0, 1], so on that scale an unrelated query sits near 0.5 and a threshold picked as though it were cosine would admit almost everything. Converting at the boundary keeps the threshold on the same scale as the measurements it is derived from. A query whose width does not match the index is an error rather than a miss: width is the only mechanically detectable form of 'wrong embedding model', and the wrong model of the same width would otherwise return confident nonsense.
Build and evaluate live in one command because both must use the same embedding model; an index and an accuracy figure produced by different models say nothing about each other. evaluate exits non-zero below --min-accuracy, which is how CI gates a corpus or model change. It also prints per-task accuracy and the first misses, since an aggregate that drops usually drops in one task and the label tells you which.
Each model takes minutes to embed the corpus, so a run interrupted partway previously reported nothing at all.
The first bake-off put chat at 26% and summarization at 36% while code was at 100%. The cause was the corpus, not the models. Dolly's summarization and information_extraction rows carry their text in a separate context field — 100% of them — and taking the instruction alone leaves 'What is a dispersive prism?' labelled summarization. Extraction now appends the passage, since extraction is defined by having something to extract from. Summarization now comes from XSum wrapped in an explicit instruction, dropping dolly's category entirely: its instructions are question-answering over a context, which is a different task. Dolly's creative_writing also holds open-ended musings and even arithmetic word problems, so it is filtered to rows that ask for text to be written. OpenAssistant conversation roots are dropped for chat. They are arbitrary user requests spanning every task in this taxonomy — 'Make a story about Rick and morty', 'Go through my emails and provide a status update' — so labelling them chat taught the classifier that any request is chat. Chat now comes from dolly's open_qa, general_qa and brainstorming rows, which are genuinely open-ended questions carrying no context. Reasoning now spans six BIG-Bench Hard configs rather than boolean_expressions alone, which was a narrow synthetic slice.
The first bake-off on the corrected corpus showed the nearest wrong training neighbour sitting at distance 0.267 while the median correct one sits at 0.430. The two distributions overlap, so no threshold can separate right answers from wrong ones — any cut that removed mistakes would remove far more correct classifications. What a threshold can do is decline queries no task covers, so an out-of-domain request routes on cost and latency instead of being forced into whichever task happens to be least distant. This measures both distributions against the same index and sweeps candidate thresholds, so the shipped value is read off the trade rather than picked.
The index is a directory of memory-mapped files and a jar entry cannot be mapped, so it travels as a zip and is expanded once into a cache directory. The cache path carries a digest of the archive, so a jar upgrade that changes the index lands in a new directory instead of silently reusing the old one. Expansion goes into a scratch directory that is then moved into place, so a reader never observes a half-written index and two processes racing cannot interleave their writes. Entries whose resolved path escapes the destination are rejected: a crafted archive could otherwise write anywhere the process can reach.
…nifest Map.copyOf and the other immutable collections randomise iteration order per JVM run — three runs of the same input gave three different orders. The index derives its document ids from that order and the manifest now derives a corpus digest from it, so both would have differed on every build from identical input. The corpus is held in a LinkedHashMap wrapped as unmodifiable instead. The manifest records corpusSha256 so a build can tell an index was left behind by a corpus edit. A prompt count alone would miss one prompt being swapped for another, which is the edit most likely to be made by hand and forgotten.
Generated against the pinned oracle (6ea215d17) with mean pooling and L2 normalization, over the unchanged probe set. The artifact digest matches the ModelJars catalog entry exactly. The runtime cannot load this model yet — its GGUF declares general.architecture=gemma-embedding, which PureJavaBackend rejects — so this is groundwork rather than a qualification. Generating it now records the oracle version that was current when the vectors were taken; deferring it would mean regenerating against whatever llama.cpp happens to be checked out later.
TaskIndexCli returns an exit status rather than calling System.exit, matching EmbeddingEquivalenceCli: its evaluate subcommand is a gate, so the verdict has to travel out through the process status.
EmbeddingGemma is the strongest task classifier in the bake-off (90.8% vs the
next best), but nothing could run it: PureJavaBackend rejected the architecture
outright. It is an encoder, and the existing pass is causal throughout.
Bidirectional attention is not a mask swapped into the decoder loop. A causal
model walks tokens one at a time because position p only needs keys it has
already visited — that is what makes a KV cache work. Here every position needs
every other position's key at the same layer, so layer L cannot start until
L-1 has finished the whole sequence. That inverts the loop nesting, which is
why this is a separate pass rather than a flag on the existing one. llama.cpp
allocates no KV cache for these architectures for the same reason.
Read out of llama.cpp at the pinned oracle rather than inferred, because four
details would each have produced a model that runs and disagrees:
- LLAMA_SWA_TYPE_SYMMETRIC masks when |p1 - p0| > n_swa/2, so the window is
centred on the query and reaches 256 either side. The causal reading of the
same n_swa=512 would look back twice as far.
- the projection head runs after pooling, not per position
- f_attention_scale is 1/sqrt(n_embd_head_k), which happens to match what the
existing code already computes
- the SWA period is 6 with no metadata key to read it from
The dense head travels with the model: EmbeddingGemma's published vectors are
its output, so skipping it gives a vector of the right width pointing somewhere
else. Loaded as a unit — a half-present head is rejected rather than partly
applied — and pooling likewise moves under the model's control, with an explicit
pooling() call now refused rather than silently ignored.
Verified against llama.cpp 6ea215d17, and each guard checked by breaking it:
8 committed probes min cosine 0.99956 (floor 0.999)
940 tokens, ad hoc 0.99939
940 tokens, window disabled 0.86032
8 probes, forced causal 0.57266
That last figure is the reason for the LlamaForwardPass guard rejecting a
bidirectional config: causal weights run through the wrong pass return a
plausible unit vector, not an error.
The committed gate cannot cover the window — its longest probe is ~25 tokens
and the window only masks past position 256, so every probe would pass
unchanged with it removed. EncoderAttentionWindowTest pins the arithmetic
instead, and records why the gate is blind to it.
TaskIndexCli passed --pooling unconditionally, which an encoder now rejects: it reads its pooling from its own metadata and refuses to be told otherwise. Applies the flag only where the caller genuinely has to supply it.
EmbeddingGemma is trained with Matryoshka Representation Learning, so each prefix of its output is itself a usable embedding. Truncating shrinks the router index proportionally, which is worth measuring before deciding how the index ships. Named matryoshkaDimensions rather than dimensions on purpose. No GGUF metadata records whether a model was trained this way, so nothing can check it, and truncating a model that was not discards trained dimensions and degrades retrieval without failing. A caller who does not know the word will not reach for the method by accident. Truncation precedes normalization: a prefix of a unit vector is shorter than one, and cosine over an unrescaled prefix is not the similarity the model was trained to produce. Measured on the router corpus, 1881 prompts, 469 held out: 768 0.9083 5.13 MB zipped 256 0.8891 1.72 MB 128 0.8486 0.87 MB The aggregate understates what happens. Going 768 -> 256 loses nine prompts in total and eight of them are extraction, which falls 0.720 -> 0.560 while the other nine tasks stay flat or improve. Extraction prompts are the long context-bearing ones, and they are the ones that need the fine dimensions to separate from chat.
Two things needed to compare storage choices for the shipped router index. The quantizer is now written to the manifest and restored by TaskIndex.open. VectorCollectionBuilder defaults to NONE, so opening a quantized index without this would have handed the reader a config that disagrees with the bytes on disk. Indexes written before the key existed read as NONE, which is what they were. The embedding cache is the difference between a comparison that gets run and one that does not. Embedding is 25 minutes of a 25-minute build and everything after it takes milliseconds, so five storage variants would have meant two hours of recomputing identical vectors. Keyed by prompt text and scoped by model and width: a cache reused across models would answer confidently with another model's vectors, so a mismatched identity discards rather than mixes. Measured with them, 1881 prompts at 768 dims, total zipped and the quantized payload alone: NONE 0.9083 5.13 MB (vectors.bin 5.51 MB raw) SQ8 0.9083 6.28 MB quantized.bin 1.38 MB FP16 0.9083 7.67 MB quantized.bin 2.76 MB SQ4 0.9083 5.75 MB quantized.bin 0.70 MB BQ 0.9190 5.33 MB quantized.bin 0.20 MB Quantization adds a compressed copy beside the full-precision vectors rather than replacing them, so it makes the artifact bigger. Accuracy is identical across NONE/SQ8/FP16/SQ4 because FLAT rescores against full precision — the quantized copy never reaches the answer. See #28.
vectors 0.1.7 adds quantizedOnly, which drops the full-precision copy instead
of writing compressed codes beside it. Until now a quantizer made this index
bigger; now it makes it eight times smaller.
Rebuilt with SQ4 against the same 1881 prompts:
accuracy zipped
float32 0.9083 5.13 MB
SQ4 0.9019 0.64 MB
vectors.bin is zero bytes and quantized.bin is the index. 0.9019 matches the
standalone measurement of 4-bit codes exactly, so the read path in vectors
agrees with scoring them directly.
Per task the loss is diffuse rather than concentrated: reasoning 0.810 ->
0.762 and chat 0.840 -> 0.820, everything else unchanged including extraction
at 0.720. That is the opposite of Matryoshka truncation to 256 dims, which was
larger (1.72 MB), less accurate (0.8891), and took nearly all of its loss out
of extraction alone.
Also bumps vectorsVersion 0.1.5 -> 0.1.7 and relocks. Two versions in one
step, so the whole build was checked rather than just models-router.
Three changes that finish the router index as a shippable artifact. Encoder positions now run concurrently. Within a layer they are independent by construction — each reads the whole sequence's keys and values and writes only its own hidden state — so the only thing keeping them sequential was sharing one set of scratch buffers. One bundle per worker fixes that without touching any arithmetic. vectors' own parallel-matmul threshold never fired here: the widest projection is under a million elements, so the whole encode ran on one core however many were free. index build 1522 s -> 300 s (5.1x) 8-probe gate 1.8 s -> 0.9 s Bit-exact, not just close: the parallel build produces a quantized.bin with the same SHA-256 as the sequential one, and the equivalence gate still reports min cosine 0.9995565 against llama.cpp. Strided rather than blocked work assignment so a sliding window, whose cost varies with position, does not leave one worker holding the expensive rows. The 0.64 MB index is now a models-router resource, with tests that open the real archive rather than a fixture — a resource is easy to ship broken because nothing reads it until a user does. They pin the corpus digest, the prompt count and the quantizer, so a rebuild against a different corpus or a silent return to full precision both fail here. TaskIndexResourceTest previously asserted that extractTo fails because no index was packaged. That state no longer exists, so it now asserts the packaged one expands. MAX_TOTAL_BYTES drops from 512 MiB to 64 MiB, set against the real artifact rather than guessed before one existed.
The prompts are third-party text under a mix of source licences — MBPP, GSM8K, dolly, XSum, OPUS-100, ARC, BIG-Bench Hard, a Hermes function-calling set. Vendoring them here means a licence question about any one benchmark blocks a release of the whole library. They now live in integrallis/model-router-corpus, and what ships here is the derived index: embeddings, not prompts. Not a size decision. The tracked corpus was 617 KB against an 11 MB .git; the 1.1 GB under corpus/ was always the gitignored HuggingFace download cache. Nothing in the build depended on the corpus — the index is a committed resource — so the only reference was TaskExemplarsTest, which already skipped when the corpus was absent. It now looks for a sibling checkout and honours -Dmodels.router.corpus, so it still runs for anyone who has one. The index pins corpusSha256, so the two cannot drift apart unnoticed: an index built from an edited corpus fails BundledTaskIndexTest rather than shipping.
The classifier had a shape gap, not a volume gap. Every farmed summarization prompt carried an XSum passage and every extraction prompt a dolly context, so a request sent without a document embedded nowhere near them. Fifteen ordinary instruction-only queries — "Summarise this quarterly report in three bullet points", "Extract the invoice number and total from this receipt" — classified as tool-use, chat, creative, sql and code. None correctly. integrallis/model-router-corpus adds 60 instruction-only prompts covering that phrasing. All fifteen now classify correctly, including two that appear nowhere in the corpus, so the fix generalises. Held-out accuracy 0.9019 -> 0.9044, extraction 0.720 -> 0.741, reasoning 0.762 -> 0.786. The aggregate barely moves because the eval split hardly sampled the region that was broken — which is how a 90% classifier could misroute every bare instruction it was given. Index 0.64 -> 0.65 MB. BundledTaskIndexTest repins the corpus digest and prompt count, so index and corpus still cannot drift apart unnoticed.
Using the router began with writing out every model's price, latency and
per-task quality by hand. Nobody knows their model's time to first token, and
the figure is hardware-specific, so anything copied from a README is a
measurement of somebody else's machine — and the router would score on it with
complete confidence.
Every value already exists in the ModelJars catalog, selector-matched to the
current CPU, core count, JDK and vector width. Nothing was reading it.
ModelRouter router = ModelRouter.discoverLocal()
.policy(RoutingPolicy.BALANCED)
.build();
ModelCatalogProvider lives in models-api because the dependency runs the other
way: ModelJars depends on models, so an SPI it implements cannot live beside
the router without inverting the tiers. DiscoveredModel is deliberately not a
routing type — a catalog reports what a model is, and something downstream
decides what that is worth.
Three failure modes it refuses to paper over. No catalog installed logs how to
add one and returns an empty fleet, because routing between hosted models only
is legitimate. A catalog that throws is skipped with a warning rather than
taking down discovery, so a second working catalog still contributes. And a
model with no performance profile for this hardware is skipped loudly rather
than defaulted: inventing a latency is exactly the failure this replaces.
Hosted pricing is the other half and is not here yet — see #30.
Apple Intelligence is present because of the hardware, not because anyone installed it. Discovering it by default would make identical code route differently on a developer's Mac than in production, and that difference surfaces as unexplained behaviour rather than as a decision. So discoverLocal() leaves it out and discoverLocal(true) includes it. Implemented as an opt-in flag on the SPI rather than an Apple special case in the router: requiresOptIn() defaults to false, and Apple is simply the first provider to set it. Anything else environment-provided gets the same treatment for free. The provider measures rather than guesses. Apple publishes no throughput figures and the framework exposes no counters, so the options were to measure, invent, or report nothing and be silently skipped by CatalogDiscovery — which would have made the opt-in a no-op that looked like it worked. It times one short generation, once, behind the opt-in. Whole-call latency stands in for time to first token because this path does not stream; that overstates TTFT, biasing the router away from this model rather than towards it, which is the safe direction for an estimate. Availability is checked, not inferred from the platform: Apple Intelligence can be switched off, unavailable in a region, or still downloading, and macOS on Apple silicon is necessary but not sufficient. Tests run everywhere. The common case is a machine with no on-device model at all, where reporting nothing must be silent and cheap, and one test asserts the services registration because a missing file there is invisible until a user hits it.
Skipping was the wrong call. A model absent from the fleet cannot be chosen at all, and the usual reason it has no profile is that it was installed recently, not that it is unusable. The estimate is calibrated rather than invented wherever it can be. Local generation reads the whole weight file per token, so it is bounded by memory bandwidth and tokensPerSecond x sizeBytes is roughly constant across models on one machine. Any measured peer therefore fixes that constant for this hardware, and a model's size predicts its rate from there — a property of how inference works rather than a guess about a particular model. A measured 4 GB model at 40 tok/s implies about 20 for an 8 GB one, which is what the test pins. DiscoveredModel gains sizeBytes for this. It is the strongest predictor available and every catalog already knows it. Median across measured peers, not mean: one model benchmarked while the machine was busy would otherwise drag every estimate with it. With nothing measured and no size known there is no signal left, and the estimate falls back to deliberately pessimistic constants. Pessimism is the safe direction — an estimate that flatters a model wins latency-sensitive routing it has never demonstrated, while one that understates it merely loses ties to models that have actually been measured. Estimated models are named in a log line saying whether the estimate was calibrated or a fallback. An estimate that scores exactly like a measurement is the thing worth being able to find out about. A measurement is never overwritten by an estimate, and a catalog reporting an absurd size yields a valid low rate rather than a failed discovery.
0.3.0 is published and immutable, and this release carries the catalog SPI that ModelJars needs to implement a provider — it cannot compile against 0.3.0, which predates com.integrallis.models.api.catalog.
A global gitignore carrying *.zip kept the index out of every commit that added it. Local builds passed throughout because they read the file straight off disk, tracked or not; CI is the first place the difference between "on disk" and "in the repo" is visible, and it failed there immediately. BundledTaskIndexTest did its job — opening the real archive rather than a fixture is what turned an artifact that would have shipped empty into a failed build. The local run was the misleading one. Force-added, with a negation in models-router/.gitignore so the same global rule cannot silently drop it again.
verifyArchitectureBoundaries forbids the ModelJars group coordinate anywhere in models source, strings included — the task is written to avoid matching its own literal, so that breadth is deliberate rather than an oversight. The install hint carried the coordinate verbatim and tripped it. The guard is right and the message was wrong. A caller who needs a catalog can find the artifact from the project name; models spelling out its coordinate is the beginning of exactly the coupling the check exists to prevent. Also the reason my local runs missed this: verifyArchitectureBoundaries is not wired into check, so it only ran in CI's compliance job.
verifyReleaseMetadata holds sixteen files to the same version, and bumping gradle.properties alone left every one of them behind. It is not wired into check, so this only failed in CI — the same reason the architecture guard did. CHANGELOG gains a dated 0.3.1 section covering the encoder architecture, the pretrained classifier and its shipped index, the catalog SPI, Apple on-device discovery, Matryoshka truncation, the parallel encoder, and the vectors 0.1.7 bump. Also updated: Antora component and attributes, docs package manifests and lockfile, landing-page pill, native crate and its lockfile, notebook defaults, and the published coordinates quoted across four READMEs. Antora's vectors-version attribute moves to 0.1.7 as well, which was simply stale.
My version bump rewrote antora.yml's `version: 'current'` to '0.3.1'. That is Antora's component version, not a display string: it decides the published path, so the site moved from site/models/current/ to site/models/0.3.1/ and verifyGeneratedSite — which looks under current/ — found none of the 1272 Javadoc attachments it expected. Only display_version and the models-version attribute track the release.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Everything needed to ship the pretrained router classifier, plus the pieces that turned out to be
prerequisites.
EmbeddingGemma runs on our own runtime
EncoderForwardPass— a separate pass, not a flag on the decoder. A causal model walks tokens oneat a time because position p only needs keys it has already visited; an encoder needs every
position's key at the same layer, so layer L cannot start until L−1 has finished the whole sequence.
Inverted loop nesting, no KV cache, matching llama.cpp.
Read out of llama.cpp at the pinned oracle rather than inferred.
LLAMA_SWA_TYPE_SYMMETRICmaskswhen
|p1 − p0| > n_swa/2, so the window is centred on the query and reaches 256 either side — thecausal reading of the same
n_swa=512would look back twice as far.Each guard checked by breaking it:
That last row is why
LlamaForwardPassnow rejects a bidirectional config: encoder weights throughthe causal pass return a plausible unit vector, not an error.
Positions within a layer are independent, so the pass parallelises: 1522 s → 300 s for a full
index build, producing a
quantized.binwith the same SHA-256 as the sequential build.The index ships at 0.65 MB
Requires vectors 0.1.7, which adds
quantizedOnly— before it, a quantizer added a compressedcopy beside the full-precision vectors and made the artifact bigger.
Matryoshka truncation to 256 dims was measured and rejected: larger (1.72 MB), less accurate
(0.8891), and it took nearly all its loss out of extraction alone.
Corpus moved out, and a shape gap closed
Prompts now live in model-router-corpus so a
licence question about one benchmark cannot block a release. The index pins the corpus digest.
Every farmed summarization prompt carried an XSum passage, so a request sent without a document
embedded nowhere near them. Fifteen ordinary instruction-only queries classified as tool-use, chat,
creative, sql and code — none correctly. After 60 instruction-only prompts, all fifteen are,
including two that appear nowhere in the corpus. Held-out accuracy barely moves (0.9019 → 0.9044)
because the eval split hardly sampled the broken region, which is how a 90% classifier misrouted
every bare instruction it saw.
Models register themselves
ModelCatalogProviderinmodels-api— ModelJars depends on models, so an SPI it implements cannotlive beside the router without inverting the tiers.
Apple Intelligence is opt-in via a general
requiresOptIn()flag: it is present because of thehardware, so discovering it by default would make identical code route differently on a Mac than in
production.
Unprofiled models are estimated, not dropped. Local generation is memory-bandwidth bound, so
tokensPerSecond × sizeBytesis roughly constant on one machine — any measured peer calibrates therest. Falls back to pessimistic constants when nothing has been measured, and never overwrites a
real measurement.
Reviewer notes
AppleFoundationModelsCataloghas only run on Linux, where it correctly returns empty. Itsmeasurement path has never executed on Apple silicon — worth a look from someone with a Mac.
window only masks past position 256.
EncoderAttentionWindowTestpins the arithmetic instead;end-to-end coverage is feat(purejava): enable Q5_K batched prefill #25.