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 cd2fbba3cff059ccc658294c5f223d2dbfd5b1e8 Author: Kristian Rickert <[email protected]> AuthorDate: Sat Oct 10 08:26:13 2026 -0400 GH-1391: Cap the token positions of one embedAll inference embedAll ran each group of same-length texts as one inference, however large the group was. Groups are now split so that no inference takes more than 16384 token positions. An input longer than that runs in a batch of its own. Groups run in the order their length first appears in the call. Signed-off-by: Kristian Rickert <[email protected]> --- .../java/opennlp/dl/vectors/SentenceVectorsDL.java | 26 ++++++++++++++++------ 1 file changed, 19 insertions(+), 7 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 f62f59f7e..a5afc055f 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 + * {@value #MAX_BATCH_TOKEN_POSITIONS} token positions is split into several inferences.</p> */ @ThreadSafe public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { @@ -277,8 +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> + * through the session with shape {@code [rows, length]}, split so that no inference exceeds + * {@value #MAX_BATCH_TOKEN_POSITIONS} token positions. A batch never pads, so every row is + * computed from the tensors its single-input call would have used.</p> */ @Override public float[][] embedAll(final List<? extends CharSequence> texts) { @@ -313,18 +316,27 @@ public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { } /** - * Splits the indices of a call into the batches that run as single inferences, one batch per - * tokenized length. + * 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. * * @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<>(); + 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); } - return new ArrayList<>(byLength.values()); + 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; } /**
