Repository navigation
OPENNLP-1877: Static embedding engine for modern distilled embedding tables (opennlp-embeddings) - #1152
OPENNLP-1877: Static embedding engine for modern distilled embedding tables (opennlp-embeddings)#1152krickert wants to merge 120 commits into
Conversation
2085efd to
7f6f67f
Compare
e60bdba to
7f6f67f
Compare
|
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. |
I'll do that now |
7f6f67f to
e60bdba
Compare
|
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. |
|
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. |
|
Thanks. Addressed in 9dd6ff7, one pushback on the marker question: Blocking:
Minor:
|
|
Moved this to draft - interfacing to match Onnx impl. Onnx's impl is better for accuracy, this one is for speed. Update incoming. |
8230a2c to
6c216b7
Compare
4f7a84e to
6184021
Compare
|
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. |
…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.
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.
Add StaticEmbeddingUsageExampleTest asserting the load-and-query workflow and point the embeddings manual section at it.
…eview conventions
…nd edge tests Annotate ModelAssembler, ModelDistiller, and TextEmbedder with @experimental like the other public embeddings types and extend the manual's note. Link Model2Vec and the randomized SVD paper at first mention. Make the usage example test mirror the manual's directory-load listing. Cover whitespace-only and supplementary-plane input in the embed path. Document the Checksum constructor.
The chapter's second programlisting (the explicit load and loadSentencePiece overloads) had no test pinning it. testExplicitOverloads loads both fixture layouts exactly as the listing shows, backed by a new minimal SentencePiece directory fixture, and the chapter now cites the test the way it already cites the directory-load listing.
…test The headline use case, ranking documents against a query by cosine similarity, existed only in the module README. The chapter gains a Semantic Search section whose listing scores each document with the public similarity method and sorts by descending score, and the new StaticEmbeddingSearchExampleTest asserts the exact ranking on the analogy fixture's known geometry.
…d pieces A distillation can now take a term list, a learned corpus vocabulary of whole words and multi-word phrases. Each term is segmented by the teacher's own tokenizer, encoded through the teacher as one sequence, and appended to the table after the subword rows, recorded as terms.txt in the model directory. The same PCA and Zipf pipeline spans all rows, and a term equal to a surviving vocabulary token is dropped as a duplicate row. At embed time the model matches text against its terms greedily longest-first over case-folded word runs (StringUtil.toLowerCase, one code point to one code point) and pools a matched term's single row in place of its words' subword pieces; text between matches tokenizes as before, and a model without a terms file embeds exactly as it did. Terms are neighbor candidates in mostSimilar like any token. The DistillModel tool gains -terms, documented in the manual. TeacherTokenizer's two Patterns (the unused-token filter and the template splitter) are replaced with cursor scans along the way.
Add a general subword tokenizer contract with original-text UTF-16 offsets, plus a dependency-free WordPiece implementation and BERT compatibility layer. Document the API and cover vocabulary validation, reference sequences, Unicode, and offset behavior. Red evidence: - A supplementary-plane word at 100 code points was rejected because UTF-16 code units were counted. - Negative piece ids and empty vocabulary pieces were accepted.
Red evidence on the prior implementation: Greek final sigma and Unicode control categories did not match the BERT reference sequence. The sibling tokenizers also disagreed on the 100-code-point limit and model callers rebuilt non-contiguous ids.
A relative path with a '..' segment is ambiguous as a teacher reference: Windows collapses '..' lexically without checking that the segment before it exists, so a misspelled hub id such as BAAI/.. silently names the working directory there while POSIX reports it as nonexistent. This is the windows-latest CI failure in HuggingFaceModelCacheTest testMalformedTeacherReferenceIsRejectedBeforeAnyRequest[14]. Red evidence on Linux: the new test fails against the current code with 'Expected java.lang.IllegalArgumentException to be thrown, but nothing was thrown', because resolve(src/..) returns the module directory.
Windows collapses '..' lexically without checking that the segment before it exists, so resolve(BAAI/..) treated the working directory as the teacher directory there while POSIX rejects the reference; this was the windows-latest CI failure in HuggingFaceModelCacheTest testMalformedTeacherReferenceIsRejectedBeforeAnyRequest[14]. The local directory shortcut now refuses relative paths containing a '..' segment, so the rejection is identical on every platform. Absolute paths and clean relative paths are unaffected, and a local directory still wins over the hub. Also rework the teacher-name escaping test, which created a directory named teacher"quoted: the quote is illegal in Windows file names, so the test could not even start there. The name-to-config plumbing is now covered with a plain directory name, and the JSON escaper behind tokenizer_name is exercised directly with quotes, backslashes, and control characters.
Three regression tests failed before enforcing the vector length across batches. The encoder now rejects inconsistent output before distillation writes model files. Add an original ONNX lookup fixture, executable distillation examples and matching manual guidance. Clean embeddings, dependency and manual verification passed; all 496 embeddings tests passed.
Correct embedding batch accounting and incremental JMH annotation processing. Add loading and allocation benchmarks, fixture tests and usage instructions.
Select model factories by provider ID, retain ONNX as the default, and allow independently owned models through the closeable TextEmbedder API. Keep model initialization out of provider discovery and preserve the existing SentenceVectorsDL inference path. Tests cover multiple providers, missing and duplicate IDs, independent model lifetimes, and ONNX output parity. The root model-path regression first failed with NullPointerException; it now rejects the invalid path. The SPI tests initially failed to compile before the API was introduced. Validation: targeted Maven verify with dependencies, apache-rat:check, and final path-validation regression tests passed.
Move the backend-free implementation into opennlp-embeddings-core and ship ONNX and CLI support in separate artifacts. Keep opennlp-embeddings as the default bundle so consumers can exclude opennlp-embeddings-onnx without changing the core API. The binary distribution retains the CLI. Discover static text models and ONNX teacher encoders through their SPIs. Preserve prepared-token distillation and pooling behavior, and document provider selection and dependency exclusions. The dependency-isolation test failed against the original packaging because ONNX was present. It now passes without ONNX or CLI, alongside existing static inference, distillation, and CLI tests. Maven verify, apache-rat:check, and final path-validation regression tests passed.
…olution The test called a static SentenceVectorsDL.encode(String, WordpieceEncoder) method that no longer exists; encoding now goes through AbstractDL's shared encodeTokens(). Coverage for that path lives in SentenceVectorsDLEmbedderTest and CreateTokenizerTest. The file was an artifact of resolving duplicate commits during the rebase onto f09b7f0 and is absent from the pre-rebase branch tip.
CreateTokenizerTest, DocumentCategorizerDLTest and NameFinderDLTest kept an unresolved second copy of tests calling static encode(String, encoder) overloads that predate the shared AbstractDL.encodeTokens() path; those overloads no longer exist. WordpieceEncoderTest kept a duplicated block around one stray piece.span() assertion. rerere misapplied an old resolution during the final conflicted commit of the rebase and left literal conflict markers plus now-dead code in the tree; this restores each file to exactly the content already verified against the pre-rebase branch tip.
The distillation script and parity harness live under dev/embeddings/ in the final layout; this stray paragraph still pointed at the old scripts/ location that the module split away.
…Distiller A merge during the rebase onto f09b7f0 left the same import listed twice with a blank line between them, which checkstyle's CustomImportOrder check rejects as extra separation within one import group.
Repeated conflict resolution on this file during the rebase onto f09b7f0 left a second, partially-overlapping copy of six test methods appended after the file's real end, three of which duplicated method signatures already present earlier in the class and failed to compile. The duplicated block, including the three methods that were not literal redefinitions, does not exist in the pre-rebase branch tip; restoring the file to that verified content fixes the build.
tokenizer.xml had accumulated four copies of the tools.tokenizer.subword section from repeated conflict resolution during the rebase, which FOP's PDF generation rejects because docbook block IDs must be unique. Restored the file to the pre-rebase branch tip and reinserted the Independent whitespace tokenizers section that main added on its own (OPENNLP-1947) and that the pre-rebase tip, based on an older main, did not have yet. embeddings.xml had a duplicated Semantic Search section from the same pattern; embeddings.xml does not exist on main, so it is restored to the pre-rebase tip as is.
6e8d843 to
b4339d4
Compare
…d priority Brings the provider SPI in line with the OPENNLP-1937 review. The API copies of TextEmbedderProvider, TextEmbedderProviders, their test, the ONNX text embedder provider and its tests are the same files as on the 1937 branch: no DEFAULT_PROVIDER constant and no getDefault(), a select(model, options) that ranks the installed, available and supporting providers by priority, get(name) ranked the same way under one name, installed() that skips a broken registration with a warning, and the system property opennlp.embedder.provider to pin a provider. StaticTextEmbedderProvider supports a directory with model.safetensors and config.json and an empty option map, so a static model directory routes to it and an ONNX file routes to the ONNX provider. The teacher encoder SPI in opennlp-embeddings-core had the same "onnx" default constant in a module that must not depend on ONNX, so it takes the same shape: TeacherEncoderProvider adds priority(), isAvailable() and supports(model); TeacherEncoderProviders offers installed(), get(name) and select(model), pinned by the system property opennlp.embeddings.teacher.provider. OnnxTeacherEncoderProvider supports *.onnx files and is available when the ONNX Runtime classes load. ModelDistiller selects the teacher provider for the teacher's ONNX file by default; the overload with a provider name pins one. Tests: StaticTextEmbedderProviderTest and OnnxTeacherEncoderProviderTest cover the capability rules, availability and registration; TeacherEncoderProvidersTest covers name lookup, priority replacement, ties, capability routing, unavailable providers, the pin property and skipped registrations; the bundle and isolation tests check the routing with and without the ONNX module. The README and the embeddings chapter describe selection and the two pin properties.
apache#1213 to apache#1215 sit on the apache#1152 head with the split module layout, the helper was regenerated without relocations, and the recurring resolution list shrinks to the AD and CoNLL-U fixes.
|
This shall be a PR / feature for |
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