This is an automated email from the ASF dual-hosted git repository. krickert pushed a commit to branch static-embeddings in repository https://gitbox.apache.org/repos/asf/opennlp.git
commit 9dd6ff7dd77e933015efc89ec8c392140a52cd03 Author: Kristian Rickert <[email protected]> AuthorDate: Fri Jul 10 16:49:27 2026 -0400 OPENNLP-1877: Address review: value-safe TensorInfo, checked IO, typed load switches, distribution wiring TensorInfo copies its shape on construction and access and compares by value, so callers cannot corrupt validated metadata and equality behaves like exported API should; elementCount() detects overflow from crafted headers. The load and read entry points now declare IOException like the rest of the project's resource loaders instead of wrapping in UncheckedIOException. The two adjacent boolean load flags became the Casing and Normalization enums, so the switches cannot be swapped silently; the directory-based load maps the model's configuration onto them and the manual, README, tests, and benchmark use the new signature. The binary distribution now ships the module: opennlp-embeddings added to the distr dependencies and its apidocs to the assembly, matching the other extension modules. JSON parsing tightened to the fail-loud contract: unicode escapes must be four hex digits (Integer.parseInt also accepted signed input) and numbers in skipped fields are held to the JSON grammar. WordPieceVocabulary gains its own test class covering the line-number ids, the duplicate rejection, the -1 sentinel, and reverse-lookup bounds, which now fail loud instead of leaking an index exception. The safetensors fixture writer that was copied across the tests is now a shared test utility, the stale cross-module javadoc reference is gone, the element-count versus byte-range guard has its negative test, and else/catch placement matches the project style. --- opennlp-distr/pom.xml | 5 ++ opennlp-distr/src/main/assembly/bin.xml | 7 ++ opennlp-docs/src/docbkx/embeddings.xml | 4 +- opennlp-extensions/opennlp-embeddings/README.md | 4 +- .../embeddings/StaticEmbeddingModelBenchmark.java | 5 +- .../java/opennlp/embeddings/FlatJsonFields.java | 25 ++---- .../main/java/opennlp/embeddings/JsonCursor.java | 76 +++++++++++++------ .../src/main/java/opennlp/embeddings/Neighbor.java | 1 + .../java/opennlp/embeddings/SafetensorsFile.java | 15 +--- .../embeddings/SafetensorsHeaderParser.java | 3 +- .../opennlp/embeddings/StaticEmbeddingModel.java | 71 ++++++++++++----- .../main/java/opennlp/embeddings/TensorInfo.java | 53 ++++++++++++- .../opennlp/embeddings/WordPieceVocabulary.java | 18 ++--- .../opennlp/embeddings/FlatJsonFieldsTest.java | 4 +- .../opennlp/embeddings/SafetensorsFileTest.java | 42 +++++++++++ .../embeddings/SafetensorsHeaderParserTest.java | 22 ++++++ .../opennlp/embeddings/SafetensorsTestFiles.java | 88 ++++++++++++++++++++++ .../StaticEmbeddingModelConcurrencyTest.java | 32 ++------ .../StaticEmbeddingModelSimilarityTest.java | 27 ++----- .../embeddings/StaticEmbeddingModelTest.java | 76 ++++++++----------- .../embeddings/WordPieceVocabularyTest.java | 81 ++++++++++++++++++++ 21 files changed, 479 insertions(+), 180 deletions(-) diff --git a/opennlp-distr/pom.xml b/opennlp-distr/pom.xml index e9092d882..a0351c248 100644 --- a/opennlp-distr/pom.xml +++ b/opennlp-distr/pom.xml @@ -92,6 +92,11 @@ <artifactId>opennlp-spellcheck</artifactId> </dependency> + <dependency> + <groupId>org.apache.opennlp</groupId> + <artifactId>opennlp-embeddings</artifactId> + </dependency> + <!-- OpenNLP Docs - Ensure that "docs" are build before distribution --> <dependency> <groupId>org.apache.opennlp</groupId> diff --git a/opennlp-distr/src/main/assembly/bin.xml b/opennlp-distr/src/main/assembly/bin.xml index 2db4eafc6..34e38c4c4 100644 --- a/opennlp-distr/src/main/assembly/bin.xml +++ b/opennlp-distr/src/main/assembly/bin.xml @@ -232,6 +232,13 @@ <outputDirectory>docs/apidocs/opennlp-morfologik</outputDirectory> </fileSet> + <fileSet> + <directory>../opennlp-extensions/opennlp-embeddings/target/reports/apidocs</directory> + <fileMode>644</fileMode> + <directoryMode>755</directoryMode> + <outputDirectory>docs/apidocs/opennlp-embeddings</outputDirectory> + </fileSet> + <fileSet> <directory>../opennlp-extensions/opennlp-spellcheck/target/reports/apidocs</directory> <fileMode>644</fileMode> diff --git a/opennlp-docs/src/docbkx/embeddings.xml b/opennlp-docs/src/docbkx/embeddings.xml index 09e2059aa..abfe59c07 100644 --- a/opennlp-docs/src/docbkx/embeddings.xml +++ b/opennlp-docs/src/docbkx/embeddings.xml @@ -69,8 +69,8 @@ List<Neighbor> analogy = model.analogy("man", "king", "woman", 1);]]> <![CDATA[ StaticEmbeddingModel model = StaticEmbeddingModel.load( Path.of("vocab.txt"), Path.of("model.safetensors"), - true, // lowerCase, from the model's do_lower_case - true); // normalize, from the model's config + StaticEmbeddingModel.Casing.UNCASED, // from the model's do_lower_case + StaticEmbeddingModel.Normalization.L2); // from the model's config ]]> </programlisting> <para> diff --git a/opennlp-extensions/opennlp-embeddings/README.md b/opennlp-extensions/opennlp-embeddings/README.md index cc924cb49..5c64565f8 100644 --- a/opennlp-extensions/opennlp-embeddings/README.md +++ b/opennlp-extensions/opennlp-embeddings/README.md @@ -41,8 +41,8 @@ For a model laid out differently, the explicit overload takes the two data files ```java StaticEmbeddingModel model = StaticEmbeddingModel.load( Path.of("vocab.txt"), Path.of("model.safetensors"), - true, // lowerCase, from the model's do_lower_case - true); // normalize, from the model's config + StaticEmbeddingModel.Casing.UNCASED, // from the model's do_lower_case + StaticEmbeddingModel.Normalization.L2); // from the model's config ``` Instances are immutable and safe for concurrent use, so one loaded model can serve every thread of an application. Texts with no in-vocabulary tokens embed to a zero vector rather than raising an error. diff --git a/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java index 03b5d7571..1ede49337 100644 --- a/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java +++ b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java @@ -45,6 +45,8 @@ import org.openjdk.jmh.infra.Blackhole; import org.openjdk.jmh.runner.Runner; import org.openjdk.jmh.runner.options.Options; import org.openjdk.jmh.runner.options.OptionsBuilder; +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; /** * JMH benchmark for {@link StaticEmbeddingModel}, the raw-lookup-throughput number the module's @@ -92,7 +94,8 @@ public class StaticEmbeddingModelBenchmark { tempDir = Files.createTempDirectory("opennlp-embeddings-jmh"); final Path vocabFile = writeVocab(tempDir); final Path safetensorsFile = writeSafetensors(tempDir); - model = StaticEmbeddingModel.load(vocabFile, safetensorsFile, true, true); + model = StaticEmbeddingModel.load(vocabFile, safetensorsFile, + Casing.UNCASED, Normalization.L2); } @TearDown(Level.Trial) diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FlatJsonFields.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FlatJsonFields.java index 6ba963dbd..81fb96967 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FlatJsonFields.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FlatJsonFields.java @@ -17,7 +17,6 @@ package opennlp.embeddings; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; @@ -44,16 +43,10 @@ final class FlatJsonFields { * {@code null} (the formats treat those the same: fall back to the default). * @throws IllegalArgumentException Thrown if the file is not a well-formed JSON object, the * field appears more than once, or its value is neither a boolean nor {@code null}. - * @throws UncheckedIOException Thrown if reading the file fails. + * @throws IOException Thrown if reading the file fails. */ - static Boolean topLevelBoolean(Path file, String field) { - final String json; - try { - json = Files.readString(file); - } - catch (IOException e) { - throw new UncheckedIOException("Unable to read " + file, e); - } + static Boolean topLevelBoolean(Path file, String field) throws IOException { + final String json = Files.readString(file); final JsonCursor cursor = new JsonCursor(json, file.getFileName().toString()); cursor.skipWhitespace(); cursor.expect('{'); @@ -62,8 +55,7 @@ final class FlatJsonFields { boolean seen = false; if (cursor.peek() == '}') { cursor.consume(); - } - else { + } else { while (true) { cursor.skipWhitespace(); final String key = cursor.parseString(); @@ -77,15 +69,12 @@ final class FlatJsonFields { seen = true; if (cursor.consumeLiteral("true")) { value = Boolean.TRUE; - } - else if (cursor.consumeLiteral("false")) { + } else if (cursor.consumeLiteral("false")) { value = Boolean.FALSE; - } - else if (!cursor.consumeLiteral("null")) { + } else if (!cursor.consumeLiteral("null")) { throw cursor.malformed("Field '" + field + "' must be a boolean or null"); } - } - else { + } else { cursor.skipValue(); } cursor.skipWhitespace(); diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java index 6947950d0..237b6e411 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java @@ -96,8 +96,7 @@ final class JsonCursor { } if (c == '\\') { value.append(parseEscape()); - } - else { + } else { value.append(c); } } @@ -128,11 +127,53 @@ final class JsonCursor { } final String hex = text.substring(position, position + 4); position += 4; - try { - return (char) Integer.parseInt(hex, 16); + // Each of the four characters must be a hex digit; Integer.parseInt alone would also + // accept a sign and silently decode the wrong character. + int value = 0; + for (int i = 0; i < 4; i++) { + final int digit = Character.digit(hex.charAt(i), 16); + if (digit < 0) { + throw malformed("Malformed \\u escape sequence: " + hex); + } + value = (value << 4) | digit; } - catch (NumberFormatException e) { - throw malformed("Malformed \\u escape sequence: " + hex); + return (char) value; + } + + // Skips one number, holding it to the JSON grammar (optional minus, digits, optional + // fraction, optional signed exponent) so malformed input fails loud even in skipped fields. + private void skipNumber() { + if (peek() == '-') { + position++; + } + if (position >= text.length() || !Character.isDigit(text.charAt(position))) { + throw malformed("Malformed number"); + } + while (position < text.length() && Character.isDigit(text.charAt(position))) { + position++; + } + if (position < text.length() && text.charAt(position) == '.') { + position++; + if (position >= text.length() || !Character.isDigit(text.charAt(position))) { + throw malformed("Malformed number: digit expected after the decimal point"); + } + while (position < text.length() && Character.isDigit(text.charAt(position))) { + position++; + } + } + if (position < text.length() + && (text.charAt(position) == 'e' || text.charAt(position) == 'E')) { + position++; + if (position < text.length() + && (text.charAt(position) == '+' || text.charAt(position) == '-')) { + position++; + } + if (position >= text.length() || !Character.isDigit(text.charAt(position))) { + throw malformed("Malformed number: digit expected in the exponent"); + } + while (position < text.length() && Character.isDigit(text.charAt(position))) { + position++; + } } } @@ -149,8 +190,7 @@ final class JsonCursor { } try { return Long.parseLong(text.substring(start, position)); - } - catch (NumberFormatException e) { + } catch (NumberFormatException e) { throw malformed("Malformed integer: " + text.substring(start, position)); } } @@ -162,8 +202,7 @@ final class JsonCursor { final char c = peek(); if (c == '"') { parseString(); - } - else if (c == '[') { + } else if (c == '[') { position++; skipWhitespace(); if (peek() != ']') { @@ -182,8 +221,7 @@ final class JsonCursor { } } position++; - } - else if (c == '{') { + } else if (c == '{') { position++; skipWhitespace(); if (peek() != '}') { @@ -205,17 +243,11 @@ final class JsonCursor { } } position++; - } - else if (c == '-' || Character.isDigit(c)) { - position++; - while (position < text.length() && "0123456789.eE+-".indexOf(text.charAt(position)) >= 0) { - position++; - } - } - else if (consumeLiteral("true") || consumeLiteral("false") || consumeLiteral("null")) { + } else if (c == '-' || Character.isDigit(c)) { + skipNumber(); + } else if (consumeLiteral("true") || consumeLiteral("false") || consumeLiteral("null")) { // consumed, nothing to record - } - else { + } else { throw malformed("Unexpected character while skipping a value: '" + c + "'"); } } diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java index 68d6f3d58..b37a9b780 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java @@ -16,6 +16,7 @@ */ package opennlp.embeddings; + /** * One vocabulary token found near a query vector by {@link StaticEmbeddingModel#mostSimilar} * or {@link StaticEmbeddingModel#analogy}, most similar first. diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java index 3fd36d2b0..36c048b16 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java @@ -17,7 +17,6 @@ package opennlp.embeddings; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.FileChannel; @@ -89,9 +88,9 @@ public final class SafetensorsFile { * file's actual length. * @throws IllegalArgumentException Thrown if {@code file} is {@code null} or missing, or the * file is malformed. - * @throws UncheckedIOException Thrown if reading the file fails. + * @throws IOException Thrown if reading the file fails. */ - public static SafetensorsFile read(Path file) { + public static SafetensorsFile read(Path file) throws IOException { if (file == null) { throw new IllegalArgumentException("File must not be null"); } @@ -139,9 +138,6 @@ public final class SafetensorsFile { return new SafetensorsFile(file, dataStart, Collections.unmodifiableMap(tensorsByName), Collections.unmodifiableMap(parsed.metadata())); } - catch (IOException e) { - throw new UncheckedIOException("Unable to read safetensors file " + file, e); - } } /** {@return the names of every tensor declared in the header, in header order} */ @@ -178,9 +174,9 @@ public final class SafetensorsFile { * this file, not declared with dtype {@code F32}, or larger than a Java array can hold. * @throws IllegalStateException Thrown if the file has been truncated since * {@link #read(Path)} validated the tensor's byte range. - * @throws UncheckedIOException Thrown if reading the file fails. + * @throws IOException Thrown if reading the file fails. */ - public float[] readFloat32(String name) { + public float[] readFloat32(String name) throws IOException { final TensorInfo info = tensorInfo(name); if (!"F32".equals(info.dtype())) { throw new IllegalArgumentException( @@ -222,9 +218,6 @@ public final class SafetensorsFile { } return values; } - catch (IOException e) { - throw new UncheckedIOException("Unable to read tensor '" + name + "' from " + file, e); - } } // Fills the buffer with bytes starting at the given file position; fails loud if the file diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java index 8a6b0989e..41a54cb90 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java @@ -76,8 +76,7 @@ final class SafetensorsHeaderParser { cursor.skipWhitespace(); if (METADATA_KEY.equals(key)) { metadata = parseStringMap(); - } - else { + } else { tensors.add(parseTensorInfo(key)); } cursor.skipWhitespace(); diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java index c10244524..99755c8fe 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java @@ -16,6 +16,7 @@ */ package opennlp.embeddings; +import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.List; @@ -55,6 +56,26 @@ import opennlp.tools.tokenize.WordpieceTokenizer; @ThreadSafe public final class StaticEmbeddingModel { + /** How the tokenizer treats letter case, matching the base model's tokenizer configuration. */ + public enum Casing { + + /** Lower-case and strip accents, the uncased BGE/BERT convention. */ + UNCASED, + + /** Preserve case and accents. */ + CASED + } + + /** Whether pooled vectors are length-normalized, matching the model's configuration. */ + public enum Normalization { + + /** L2-normalize each pooled vector. */ + L2, + + /** Leave pooled vectors unnormalized. */ + NONE + } + private static final float NORMALIZE_EPSILON = 1e-12f; private static final String WEIGHTS_TENSOR_NAME = "weights"; private static final String VOCABULARY_FILE_NAME = "vocab.txt"; @@ -104,7 +125,7 @@ public final class StaticEmbeddingModel { * * <p>A {@code strip_accents} that is absent or JSON {@code null} follows the BERT convention * of stripping accents exactly when lower-casing, which is what the single lower-case switch - * of {@link #load(Path, Path, boolean, boolean)} does. A model that explicitly sets + * of {@link #load(Path, Path, Casing, Normalization)} does. A model that explicitly sets * {@code strip_accents} against its {@code do_lower_case} value cannot be represented by * that switch, so it is rejected rather than silently mis-tokenized.</p> * @@ -116,9 +137,9 @@ public final class StaticEmbeddingModel { * field, the accent handling is not representable, or the vocabulary and the embedding * matrix disagree; the message names the explicit overload as the fallback for * differently laid-out models. - * @throws java.io.UncheckedIOException Thrown if reading a file fails. + * @throws IOException Thrown if reading a file fails. */ - public static StaticEmbeddingModel load(Path modelDirectory) { + public static StaticEmbeddingModel load(Path modelDirectory) throws IOException { if (modelDirectory == null) { throw new IllegalArgumentException("ModelDirectory must not be null"); } @@ -133,15 +154,15 @@ public final class StaticEmbeddingModel { final Boolean normalize = FlatJsonFields.topLevelBoolean(configFile, "normalize"); if (normalize == null) { throw new IllegalArgumentException(configFile + " has no boolean 'normalize' field; " - + "use load(vocabularyFile, safetensorsFile, lowerCase, normalize) and choose " + + "use load(vocabularyFile, safetensorsFile, casing, normalization) and choose " + "explicitly"); } final Boolean lowerCase = FlatJsonFields.topLevelBoolean(tokenizerConfigFile, "do_lower_case"); if (lowerCase == null) { throw new IllegalArgumentException(tokenizerConfigFile + " has no boolean " - + "'do_lower_case' field; use load(vocabularyFile, safetensorsFile, lowerCase, " - + "normalize) and choose explicitly"); + + "'do_lower_case' field; use load(vocabularyFile, safetensorsFile, casing, " + + "normalization) and choose explicitly"); } final Boolean stripAccents = FlatJsonFields.topLevelBoolean(tokenizerConfigFile, "strip_accents"); @@ -149,10 +170,12 @@ public final class StaticEmbeddingModel { throw new IllegalArgumentException(tokenizerConfigFile + " sets strip_accents=" + stripAccents + " against do_lower_case=" + lowerCase + "; the single lower-case " + "switch strips accents exactly when lower-casing, so this model must be loaded " - + "with load(vocabularyFile, safetensorsFile, lowerCase, normalize) after choosing " + + "with load(vocabularyFile, safetensorsFile, casing, normalization) after choosing " + "deliberately"); } - return load(vocabularyFile, safetensorsFile, lowerCase, normalize); + return load(vocabularyFile, safetensorsFile, + lowerCase ? Casing.UNCASED : Casing.CASED, + normalize ? Normalization.L2 : Normalization.NONE); } private static Path requiredFile(Path modelDirectory, String name) { @@ -160,7 +183,7 @@ public final class StaticEmbeddingModel { if (!Files.isRegularFile(file)) { throw new IllegalArgumentException("Model directory " + modelDirectory + " has no " + name + "; for a different layout, use load(vocabularyFile, safetensorsFile, " - + "lowerCase, normalize)"); + + "casing, normalization)"); } return file; } @@ -179,23 +202,35 @@ public final class StaticEmbeddingModel { * An optional 1-D {@code F32} tensor named {@code "weights"}, one * scalar per vocabulary row, is used as a per-token pooling weight * when present. - * @param lowerCase Whether the tokenizer should lower-case and strip accents, matching - * the base model's tokenizer configuration ({@code true} for the - * uncased BGE/BERT family this module targets). - * @param normalize Whether {@link #embed(String)} L2-normalizes its result, matching - * the source model's {@code config.json} {@code normalize} field. + * @param casing Whether the tokenizer lower-cases and strips accents + * ({@link Casing#UNCASED}, matching the uncased BGE/BERT family this + * module targets) or preserves case ({@link Casing#CASED}), matching + * the base model's tokenizer configuration. + * @param normalization Whether {@link #embed(String)} L2-normalizes its result + * ({@link Normalization#L2}), matching the source model's + * {@code config.json} {@code normalize} field. * @return The loaded model. * @throws IllegalArgumentException Thrown if an argument is {@code null}, a file is missing * or malformed, or the vocabulary size and the embedding matrix's row count disagree. + * @throws IOException Thrown if reading a file fails. */ public static StaticEmbeddingModel load(Path vocabularyFile, Path safetensorsFile, - boolean lowerCase, boolean normalize) { + Casing casing, Normalization normalization) + throws IOException { if (vocabularyFile == null) { throw new IllegalArgumentException("VocabularyFile must not be null"); } if (safetensorsFile == null) { throw new IllegalArgumentException("SafetensorsFile must not be null"); } + if (casing == null) { + throw new IllegalArgumentException("Casing must not be null"); + } + if (normalization == null) { + throw new IllegalArgumentException("Normalization must not be null"); + } + final boolean lowerCase = casing == Casing.UNCASED; + final boolean normalize = normalization == Normalization.L2; final WordPieceVocabulary vocabulary = WordPieceVocabulary.read(vocabularyFile); final SafetensorsFile tensors = SafetensorsFile.read(safetensorsFile); @@ -276,8 +311,7 @@ public final class StaticEmbeddingModel { for (int d = 0; d < dimension; d++) { sum[d] += embeddings[base + d]; } - } - else { + } else { final float weight = weights[row]; for (int d = 0; d < dimension; d++) { sum[d] += embeddings[base + d] * weight; @@ -529,8 +563,7 @@ public final class StaticEmbeddingModel { swap(parent, i); i = parent; } - } - else if (similarity > similarities[0]) { + } else if (similarity > similarities[0]) { similarities[0] = similarity; rows[0] = row; siftDown(); diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.java index 64c0bb8d8..e07c7c212 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.java @@ -16,6 +16,9 @@ */ package opennlp.embeddings; +import java.util.Arrays; + + /** * Header metadata for one tensor in a safetensors file, as declared by the file's own JSON * header. Carries no data; {@link SafetensorsFile#readFloat32(String)} resolves the bytes. @@ -32,14 +35,62 @@ package opennlp.embeddings; public record TensorInfo(String name, String dtype, int[] shape, long dataOffsetBegin, long dataOffsetEnd) { + /** + * Creates the metadata, copying {@code shape} so later mutation of the caller's array cannot + * corrupt the validated state. + */ + public TensorInfo { + shape = shape.clone(); + } + + /** + * @return The tensor's dimensions, outermost first, as a copy; mutating it does not affect + * this record. + */ + @Override + public int[] shape() { + return shape.clone(); + } + /** * @return The number of elements the tensor holds, the product of {@link #shape()}. + * @throws IllegalArgumentException Thrown if the product overflows a {@code long}, which only + * a crafted header can produce. */ public long elementCount() { long count = 1; for (int dimension : shape) { - count *= dimension; + try { + count = Math.multiplyExact(count, dimension); + } catch (ArithmeticException e) { + throw new IllegalArgumentException("Tensor '" + name + "' declares a shape " + + Arrays.toString(shape) + " whose element count overflows a long", e); + } } return count; } + + @Override + public boolean equals(Object other) { + return other instanceof TensorInfo that + && name.equals(that.name) && dtype.equals(that.dtype) + && Arrays.equals(shape, that.shape) + && dataOffsetBegin == that.dataOffsetBegin && dataOffsetEnd == that.dataOffsetEnd; + } + + @Override + public int hashCode() { + int result = name.hashCode(); + result = 31 * result + dtype.hashCode(); + result = 31 * result + Arrays.hashCode(shape); + result = 31 * result + Long.hashCode(dataOffsetBegin); + result = 31 * result + Long.hashCode(dataOffsetEnd); + return result; + } + + @Override + public String toString() { + return "TensorInfo[name=" + name + ", dtype=" + dtype + ", shape=" + Arrays.toString(shape) + + ", dataOffsetBegin=" + dataOffsetBegin + ", dataOffsetEnd=" + dataOffsetEnd + "]"; + } } diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/WordPieceVocabulary.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/WordPieceVocabulary.java index cdff03c83..7cfd65d10 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/WordPieceVocabulary.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/WordPieceVocabulary.java @@ -17,7 +17,6 @@ package opennlp.embeddings; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Collections; @@ -55,23 +54,16 @@ final class WordPieceVocabulary { * @return The parsed vocabulary. * @throws IllegalArgumentException Thrown if {@code file} is {@code null}, missing, or * contains a duplicate token. - * @throws UncheckedIOException Thrown if reading the file fails. + * @throws IOException Thrown if reading the file fails. */ - static WordPieceVocabulary read(Path file) { + static WordPieceVocabulary read(Path file) throws IOException { if (file == null) { throw new IllegalArgumentException("File must not be null"); } if (!Files.isRegularFile(file)) { throw new IllegalArgumentException("File does not exist or is not a regular file: " + file); } - final List<String> lines; - try { - lines = Files.readAllLines(file); - } - catch (IOException e) { - throw new UncheckedIOException("Unable to read vocabulary file " + file, e); - } - return fromLines(lines, file.toString()); + return fromLines(Files.readAllLines(file), file.toString()); } // Package-private so tests can build a vocabulary from in-memory lines without a temp file. @@ -121,6 +113,10 @@ final class WordPieceVocabulary { * @return The token at that id. */ String token(int id) { + if (id < 0 || id >= tokenById.size()) { + throw new IllegalArgumentException( + "Id " + id + " is outside [0, " + tokenById.size() + ")"); + } return tokenById.get(id); } } diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/FlatJsonFieldsTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/FlatJsonFieldsTest.java index 27fee6954..c37bb51ad 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/FlatJsonFieldsTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/FlatJsonFieldsTest.java @@ -17,7 +17,7 @@ package opennlp.embeddings; import java.io.IOException; -import java.io.UncheckedIOException; +import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; @@ -113,7 +113,7 @@ class FlatJsonFieldsTest { @Test void testMissingFileFailsAsAnIoProblem(@TempDir Path dir) { - assertThrows(UncheckedIOException.class, + assertThrows(IOException.class, () -> FlatJsonFields.topLevelBoolean(dir.resolve("absent.json"), "normalize")); } } diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java index 21a6cb30d..1595fd2c2 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java @@ -291,4 +291,46 @@ class SafetensorsFileTest { assertThrows(IllegalStateException.class, () -> parsed.readFloat32("w")); assertTrue(e.getMessage().contains("truncated")); } + + @Test + void testReadFloat32RejectsElementCountByteRangeMismatch(@TempDir Path dir) throws IOException { + // Shape [2] declares two F32 elements (8 bytes) but the data range holds only one. + final byte[] data = floatsToLittleEndianBytes(1f); + final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[2]," + + "\"data_offsets\":[0," + data.length + "]}}"; + final Path file = writeFile(dir, "model.safetensors", header, data); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + final IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> parsed.readFloat32("w")); + assertTrue(e.getMessage().contains("2 F32 elements"), e.getMessage()); + } + + @Test + void testTensorInfoShapeIsDefensivelyCopied() { + final int[] shape = {2, 3}; + final TensorInfo info = new TensorInfo("t", "F32", shape, 0, 24); + shape[0] = 99; + assertEquals(2, info.shape()[0], "construction must copy the caller's array"); + info.shape()[0] = 99; + assertEquals(2, info.shape()[0], "the accessor must return a copy"); + assertEquals(6, info.elementCount()); + } + + @Test + void testTensorInfoEqualsByValue() { + final TensorInfo a = new TensorInfo("t", "F32", new int[] {2, 3}, 0, 24); + final TensorInfo b = new TensorInfo("t", "F32", new int[] {2, 3}, 0, 24); + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + } + + @Test + void testTensorInfoElementCountOverflowFailsLoudly() { + final TensorInfo crafted = new TensorInfo("t", "F32", + new int[] {Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE}, 0, 8); + final IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, crafted::elementCount); + assertTrue(e.getMessage().contains("overflows"), e.getMessage()); + } } diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java index e9ebf7153..966de9e19 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java @@ -155,4 +155,26 @@ class SafetensorsHeaderParserTest { assertTrue(e.getMessage().contains("Malformed safetensors header at offset"), () -> "Message should carry the offset, got: " + e.getMessage()); } + @Test + void testSignedUnicodeEscapeFailsLoudly() { + // Integer.parseInt would accept "-0FF" and decode the wrong character; the parser must not. + final String header = "{\"__metadata__\":{\"note\":\"a\\u-0FFb\"}," + + "\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}}"; + assertThrows(IllegalArgumentException.class, () -> SafetensorsHeaderParser.parse(header)); + } + + @Test + void testMalformedNumberInSkippedFieldFailsLoudly() { + // Skipped unknown fields still hold values to the JSON grammar; "1e++--..5" is not a number. + final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4],\"unknown\":1e++--..5}}"; + assertThrows(IllegalArgumentException.class, () -> SafetensorsHeaderParser.parse(header)); + } + + @Test + void testWellFormedNumbersInSkippedFieldsAreAccepted() { + final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4],\"a\":-1.5e+10,\"b\":0.25,\"c\":3}}"; + assertEquals(1, SafetensorsHeaderParser.parse(header).tensors().size()); + } } diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsTestFiles.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsTestFiles.java new file mode 100644 index 000000000..f590a4edb --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsTestFiles.java @@ -0,0 +1,88 @@ +/* + * 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.embeddings; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.StringJoiner; + +/** + * Writes small well-formed safetensors fixtures for tests and benchmarks, replacing the writer + * that used to be copied into every test class. Negative tests that need deliberately malformed + * bytes still hand-roll them. + */ +final class SafetensorsTestFiles { + + private SafetensorsTestFiles() { + } + + /** One F32 tensor to write: a name, a shape, and the values in row-major order. */ + record Tensor(String name, int[] shape, float[] values) { + } + + static Tensor matrix(String name, float[][] rows) { + final int dimension = rows[0].length; + final float[] values = new float[rows.length * dimension]; + for (int r = 0; r < rows.length; r++) { + System.arraycopy(rows[r], 0, values, r * dimension, dimension); + } + return new Tensor(name, new int[] {rows.length, dimension}, values); + } + + static Tensor vector(String name, float[] values) { + return new Tensor(name, new int[] {values.length}, values); + } + + /** + * Writes a safetensors file holding the given F32 tensors, header first, data in declaration + * order. + */ + static void write(Path file, Tensor... tensors) throws IOException { + final ByteArrayOutputStream data = new ByteArrayOutputStream(); + final StringJoiner header = new StringJoiner(",", "{", "}"); + int offset = 0; + for (final Tensor tensor : tensors) { + final ByteBuffer buffer = + ByteBuffer.allocate(tensor.values().length * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + for (final float value : tensor.values()) { + buffer.putFloat(value); + } + data.writeBytes(buffer.array()); + final StringJoiner shape = new StringJoiner(",", "[", "]"); + for (final int dimension : tensor.shape()) { + shape.add(Integer.toString(dimension)); + } + final int end = offset + tensor.values().length * Float.BYTES; + header.add("\"" + tensor.name() + "\":{\"dtype\":\"F32\",\"shape\":" + shape + + ",\"data_offsets\":[" + offset + "," + end + "]}"); + offset = end; + } + final byte[] headerBytes = header.toString().getBytes(StandardCharsets.UTF_8); + final ByteBuffer out = ByteBuffer.allocate(8 + headerBytes.length + data.size()) + .order(ByteOrder.LITTLE_ENDIAN); + out.putLong(headerBytes.length); + out.put(headerBytes); + out.put(data.toByteArray()); + Files.write(file, out.array()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelConcurrencyTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelConcurrencyTest.java index 1ccc23abd..840ff58e0 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelConcurrencyTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelConcurrencyTest.java @@ -16,11 +16,7 @@ */ package opennlp.embeddings; -import java.io.ByteArrayOutputStream; import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; @@ -35,14 +31,17 @@ import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; + import static org.junit.jupiter.api.Assertions.assertTrue; /** * A concurrency smoke test for the {@code @ThreadSafe} claim on {@link StaticEmbeddingModel}: * one shared instance, many threads, every concurrent result compared against the * single-threaded reference computed up front. All operations are deterministic, so any - * deviation under concurrency is a thread-safety defect by definition. Mirrors the - * {@code LexiconConcurrencyTest} pattern from the opennlp-wordnet module. + * deviation under concurrency is a thread-safety defect by definition: one shared instance, + * reference results computed single-threaded first, then compared under contention. */ class StaticEmbeddingModelConcurrencyTest { @@ -57,25 +56,10 @@ class StaticEmbeddingModelConcurrencyTest { {0f, 0f}, {0f, 0f}, {0f, 0f}, {3f, 3f}, {2f, 4f}, {2f, 1f}, {1f, 2f}, {-3f, -1f}, }; - final ByteBuffer buffer = ByteBuffer.allocate(rows.length * 2 * 4) - .order(ByteOrder.LITTLE_ENDIAN); - for (final float[] row : rows) { - for (final float value : row) { - buffer.putFloat(value); - } - } - final byte[] data = buffer.array(); - final String header = "{\"embeddings\":{\"dtype\":\"F32\",\"shape\":[" + rows.length - + ",2],\"data_offsets\":[0," + data.length + "]}}"; - final byte[] headerBytes = header.getBytes(StandardCharsets.UTF_8); - final ByteArrayOutputStream out = new ByteArrayOutputStream(); - out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) - .putLong(headerBytes.length).array()); - out.write(headerBytes); - out.write(data); final Path tensors = dir.resolve("model.safetensors"); - Files.write(tensors, out.toByteArray()); - return StaticEmbeddingModel.load(vocab, tensors, true, true); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", rows)); + return StaticEmbeddingModel.load(vocab, tensors, + Casing.UNCASED, Normalization.L2); } @Test diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java index b208de352..8a553d7c0 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java @@ -28,6 +28,9 @@ import java.util.List; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; + import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -66,29 +69,14 @@ class StaticEmbeddingModelSimilarityTest { } private static Path writeSafetensors(Path dir) throws IOException { - final ByteBuffer buffer = - ByteBuffer.allocate(ROWS.length * DIMENSION * 4).order(ByteOrder.LITTLE_ENDIAN); - for (final float[] row : ROWS) { - for (final float value : row) { - buffer.putFloat(value); - } - } - final byte[] data = buffer.array(); - final String header = "{\"embeddings\":{\"dtype\":\"F32\",\"shape\":[" + ROWS.length + "," - + DIMENSION + "],\"data_offsets\":[0," + data.length + "]}}"; - final byte[] headerBytes = header.getBytes(StandardCharsets.UTF_8); - final ByteArrayOutputStream out = new ByteArrayOutputStream(); - out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) - .putLong(headerBytes.length).array()); - out.write(headerBytes); - out.write(data); final Path file = dir.resolve("model.safetensors"); - Files.write(file, out.toByteArray()); + SafetensorsTestFiles.write(file, SafetensorsTestFiles.matrix("embeddings", ROWS)); return file; } private static StaticEmbeddingModel load(Path dir) throws IOException { - return StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir), true, false); + return StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir), + Casing.UNCASED, Normalization.NONE); } @Test @@ -237,7 +225,8 @@ class StaticEmbeddingModelSimilarityTest { out.write(data); final Path tensors = dir.resolve("zero-model.safetensors"); Files.write(tensors, out.toByteArray()); - final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocab, tensors, true, false); + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(vocab, tensors, Casing.UNCASED, Normalization.NONE); final List<Neighbor> result = model.mostSimilar("a", 5); diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java index 31c7becdf..29b1062cc 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java @@ -28,6 +28,9 @@ import java.util.List; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; + import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -57,51 +60,24 @@ class StaticEmbeddingModelTest { } private static Path writeSafetensors(Path dir, boolean withWeights) throws IOException { - final ByteArrayOutputStream data = new ByteArrayOutputStream(); - final ByteBuffer embeddingBuffer = - ByteBuffer.allocate(ROWS.length * DIMENSION * 4).order(ByteOrder.LITTLE_ENDIAN); - for (final float[] row : ROWS) { - for (final float value : row) { - embeddingBuffer.putFloat(value); - } - } - final byte[] embeddingBytes = embeddingBuffer.array(); - data.write(embeddingBytes); - - String header = "{\"embeddings\":{\"dtype\":\"F32\",\"shape\":[" + ROWS.length + "," - + DIMENSION + "],\"data_offsets\":[0," + embeddingBytes.length + "]}"; + final Path file = dir.resolve("model.safetensors"); if (withWeights) { // Weight per row: [1, 1, 1, 2, 1, 1] so "hello" (row 3) counts double in the sum but not // in the pooling denominator, which is the exact behavior being pinned. - final float[] weightValues = {1f, 1f, 1f, 2f, 1f, 1f}; - final ByteBuffer weightBuffer = - ByteBuffer.allocate(weightValues.length * 4).order(ByteOrder.LITTLE_ENDIAN); - for (final float value : weightValues) { - weightBuffer.putFloat(value); - } - final byte[] weightBytes = weightBuffer.array(); - final int start = embeddingBytes.length; - data.write(weightBytes); - header += ",\"weights\":{\"dtype\":\"F32\",\"shape\":[" + weightValues.length - + "],\"data_offsets\":[" + start + "," + (start + weightBytes.length) + "]}"; + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.matrix("embeddings", ROWS), + SafetensorsTestFiles.vector("weights", new float[] {1f, 1f, 1f, 2f, 1f, 1f})); + } else { + SafetensorsTestFiles.write(file, SafetensorsTestFiles.matrix("embeddings", ROWS)); } - header += "}"; - - final byte[] headerBytes = header.getBytes(StandardCharsets.UTF_8); - final ByteArrayOutputStream out = new ByteArrayOutputStream(); - out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) - .putLong(headerBytes.length).array()); - out.write(headerBytes); - out.write(data.toByteArray()); - final Path file = dir.resolve("model.safetensors"); - Files.write(file, out.toByteArray()); return file; } @Test void testEmbedMeanPoolsWithoutWeights(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); final float[] result = model.embed("hello world"); @@ -113,7 +89,8 @@ class StaticEmbeddingModelTest { void testEmbedAppliesPerTokenWeightsButDividesByTokenCount(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, true), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, true), + Casing.UNCASED, Normalization.NONE); final float[] result = model.embed("hello world"); @@ -125,7 +102,8 @@ class StaticEmbeddingModelTest { @Test void testEmbedNormalizesToUnitLength(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, true); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.L2); final float[] result = model.embed("cat"); @@ -142,7 +120,8 @@ class StaticEmbeddingModelTest { @Test void testEmbedSkipsUnknownTokens(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); // "xyzzy" cannot be represented by any vocabulary piece, so it becomes [UNK] and must be // excluded from both the sum and the pooling denominator, leaving just "cat". @@ -154,7 +133,8 @@ class StaticEmbeddingModelTest { @Test void testEmbedOfTextWithNoInVocabularyTokensIsZeroVector(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed("xyzzy"), 1e-5f); } @@ -162,7 +142,8 @@ class StaticEmbeddingModelTest { @Test void testEmbedOfEmptyTextIsZeroVectorNotAnError(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, true); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.L2); assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed(""), 1e-5f); } @@ -170,7 +151,8 @@ class StaticEmbeddingModelTest { @Test void testDimensionAndVocabularySizeAccessors(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); assertEquals(DIMENSION, model.dimension()); assertEquals(VOCAB_TOKENS.size(), model.vocabularySize()); @@ -179,7 +161,8 @@ class StaticEmbeddingModelTest { @Test void testEmbedRejectsNullText(@TempDir Path dir) throws IOException { final StaticEmbeddingModel model = - StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), true, false); + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); assertThrows(IllegalArgumentException.class, () -> model.embed(null)); } @@ -190,9 +173,9 @@ class StaticEmbeddingModelTest { final Path tensors = writeSafetensors(dir, false); assertThrows(IllegalArgumentException.class, - () -> StaticEmbeddingModel.load(null, tensors, true, false)); + () -> StaticEmbeddingModel.load(null, tensors, Casing.UNCASED, Normalization.NONE)); assertThrows(IllegalArgumentException.class, - () -> StaticEmbeddingModel.load(vocab, null, true, false)); + () -> StaticEmbeddingModel.load(vocab, null, Casing.UNCASED, Normalization.NONE)); } @Test @@ -201,7 +184,8 @@ class StaticEmbeddingModelTest { Files.write(shortVocab, List.of("[CLS]", "[SEP]", "[UNK]")); final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> StaticEmbeddingModel.load(shortVocab, writeSafetensors(dir, false), true, false)); + () -> StaticEmbeddingModel.load(shortVocab, writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE)); assertTrue(e.getMessage().contains("rows")); } @@ -235,7 +219,7 @@ class StaticEmbeddingModelTest { Files.write(file, out.toByteArray()); final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> StaticEmbeddingModel.load(writeVocab(dir), file, true, false)); + () -> StaticEmbeddingModel.load(writeVocab(dir), file, Casing.UNCASED, Normalization.NONE)); assertTrue(e.getMessage().contains("weights")); } diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/WordPieceVocabularyTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/WordPieceVocabularyTest.java new file mode 100644 index 000000000..4ae76c1ba --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/WordPieceVocabularyTest.java @@ -0,0 +1,81 @@ +/* + * 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.embeddings; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The vocabulary contract: line number is the id, duplicates fail loud, the id lookup uses a + * {@code -1} sentinel, and the reverse lookup enforces its bounds. + */ +class WordPieceVocabularyTest { + + @Test + void testLineNumberIsTheTokenId() { + final WordPieceVocabulary vocabulary = + WordPieceVocabulary.fromLines(List.of("[CLS]", "[SEP]", "hello", "world"), "test"); + assertEquals(4, vocabulary.size()); + assertEquals(0, vocabulary.id("[CLS]")); + assertEquals(2, vocabulary.id("hello")); + assertEquals("world", vocabulary.token(3)); + assertTrue(vocabulary.tokens().contains("hello")); + } + + @Test + void testUnknownTokenIdIsTheSentinel() { + final WordPieceVocabulary vocabulary = + WordPieceVocabulary.fromLines(List.of("hello"), "test"); + assertEquals(-1, vocabulary.id("missing")); + assertThrows(IllegalArgumentException.class, () -> vocabulary.id(null)); + } + + @Test + void testDuplicateTokenFailsLoudlyNamingBothLines() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> WordPieceVocabulary.fromLines(List.of("hello", "world", "hello"), "test")); + assertTrue(e.getMessage().contains("hello"), e.getMessage()); + assertTrue(e.getMessage().contains("0") && e.getMessage().contains("2"), e.getMessage()); + } + + @Test + void testReverseLookupEnforcesBounds() { + final WordPieceVocabulary vocabulary = + WordPieceVocabulary.fromLines(List.of("hello"), "test"); + assertEquals("hello", vocabulary.token(0)); + assertThrows(IllegalArgumentException.class, () -> vocabulary.token(-1)); + assertThrows(IllegalArgumentException.class, () -> vocabulary.token(1)); + } + + @Test + void testReadFromFileMatchesInMemoryLines(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("vocab.txt"); + Files.write(file, List.of("[CLS]", "token")); + final WordPieceVocabulary read = WordPieceVocabulary.read(file); + assertEquals(2, read.size()); + assertEquals(1, read.id("token")); + } +}
