Skip to content

OPENNLP-1910: Add bounded in-memory vector indexes for static embeddings - #1214

Draft
krickert wants to merge 148 commits into
apache:mainfrom
ai-pipestream:OPENNLP-1910-in-memory-vector-index
Draft

krickert wants to merge 148 commits into
apache:mainfrom
ai-pipestream:OPENNLP-1910-in-memory-vector-index

Conversation

@krickert

Copy link
Copy Markdown
Contributor

Summary

Depends on #1213, which in turn depends on #1152. Please review this after its parents.

This introduces a bounded, dependency-free VectorIndex contract for static embedding similarity search inside one JVM:

  • a build-once, read-many add, freeze, and top-k lifecycle
  • caller-provided identifiers and fixed vector dimensionality
  • common validation, cosine-score semantics, and deterministic result ordering
  • an exact full-precision flat index
  • a TurboQuant-backed index using OPENNLP-1895
  • concurrent queries after a frozen index has been safely published
  • parameterized contract tests shared by both shipped implementations
  • manual and executable usage-example coverage

The intended scope is document-local or bounded-corpus search. This is not a distributed search system, a mutable disk index, or a replacement for Lucene, Solr, or OpenSearch.

Verification

The branch passed its focused index contract and usage tests. It also passed the complete 15-project compilation, packaging, Checkstyle, forbidden-API, Javadoc, and RAT reactor gate with opennlp.forkCount=1.

The unrestricted upstream reactor test command additionally encountered unrelated model-download failures in inherited runtime and formats tests. Those failures could not obtain public tokenizer and sentence models; 1,682 runtime unit tests passed before the download-dependent failures.

JIRA

https://issues.apache.org/jira/browse/OPENNLP-1910

krickert added a commit that referenced this pull request Aug 16, 2026
krickert added a commit to ai-pipestream/opennlp that referenced this pull request Aug 22, 2026
krickert added a commit that referenced this pull request Aug 30, 2026
krickert and others added 26 commits September 16, 2026 02:13
…with exact original-text spans

New opennlp-extensions module implementing SentencePiece model inference
without native code: the ModelProto reader, the model-embedded normalizer
(precompiled character map over a Darts-clone double-array trie, whitespace
collapsing and escaping, the dummy word-boundary marker), unigram best-path
segmentation, BPE agenda merging, byte fallback, and user-defined symbol
handling. The public contract is SubwordTokenizer/SubwordPiece; every piece
reports the exact UTF-16 span of the caller's original text it came from,
and the model normalizer is also exposed as an OffsetAwareNormalizer
producing AlignedText.

Parity with the reference implementation is asserted, not assumed: five
tiny bundled models (unigram, unigram with byte fallback, BPE, identity
normalization, whitespace-as-suffix) carry fixtures generated by the
sentencepiece Python package over 40 inputs each, checked piece for piece,
id for id, span for span, plus each model's embedded self-test samples.
An opt-in test (-Dopennlp.subword.eval.dir) runs the same assertions
against real downloaded models; T5-small and ALBERT-base-v2 pass exactly,
including mixed scripts, emoji ZWJ sequences, BOM, and CRLF inputs.
…step

The vocabulary trie dispatches wide nodes (the root and first level of a
real vocabulary) through a 256-entry direct table, one load per byte, and
scans narrow nodes' short label slices linearly instead of binary
searching; a randomized differential test holds both layouts against a
map-backed reference, and moving the duplicate-piece detection into the
counting pass fixes the index error it previously produced. Non-unknown
segments reuse the vocabulary's piece string instead of decoding their
bytes, since the trie match means the bytes are identical.

The normalizer precomputes, per possible first byte, whether any
character-map rule or user-defined symbol starts with it; a clear bit
proves the prefix machinery would pass the byte through raw, so plain
ASCII text skips it entirely. The per-chunk record became a per-call
scratch, the input view keeps its oversized buffers with an explicit
length instead of trimming (pure-ASCII text gets an identity offset map
and no map array at all), the Viterbi scratch is one interleaved array
with scores as raw float bits, and the character-map trie walk relies on
the JVM's own bounds checks with the fail-loud translation on the cold
path.

All 37 bundled parity tests and the T5-small and ALBERT real-model
fixtures pass byte-identically. Single-thread throughput on the T5-small
vocabulary goes from 2.83M to 6.47M pieces per second, from 0.62x to
1.42x of the reference implementation measured through its Python
binding.
SubwordTokenizer and SubwordPiece move to opennlp.tools.tokenize, next to
Tokenizer and WordpieceTokenizer, matching where every other seam of this
round lives. The opennlp-subword module keeps only the SentencePiece
implementation.
…zer into it

WordpieceEncoder in opennlp-api runs the full BERT tokenization pipeline
as a SubwordTokenizer: every piece carries its vocabulary id and the span
of the original text, surviving the normalization steps that change,
insert, and remove characters. Content is computed with the same library
calls the previous pipeline made; offsets come from a per-code-point
rerun, with contextual case mappings (Greek final sigma) falling back to
word-wide spans that widen but never misplace. List and map constructors
cover line-number and explicit-id vocabularies.

BertTokenizer, unreleased and superseded, is removed. The dl tokenizer
creation builds on the encoder behind the existing Tokenizer plumbing via
a package-private adapter with unchanged special-token selection, and a
vocabulary missing its special tokens now fails at construction instead
of at the first id mapping, pinned by a test.

Parity is enforced twice: a differential suite against the reference
pipeline (kept test-only as ReferenceBertPipeline) over a curated corpus
plus 800 randomized inputs, and the removed class's reference token
sequences ported case for case. WordpieceTokenizer is untouched.
…verrides

Applies the review conventions from the OPENNLP-1869 review: class javadoc states the contracts instead of design narrative, every override carries inheritDoc with its null contract, and the private helpers are documented.
The tokenizer is Serializable through the OffsetAwareNormalizer contract but declared no serialVersionUID, which the compiler warns about. Added the serialver-computed value so it matches the convention used across the normalizer classes.
Adds a Subword Tokenization section to the Tokenizer chapter: the SubwordTokenizer
contract and its original-text span guarantee, loading and using a SentencePiece
model including the OffsetAwareNormalizer face, and the WordpieceEncoder pipeline
with its vocab.txt construction and special-token framing.
…s, name the format constants, document every helper
…bjectInputFilter

SentencePieceTokenizer gains serialize(OutputStream) and deserialize(InputStream)
methods. Reads are filtered through an ObjectInputFilter that allow-lists only the
classes reachable from a legitimate tokenizer graph and bounds graph depth,
references, and array length; foreign payloads are rejected with
InvalidClassException before being materialised. Limits are adjustable through a
DeserializationLimits record for unusually large vocabularies; the allow-list is
not configurable. The serialVersionUID is recomputed for the new public methods.
Add SentencePieceUsageExampleTest asserting the load-and-encode workflow and
point the tokenizer manual section at it.
…ixtures, thread safety wording

- Normalize the argument validation messages to the project style, naming the
  offending parameter and dropping the leading article and the trailing period, in
  WordpieceEncoder, SentencePieceTokenizer, ModelProtoReader, BpeEncoder,
  UnigramEncoder and the SubwordPiece compact constructor.
- Stop promising thread safety in the SubwordTokenizer contract and state that it is
  implementation specific instead; the manual now records that both shipped
  implementations are immutable and therefore safe for concurrent use.
- Drop the @throws IllegalArgumentException tags that only repeated the inherited
  contract on the normalize and normalizeAligned overrides, leaving a plain
  {@inheritdoc} as the rest of the class does.
- Remove commentary about release history rather than about the code: the pointer to
  the BertTokenizer class of the 3.0.0 milestone builds in WordpieceTokenizer, and
  the "frozen" qualifier on the ReferenceBertPipeline baseline.
- Move the bundled model loading and the fixture file reading out of
  SentencePieceParityTest into SentencePieceFixtures, so the alignment, validation
  and serialization tests no longer reach into another test class for a tokenizer.
- Extract MODEL_SUFFIX and FIXTURES_SUFFIX constants on SentencePieceFixtures and use
  them in SentencePieceRealModelEvalTest when deriving a fixture path from a model
  path, instead of repeating the two literals.
- Fold the five duplicated @valuesource model lists into a single
  SentencePieceFixtures#models @MethodSource, so adding a bundled model stays a one
  line change.
- Correct the parity test javadoc, which credited a nonexistent gen_fixtures.tsv
  sibling script instead of the gen_fixtures.py script in the test resources.
- Pin accessors that had no coverage: every score is finite and out of range ids are
  rejected, byte pieces occur only in byte fallback models and always render in the
  <0x..> form, and isByte rejects negative ids.
- Assert SubwordPiece.span() next to start and end in WordpieceEncoderTest so the
  derived span stays covered by the piece assertions.
- Document the IOException of the serialized helper in the serialization test and
  fully qualify the OutputStream javadoc link now that the import is gone.
… per review

Applies the review: the old entry point stays through one stable release
instead of being removed, and the DL extension point keeps its descriptor.

- Recreate BertTokenizer in opennlp-api as a thin shim, deprecated since
  3.0.0 forRemoval, with the original three Set based constructors and the
  original tokenizePos message. tokenize delegates to encodeToPieces; ids
  are synthesized from the set order because the tokenize path never reads
  them. Null contract follows this branch's reviewed convention,
  IllegalArgumentException, documented in the throws clauses.
- Delete the package-private EncoderTokenizer; the adapter now lives in
  opennlp-api where downstream code can reach it. AbstractDL's protected
  createTokenizer returns BertTokenizer again, restoring the override
  descriptor so an already compiled subclass keeps overriding at runtime,
  and createPipelineTokenizer hands back the shim.
- Delete ReferenceBertPipeline and point the curated and randomized
  differential tests in WordpieceEncoderTest at the shim, pinning shim and
  encoder to one sequence. Add BertTokenizerTest covering each constructor's
  argument validation, the default special token chain, and the exact
  tokenizePos message. Independent expected sequences continue to live in
  WordpieceEncoderReferenceSequencesTest.
- Fix a real divergence the compatibility check surfaced: the encoder kept
  U+2028 and U+2029 inside words while the old pipeline split on them, so a
  word carrying a line or paragraph separator became the unknown piece.
  cleanAndIsolateCjk now maps Zl and Zp to a space, with a span asserting
  regression test.
- Manual: the WordPiece section describes the deprecation and the migration,
  including the Set to List vocabulary change.
Add a README for regenerating the bundled SentencePiece parity fixtures
and clarify that Utf8Text is the encode-path span bridge behind
SubwordPiece offsets.
State that the reader is an independent re-implementation of the
serialized format, cite Aoe (1989), Yata et al. (2007), and Kanda et
al. (2023) in the class javadoc, and decode the bit-9 offset extension
with a plain conditional instead of the branchless form.
…ink it from the manual

State that the reference implementation produces the expected fixture
outputs, add the end-to-end validation steps for the bundled and real
models, and point the manual's SentencePiece section at the README.
The absolute GitHub URL 404s until merge and pins the branch layout.
…ontains

The manual states which parts are in opennlp-api and which in the
opennlp-subword artifact, that no model is bundled, what the parts of a
SubwordPiece are, and which tests pin the span rules and reference parity.
krickert and others added 29 commits September 16, 2026 06:37
…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.
…t quantization, expand tests

Fail loud when a directory holds both model.quantized and model.safetensors
instead of preferring one silently. Widen the QuantizeModel seed argument to
long so the full seed space is expressible. Add a Quantized Models manual
section pointing at the workflow tests. Convert reconstruction and embed-parity
tests to parameterized bit-width cases, add the SentencePiece quantized path
and a both-files-present rejection, and share the cosine and SentencePiece
fixture helpers.
A quantized file header declaring a dimension near 2^29 made
paddedDimension * bits overflow a signed int, so the per-row byte count went
negative and reading the file crashed with an undocumented
NegativeArraySizeException instead of a clean rejection; the same overflow
would corrupt the bit addressing in readCode/writeCode. Compute the row byte
count in long arithmetic through a single range-checked helper used by
quantize, read, and the constructor. Add edge-case tests across non-power-of-two
dimensions, a one-dimensional matrix, and adversarial rows the rotation must
still reconstruct.
…eption

Content errors in read now throw the checked loader exception instead of
IllegalArgumentException, with a size plausibility guard before any allocation.
The ambiguous matrix source and the quantizer's missing safetensors follow the
same contract. The CLI pin includes QuantizeModel, which this branch registers.
…d table path

The row count disagreement and the wordpiece unknown-token check now throw
InvalidFormatException like the rest of the module after the OPENNLP-1877
conversion.
…path and a size bound before allocation (failing tests)

The loader contract says malformed file content fails with the checked
InvalidFormatException, but the decoded-norm, pooling-weight, stored-grid,
and trailing-byte rejections still throw IllegalArgumentException, and a
small hostile file declaring huge dimensions reaches per-row allocation and
dies with EOFException instead of a format error. These tests pin the
intended behavior and fail against the current reader:

- testReadRejectsForeignAndTruncatedFiles now expects
  InvalidFormatException for trailing bytes (was pinning
  IllegalArgumentException, the wrong contract)
- testDeclaredPayloadBeyondFileSizeFailsBeforeAllocating: a 1.1 MB file
  declaring 1,000,000 rows of 512 dims at 4 bits must fail fast, before
  allocating 256 MB of codes plus 8 MB of scales and norms
- non-finite stored grid levels, decoded norms, and pooling weights, and a
  pooling-weight flag with no weights present, must all fail with
  InvalidFormatException
…eption and bound declared sizes before allocating

read(Path) promises the checked InvalidFormatException for malformed
content, but four rejections still threw IllegalArgumentException and
escaped every catch (IOException): an invalid decoded norm, a non-finite
pooling weight, trailing bytes after the declared content, and the
row-byte-count and storable-size checks reached from a hostile header. A
stored grid that fromLevels rejects also surfaced as
IllegalArgumentException. All of these now throw InvalidFormatException
naming the file and the offending field; constructor validation for
programmatic callers stays IllegalArgumentException.

The reader also only checked rowCount and dimension individually against
the file size, so a 1.1 MB file declaring 1,000,000 rows of 512 dims at 4
bits forced a 256 MB code allocation before any content check. The header
fully determines the file size, so the declared total (28 fixed bytes, the
grid levels, a scale and a decoded norm per row, the flag byte, and the
packed codes, plus the per-row pooling weights when flagged) is now held
against the actual file size and rejected before anything row-sized is
allocated.
…ding

The class comment and the manual claimed a 500,000-row, 300-dimension
table drops to 77 MB at 4 bits. The Hadamard rotation pads rows to the
next power of two, so 300 dimensions store 512 codes per row, and each row
also carries two floats (the fitted scale and the decoded norm), not one.
Measured from a written file, that is 264 bytes per row: 132 MB against
600 MB of float32, 4.5 times smaller, not 7.8. The overall shrink range
becomes roughly 4 to 16 times depending on bit width and padding distance.
The padding behavior itself is unchanged.
Red tests covered numeric extremes, model shapes, format validation, and tokenizer paths. Green affected reactor: 3,691 tests with two expected skips; package, documentation, and policy checks passed.
A quantized file header declaring a dimension near 2^29 made
paddedDimension * bits overflow a signed int, so the per-row byte count went
negative and reading the file crashed with an undocumented
NegativeArraySizeException instead of a clean rejection; the same overflow
would corrupt the bit addressing in readCode/writeCode. Compute the row byte
count in long arithmetic through a single range-checked helper used by
quantize, read, and the constructor. Add edge-case tests across non-power-of-two
dimensions, a one-dimensional matrix, and adversarial rows the rotation must
still reconstruct.
…eption

Content errors in read now throw the checked loader exception instead of
IllegalArgumentException, with a size plausibility guard before any allocation.
The ambiguous matrix source and the quantizer's missing safetensors follow the
same contract. The CLI pin includes QuantizeModel, which this branch registers.
…ding

The class comment and the manual claimed a 500,000-row, 300-dimension
table drops to 77 MB at 4 bits. The Hadamard rotation pads rows to the
next power of two, so 300 dimensions store 512 codes per row, and each row
also carries two floats (the fitted scale and the decoded norm), not one.
Measured from a written file, that is 264 bytes per row: 132 MB against
600 MB of float32, 4.5 times smaller, not 7.8. The overall shrink range
becomes roughly 4 to 16 times depending on bit width and padding distance.
The padding behavior itself is unchanged.
Six regression cases failed when float midpoint rounding selected a non-nearest level. Preserve double precision for encoding thresholds and test exact midpoint selection.

Add independent ONQ2 fixtures for 2-, 3-, and 4-bit matrices, including saved bytes, decoded coordinates, scoring, pooling, truncation, and trailing data checks.
QuantizedEmbeddingMatrix still had a literal merge conflict marker
block around its class Javadoc. QuantizedEmbeddingMatrixTest had the
overflow-dimension test defined twice with slightly different wording,
which does not compile. StaticEmbeddingModelSentencePieceQuantizedTest
had the same InvalidFormatException import listed twice. Each fix
matches the content already checked against the pre-rebase
OPENNLP-1895 branch tip for that file.
Red evidence: test compilation fails because VectorIndex, FlatFloatIndex, and TurboQuantIndex do not exist.
Implements the contract pinned in 86e616d. The focused suite now passes 37 tests across the exact and quantized implementations.
Document the exact and TurboQuant choices, their bounded single-JVM scope, the build-freeze-query lifecycle, concurrency contract, and TurboQuant persistence. A mirrored usage test passes with the focused 38-test index suite.
A frozen exact index writes its full-precision row-major floats and its ids in
row order as a directory, and reads them back into a frozen index that scores
identically, mirroring TurboQuantIndex. Malformed, truncated, or mismatched
files fail as InvalidFormatException.

(cherry picked from commit ff73e3a991da54c1c8d386b7435ef241c8e4c24e)
Red tests covered tiny vectors, non-finite queries, tie ordering, overflow, and malformed persisted indexes. Green full affected reactor, documentation, and policy checks passed, including 514 embeddings tests.
Red evidence: test compilation fails because VectorIndex, FlatFloatIndex, and TurboQuantIndex do not exist.
Implements the contract pinned in 86e616d. The focused suite now passes 37 tests across the exact and quantized implementations.
Document the exact and TurboQuant choices, their bounded single-JVM scope, the build-freeze-query lifecycle, concurrency contract, and TurboQuant persistence. A mirrored usage test passes with the focused 38-test index suite.
10 file-size tests failed with the 4-byte metadata estimate; 13 truncated-header tests failed with EOFException. Count the 8-byte scale and norm, preserve the read failure cause, and test saved-index queries.

Update the manual and persistence example. Index file content and scoring are unchanged.
Conflict resolution during the OPENNLP-1910 rebase left two copies of
the tools.embeddings.vector-index section, which FOP's PDF generation
rejects because docbook block IDs must be unique. The kept copy matches
the content already checked against the pre-rebase OPENNLP-1910 branch
tip.
@krickert
krickert force-pushed the OPENNLP-1910-in-memory-vector-index branch from 8096837 to a5f9afc Compare September 16, 2026 17:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant