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 82156b04f14d96bd5852778138253b0c5756a222 Author: Kristian Rickert <[email protected]> AuthorDate: Sun Oct 11 01:48:55 2026 -0400 GH-1391: Share the batch test fixtures and drop duplicate tests The batch tests use TinyVectorModels and list the expected batch sizes directly. The cap tests added to SentenceVectorsDLEmbedderTest repeated the cases of SentenceVectorsDLBatchLimitTest and are removed. Signed-off-by: Kristian Rickert <[email protected]> --- .../vectors/SentenceVectorsDLBatchLimitTest.java | 255 ++++++++++----------- .../dl/vectors/SentenceVectorsDLEmbedderTest.java | 57 ----- 2 files changed, 120 insertions(+), 192 deletions(-) diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java index 6705d8e04..01c5b03c2 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java @@ -14,47 +14,40 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package opennlp.dl.vectors; import java.io.File; -import java.io.IOException; -import java.io.InputStream; -import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.StandardCopyOption; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; -import java.util.Objects; import java.util.stream.Stream; -import ai.onnxruntime.OrtEnvironment; -import ai.onnxruntime.OrtException; -import ai.onnxruntime.OrtProvider; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; -import opennlp.dl.InferenceOptions; import opennlp.dl.Tokens; +import static opennlp.dl.vectors.TinyVectorModels.TINY_POOLED; +import static opennlp.dl.vectors.TinyVectorModels.TINY_VECTORS; +import static opennlp.dl.vectors.TinyVectorModels.gpuOrAbort; +import static opennlp.dl.vectors.TinyVectorModels.model; +import static opennlp.dl.vectors.TinyVectorModels.scale; +import static opennlp.dl.vectors.TinyVectorModels.vocab; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.api.Assumptions.abort; -import static org.junit.jupiter.api.Assumptions.assumeTrue; /** * Tests the split of {@link SentenceVectorsDL#embedAll(List)} into inferences of at most * {@value SentenceVectorsDL#MAX_BATCH_TOKEN_POSITIONS} token positions: the boundaries of the - * plan, and the vectors that come back through a split group. - * - * <p>The test models are the ones of {@link SentenceVectorsDLEmbedderTest}. With the vocabulary - * below, "hello" encodes to 3 tokens, "hello world" to 4 and the empty text to 2.</p> + * batches, and the vectors that come back through a split group. "hello" encodes to 3 tokens and + * "hello world" to 4. */ class SentenceVectorsDLBatchLimitTest { @@ -66,110 +59,129 @@ class SentenceVectorsDLBatchLimitTest { private static final String HELLO_WORLD = "hello world"; - private static float[] scale(final float factor) { - return new float[] {0.5f * factor, -1f * factor, 2f * factor}; - } - private static final float[] UNIT = scale(1 / (float) Math.sqrt(5.25)); - private static File model(final Path dir, final String name) throws IOException { - final Path file = dir.resolve(name); - try (InputStream is = Objects.requireNonNull(SentenceVectorsDLBatchLimitTest.class - .getResourceAsStream("/opennlp/dl/vectors/" + name))) { - Files.copy(is, file, StandardCopyOption.REPLACE_EXISTING); - } - return file.toFile(); - } - - private static File vocab(final Path dir) throws IOException { - final Path file = dir.resolve("vocab.txt"); - Files.write(file, List.of("[PAD]", "unused1", "[UNK]", "[SEP]", "hello", "world", - "unused2", "[CLS]")); - return file.toFile(); - } - + /** + * {@return a session on the CPU over the given tiny model} + * + * @param dir The directory to copy the model into. + * @param model {@link TinyVectorModels#TINY_VECTORS} or {@link TinyVectorModels#TINY_POOLED}. + * @param pooling The pooling. + * @param normalize {@code true} to scale every vector to unit length. + * @param maxLength The maximum number of tokens per input. + */ private static SentenceVectorsDL vectors(final Path dir, final String model, final Pooling pooling, final boolean normalize, final int maxLength) throws Exception { return new SentenceVectorsDL(model(dir, model), vocab(dir), true, pooling, normalize, maxLength); } - private static Tokens[] tokens(final int length, final int count) { - return Collections.nCopies(count, new Tokens(new String[length], new long[length], - new long[length], new long[length])).toArray(new Tokens[0]); + /** + * {@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]); } - /** The sizes of the batches of {@code count} encodings of {@code length} tokens. */ + /** + * {@return the sizes of the batches of {@code count} encodings of {@code length} tokens} + * + * @param length The number of tokens of every encoding. + * @param count The number of encodings. + */ private static List<Integer> sizes(final int length, final int count) { + final Tokens[] encoded = Collections.nCopies(count, tokens(length)).toArray(new Tokens[0]); final List<Integer> sizes = new ArrayList<>(); - for (final List<Integer> batch : SentenceVectorsDL.batches(tokens(length, count))) { + for (final List<Integer> batch : SentenceVectorsDL.batches(encoded)) { sizes.add(batch.size()); } return sizes; } /** - * Checks the plan at the edges of the row count: a group that fills one inference exactly, - * one row more, one row fewer, and lengths that do or do not divide the cap evenly. + * {@return {@code pairs} times "hello world" followed by "hello"} + * + * @param pairs The number of pairs. */ - @ParameterizedTest - @CsvSource({ - "4, 4096, 4096", - "4, 4097, 4096;1", - "4, 4095, 4095", - "2, 8192, 8192", - "2, 8193, 8192;1", - "3, 5461, 5461", - "3, 5462, 5461;1", - "3, 10922, 5461;5461", - "3, 10923, 5461;5461;1", - "1, 16384, 16384", - "1, 16385, 16384;1"}) - void testRowCountBoundaries(final int length, final int count, final String expected) { - final List<Integer> sizes = new ArrayList<>(); - for (final String size : expected.split(";")) { - sizes.add(Integer.parseInt(size.trim())); + private static List<String> alternating(final int pairs) { + final List<String> texts = new ArrayList<>(2 * pairs); + for (int i = 0; i < pairs; i++) { + texts.add(HELLO_WORLD); + texts.add(HELLO); } - assertEquals(sizes, sizes(length, count)); + return texts; } - /** A length equal to the cap fits one row per inference, and one more token the same. */ + /** + * {@return the length, the count and the expected batch sizes at the edges of the row count} + */ + static Stream<Arguments> rowCountBoundaries() { + return Stream.of( + Arguments.of(4, 4096, List.of(4096)), + Arguments.of(4, 4097, List.of(4096, 1)), + Arguments.of(4, 4095, List.of(4095)), + Arguments.of(2, 8192, List.of(8192)), + Arguments.of(2, 8193, List.of(8192, 1)), + Arguments.of(3, 5461, List.of(5461)), + Arguments.of(3, 5462, List.of(5461, 1)), + Arguments.of(3, 10922, List.of(5461, 5461)), + Arguments.of(3, 10923, List.of(5461, 5461, 1)), + Arguments.of(1, 16384, List.of(16384)), + Arguments.of(1, 16385, List.of(16384, 1))); + } + + /** + * Checks a group that fills one inference exactly, one row more, one row fewer, and lengths + * that do or do not divide the cap evenly. + */ @ParameterizedTest - @CsvSource({"16384", "16385"}) + @MethodSource("rowCountBoundaries") + void testRowCountBoundaries(final int length, final int count, final List<Integer> expected) { + assertEquals(expected, sizes(length, count)); + } + + /** An input as long as the cap, or longer, runs in a batch of its own. */ + @ParameterizedTest + @ValueSource(ints = {LIMIT, LIMIT + 1}) void testLengthAtAndOverTheLimitRunsOneRowPerBatch(final int length) { assertEquals(List.of(1, 1, 1), sizes(length, 3)); } + /** An empty call has no batches. */ @Test void testEmptyCallHasNoBatches() { assertTrue(SentenceVectorsDL.batches(new Tokens[0]).isEmpty()); } + /** Inputs of different lengths run one batch per length, in the order each length appears. */ + @Test + void testGroupsByLengthInCallOrder() { + 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)})); + } + /** - * Checks that only the group that is too large is split: 4-token inputs need two inferences, - * the 3-token inputs between them stay in one, and the groups keep the call order. + * Checks that only the group that is too large is split: the 4-token inputs need two + * inferences, the 3-token inputs between them stay in one. */ @Test void testOnlyTheLargeGroupIsSplit() { - final int wide = LIMIT / 4 + 1; - final Tokens[] plan = new Tokens[wide + 2]; - System.arraycopy(tokens(4, 1), 0, plan, 0, 1); - plan[1] = tokens(3, 1)[0]; - plan[2] = tokens(3, 1)[0]; - for (int i = 3; i < plan.length; i++) { - plan[i] = tokens(4, 1)[0]; - } - final List<List<Integer>> batches = SentenceVectorsDL.batches(plan); + final Tokens[] encoded = new Tokens[LIMIT / 4 + 3]; + Arrays.fill(encoded, tokens(4)); + encoded[1] = tokens(3); + encoded[2] = tokens(3); + final List<List<Integer>> batches = SentenceVectorsDL.batches(encoded); assertEquals(3, batches.size()); assertEquals(LIMIT / 4, batches.get(0).size()); + assertEquals(List.of(0, 3), batches.get(0).subList(0, 2)); assertEquals(1, batches.get(1).size()); assertEquals(List.of(1, 2), batches.get(2)); - assertEquals(0, batches.get(0).get(0)); - assertEquals(3, batches.get(0).get(1)); } /** - * The same mix through the session: one group splits, the other does not, and every row + * Runs the same mix through the session: one group splits, the other does not, and every row * equals the vector of its own text. */ @Test @@ -177,7 +189,7 @@ class SentenceVectorsDLBatchLimitTest { final List<String> texts = new ArrayList<>(Collections.nCopies(LIMIT / 4 + 1, HELLO_WORLD)); texts.add(1, HELLO); texts.add(7, HELLO); - try (SentenceVectorsDL v = vectors(dir, "tiny-vectors.onnx", Pooling.MEAN, false, + try (SentenceVectorsDL v = vectors(dir, TINY_VECTORS, Pooling.MEAN, false, SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { final float[][] batch = v.embedAll(texts); assertEquals(texts.size(), batch.length); @@ -187,14 +199,11 @@ class SentenceVectorsDLBatchLimitTest { } } - /** - * A single input as wide as the cap runs next to a short group, and comes back in call - * order. - */ + /** An input as wide as the cap runs next to short inputs and comes back in call order. */ @Test void testInputAtTheLimitThroughTheSession(@TempDir final Path dir) throws Exception { final String wide = String.join(" ", Collections.nCopies(LIMIT - 2, HELLO)); - try (SentenceVectorsDL v = vectors(dir, "tiny-vectors.onnx", Pooling.MEAN, false, LIMIT)) { + try (SentenceVectorsDL v = vectors(dir, TINY_VECTORS, Pooling.MEAN, false, LIMIT)) { final List<String> texts = List.of(HELLO, wide, HELLO_WORLD, wide); final float[][] batch = v.embedAll(texts); for (int i = 0; i < texts.size(); i++) { @@ -203,47 +212,27 @@ class SentenceVectorsDLBatchLimitTest { } } - /** Alternating lengths, both groups split: results are in input order. */ - @Test - void testAlternatingLengthsKeepInputOrder(@TempDir final Path dir) throws Exception { - final int pairs = LIMIT / 3 + 1; - final List<String> texts = alternating(pairs); - try (SentenceVectorsDL v = vectors(dir, "tiny-vectors.onnx", Pooling.MEAN, false, - SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { - final float[][] batch = v.embedAll(texts); - assertEquals(texts.size(), batch.length); - for (int i = 0; i < batch.length; i++) { - assertArrayEquals(scale(i % 2 == 0 ? 19 / 4f : 14 / 3f), batch[i], DELTA, "row " + i); - } - } - } - - private static List<String> alternating(final int pairs) { - final List<String> texts = new ArrayList<>(2 * pairs); - for (int i = 0; i < pairs; i++) { - texts.add(HELLO_WORLD); - texts.add(HELLO); - } - return texts; - } - + /** + * {@return the model, pooling, unit length, and the expected vectors for "hello world" and + * "hello"} + */ static Stream<Arguments> pathCases() { return Stream.of( - // model, pooling, normalize, expected for "hello world", expected for "hello" - Arguments.of("tiny-vectors.onnx", Pooling.CLS, false, scale(7), scale(7)), - Arguments.of("tiny-vectors.onnx", Pooling.CLS, true, UNIT, UNIT), - Arguments.of("tiny-vectors.onnx", Pooling.MEAN, true, UNIT, UNIT), - Arguments.of("tiny-pooled.onnx", Pooling.CLS, false, scale(19), scale(14)), - Arguments.of("tiny-pooled.onnx", Pooling.MEAN, true, UNIT, UNIT)); + Arguments.of(TINY_VECTORS, Pooling.MEAN, false, scale(19 / 4f), scale(14 / 3f)), + Arguments.of(TINY_VECTORS, Pooling.CLS, false, scale(7), scale(7)), + Arguments.of(TINY_VECTORS, Pooling.CLS, true, UNIT, UNIT), + Arguments.of(TINY_VECTORS, Pooling.MEAN, true, UNIT, UNIT), + Arguments.of(TINY_POOLED, Pooling.CLS, false, scale(19), scale(14)), + Arguments.of(TINY_POOLED, Pooling.MEAN, true, UNIT, UNIT)); } /** - * Checks the pooled-output model, CLS pooling and unit length scaling through a split group, - * since each takes its own branch after the inference. + * Checks every path after the inference through split groups of alternating lengths, with the + * results in input order. */ @ParameterizedTest @MethodSource("pathCases") - void testPoolingAndNormalizeThroughASplitGroup(final String model, final Pooling pooling, + void testPathsThroughASplitGroup(final String model, final Pooling pooling, final boolean normalize, final float[] helloWorld, final float[] hello, @TempDir final Path dir) throws Exception { final List<String> texts = alternating(LIMIT / 3 + 1); @@ -257,31 +246,27 @@ class SentenceVectorsDLBatchLimitTest { } } + /** + * {@return the model, pooling and unit length combinations, one per path after the inference} + */ + static Stream<Arguments> gpuCases() { + return pathCases().map(arguments -> Arguments.of(Arrays.copyOf(arguments.get(), 3))); + } + /** * Checks that a split group gives the same vectors on CUDA device {@code 0} as on the CPU, * where a CUDA device exists; skipped elsewhere. */ @ParameterizedTest - @MethodSource("pathCases") + @MethodSource("gpuCases") void testGpuMatchesTheCpuForASplitGroup(final String model, final Pooling pooling, - final boolean normalize, final float[] ignored, final float[] alsoIgnored, - @TempDir final Path dir) throws Exception { - assumeTrue(OrtEnvironment.getAvailableProviders().contains(OrtProvider.CUDA), - "the ONNX Runtime on the classpath has no CUDA execution provider"); - final InferenceOptions options = new InferenceOptions(); - options.setGpu(true); - options.setGpuDeviceId(0); - final SentenceVectorsDL gpu; - try { - gpu = new SentenceVectorsDL(model(dir, model), vocab(dir), true, pooling, normalize, - SentenceVectorsDL.DEFAULT_MAX_LENGTH, options); - } catch (OrtException e) { - abort("no usable CUDA device 0: " + e.getMessage()); - return; - } + final boolean normalize, @TempDir final Path dir) throws Exception { + final File modelFile = model(dir, model); + final File vocab = vocab(dir); final List<String> texts = alternating(LIMIT / 3 + 1); - try (gpu; SentenceVectorsDL cpu = vectors(dir, model, pooling, normalize, - SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { + try (SentenceVectorsDL gpu = gpuOrAbort(modelFile, vocab, pooling, normalize); + SentenceVectorsDL cpu = new SentenceVectorsDL(modelFile, vocab, true, pooling, + normalize, SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { final float[][] expected = cpu.embedAll(texts); final float[][] actual = gpu.embedAll(texts); assertEquals(expected.length, actual.length); 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 5adab3fbf..d12d97878 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,14 +22,12 @@ 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; @@ -250,61 +248,6 @@ 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 a quarter of the cap in rows. - */ - @Test - void testEmbedAllCapsTokenPositionsPerInference(@TempDir final Path dir) throws Exception { - final int helloWorldTokens = 4; - final int rowsPerInference = SentenceVectorsDL.MAX_BATCH_TOKEN_POSITIONS / helloWorldTokens; - final int count = rowsPerInference + rowsPerInference / 4; - try (SentenceVectorsDL vectors = meanVectors(dir, SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { - final List<String> texts = Collections.nCopies(count, "hello world"); - final List<List<Integer>> batches = SentenceVectorsDL.batches( - Collections.nCopies(count, tokens(helloWorldTokens)).toArray(new Tokens[0])); - assertEquals(2, batches.size()); - assertEquals(rowsPerInference, batches.get(0).size()); - assertEquals(count - rowsPerInference, 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.
