OPENNLP-1877: Static embedding engine for modern distilled embedding tables (opennlp-embeddings) - #1152
Draft
krickert wants to merge 51 commits into
Draft
OPENNLP-1877: Static embedding engine for modern distilled embedding tables (opennlp-embeddings)#1152krickert wants to merge 51 commits into
krickert wants to merge 51 commits into
Conversation
mawiesne
marked this pull request as draft
July 8, 2026 04:32
krickert
force-pushed
the
static-embeddings
branch
2 times, most recently
from
July 8, 2026 14:44
2085efd to
7f6f67f
Compare
krickert
marked this pull request as ready for review
July 8, 2026 15:00
krickert
force-pushed
the
static-embeddings
branch
2 times, most recently
from
July 8, 2026 19:05
e60bdba to
7f6f67f
Compare
Contributor
|
This probably requires an update to the docs? Probably doesn't need to be extensive - just a description of when to use it and how. |
jzonthemtn
reviewed
Jul 9, 2026
Contributor
Author
I'll do that now |
krickert
force-pushed
the
static-embeddings
branch
from
July 9, 2026 20:02
7f6f67f to
e60bdba
Compare
Contributor
Author
|
Done in 76d58c5: a Static Embeddings chapter in the Dev Manual (when to use a static table over a contextual model, both load calls, thread safety, the no-bundled-model license note) plus a module README with the same ground covered for people landing in the source tree. |
Contributor
|
Hi @krickert - as mentioned on Slack I currently dont have the time for a manual review but I just let Fable do a comprehensive review on this PR. Here is the result: Blocking / should be addressed before merge:
Minor:
Human review will follow. |
Contributor
Author
|
Thanks. Addressed in 9dd6ff7, one pushback on the marker question: Blocking:
Minor:
|
krickert
marked this pull request as draft
July 11, 2026 13:22
Contributor
Author
|
Moved this to draft - interfacing to match Onnx impl. Onnx's impl is better for accuracy, this one is for speed. Update incoming. |
krickert
force-pushed
the
static-embeddings
branch
from
July 13, 2026 00:51
8230a2c to
6c216b7
Compare
krickert
force-pushed
the
static-embeddings
branch
from
July 13, 2026 08:18
4f7a84e to
6184021
Compare
Contributor
Author
|
Before this leaves draft: re-verify the two embed() throughput rows in the README performance table against the original JMH runs. The surrounding prose was reworked, but the numbers were deliberately left untouched pending that check. |
…ighbor scan, thread-safety hardening
Two analogy() bugs fixed. Passing equal terms crashed with
IllegalArgumentException("duplicate element") from Set.of; and the
exclusion compared raw input strings against vocabulary tokens, so on
an uncased model analogy("Man", "King", "Woman", k) handed "king"
straight back as a result. Exclusion now folds the terms through the
model's own tokenizer and excludes the resulting vocabulary rows,
which makes it case- and accent-consistent with embed() and tolerant
of equal or multiword terms. Both are pinned by new tests.
Nearest-neighbor scan reworked around three observations: per-row L2
norms are constants of the model, so they are precomputed at load
instead of recomputed (with a sqrt) for every row on every query; the
top-K selection now uses a bounded min-heap over primitive parallel
arrays instead of materializing and fully sorting one record per
vocabulary row per query; and the special-token check is a
precomputed boolean mask instead of per-row string hashing. The dot
loop uses four accumulators because the JIT must not reorder
floating-point additions and so cannot unroll the reduction itself.
The zero-norm-row NaN guard is preserved and now has its own test.
embed() drops an OptionalInt allocation per token (primitive -1
sentinel) and hoists the weight branch out of the accumulation loop.
Forked JMH, same configuration and fixture as the recorded baseline:
embed: 999,222 -> 1,041,654 ops/s (+4.2%)
mostSimilarTop10: 3,292 -> 9,173 ops/s (2.79x)
Thread safety reviewed and hardened: @threadsafe on
StaticEmbeddingModel, SafetensorsFile, and WordPieceVocabulary; the
class javadoc now documents why the one piece of global mutable state
in the tokenizer chain (WhitespaceTokenizer.INSTANCE's keepNewLines
flag) cannot affect results, since BERT basic tokenization replaces
all whitespace with plain spaces before that split runs; and a new
concurrency test runs 8 threads against one shared instance comparing
every result to the single-threaded reference.
Replaces the whole-file byte[] (capped at 2 GB by Java's int-indexed arrays, and failing as an opaque OutOfMemoryError beyond it) with positional FileChannel reads: the header is read eagerly, tensor data streams straight into the caller's float[] through a reused 1 MB chunk. File size is now unlimited; the remaining ceiling is per decoded tensor (a float[] holds at most ~2.1 billion elements) and is checked with a clear message. Peak load memory drops since file bytes and the decoded array no longer coexist. A file truncated between read() and readFloat32() fails loud instead of returning partial data.
…rbage Writing the direct tests surfaced one gap: parseTop stopped at the closing brace and silently ignored anything after it. Trailing whitespace stays legal (writers space-pad the header to align the data section), any other trailing content now fails loud.
… the model configs Reads normalize from config.json and do_lower_case from tokenizer_config.json (field names verified against published Model2Vec-family releases), so callers no longer have to know a model's switches to load it. The JSON scanning primitives move from SafetensorsHeaderParser into a shared JsonCursor; FlatJsonFields reads just the top-level booleans and skips everything else structurally. A strip_accents explicitly set against do_lower_case is rejected rather than silently mis-tokenized: the single lower-case switch follows the BERT convention of stripping accents exactly when lower-casing, and the error points at the explicit overload for deliberate choices.
Covers when to use a static table over a contextual model, the one-call model-directory load and the explicit overload, thread safety, the no-bundled-model license posture, and the safetensors reader's guarantees.
…d load switches, distribution wiring TensorInfo copies its shape on construction and access and compares by value, so callers cannot corrupt validated metadata and equality behaves like exported API should; elementCount() detects overflow from crafted headers. The load and read entry points now declare IOException like the rest of the project's resource loaders instead of wrapping in UncheckedIOException. The two adjacent boolean load flags became the Casing and Normalization enums, so the switches cannot be swapped silently; the directory-based load maps the model's configuration onto them and the manual, README, tests, and benchmark use the new signature. The binary distribution now ships the module: opennlp-embeddings added to the distr dependencies and its apidocs to the assembly, matching the other extension modules. JSON parsing tightened to the fail-loud contract: unicode escapes must be four hex digits (Integer.parseInt also accepted signed input) and numbers in skipped fields are held to the JSON grammar. WordPieceVocabulary gains its own test class covering the line-number ids, the duplicate rejection, the -1 sentinel, and reverse-lookup bounds, which now fail loud instead of leaking an index exception. The safetensors fixture writer that was copied across the tests is now a shared test utility, the stale cross-module javadoc reference is gone, the element-count versus byte-range guard has its negative test, and else/catch placement matches the project style.
…ement The distribution dependency added for the review round needs the managed version like the other extension modules; without it the distr module fails to resolve the artifact.
Matches the casing of the existing public WordpieceTokenizer.
TextEmbedder in opennlp-api is the text-level embedding contract: embed(CharSequence), embedAll (default loops one at a time; runtimes that batch efficiently should override), and dimension(). It is the text-level counterpart of the word-level WordVectorTable, and the javadoc states the layer difference. StaticEmbeddingModel implements it with a thin CharSequence overload; the existing embed(String) hot path is untouched. SentenceVectorsDL implements it as the contextual tier: getVectors stays the primary entry point unchanged, embed adapts it (OrtException wrapped unchecked, the seam is runtime-neutral), and dimension() reads the model's declared output metadata with a one-time probe fallback for dynamic shapes. A batched embedAll override remains open as a follow-up. The adapter test drives a real ONNX session: a committed 373-byte model (generation script alongside) computes token_id times a fixed weight row, so every expected vector is hand-computable. The gitignored *.onnx pattern is force-added for this one bundled fixture; rat-excludes covers the binary.
WordpiecePipeline is a module-internal copy of the five-stage BERT pipeline over WordpieceTokenizer, pinned by the reference token sequences. StaticEmbeddingModel uses it instead of a shared class, so the hot path is unchanged operation for operation and the module no longer depends on any tokenizer being reworked elsewhere. The SentenceVectorsDL class javadoc is reduced to the factual contract.
SentenceVectorsDLEmbedderTest resolved the ONNX model with new File(url.toURI()), which only works when the resource sits on the filesystem. When the test runs from the opennlp-dl test-jar (as it does in opennlp-dl-gpu) the resource URI points inside a jar and is not hierarchical, so new File(uri) threw IllegalArgumentException and failed the opennlp-dl-gpu build. Copy the model out of the classpath into the JUnit temp dir, matching the existing pattern in LoadVocabTest.
… align the Dev Manual chapter
… TextEmbedder and both impls
…er seam Static embedding tables distilled from SentencePiece teachers (bge-m3 and the XLM-RoBERTa family) now load and embed, multilingual vectors included. The module's own WordPiece pipeline is replaced by the SubwordTokenizer seam: WordpieceEncoder from opennlp-api for WordPiece models, the pure-JVM SentencePieceTokenizer from opennlp-subword for SentencePiece models. Matrix rows are resolved by piece string, never by tokenizer id, because the trained .model file and the matrix vocabulary routinely order and offset ids differently; a load-time sweep verifies every poolable piece has a row, so a wrong file pairing fails loud at load, not at query time. The SentencePiece row order comes from the Unigram model.vocab list of tokenizer.json (with added_tokens overlaid), read by a purpose-built parser in the package's existing JsonCursor style. The directory loader detects the tokenizer family from the files present. Verified for exact output parity against the reference implementation on a real distilled multilingual model, including CJK input.
Applies the review conventions from the OPENNLP-1869 review to the embeddings module: private constructors and helpers gain javadoc, the TensorInfo record and FlatJsonFields validate their arguments at the boundary, and the hot-path and test-history commentary shrinks to what the code does.
…ories The benchmark took only a synthetic fixture. A modelDir parameter now selects the table: the default 'synthetic' keeps the offline fixture, and passing model directories (-p modelDir=dirA,dirB) benchmarks real tables and reports one row each, which is how the potion-vs-bge-m3 comparison is produced.
…k numbers TRAINING.md walks through distilling a table from a sentence-transformer teacher, with the multilingual bge-m3 SentencePiece model as the worked example: dimension guidance, assembling the model directory, loading and verifying parity in the JVM, and the WordPiece variant. The README performance section carries real potion-vs-bge-m3 throughput split into the embed and nearest-neighbor cost drivers.
Model2Vec mean-pools content pieces and never frames, so it removes [CLS]/[SEP] from the distilled vocabulary, keeping only [PAD]/[UNK]. Such tables could not load because the WordPiece encoder requires the frame tokens. The loader now caches the absent frame tokens onto the unknown row: the encoder still frames, and pooling skips the frame by id exactly as before, so which pieces are pooled does not change. [PAD] and [MASK] join the neighbor-exclusion set, since a distilled table keeps rows for them that text never tokenizes to.
…led directories A Model2Vec distillation writes model.safetensors, tokenizer.json, and config.json, but not the vocab.txt and tokenizer_config.json a WordPiece model needs, nor the trained SentencePiece .model file. AssembleModel completes a WordPiece directory by deriving the two missing files from tokenizer.json (the vocabulary in id order, the casing from the normalizer's lowercase flag) and reports the SentencePiece .model file it cannot fabricate, then loads the result to verify it. Wired as its own opennlp-embeddings command with a launcher script and distribution entry, mirroring the spellcheck module.
scripts/distill_bge_m3.py is the runnable form of the TRAINING.md worked example. scripts/parity holds the reproducible comparison against the model2vec Python reference: the same model and the same multilingual sentences on both sides, the two vector sets checked against each other, and both single-thread throughputs measured with the same fixed-duration methodology. A run passes only when the vectors agree within float tolerance, so the two speeds it prints are for implementations producing the same answer.
…eplace unsupported doc claims
Add StaticEmbeddingUsageExampleTest asserting the load-and-query workflow and point the embeddings manual section at it.
…eview conventions
Distillation, so producing a table no longer needs a Python environment:
- Add ModelDistiller, reproducing Model2Vec in Java: clean the teacher's vocabulary,
run every surviving token through the teacher's ONNX graph as [bos, token, eos] and
mean-pool the last hidden states, project onto the top principal components, then
scale each row by its Zipf weight sif / (sif + p).
- Add OnnxTeacherEncoder for that forward pass, feeding the graph exactly the inputs it
declares (input_ids, attention_mask, plus a zero token_type_ids for the BERT-family
graphs that ask for one), and declare the onnxruntime dependency in the module at the
root-managed version, the same engine opennlp-dl already runs.
- Add TeacherTokenizer, which decides which vocabulary rows survive into the table and
rewrites tokenizer.json to describe it, copying every field it does not change byte
for byte so the result stays a faithful fast-tokenizer description.
- Add RandomizedPca (Halko, Martinsson, Tropp), because a dense SVD of a 250k-row
multilingual matrix is not practical in pure Java; component signs are fixed the way
scikit-learn's svd_flip fixes them, so two runs are comparable rather than mirrored.
- Add SafetensorsWriter, the write side of the format SafetensorsFile reads, streaming
the matrix in chunks so its overhead beyond the caller's array is constant.
- Add HuggingFaceModelCache so -teacher accepts a hub id and its files download once
into a local cache; an optional file the repository does not have (a WordPiece
teacher has no SentencePiece model) is reported absent, not an error.
- Register the DistillModel tool in the module CLI. It ends by assembling and verifying
its own output directory, so a run that prints a summary is a directory that loads.
Documentation:
- Rewrite TRAINING.md around the DistillModel command instead of the Python uv and
model2vec setup, keep the Python flow only as the parity reference, and record that
two tables distilled independently from one teacher agree on their pairwise geometry
but not axis by axis.
- Point the README "Getting a model" section at DistillModel as an alternative to
downloading an already released table.
Folded duplication:
- Move firstRegularFile into ModelFileNames, next to the name lists it scans;
StaticEmbeddingModel and ModelAssembler each carried a copy.
- Fold the boolean and string readers of FlatJsonFields onto one top-level walker with
a ValueReader seam, so the object grammar, the duplicate-field check, and the
trailing-content check exist once.
- Collapse the two identical file checks in EmbeddingVocabulary into requireRegularFile.
- Reduce SentenceVectorsDL.declaredOutputDimension to the single first-output read the
loop actually performed, since every path returned on the first output anyway.
Javadoc and commentary:
- Document JsonCursor.consumeLiteral, requireEnd, and the new position(), and the two
writer overloads in SafetensorsTestFiles.
- Fix TensorInfo's stale link to SafetensorsFile.readFloat32, which is readFloats now,
and move its element-count description into the {@return} form.
- Describe Neighbor.token as one subword piece of the model's tokenizer rather than a
WordPiece, now that SentencePiece tables load too.
- State on TextEmbedder that thread safety is implementation specific, which is what
the two implementations actually promise, and mark up null in its tags.
- Drop the commentary that narrates history rather than behavior: the benchmark's
design-doc justification, the "before the fix" and "used to crash" notes in the
similarity tests, and the "untouched by the interface adoption" note in the DL test.
Tests:
- Add EmbeddingTestFixtures holding the king/queen/man/woman analogy table and the JSON
string quoter, and point the similarity, concurrency, SentencePiece, usage-example,
and assembler tests at it instead of their own copies of each.
- Build the size-mismatch and zero-vector fixtures with SafetensorsTestFiles rather
than hand-rolling the header and the little-endian payload in the test.
- Pin FlatJsonFields.topLevelString: escapes, absent versus explicit null, nested names
not matching, non-string values, duplicates, and null arguments.
- Add ModelDistillerTest, RandomizedPcaTest, and TeacherTokenizerTest over the
distiller's pure pieces, and drop a duplicate IOException import in
FlatJsonFieldsTest.
Build:
- Sort opennlp-spellcheck before opennlp-subword in the extensions module list.
…urface A review of the model distiller, which had not been reviewed before, found four defects that produce a wrong result rather than an error, and a set of classes with no tests at all. Correctness: - RandomizedPca floored the CholeskyQR jitter at an absolute value, so the decomposition was not scale invariant. Scaling the input down by 1e-6 turned an exact rank-6 recovery into an explained variance of 0.0017 with pairwise dot products wrong by three orders of magnitude, while still returning a plausible table. The jitter is now relative to the Gram trace. - RandomizedPca returned NaN instead of failing when the centered matrix was exactly zero, and a single non-finite input value poisoned its column mean and turned the whole result NaN. Both are now rejected. - The distiller guard named nanToZero tested only Float.isNaN, so an infinity from the teacher passed through into the PCA. Renamed to nonFiniteToZero and switched to Float.isFinite, with the divergence from numpy nan_to_num noted. - The PCA-skip branch kept the un-reduced matrix but left the requested dimension as the row stride, so the Zipf loop scaled the wrong cells and the writer rejected the array. Reachable with any vocabulary smaller than pcaDims, for example 100 tokens at the default 256. Tests, all previously absent: - SafetensorsWriter had no test file and its one case wrote six floats, so the 1 MiB streaming loop never crossed a chunk boundary. Added a 400x1024 round trip along with the header layout, alignment and reject paths. - ModelDistillerTest never called distill, leaving every argument check unverified. It now covers all five reject paths. - Added first tests for ModelFileNames, HuggingFaceModelCache, OnnxTeacherEncoder and the module CLI. - RandomizedPcaTest pins scale invariance, the non-finite rejection and determinism across pool sizes. Also corrected the TRAINING.md WordPiece section, which told the reader to run AssembleModel on output the distiller has already assembled. Module tests go from 156 to 273, none skipped.
The distiller fetched teacher weights from a moving ref and executed the ONNX graph without checking anything. DownloadUtil has always refused a model whose sha512 sidecar it cannot read, so this brings the hub path to the same posture. - Resolve the ref to a commit once, then request every file at that sha, so a force-push midway cannot mix two revisions into one cache directory. A teacher may now name a revision as org/model@revision. - Verify every file against the digest the hub publishes in x-linked-etag, choosing the algorithm by hex length: 40 is the git blob SHA-1 over "blob <len>\0" and the bytes, 64 is the SHA-256 of the content. The digest is computed over what was written to disk. A mismatch deletes the partial file and throws, in DownloadUtil's wording. - Refuse a file whose digest the hub does not publish, rather than accept something unverifiable. This is the point of the change. - Record the resolved commit in .opennlp-revision, which also marks the directory as a complete snapshot, and carry it into the distilled model's config.json as teacher_revision so a table can name the teacher it came from. - Make the cache directory name injective. It replaced '/', '.' and '@' without distinguishing them, so acme/model@v1, acme/model.v1 and acme/model_v1 shared one directory, and the cached path answers without contacting the hub, which would have served one teacher another's files. Headers are read by walking HttpResponse.previousResponse(): with Redirect.NORMAL the final CDN response carries neither header, while the redirecting hub response carries both. A test fails if that walk is removed. Tests need no network. A loopback HttpServer serves canned replies and records the requests made, covering both digest forms, a corrupted body, a missing and six malformed etags, optional and required 404s, the redirect path, the zero request cached path, revision pinning and the directory collision. Module tests go from 273 to 310, none skipped.
krickert
force-pushed
the
static-embeddings
branch
from
July 31, 2026 10:38
558ebd3 to
ce7090c
Compare
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
3 participants
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.
Adds a new
opennlp-extensions/opennlp-embeddingsmodule: a pure-JVM engine for modern static embedding tables (Model2Vec-family distillations, the 2024/2025 successors to word2vec/GloVe: same flat per-token table shape, sentence-transformer semantics, inference is pure lookup).Positioning: this is the speed tier of text embedding, not a replacement for anything. The ONNX
SentenceVectorsDLin opennlp-dl runs the actual transformer and produces contextual vectors, so this isn't made to be a replacement as that's the OOTB accuracy tier.This module trades that context for a static table: no native runtime, no GPU, and throughput in the hundreds of thousands of texts per second per core. Different pipeline stages want different points on that curve. Matches the philosophy of this style of embedding.
What's in it:
TextEmbedder(new interface inopennlp-api, packageopennlp.tools.embeddings): the text-level embedding contract.embed(CharSequence),embedAll(List)(default implementation loops; runtimes that batch efficiently should override), anddimension(). It is the text-level counterpart of the existing word-levelWordVectorTable, and the javadoc states that layer difference explicitly.StaticEmbeddingModelis the first implementation.SentenceVectorsDLis the natural second one: adopting the interface is purely additive (its existing constructors andgetVectorsare untouched), and its ONNX runtime is exactly what the overridableembedAllbatch method exists for. That adoption is a separate discussion, not part of this PR.SafetensorsFile: reads the safetensors format with a purpose-built cursor parser for the JSON header (no third-party JSON dependency; decodes F32, F16, and BF16, widening the 16-bit types to float). safetensors carries no executable content, unlike pickle-based checkpoints, so loading is safe by construction. The embedding matrix is auto-detected as the single 2-D float tensor, failing loud and listing candidates on ambiguity rather than guessing a key-name convention.WordpieceVocabulary: BERT-stylevocab.txt, line number = embedding row id. (Named to match the existingWordpieceTokenizercasing.)StaticEmbeddingModel: embeds through the existingBertTokenizer/WordpieceTokenizer, reused unchanged. The pooling formula is verified against the Moust reference implementations:[CLS]/[SEP]never pooled, unknown tokens dropped from sum and denominator, optional per-rowweightstensor, token-count denominator, epsilon-floored normalization.similarity,mostSimilar(bounded top-K over precomputed row norms),analogy(inputterms excluded by folding through the mode
Posture: code only, bring your own tab fetched at build or run time, no new dependencies.
Thread safety: immutable,
@ThreadSafe, with an 8-thread concurrency test comparing every result against thesingle-threaded reference.
Measured (JMH, opt-in
jmhprofile matching opennlp-runtime's pattern; fixture at real published-table scale, 29,528 x 256):embed~766k short sentences/s on one core (1.04M ops/s of 5 sentences at 32 threads); full-vocabulary top-10 scan 649/s per core, ~9.2k/s at 32 threads.Verification: opennlp-embeddings 43/0 plus the new interface contract test,
mvn verifygreen including checkstyle and forbiddenapis.Follow-ups (deliberately out of this PR):
TextEmbedderwith a real batchedembedAll(pending discussion with its author), the gRPC backend in opennlp-sandbox, a concurrent-load comparison against a Python baseline, an ANN index formostSimilar, and bundled-default-model license diligence.https://issues.apache.org/jira/browse/OPEN-1877