This is an automated email from the ASF dual-hosted git repository.
mawiesne pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/opennlp.git
The following commit(s) were added to refs/heads/main by this push:
new 95e990636 GH-1391: Cap the token positions of one
SentenceVectorsDL.embedAll inference (#1468)
95e990636 is described below
commit 95e9906364baba21625234dc30f198c7140bbc0c
Author: Kristian Rickert <[email protected]>
AuthorDate: Wed Oct 21 15:24:24 2026 -0400
GH-1391: Cap the token positions of one SentenceVectorsDL.embedAll
inference (#1468)
TextEmbedder.embedAll and the manual say the implementation forms its
own batches, so a caller may pass a list of any size. SentenceVectorsDL
grouped the texts by tokenized length and ran each group as one
inference, however many texts it had. Each group is now split so that no
inference takes more than 16384 token positions (rows times row length),
32 rows at the BERT length of 512. An input longer than that runs in a
batch of its own. Groups run in the
order their length first appears in the call.
Vectors do not change: batches are not padded, so every row is computed
from the same tensors as its single-input call.
The grouping moved into a package-private batches method so tests can
check the inference plan. SentenceVectorsDLBatchLimitTest covers row
count boundaries (exact fit, one over, one under, lengths that do and do
not divide the cap), a length at and over the cap, an empty call, a mix
where only one group splits, input order with alternating lengths, CLS
and mean pooling, the pooled-output model and unit length scaling
through a split group. The GPU parity test also runs a split group, and
the tiny model texts and expected vectors are shared in TinyVectorModels.
Signed-off-by: Kristian Rickert <[email protected]>
Co-authored-by: Claude Code 2.1.296 (Opus 5.5) <[email protected]>
---
.../java/opennlp/dl/vectors/SentenceVectorsDL.java | 49 ++++-
.../vectors/SentenceVectorsDLBatchLimitTest.java | 240 +++++++++++++++++++++
.../dl/vectors/SentenceVectorsDLEmbedderTest.java | 17 +-
.../SentenceVectorsDLInferenceOptionsTest.java | 20 +-
.../java/opennlp/dl/vectors/TinyVectorModels.java | 40 ++++
opennlp-docs/src/docbkx/machine-learning.xml | 4 +-
6 files changed, 343 insertions(+), 27 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 08e114d83..fa4902fe3 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
@@ -23,6 +23,7 @@ import java.io.UncheckedIOException;
import java.nio.LongBuffer;
import java.util.ArrayList;
import java.util.HashMap;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -74,7 +75,8 @@ import opennlp.tools.util.ParamChecks;
* {@code close()} with inference methods.</p>
*
* <p>{@link #embedAll(List)} runs one batched session per distinct tokenized
length, so a batch
- * of same-length inputs costs one inference instead of one per input.</p>
+ * of same-length inputs costs one inference instead of one per input. A group
larger than
+ * 16384 token positions is split into several inferences.</p>
*/
@ThreadSafe
public class SentenceVectorsDL extends AbstractDL implements TextEmbedder {
@@ -82,6 +84,15 @@ 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. The bound is 32 rows at the BERT length of
+ * {@value #DEFAULT_MAX_LENGTH}.
+ */
+ 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;
@@ -268,9 +279,9 @@ public class SentenceVectorsDL extends AbstractDL
implements TextEmbedder {
/**
* {@inheritDoc}
*
- * <p>The inputs are tokenized up front, grouped by tokenized length, and
each group runs
- * through the session once with shape {@code [group size, length]}. A batch
never pads, so
- * every row is computed from the tensors its single-input call would have
used.</p>
+ * <p>The inputs are grouped by tokenized length, never padded, and each
group runs in
+ * inferences of at most 16384 token positions. Every row is computed from
the tensors its
+ * single-input call would have used.</p>
*/
@Override
public float[][] embedAll(final List<? extends CharSequence> texts) {
@@ -282,15 +293,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 +315,32 @@ public class SentenceVectorsDL extends AbstractDL
implements TextEmbedder {
return vectors;
}
+ /**
+ * Splits the indices of a call into the batches that run as single
inferences. A batch contains
+ * indices of one tokenized length, in call order, and no more rows than fit
in
+ * {@value #MAX_BATCH_TOKEN_POSITIONS} token positions, but always at least
one, so a row
+ * wider than the bound runs alone.
+ *
+ * @param encoded The encodings of the call, in call order. An encoding
contains at least
+ * the start and end tokens, so no length is zero.
+ * @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 LinkedHashMap<>();
+ for (int i = 0; i < encoded.length; i++) {
+ byLength.computeIfAbsent(encoded[i].ids().length, length -> new
ArrayList<>()).add(i);
+ }
+ final List<List<Integer>> batches = new ArrayList<>();
+ for (final Map.Entry<Integer, List<Integer>> group : byLength.entrySet()) {
+ final int rows = Math.max(1, MAX_BATCH_TOKEN_POSITIONS / group.getKey());
+ final List<Integer> indices = group.getValue();
+ for (int from = 0; from < indices.size(); from += rows) {
+ batches.add(indices.subList(from, Math.min(from + rows,
indices.size())));
+ }
+ }
+ return batches;
+ }
+
/**
* {@inheritDoc}
*
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..e55fc0e89
--- /dev/null
+++
b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLBatchLimitTest.java
@@ -0,0 +1,240 @@
+/*
+ * 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.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+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.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import opennlp.dl.Tokens;
+
+import static opennlp.dl.vectors.TinyVectorModels.CLS_VECTOR;
+import static opennlp.dl.vectors.TinyVectorModels.DELTA;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_MEAN;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_SUM;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD_MEAN;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD_SUM;
+import static opennlp.dl.vectors.TinyVectorModels.TINY_POOLED;
+import static opennlp.dl.vectors.TinyVectorModels.TINY_VECTORS;
+import static opennlp.dl.vectors.TinyVectorModels.UNIT_VECTOR;
+import static opennlp.dl.vectors.TinyVectorModels.alternating;
+import static opennlp.dl.vectors.TinyVectorModels.model;
+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;
+
+/**
+ * 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
+ * batches, and the vectors that come back through a split group. "hello"
encodes to 3 tokens and
+ * "hello world" to 4.
+ */
+class SentenceVectorsDLBatchLimitTest {
+
+ private static final int LIMIT = SentenceVectorsDL.MAX_BATCH_TOKEN_POSITIONS;
+
+ /**
+ * {@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);
+ }
+
+ /**
+ * {@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]);
+ }
+
+ /**
+ * {@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(encoded)) {
+ sizes.add(batch.size());
+ }
+ return sizes;
+ }
+
+ /**
+ * {@return the length, the count and the expected batch sizes at the edges
of the row count}
+ */
+ static Stream<Arguments> rowCountBoundaries() {
+ final int rows4 = LIMIT / 4;
+ final int rows2 = LIMIT / 2;
+ final int rows3 = LIMIT / 3;
+ return Stream.of(
+ Arguments.of(4, rows4, List.of(rows4)),
+ Arguments.of(4, rows4 + 1, List.of(rows4, 1)),
+ Arguments.of(4, rows4 - 1, List.of(rows4 - 1)),
+ Arguments.of(2, rows2, List.of(rows2)),
+ Arguments.of(2, rows2 + 1, List.of(rows2, 1)),
+ Arguments.of(3, rows3, List.of(rows3)),
+ Arguments.of(3, rows3 + 1, List.of(rows3, 1)),
+ Arguments.of(3, 2 * rows3, List.of(rows3, rows3)),
+ Arguments.of(3, 2 * rows3 + 1, List.of(rows3, rows3, 1)),
+ Arguments.of(1, LIMIT, List.of(LIMIT)),
+ Arguments.of(1, LIMIT + 1, List.of(LIMIT, 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
+ @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: the 4-token inputs
need two
+ * inferences, the 3-token inputs between them stay in one.
+ */
+ @Test
+ void testOnlyTheLargeGroupIsSplit() {
+ 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));
+ }
+
+ /**
+ * Runs the same mix through the session: one group splits, the other does
not. The rows at
+ * the start and end of each inference, and the two short texts, equal the
vector of their 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, Pooling.MEAN, false,
+ SentenceVectorsDL.DEFAULT_MAX_LENGTH)) {
+ final float[][] batch = v.embedAll(texts);
+ assertEquals(texts.size(), batch.length);
+ // The first inference ends with the (LIMIT / 4)-th "hello world", at
index LIMIT / 4 + 1.
+ final int lastOfFirst = LIMIT / 4 + 1;
+ for (final int i : new int[] {0, 1, 2, 7, lastOfFirst, lastOfFirst + 1,
texts.size() - 1}) {
+ assertArrayEquals(v.embed(texts.get(i)), batch[i], DELTA, "row " + i);
+ }
+ }
+ }
+
+ /** 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, 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);
+ }
+ }
+ }
+
+ /**
+ * {@return the model, pooling, unit length, and the expected vectors for
"hello world" and
+ * "hello"}
+ */
+ static Stream<Arguments> pathCases() {
+ return Stream.of(
+ Arguments.of(TINY_VECTORS, Pooling.MEAN, false, HELLO_WORLD_MEAN,
HELLO_MEAN),
+ Arguments.of(TINY_VECTORS, Pooling.CLS, false, CLS_VECTOR, CLS_VECTOR),
+ Arguments.of(TINY_VECTORS, Pooling.CLS, true, UNIT_VECTOR,
UNIT_VECTOR),
+ Arguments.of(TINY_VECTORS, Pooling.MEAN, true, UNIT_VECTOR,
UNIT_VECTOR),
+ Arguments.of(TINY_POOLED, Pooling.CLS, false, HELLO_WORLD_SUM,
HELLO_SUM),
+ Arguments.of(TINY_POOLED, Pooling.MEAN, true, UNIT_VECTOR,
UNIT_VECTOR));
+ }
+
+ /**
+ * Checks every path after the inference through split groups of alternating
lengths, with the
+ * results in input order.
+ */
+ @ParameterizedTest
+ @MethodSource("pathCases")
+ 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);
+ 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);
+ }
+ }
+ }
+}
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 01d2fe470..6d3356330 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
@@ -39,9 +39,14 @@ import static
opennlp.dl.vectors.OnnxTextEmbedderProvider.MAX_LENGTH_OPTION;
import static opennlp.dl.vectors.OnnxTextEmbedderProvider.NORMALIZE_OPTION;
import static opennlp.dl.vectors.OnnxTextEmbedderProvider.POOLING_OPTION;
import static opennlp.dl.vectors.OnnxTextEmbedderProvider.VOCABULARY_OPTION;
+import static opennlp.dl.vectors.TinyVectorModels.CLS_VECTOR;
import static opennlp.dl.vectors.TinyVectorModels.DELTA;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_MEAN;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD_MEAN;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD_SUM;
import static opennlp.dl.vectors.TinyVectorModels.TINY_POOLED;
import static opennlp.dl.vectors.TinyVectorModels.TINY_VECTORS;
+import static opennlp.dl.vectors.TinyVectorModels.UNIT_VECTOR;
import static opennlp.dl.vectors.TinyVectorModels.model;
import static opennlp.dl.vectors.TinyVectorModels.scale;
import static opennlp.dl.vectors.TinyVectorModels.vocab;
@@ -63,18 +68,6 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
*/
class SentenceVectorsDLEmbedderTest {
- // W / |W|, the unit-length vector of every input of the tiny models.
- private static final float[] UNIT_VECTOR = scale(1 / (float)
Math.sqrt(5.25));
-
- // "hello world" = [CLS]=7 hello=4 world=5 [SEP]=3
- private static final float[] HELLO_WORLD_MEAN = scale(19 / 4f);
- private static final float[] HELLO_WORLD_SUM = scale(19);
-
- // "hello" = [CLS]=7 hello=4 [SEP]=3
- private static final float[] HELLO_MEAN = scale(14 / 3f);
-
- private static final float[] CLS_VECTOR = scale(7);
-
// The vocabulary above as a tokenizer.json of an uncased WordPiece model.
private static File tokenizerJson(final Path dir, final boolean lowercase)
throws IOException {
final Path file = dir.resolve("tokenizer.json");
diff --git
a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLInferenceOptionsTest.java
b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLInferenceOptionsTest.java
index da8cb3a70..a19b4fc88 100644
---
a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLInferenceOptionsTest.java
+++
b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLInferenceOptionsTest.java
@@ -35,8 +35,10 @@ import opennlp.dl.InferenceOptions;
import static opennlp.dl.OnnxTestSupport.UNUSABLE_DEVICE_ID;
import static opennlp.dl.OnnxTestSupport.gpu;
import static opennlp.dl.vectors.TinyVectorModels.DELTA;
+import static opennlp.dl.vectors.TinyVectorModels.HELLO_WORLD;
import static opennlp.dl.vectors.TinyVectorModels.TINY_POOLED;
import static opennlp.dl.vectors.TinyVectorModels.TINY_VECTORS;
+import static opennlp.dl.vectors.TinyVectorModels.alternating;
import static opennlp.dl.vectors.TinyVectorModels.gpuOrAbort;
import static opennlp.dl.vectors.TinyVectorModels.model;
import static opennlp.dl.vectors.TinyVectorModels.scale;
@@ -50,10 +52,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
*/
class SentenceVectorsDLInferenceOptionsTest {
- private static final String HELLO = "hello";
-
- private static final String TEXT = "hello world";
-
private static final String LONGER_TEXT = "world hello world";
/**
@@ -83,13 +81,21 @@ class SentenceVectorsDLInferenceOptionsTest {
SentenceVectorsDL cpu = new SentenceVectorsDL(model, vocab, true,
pooling, normalize,
SentenceVectorsDL.DEFAULT_MAX_LENGTH, new InferenceOptions())) {
assertEquals(cpu.dimension(), gpu.dimension());
- assertArrayEquals(cpu.embed(TEXT), gpu.embed(TEXT), DELTA);
- final List<String> texts = List.of(TEXT, LONGER_TEXT, "");
+ assertArrayEquals(cpu.embed(HELLO_WORLD), gpu.embed(HELLO_WORLD), DELTA);
+ final List<String> texts = List.of(HELLO_WORLD, LONGER_TEXT, "");
final float[][] cpuBatch = cpu.embedAll(texts);
final float[][] gpuBatch = gpu.embedAll(texts);
for (int i = 0; i < cpuBatch.length; i++) {
assertArrayEquals(cpuBatch[i], gpuBatch[i], DELTA, "row " + i);
}
+ // Both length groups exceed the token position cap and are split.
+ final List<String> split =
alternating(SentenceVectorsDL.MAX_BATCH_TOKEN_POSITIONS / 3 + 1);
+ final float[][] cpuSplit = cpu.embedAll(split);
+ final float[][] gpuSplit = gpu.embedAll(split);
+ assertEquals(split.size(), gpuSplit.length);
+ for (int i = 0; i < cpuSplit.length; i++) {
+ assertArrayEquals(cpuSplit[i], gpuSplit[i], DELTA, "split row " + i);
+ }
}
}
@@ -170,7 +176,7 @@ class SentenceVectorsDLInferenceOptionsTest {
SentenceVectorsDL.DEFAULT_MAX_LENGTH);
SentenceVectorsDL withOptions = new SentenceVectorsDL(model, vocab,
true, Pooling.MEAN,
false, SentenceVectorsDL.DEFAULT_MAX_LENGTH, new
InferenceOptions())) {
- assertArrayEquals(plain.embed(TEXT), withOptions.embed(TEXT), DELTA);
+ assertArrayEquals(plain.embed(HELLO_WORLD),
withOptions.embed(HELLO_WORLD), DELTA);
}
}
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 d48d4ae16..9fe03f057 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
@@ -20,6 +20,7 @@ import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
+import java.util.ArrayList;
import java.util.List;
import ai.onnxruntime.OrtException;
@@ -47,6 +48,30 @@ public final class TinyVectorModels {
/** The tolerance for comparing vectors of the tiny models. */
public static final float DELTA = 1e-5f;
+ /** A text that encodes to 3 tokens: {@code [CLS]=7 hello=4 [SEP]=3}. */
+ public static final String HELLO = "hello";
+
+ /** A text that encodes to 4 tokens: {@code [CLS]=7 hello=4 world=5
[SEP]=3}. */
+ public static final String HELLO_WORLD = "hello world";
+
+ /** {@code W / |W|}, the unit-length vector of every input. */
+ public static final float[] UNIT_VECTOR = scale(1 / (float) Math.sqrt(5.25));
+
+ /** The mean pooled vector of {@link #HELLO_WORLD}. */
+ public static final float[] HELLO_WORLD_MEAN = scale(19 / 4f);
+
+ /** The summed vector of {@link #HELLO_WORLD}, the {@code
sentence_embedding} of the pooled model. */
+ public static final float[] HELLO_WORLD_SUM = scale(19);
+
+ /** The mean pooled vector of {@link #HELLO}. */
+ public static final float[] HELLO_MEAN = scale(14 / 3f);
+
+ /** The summed vector of {@link #HELLO}, the {@code sentence_embedding} of
the pooled model. */
+ public static final float[] HELLO_SUM = scale(14);
+
+ /** The {@code [CLS]} pooled vector of every input. */
+ public static final float[] CLS_VECTOR = scale(7);
+
private static final String VECTORS_RESOURCES = "/opennlp/dl/vectors/";
private TinyVectorModels() {
@@ -88,6 +113,21 @@ public final class TinyVectorModels {
return new float[] {0.5f * factor, -1f * factor, 2f * factor};
}
+ /**
+ * {@return {@code pairs} times {@link #HELLO_WORLD} followed by {@link
#HELLO}, so the two
+ * lengths alternate}
+ *
+ * @param pairs The number of pairs.
+ */
+ public 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;
+ }
+
/**
* Opens a session on CUDA device {@code 0}, or aborts the calling test
where the ONNX Runtime
* on the classpath has no CUDA execution provider or the device is not
usable.
diff --git a/opennlp-docs/src/docbkx/machine-learning.xml
b/opennlp-docs/src/docbkx/machine-learning.xml
index 69c3f72d3..dc1f0f07a 100644
--- a/opennlp-docs/src/docbkx/machine-learning.xml
+++ b/opennlp-docs/src/docbkx/machine-learning.xml
@@ -157,7 +157,9 @@ org-start wc=ic=1.0 w&c=apache,ic=2.0 p1wc=lc=3.0]]>
<para>
An embedder is safe for concurrent use until it is
closed. <code>embedAll</code> checks
every text before embedding any of them and forms its
own batches, so it accepts a list
- of any size. Inference failures are reported as
<code>EmbeddingException</code>.
+ of any size. No inference of
<code>SentenceVectorsDL</code> exceeds 16,384 token
+ positions, except that a single text whose encoding is
longer runs as an inference of
+ its own. Inference failures are reported as
<code>EmbeddingException</code>.
</para>
<programlisting language="java"><