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 8c7c439df41c0561a56a95bf3b9a89986d3ab152 Author: Kristian Rickert <[email protected]> AuthorDate: Sun Oct 11 02:00:42 2026 -0400 GH-1392: Pad to the longest text by default on a GPU On a GPU, embedAll without a padding strategy grouped texts by exact length, so a call of mixed lengths ran nearly one inference per text and was no faster than calling embed in a loop. A session built without a PaddingStrategy now uses LONGEST on a GPU when the vocabulary has a padding token, and EXACT_LENGTH otherwise. An explicit strategy is kept as given, and the CPU default is unchanged. Before the change the new tests fail: testDevicePadding[3]: expected: <LONGEST> but was: <EXACT_LENGTH> testGpuBatchesMixedLengthsByDefault: expected: <[[1, 0]]> but was: <[[0], [1]]> SentenceVectorsDLBenchmark on all-MiniLM-L6-v2, RTX 4080 SUPER, texts of 4 to 40 words, embedAll time per call before and after: 8 texts: 3603 us -> 862 us 32 texts: 11431 us -> 2152 us 128 texts: 23815 us -> 7443 us Signed-off-by: Kristian Rickert <[email protected]> --- .../src/main/java/opennlp/dl/AbstractDL.java | 9 +++ .../java/opennlp/dl/vectors/PaddingStrategy.java | 4 +- .../java/opennlp/dl/vectors/SentenceVectorsDL.java | 64 ++++++++++++++++++---- .../dl/vectors/SentenceVectorsDLEmbedderTest.java | 54 ++++++++++++++++++ .../java/opennlp/dl/vectors/TinyVectorModels.java | 14 +++-- opennlp-docs/src/docbkx/machine-learning.xml | 4 +- 6 files changed, 130 insertions(+), 19 deletions(-) diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java index 63586d683..0ccfcfe1a 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java @@ -429,6 +429,15 @@ public abstract class AbstractDL implements AutoCloseable { return id; } + /** + * {@return whether a vocabulary has the padding token that {@link #padTokenId(Map)} reads} + * + * @param vocab The vocabulary. + */ + protected static boolean hasPadToken(final Map<String, Integer> vocab) { + return vocab.containsKey(specialTokens(vocab).pad()); + } + /** * The classification, separator, unknown and padding tokens of a vocabulary. * diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/PaddingStrategy.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/PaddingStrategy.java index fba022ec5..78d00f56b 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/PaddingStrategy.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/PaddingStrategy.java @@ -35,14 +35,14 @@ public enum PaddingStrategy { /** * No padding. Texts are grouped by their exact tokenized length and each group runs on its own. - * This is the default. On natural text most groups hold one or two texts. + * This is the default on the CPU. On natural text most groups hold one or two texts. */ EXACT_LENGTH, /** * Every row of a batch is padded to the longest row of that batch. Texts are ordered by length * before batching, so one long text does not pad out a batch of short ones. Requires a padding - * token in the vocabulary. + * token in the vocabulary. This is the default on a GPU if the vocabulary has one. */ LONGEST, 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 031ef2b51..19e1a72f7 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 @@ -77,9 +77,10 @@ import opennlp.tools.util.ParamChecks; * {@code close()} with inference methods.</p> * * <p>{@link #embedAll(List)} batches its texts according to the configured - * {@link PaddingStrategy}. By default it does not pad and runs one inference per distinct - * tokenized length; {@link PaddingStrategy#LONGEST} runs texts of mixed lengths together. No - * inference exceeds {@value #MAX_BATCH_TOKEN_POSITIONS} token positions.</p> + * {@link PaddingStrategy}. By default it does not pad on the CPU and runs one inference per + * distinct tokenized length; on a GPU it uses {@link PaddingStrategy#LONGEST}, which runs texts of + * mixed lengths together. No inference exceeds {@value #MAX_BATCH_TOKEN_POSITIONS} token + * positions.</p> */ @ThreadSafe public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { @@ -95,9 +96,15 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { */ static final int MAX_BATCH_TOKEN_POSITIONS = 16384; - /** The {@link PaddingStrategy} used if none is given. */ + /** The {@link PaddingStrategy} on the CPU if none is given. */ public static final PaddingStrategy DEFAULT_PADDING = PaddingStrategy.EXACT_LENGTH; + /** + * The {@link PaddingStrategy} on a GPU if none is given and the vocabulary has a padding token. + * A GPU runs one wide inference faster than many narrow ones. + */ + public static final PaddingStrategy DEFAULT_GPU_PADDING = PaddingStrategy.LONGEST; + private static final String SENTENCE_EMBEDDING = "sentence_embedding"; private static final int POOLED_RANK = 2; private static final int TOKEN_RANK = 3; @@ -204,7 +211,8 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { * included; longer input is truncated. Must be at least {@code 2}. * @param inferenceOptions The {@link InferenceOptions}, of which the GPU settings and the lower * case setting are used. Must not be {@code null}. {@link #embedAll(List)} batches by - * {@link #DEFAULT_PADDING}. + * {@link #DEFAULT_PADDING} on the CPU, and by {@link #DEFAULT_GPU_PADDING} on a GPU if the + * vocabulary has a padding token. * * @throws IllegalArgumentException Thrown if {@code pooling} or {@code inferenceOptions} is * {@code null}, if {@code inferenceOptions} holds invalid split options, if @@ -224,8 +232,7 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { final InferenceOptions inferenceOptions) throws OrtException, IOException { - this(model, vocabulary, lowerCase, pooling, normalize, maxLength, DEFAULT_PADDING, - inferenceOptions); + this(model, vocabulary, lowerCase, pooling, normalize, maxLength, inferenceOptions, null); } @@ -266,19 +273,45 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { final PaddingStrategy padding, final InferenceOptions inferenceOptions) throws OrtException, IOException { + this(model, vocabulary, lowerCase, pooling, normalize, maxLength, inferenceOptions, + ParamChecks.requireNonNullArg(padding, "padding")); + + } + + /** + * Instantiates a {@link SentenceVectorsDL sentence vector generator} using ONNX models. + * + * @param model The file name of a sentence vectors ONNX model. + * @param vocabulary The file name of the vocabulary file for the model. + * @param lowerCase {@code true} for uncased models, {@code false} for cased models. + * @param pooling How token vectors are pooled. Must not be {@code null}. + * @param normalize {@code true} to scale every vector to unit length. + * @param maxLength The maximum number of tokens per input. Must be at least {@code 2}. + * @param inferenceOptions The {@link InferenceOptions}. Must not be {@code null}. + * @param padding The {@link PaddingStrategy}, or {@code null} for the one the device picks, + * see {@link #devicePadding(boolean, Map)}. + * + * @throws IllegalArgumentException Thrown if an argument is invalid. + * @throws OrtException Thrown if the {@code model} cannot be loaded or the GPU is unavailable. + * @throws IOException Thrown if errors occurred loading the {@code model} or {@code vocabulary}. + */ + private SentenceVectorsDL(final File model, final File vocabulary, final boolean lowerCase, + final Pooling pooling, final boolean normalize, final int maxLength, + final InferenceOptions inferenceOptions, final PaddingStrategy padding) + throws OrtException, IOException { + super(model, vocabulary, sessionOptions(inferenceOptions), resolveLowerCase(inferenceOptions, lowerCase)); try { ParamChecks.requireNonNullArg(pooling, "pooling"); - ParamChecks.requireNonNullArg(padding, "padding"); if (maxLength < MIN_LENGTH) { throw new IllegalArgumentException("maxLength must be at least " + MIN_LENGTH); } this.pooling = pooling; this.normalize = normalize; this.maxLength = maxLength; - this.padding = padding; - this.padTokenId = padding == PaddingStrategy.EXACT_LENGTH ? 0 : padTokenId(vocab); + this.padding = padding != null ? padding : devicePadding(inferenceOptions.isGpu(), vocab); + this.padTokenId = this.padding == PaddingStrategy.EXACT_LENGTH ? 0 : padTokenId(vocab); final Map<String, NodeInfo> outputs = session.getOutputInfo(); this.outputName = selectOutput(outputs); final long[] shape = ((TensorInfo) outputs.get(outputName).getInfo()).getShape(); @@ -297,6 +330,17 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { } + /** + * {@return the padding of a session built without a padding strategy: {@link #DEFAULT_PADDING} + * on the CPU, {@link #DEFAULT_GPU_PADDING} on a GPU if the vocabulary has a padding token} + * + * @param gpu {@code true} if the session runs on a GPU. + * @param vocab The vocabulary of the session. + */ + static PaddingStrategy devicePadding(final boolean gpu, final Map<String, Integer> vocab) { + return gpu && hasPadToken(vocab) ? DEFAULT_GPU_PADDING : DEFAULT_PADDING; + } + /** * Generates vectors given a sentence. * 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 61198e668..8e0674bf8 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 @@ -24,13 +24,17 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.stream.Stream; 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.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; import opennlp.dl.InferenceOptions; import opennlp.dl.Tokens; @@ -48,6 +52,7 @@ import static opennlp.dl.vectors.OnnxTextEmbedderProvider.POOLING_OPTION; import static opennlp.dl.vectors.OnnxTextEmbedderProvider.VOCABULARY_OPTION; 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.tokens; @@ -401,6 +406,55 @@ class SentenceVectorsDLEmbedderTest { () -> meanVectors(dir, 8, null)).getMessage()); } + /** + * {@return whether the session runs on a GPU, whether the vocabulary has a padding token, and + * the padding the device picks} + */ + static Stream<Arguments> devicePaddingCases() { + return Stream.of( + Arguments.of(false, true, PaddingStrategy.EXACT_LENGTH), + Arguments.of(false, false, PaddingStrategy.EXACT_LENGTH), + Arguments.of(true, true, PaddingStrategy.LONGEST), + Arguments.of(true, false, PaddingStrategy.EXACT_LENGTH)); + } + + /** + * Asserts that a session on a GPU pads to the longest row by default, and falls back to no + * padding when the vocabulary has no padding token. + */ + @ParameterizedTest + @MethodSource("devicePaddingCases") + void testDevicePadding(final boolean gpu, final boolean padToken, + final PaddingStrategy expected) { + final Map<String, Integer> vocab = new HashMap<>(Map.of("[UNK]", 2, "[SEP]", 3, "[CLS]", 7)); + if (padToken) { + vocab.put("[PAD]", 0); + } + assertEquals(expected, SentenceVectorsDL.devicePadding(gpu, vocab)); + } + + /** + * Asserts that a GPU session from the constructor without a padding strategy runs texts of + * different lengths in one batch, with the vectors of the CPU session; skipped without a CUDA + * device. + */ + @Test + void testGpuBatchesMixedLengthsByDefault(@TempDir final Path dir) throws Exception { + final File model = model(dir, TINY_VECTORS); + final File vocabulary = vocab(dir); + final List<String> texts = List.of("hello world", "hello", "world hello world"); + try (SentenceVectorsDL gpu = gpuOrAbort(model, vocabulary, Pooling.MEAN, false); + SentenceVectorsDL cpu = new SentenceVectorsDL(model, vocabulary, true, Pooling.MEAN, + false, SentenceVectorsDL.DEFAULT_MAX_LENGTH, new InferenceOptions())) { + assertEquals(List.of(List.of(1, 0)), gpu.batches(new Tokens[] {tokens(4), tokens(3)})); + final float[][] expected = cpu.embedAll(texts); + final float[][] actual = gpu.embedAll(texts); + for (int i = 0; i < expected.length; i++) { + assertArrayEquals(expected[i], actual[i], DELTA, "row " + i); + } + } + } + /** * 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. diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/TinyVectorModels.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/TinyVectorModels.java index 8be2defb4..76e54f833 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/TinyVectorModels.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/TinyVectorModels.java @@ -151,8 +151,8 @@ public final class TinyVectorModels { } /** - * Opens a session on CUDA device {@code 0} with the default maximum length and padding, or - * aborts the calling test where there is no usable CUDA device. + * Opens a session on CUDA device {@code 0} with the default maximum length and the padding the + * device picks, or aborts the calling test where there is no usable CUDA device. * * @param model The model file. * @param vocab The vocabulary file. @@ -164,7 +164,7 @@ public final class TinyVectorModels { public static SentenceVectorsDL gpuOrAbort(final File model, final File vocab, final Pooling pooling, final boolean normalize) throws IOException { return gpuOrAbort(model, vocab, pooling, normalize, SentenceVectorsDL.DEFAULT_MAX_LENGTH, - SentenceVectorsDL.DEFAULT_PADDING); + null); } /** @@ -176,7 +176,7 @@ public final class TinyVectorModels { * @param pooling The pooling. * @param normalize {@code true} to scale every vector to unit length. * @param maxLength The maximum number of tokens per input. - * @param padding The padding strategy. + * @param padding The padding strategy, or {@code null} for the one the device picks. * @return The session on the GPU. * @throws IOException Thrown if the model or vocabulary cannot be read. */ @@ -186,8 +186,10 @@ public final class TinyVectorModels { assumeTrue(OrtEnvironment.getAvailableProviders().contains(OrtProvider.CUDA), "the ONNX Runtime on the classpath has no CUDA execution provider"); try { - return new SentenceVectorsDL(model, vocab, true, pooling, normalize, maxLength, padding, - gpu(0)); + return padding == null + ? new SentenceVectorsDL(model, vocab, true, pooling, normalize, maxLength, gpu(0)) + : new SentenceVectorsDL(model, vocab, true, pooling, normalize, maxLength, padding, + gpu(0)); } catch (OrtException e) { return abort("no usable CUDA device 0: " + e.getMessage()); } diff --git a/opennlp-docs/src/docbkx/machine-learning.xml b/opennlp-docs/src/docbkx/machine-learning.xml index 00831b676..415098b19 100644 --- a/opennlp-docs/src/docbkx/machine-learning.xml +++ b/opennlp-docs/src/docbkx/machine-learning.xml @@ -196,7 +196,9 @@ try (TextEmbedder embedder = providers.select(spec).create(spec)) { ONNX Runtime release requires. If CUDA or the configured device is not available, the constructor throws an <code>OrtException</code>; ONNX Runtime does not fall back to the CPU. A lower case setting in the options takes precedence over the <code>lowerCase</code> - argument. + argument. On a GPU, <code>embedAll</code> pads each batch to its longest text by default, + if the vocabulary has a padding token, because one wide inference is faster there than + many narrow ones. The constructor that takes a <code>PaddingStrategy</code> overrides this. </para> <programlisting language="java"><