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 8592377fb41b601d57b81c1ead7e1a8e8b51ecea Author: Kristian Rickert <[email protected]> AuthorDate: Sat Oct 10 08:34:08 2026 -0400 GH-1391: Test the batch boundaries and split groups of embedAll --- .../vectors/SentenceVectorsDLBatchLimitTest.java | 293 +++++++++++++++++++++ 1 file changed, 293 insertions(+) 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 new file mode 100644 index 000000000..6705d8e04 --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java @@ -0,0 +1,293 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * 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.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 opennlp.dl.InferenceOptions; +import opennlp.dl.Tokens; + +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> + */ +class SentenceVectorsDLBatchLimitTest { + + private static final int LIMIT = SentenceVectorsDL.MAX_BATCH_TOKEN_POSITIONS; + + private static final float DELTA = 1e-5f; + + private static final String HELLO = "hello"; + + 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(); + } + + 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]); + } + + /** The sizes of the batches of {@code count} encodings of {@code length} tokens. */ + private static List<Integer> sizes(final int length, final int count) { + final List<Integer> sizes = new ArrayList<>(); + for (final List<Integer> batch : SentenceVectorsDL.batches(tokens(length, count))) { + 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. + */ + @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())); + } + assertEquals(sizes, sizes(length, count)); + } + + /** A length equal to the cap fits one row per inference, and one more token the same. */ + @ParameterizedTest + @CsvSource({"16384", "16385"}) + void testLengthAtAndOverTheLimitRunsOneRowPerBatch(final int length) { + assertEquals(List.of(1, 1, 1), sizes(length, 3)); + } + + @Test + void testEmptyCallHasNoBatches() { + assertTrue(SentenceVectorsDL.batches(new Tokens[0]).isEmpty()); + } + + /** + * 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. + */ + @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); + assertEquals(3, batches.size()); + assertEquals(LIMIT / 4, batches.get(0).size()); + 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 + * equals the vector of its own text. + */ + @Test + void testMixedGroupsThroughTheSession(@TempDir final Path dir) throws Exception { + 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, + SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { + final float[][] batch = v.embedAll(texts); + assertEquals(texts.size(), batch.length); + for (int i = 0; i < texts.size(); i++) { + assertArrayEquals(v.embed(texts.get(i)), batch[i], DELTA, "row " + i); + } + } + } + + /** + * A single input as wide as the cap runs next to a short group, 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)) { + 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++) { + assertArrayEquals(v.embed(texts.get(i)), batch[i], DELTA, "row " + i); + } + } + } + + /** 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; + } + + 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)); + } + + /** + * Checks the pooled-output model, CLS pooling and unit length scaling through a split group, + * since each takes its own branch after the inference. + */ + @ParameterizedTest + @MethodSource("pathCases") + void testPoolingAndNormalizeThroughASplitGroup(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); + try (SentenceVectorsDL v = vectors(dir, model, pooling, normalize, + SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { + final float[][] batch = v.embedAll(texts); + assertEquals(texts.size(), batch.length); + for (int i = 0; i < batch.length; i++) { + assertArrayEquals(i % 2 == 0 ? helloWorld : hello, batch[i], DELTA, "row " + i); + } + } + } + + /** + * 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") + 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 List<String> texts = alternating(LIMIT / 3 + 1); + try (gpu; SentenceVectorsDL cpu = vectors(dir, model, pooling, normalize, + SentenceVectorsDL.DEFAULT_MAX_LENGTH)) { + final float[][] expected = cpu.embedAll(texts); + final float[][] actual = gpu.embedAll(texts); + assertEquals(expected.length, actual.length); + for (int i = 0; i < expected.length; i++) { + assertArrayEquals(expected[i], actual[i], DELTA, "row " + i); + } + } + } +}
