This is an automated email from the ASF dual-hosted git repository. krickert pushed a commit to branch feat/GH-1392-embed-padding in repository https://gitbox.apache.org/repos/asf/opennlp.git
commit bcdd68d904a019e3b273458a2a20835b9d6ef7bd Author: Kristian Rickert <[email protected]> AuthorDate: Sat Oct 10 08:26:13 2026 -0400 GH-1391: Test that embedAll caps the token positions of one inference Moves the grouping of embedAll into a package-private batches method, unchanged, so tests can assert the inference plan, and names the bound the fix will apply. The new tests fail on this commit: testEmbedAllCapsTokenPositionsPerInference: expected: <2> but was: <1> testEmbedAllGroupsByLengthInCallOrder: expected: <[[0, 3], [1, 2, 4]]> but was: <[[1, 2, 4], [0, 3]]> testEmbedAllRunsAnInputWiderThanTheCapInItsOwnBatch: expected: <[[0], [1]]> but was: <[[0, 1]]> Signed-off-by: Kristian Rickert <[email protected]> --- .../java/opennlp/dl/vectors/SentenceVectorsDL.java | 27 +++++++++-- .../dl/vectors/SentenceVectorsDLEmbedderTest.java | 54 ++++++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java index 3a1dea3fb..f62f59f7e 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java @@ -82,6 +82,14 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { /** The maximum number of tokens per input if none is given, the BERT position limit. */ public static final int DEFAULT_MAX_LENGTH = 512; + /** + * The upper bound on the token positions of one inference, its row count times its row length. + * {@link #embedAll(List)} splits a larger group into consecutive inferences, so the tensors of + * one inference stay bounded however many texts a call has. A single input longer than the + * bound still runs on its own. + */ + static final int MAX_BATCH_TOKEN_POSITIONS = 16384; + private static final String SENTENCE_EMBEDDING = "sentence_embedding"; private static final int POOLED_RANK = 2; private static final int TOKEN_RANK = 3; @@ -282,15 +290,13 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { throw new IllegalArgumentException("texts[" + i + "] must not be null"); } } - final Map<Integer, List<Integer>> byLength = new HashMap<>(); final Tokens[] encoded = new Tokens[checked.length]; for (int i = 0; i < checked.length; i++) { encoded[i] = encode(checked[i]); - byLength.computeIfAbsent(encoded[i].ids().length, length -> new ArrayList<>()).add(i); } final float[][] vectors = new float[checked.length][]; try { - for (final List<Integer> group : byLength.values()) { + for (final List<Integer> group : batches(encoded)) { final Tokens[] batch = new Tokens[group.size()]; for (int b = 0; b < batch.length; b++) { batch[b] = encoded[group.get(b)]; @@ -306,6 +312,21 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { return vectors; } + /** + * Splits the indices of a call into the batches that run as single inferences, one batch per + * tokenized length. + * + * @param encoded The encodings of the call, in call order. + * @return The batches, each a list of indices into {@code encoded}. + */ + static List<List<Integer>> batches(final Tokens[] encoded) { + final Map<Integer, List<Integer>> byLength = new HashMap<>(); + for (int i = 0; i < encoded.length; i++) { + byLength.computeIfAbsent(encoded[i].ids().length, length -> new ArrayList<>()).add(i); + } + return new ArrayList<>(byLength.values()); + } + /** * {@inheritDoc} * diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java index d12d97878..f944ae208 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java @@ -22,12 +22,14 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import opennlp.dl.Tokens; import opennlp.tools.embeddings.TextEmbedder; import opennlp.tools.embeddings.TextEmbedderProvider; import opennlp.tools.util.InvalidFormatException; @@ -248,6 +250,58 @@ class SentenceVectorsDLEmbedderTest { } } + /** + * Asserts that {@code embedAll} splits a large group of same-length inputs into inferences of + * at most {@value SentenceVectorsDL#MAX_BATCH_TOKEN_POSITIONS} token positions, and that every + * row still equals its single-input vector. "hello world" encodes to 4 tokens, so one + * inference takes at most 4096 rows. + */ + @Test + void testEmbedAllCapsTokenPositionsPerInference(@TempDir final Path dir) throws Exception { + try (SentenceVectorsDL vectors = meanVectors(dir, SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { + final List<String> texts = Collections.nCopies(5000, "hello world"); + final List<List<Integer>> batches = + SentenceVectorsDL.batches(Collections.nCopies(5000, tokens(4)).toArray(new Tokens[0])); + assertEquals(2, batches.size()); + assertEquals(4096, batches.get(0).size()); + assertEquals(904, batches.get(1).size()); + final float[][] batch = vectors.embedAll(texts); + assertEquals(texts.size(), batch.length); + for (final float[] row : batch) { + assertArrayEquals(HELLO_WORLD_MEAN, row, DELTA); + } + } + } + + /** + * Asserts that inputs of different lengths still run one inference per length, in the order + * each length first appears in the call. + */ + @Test + void testEmbedAllGroupsByLengthInCallOrder() { + assertEquals(List.of(List.of(0, 3), List.of(1, 2, 4)), SentenceVectorsDL.batches( + new Tokens[] {tokens(4), tokens(3), tokens(3), tokens(4), tokens(3)})); + } + + /** + * Asserts that an input wider than the cap on its own still runs, as a batch of one. + */ + @Test + void testEmbedAllRunsAnInputWiderThanTheCapInItsOwnBatch() { + final Tokens wide = tokens(SentenceVectorsDL.MAX_BATCH_TOKEN_POSITIONS + 1); + assertEquals(List.of(List.of(0), List.of(1)), + SentenceVectorsDL.batches(new Tokens[] {wide, wide})); + } + + /** + * {@return an encoding of {@code length} tokens; only its length matters for batching} + * + * @param length The number of tokens. + */ + private static Tokens tokens(final int length) { + return new Tokens(new String[length], new long[length], new long[length], new long[length]); + } + /** * Asserts the batch contract edges: an empty input yields an empty batch, and a * {@code null} element is rejected rather than failing later inside the session.
