(opennlp) 01/07: GH-1391: Test that embedAll caps the token positions of one inference

[kr...@apache.org]

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 <kr...@apache.org>
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 <kr...@apache.org>
---
 .../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.