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 d1d3067fc7a0c022d47e8be3afeaeef2631fd091 Author: Kristian Rickert <[email protected]> AuthorDate: Thu Jul 16 06:17:31 2026 -0400 OPENNLP-1877: Name format constants, document throws, and source or replace unsupported doc claims --- opennlp-extensions/opennlp-embeddings/README.md | 18 ++- opennlp-extensions/opennlp-embeddings/TRAINING.md | 6 +- .../opennlp/embeddings/EmbeddingVocabulary.java | 2 + .../java/opennlp/embeddings/FlatJsonFields.java | 1 + .../main/java/opennlp/embeddings/JsonCursor.java | 2 + .../java/opennlp/embeddings/ModelAssembler.java | 50 +++---- .../java/opennlp/embeddings/ModelFileNames.java | 52 +++++++ .../java/opennlp/embeddings/SafetensorsFile.java | 32 +++-- .../opennlp/embeddings/StaticEmbeddingModel.java | 61 ++++----- .../main/java/opennlp/embeddings/TensorInfo.java | 4 +- .../opennlp/embeddings/TokenizerJsonVocab.java | 13 +- .../main/java/opennlp/embeddings/cmdline/CLI.java | 14 +- .../opennlp/embeddings/SafetensorsFileTest.java | 150 +++++++++------------ .../opennlp/embeddings/TokenizerJsonVocabTest.java | 7 + 14 files changed, 239 insertions(+), 173 deletions(-) diff --git a/opennlp-extensions/opennlp-embeddings/README.md b/opennlp-extensions/opennlp-embeddings/README.md index 873d575cc..4481edf0f 100644 --- a/opennlp-extensions/opennlp-embeddings/README.md +++ b/opennlp-extensions/opennlp-embeddings/README.md @@ -17,11 +17,9 @@ # OpenNLP Static Embeddings -Embeddings have become an essential part of AI workloads. As such, OpenNLP introduces a pure-JVM approach to embeddings with a modern Model2Vec engine. +Turn text into embedding vectors from a static (non-contextual) table: a per-token vector matrix plus subword tokenization, WordPiece or SentencePiece. It uses the same lookup-table approach as [word2vec](https://code.google.com/archive/p/word2vec/) and [GloVe](https://nlp.stanford.edu/projects/glove/). Distillation tools can compress a sentence-transformer into such a flat table (the [Model2Vec](https://github.com/MinishLab/model2vec) family is the primary target), and looking a sentenc [...] -Turn text into embedding vectors from a static (non-contextual) table: a per-token vector matrix plus subword tokenization, WordPiece or SentencePiece. It is the modern successor to the word2vec and GloVe workflow. Distillation tools can compress a sentence-transformer into such a flat table (the Model2Vec family is the primary target), and looking a sentence up in the table approximates the transformer's semantics at a fraction of the cost. Because SentencePiece models are supported, th [...] - -OpenNLP also supports ONNX models, which are inherently more accurate. Model2Vec sacrifices some accuracy for a large speed gain, and OpenNLP recognizes that trade-off, so both embedding methods are supported and share the same `TextEmbedder` seam. +OpenNLP also supports ONNX models, which are inherently more accurate. Model2Vec sacrifices some accuracy for a large speed gain, and OpenNLP recognizes that trade-off, so both embedding methods are supported and implement the same `TextEmbedder` interface. ## Quickstart @@ -60,7 +58,7 @@ flowchart LR E --> F["float[] vector"] ``` -1. **Tokenize.** The model's own subword tokenizer splits the text into pieces: WordPiece with the model's casing rule, or a trained SentencePiece model that carries its own text normalizer. Special pieces (the WordPiece `[CLS]`/`[SEP]`/`[UNK]` frame, a SentencePiece model's control and unknown pieces) never contribute to the pooled vector. +1. **Tokenize.** The model's own subword tokenizer splits the text into pieces: WordPiece with the model's casing rule, or a trained SentencePiece model that carries its own text normalizer. Special pieces (the WordPiece `[CLS]`, `[SEP]`, and `[UNK]` tokens, a SentencePiece model's control and unknown pieces) never contribute to the pooled vector. 2. **Gather.** Each piece contributes its matrix row, found by the piece *string* rather than the tokenizer's numeric id. The two files of a SentencePiece model routinely order and offset their ids differently (the fairseq convention shifts them by one, and distillation tools reorder the vocabulary outright), so string lookup is what keeps the pairing robust; a poolable piece with no matrix row fails loud at load time, not at query time. Unknown pieces are dropped, and a text with no in- [...] 3. **Weight and pool.** Per-token weights (when the model carries them) multiply into the running sum, and the sum is divided by the plain token count. This mean-pool matches the reference implementation of the targeted model family exactly, verified against it rather than assumed. 4. **Normalize.** The pooled vector is L2-normalized by default so cosine similarity is a dot product. Normalization can be turned off for models that expect raw pooled vectors. @@ -87,7 +85,7 @@ flowchart TD MAT --> M ``` -The weights are read with a purpose-built **safetensors** reader. Unlike pickle-based checkpoint formats, safetensors carries no executable content, so loading a downloaded file cannot execute arbitrary code. Tensor data streams directly into the decoded array, so the file size is not bound by Java's int-indexed arrays; a single decoded tensor is capped at the maximum Java array length (about 2.1 billion float elements), checked explicitly. +The weights are read with a purpose-built [safetensors](https://github.com/huggingface/safetensors) reader. Unlike pickle-based checkpoint formats, safetensors carries no executable content, so loading a downloaded file cannot execute arbitrary code. Tensor data streams directly into the decoded array, so the file size is not bound by Java's int-indexed arrays; a single decoded tensor is capped at the maximum Java array length (about 2.1 billion float elements), checked explicitly. ## Architecture @@ -106,11 +104,11 @@ flowchart TD DL["SentenceVectorsDL<br/>(opennlp-dl, ONNX)"] -. implements .-> TE ``` -Two seams keep the module small. `SubwordTokenizer` is the tokenization seam: the WordPiece encoder from `opennlp-api` and the pure-JVM SentencePiece implementation from `opennlp-subword` both produce the same piece stream, so the pooling code has exactly one path. `TextEmbedder` is the embedding seam: the static path here and the contextual ONNX path in `opennlp-dl` both implement it, so callers can swap one for the other without touching their code. +Two interfaces keep the module small. `SubwordTokenizer` is the tokenization interface: the WordPiece encoder from `opennlp-api` and the pure-JVM SentencePiece implementation from `opennlp-subword` both produce the same piece stream, so the pooling code has exactly one path. `TextEmbedder` is the embedding interface: the static path here and the contextual ONNX path in `opennlp-dl` both implement it, so callers can swap one for the other without touching their code. ## Performance -A static table wins on speed and footprint because there is no model forward pass: the hot path is a vocabulary lookup, a handful of vector adds, and one normalization. The module ships a JMH benchmark (`StaticEmbeddingModelBenchmark`) that measures `embed()` and `mostSimilar()` throughput on a real model directory (`-p modelDir=/path/to/model`), so you can reproduce numbers on your own hardware and model. +A static table wins on speed and footprint because there is no model forward pass: the hot path is a vocabulary lookup, a handful of vector adds, and one normalization. The module ships a Java Microbenchmark Harness (JMH) benchmark (`StaticEmbeddingModelBenchmark`) that measures `embed()` and `mostSimilar()` throughput on a real model directory (`-p modelDir=/path/to/model`), so you can reproduce numbers on your own hardware and model. Two things drive the numbers, and the benchmark separates them. `embed()` is tokenize-and-pool, so its cost tracks the text and the tokenizer, not the table size. `mostSimilar()` is a brute-force scan over every row, so its cost tracks the vocabulary size directly. A run comparing a small WordPiece table against the large multilingual SentencePiece table makes the split visible (throughput across all cores, one machine, indicative not publishable): @@ -119,7 +117,7 @@ Two things drive the numbers, and the benchmark separates them. `embed()` is tok | potion-base-8M | WordPiece, 29.5k | ~295k ops/s | ~9,000 ops/s | | bge-m3 (distilled) | SentencePiece, 250k | ~1.47M ops/s | ~550 ops/s | -So a large multilingual vocabulary is free for embedding and expensive for a full nearest-neighbor scan; that scan is where an approximate index earns its place once the table is large. Separately, on the potion-base-8M table the JVM path ran roughly an order of magnitude faster single-threaded than the model2vec Python reference at around a fifth of the resident memory, with output vectors matching the reference within floating-point tolerance, so the speed is not bought with accuracy. [...] +So a large multilingual vocabulary is free for embedding and expensive for a full nearest-neighbor scan; that scan is where an approximate index earns its place once the table is large. Separately, `scripts/parity/` holds a harness that reruns the single-thread speed comparison against the Model2Vec Python reference and checks that the output vectors match it within floating-point tolerance, so a cross-runtime comparison is something you reproduce on your own hardware rather than quote. [...] ## Usage @@ -176,7 +174,7 @@ IntStream.range(0, docs.size()) .forEach(i -> System.out.println(docs.get(i))); ``` -Here `dot` is any dot product over two float arrays. For a full RAG-style retriever, keep the document vectors in whatever index you already use and score queries the same way. This applies to most modern search engines, since they tend to decouple the HNSW lookups from the vectors you feed them. +Here `dot` is any dot product over two float arrays. For a full retrieval-augmented generation (RAG) retriever, keep the document vectors in whatever index you already use and score queries the same way. A vector index that stores and searches precomputed vectors, such as a Hierarchical Navigable Small World (HNSW) index, does not care how those vectors were produced, so these embeddings can feed it directly. ## Getting a model diff --git a/opennlp-extensions/opennlp-embeddings/TRAINING.md b/opennlp-extensions/opennlp-embeddings/TRAINING.md index c229c3d77..e07c0a2c1 100644 --- a/opennlp-extensions/opennlp-embeddings/TRAINING.md +++ b/opennlp-extensions/opennlp-embeddings/TRAINING.md @@ -19,7 +19,7 @@ This module loads static embedding tables; it does not produce them. A table is distilled once from a sentence-transformer teacher, offline, in Python, and then loaded in the JVM as many times as you like. This walks through distilling one and assembling the directory `StaticEmbeddingModel.load` expects, using a multilingual SentencePiece model (bge-m3) as the worked example. -The distillation tool is [Model2Vec](https://github.com/MinishLab/model2vec). It runs the teacher over its own vocabulary once, applies PCA and a Zipf weighting, and writes a flat per-token matrix. There is no training loop and no labelled data; a distillation is minutes on CPU, not hours on a GPU. +The distillation tool is [Model2Vec](https://github.com/MinishLab/model2vec). It runs the teacher over its own vocabulary once, applies principal component analysis (PCA) and a Zipf weighting (frequent tokens are down-weighted, after Zipf's law of word frequency), and writes a flat per-token matrix. There is no training loop and no labelled data; a distillation is minutes on CPU, not hours on a GPU. ## 1. Set up the distiller @@ -30,7 +30,7 @@ uv pip install --python .venv-distill "model2vec[distill]" ## 2. Distill the teacher -bge-m3 is an XLM-RoBERTa/SentencePiece model with a 250k multilingual vocabulary, native dimension 1024. +bge-m3 is an [XLM-RoBERTa](https://arxiv.org/abs/1911.02116)/SentencePiece model with a 250k multilingual vocabulary, native dimension 1024. ```python # distill_bge_m3.py @@ -106,4 +106,4 @@ A WordPiece teacher (a BERT-family model such as bge-large-en) distills the same ## Where a table's license comes from -Distillation carries the teacher's license onto the table. bge-m3 is MIT, so its distillation is freely redistributable; a table distilled from a non-commercial or share-alike teacher inherits those terms. Check the teacher before publishing a table. +Distillation carries the teacher's license onto the table. bge-m3 is published under the MIT license per its [model card](https://huggingface.co/BAAI/bge-m3) (verify at download time), so its distillation is freely redistributable; a table distilled from a non-commercial or share-alike teacher inherits those terms. Check the teacher before publishing a table. diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java index 1b63e57d1..86bd03a58 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java @@ -124,6 +124,7 @@ final class EmbeddingVocabulary { * * @param token The token to look up. Must not be {@code null}. * @return The token's id, or {@code -1} when the token is not in this vocabulary. + * @throws IllegalArgumentException Thrown if {@code token} is {@code null}. */ int id(String token) { if (token == null) { @@ -143,6 +144,7 @@ final class EmbeddingVocabulary { * * @param id The row id. Must be within {@code [0, size())}. * @return The token at that id. + * @throws IllegalArgumentException Thrown if {@code id} is outside {@code [0, size())}. */ String token(int id) { if (id < 0 || id >= tokenById.size()) { 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 5a6ec7a24..4dd4dc83a 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 @@ -28,6 +28,7 @@ import java.nio.file.Path; */ final class FlatJsonFields { + /** Not instantiable. */ private FlatJsonFields() { } 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 9e76ee23d..294595276 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 @@ -31,6 +31,8 @@ final class JsonCursor { private int position; /** + * Creates a cursor positioned at the start of the given JSON text. + * * @param text The JSON text to scan. Must not be {@code null}. * @param inputName What the text is (for error messages), e.g. {@code "safetensors header"} * or a file name. diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java index 06fc0997c..b92ec01ef 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java @@ -43,14 +43,16 @@ import java.util.Map; */ public final class ModelAssembler { - private static final String SAFETENSORS_FILE_NAME = "model.safetensors"; - private static final String TOKENIZER_JSON_FILE_NAME = "tokenizer.json"; - private static final String CONFIG_FILE_NAME = "config.json"; - private static final String VOCABULARY_FILE_NAME = "vocab.txt"; - private static final String TOKENIZER_CONFIG_FILE_NAME = "tokenizer_config.json"; - private static final List<String> SENTENCEPIECE_MODEL_FILE_NAMES = - List.of("sentencepiece.bpe.model", "spiece.model", "tokenizer.model"); + /** The WordPiece tokenizer family, the {@code model.type} of a BERT-style distillation. */ + private static final String FAMILY_WORDPIECE = "WordPiece"; + /** The Unigram {@code model.type} a SentencePiece distillation's {@code tokenizer.json} uses. */ + private static final String FAMILY_UNIGRAM = "Unigram"; + + /** The SentencePiece tokenizer family, reported for an assembled Unigram directory. */ + private static final String FAMILY_SENTENCEPIECE = "SentencePiece"; + + /** Not instantiable. */ private ModelAssembler() { } @@ -88,17 +90,17 @@ public final class ModelAssembler { throw new IllegalArgumentException( "Model directory does not exist or is not a directory: " + modelDirectory); } - requireFile(modelDirectory, SAFETENSORS_FILE_NAME); - requireFile(modelDirectory, CONFIG_FILE_NAME); - final Path tokenizerJson = requireFile(modelDirectory, TOKENIZER_JSON_FILE_NAME); + requireFile(modelDirectory, ModelFileNames.SAFETENSORS); + requireFile(modelDirectory, ModelFileNames.CONFIG); + final Path tokenizerJson = requireFile(modelDirectory, ModelFileNames.TOKENIZER_JSON); final TokenizerJson tokenizer = readTokenizerJson(tokenizerJson); return switch (tokenizer.modelType()) { - case "WordPiece" -> assembleWordpiece(modelDirectory, tokenizer); - case "Unigram" -> assembleSentencePiece(modelDirectory); + case FAMILY_WORDPIECE -> assembleWordpiece(modelDirectory, tokenizer); + case FAMILY_UNIGRAM -> assembleSentencePiece(modelDirectory); default -> throw new IllegalArgumentException(tokenizerJson + " has a '" - + tokenizer.modelType() + "' tokenizer model; only WordPiece and Unigram " - + "(SentencePiece) distillations are supported"); + + tokenizer.modelType() + "' tokenizer model; only " + FAMILY_WORDPIECE + " and " + + FAMILY_UNIGRAM + " (" + FAMILY_SENTENCEPIECE + ") distillations are supported"); }; } @@ -113,17 +115,17 @@ public final class ModelAssembler { */ private static Result assembleWordpiece(Path modelDirectory, TokenizerJson tokenizer) throws IOException { - final Path vocabularyFile = modelDirectory.resolve(VOCABULARY_FILE_NAME); + final Path vocabularyFile = modelDirectory.resolve(ModelFileNames.VOCABULARY); boolean wroteVocabulary = false; if (!Files.exists(vocabularyFile)) { if (tokenizer.orderedVocabulary() == null) { throw new IllegalArgumentException("tokenizer.json in " + modelDirectory - + " has no model.vocab dictionary; cannot derive " + VOCABULARY_FILE_NAME); + + " has no model.vocab dictionary; cannot derive " + ModelFileNames.VOCABULARY); } Files.write(vocabularyFile, tokenizer.orderedVocabulary()); wroteVocabulary = true; } - final Path tokenizerConfigFile = modelDirectory.resolve(TOKENIZER_CONFIG_FILE_NAME); + final Path tokenizerConfigFile = modelDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG); boolean wroteTokenizerConfig = false; if (!Files.exists(tokenizerConfigFile)) { // The BERT normalizer's lowercase flag is the casing; default to lower-casing (the uncased @@ -134,7 +136,7 @@ public final class ModelAssembler { wroteTokenizerConfig = true; } final StaticEmbeddingModel model = load(modelDirectory); - return new Result("WordPiece", model.dimension(), model.vocabularySize(), + return new Result(FAMILY_WORDPIECE, model.dimension(), model.vocabularySize(), wroteVocabulary, wroteTokenizerConfig); } @@ -147,14 +149,16 @@ public final class ModelAssembler { * @throws IOException Thrown if loading fails to read a file. */ private static Result assembleSentencePiece(Path modelDirectory) throws IOException { - if (firstExisting(modelDirectory, SENTENCEPIECE_MODEL_FILE_NAMES) == null) { + if (firstExisting(modelDirectory, ModelFileNames.SENTENCEPIECE_MODELS) == null) { throw new IllegalArgumentException("Model directory " + modelDirectory + " is a " - + "SentencePiece model but has no trained model file (one of " - + String.join(", ", SENTENCEPIECE_MODEL_FILE_NAMES) + "); copy it from the teacher " - + "model's repository (it is named sentencepiece.bpe.model there) into this directory"); + + FAMILY_SENTENCEPIECE + " model but has no trained model file (one of " + + String.join(", ", ModelFileNames.SENTENCEPIECE_MODELS) + "); copy it from the " + + "teacher model's repository (it is named sentencepiece.bpe.model there) into this " + + "directory"); } final StaticEmbeddingModel model = load(modelDirectory); - return new Result("SentencePiece", model.dimension(), model.vocabularySize(), false, false); + return new Result(FAMILY_SENTENCEPIECE, model.dimension(), model.vocabularySize(), + false, false); } /** diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelFileNames.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelFileNames.java new file mode 100644 index 000000000..d0d2e1d54 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelFileNames.java @@ -0,0 +1,52 @@ +/* + * 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.util.List; + +/** + * The file names of a static embedding model directory, shared by + * {@link StaticEmbeddingModel}'s loader and {@link ModelAssembler}. A WordPiece directory holds + * {@link #SAFETENSORS}, {@link #CONFIG}, {@link #VOCABULARY}, and {@link #TOKENIZER_CONFIG}; a + * SentencePiece directory holds {@link #SAFETENSORS}, {@link #CONFIG}, {@link #TOKENIZER_JSON}, + * and one of {@link #SENTENCEPIECE_MODELS}. + */ +final class ModelFileNames { + + /** The safetensors file holding the embedding matrix and optional per-token weights. */ + static final String SAFETENSORS = "model.safetensors"; + + /** The tokenizer description whose Unigram {@code model.vocab} order names the matrix rows. */ + static final String TOKENIZER_JSON = "tokenizer.json"; + + /** The model configuration carrying the {@code normalize} pooling switch. */ + static final String CONFIG = "config.json"; + + /** The BERT-style vocabulary of a WordPiece model, one token per line in row order. */ + static final String VOCABULARY = "vocab.txt"; + + /** The tokenizer configuration carrying the WordPiece {@code do_lower_case} switch. */ + static final String TOKENIZER_CONFIG = "tokenizer_config.json"; + + /** The file names SentencePiece models ship their trained {@code .model} under, in try order. */ + static final List<String> SENTENCEPIECE_MODELS = + List.of("sentencepiece.bpe.model", "spiece.model", "tokenizer.model"); + + /** Not instantiable. */ + private ModelFileNames() { + } +} 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 58e7e44b5..316c71770 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 @@ -53,6 +53,15 @@ public final class SafetensorsFile { private static final int HEADER_LENGTH_PREFIX_BYTES = 8; + /** The header's dtype marker for 32-bit IEEE floats. */ + private static final String DTYPE_F32 = "F32"; + + /** The header's dtype marker for 16-bit IEEE half floats. */ + private static final String DTYPE_F16 = "F16"; + + /** The header's dtype marker for 16-bit bfloat16 floats. */ + private static final String DTYPE_BF16 = "BF16"; + // Positional-read chunk size, a multiple of Float.BYTES so every filled chunk decodes to // whole floats. private static final int READ_CHUNK_BYTES = 1 << 20; @@ -163,7 +172,7 @@ public final class SafetensorsFile { /** * Decodes a floating-point tensor's data to {@code float[]}, streaming it from the file. * Accepts the {@code F32}, {@code F16} (IEEE half) and {@code BF16} (bfloat16) dtypes; the two - * 16-bit types are widened to {@code float} as they are read. {@code F16} is model2vec's + * 16-bit types are widened to {@code float} as they are read. {@code F16} is Model2Vec's * default output dtype, so this is the common case for downloaded distilled tables. * * @param name The tensor's name. Must not be {@code null}. @@ -226,9 +235,9 @@ public final class SafetensorsFile { */ public float[] readFloat32(String name) throws IOException { final TensorInfo info = tensorInfo(name); - if (!"F32".equals(info.dtype())) { + if (!DTYPE_F32.equals(info.dtype())) { throw new IllegalArgumentException( - "Tensor '" + name + "' has dtype " + info.dtype() + ", not F32"); + "Tensor '" + name + "' has dtype " + info.dtype() + ", not " + DTYPE_F32); } return readFloats(name); } @@ -245,14 +254,14 @@ public final class SafetensorsFile { private static void decodeInto(ByteBuffer chunk, String dtype, float[] out, int offset, int count) { switch (dtype) { - case "F32" -> chunk.asFloatBuffer().get(out, offset, count); - case "F16" -> { + case DTYPE_F32 -> chunk.asFloatBuffer().get(out, offset, count); + case DTYPE_F16 -> { final ShortBuffer shorts = chunk.asShortBuffer(); for (int i = 0; i < count; i++) { out[offset + i] = Float.float16ToFloat(shorts.get()); } } - case "BF16" -> { + case DTYPE_BF16 -> { // bfloat16 is the high 16 bits of a float32: shift back up and reinterpret. final ShortBuffer shorts = chunk.asShortBuffer(); for (int i = 0; i < count; i++) { @@ -268,20 +277,21 @@ public final class SafetensorsFile { * * @param dtype The tensor dtype. * @param tensorName The tensor's name, for the error message. - * @throws IllegalArgumentException if {@code dtype} is not a supported float type. + * @throws IllegalArgumentException Thrown if {@code dtype} is not a supported float type. */ private static int floatElementBytes(String dtype, String tensorName) { return switch (dtype) { - case "F32" -> Float.BYTES; - case "F16", "BF16" -> Short.BYTES; + case DTYPE_F32 -> Float.BYTES; + case DTYPE_F16, DTYPE_BF16 -> Short.BYTES; default -> throw new IllegalArgumentException("Tensor '" + tensorName + "' has dtype " - + dtype + ", not a supported float type (F32, F16, BF16)"); + + dtype + ", not a supported float type (" + DTYPE_F32 + ", " + DTYPE_F16 + ", " + + DTYPE_BF16 + ")"); }; } /** {@return whether {@code dtype} is a float type this reader decodes} */ private static boolean isFloatDtype(String dtype) { - return "F32".equals(dtype) || "F16".equals(dtype) || "BF16".equals(dtype); + return DTYPE_F32.equals(dtype) || DTYPE_F16.equals(dtype) || DTYPE_BF16.equals(dtype); } /** 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 c5bf73ba3..973f105c2 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 @@ -49,7 +49,7 @@ import opennlp.tools.tokenize.WordpieceTokenizer; * never by tokenizer id, so the two files may order or offset their ids differently without * corrupting lookups; a piece the matrix does not carry fails loud at load time.</p> * - * <p>Special pieces (the WordPiece {@code [CLS]}/{@code [SEP]}/{@code [UNK]} frame, a + * <p>Special pieces (the WordPiece {@code [CLS]}, {@code [SEP]}, and {@code [UNK]} tokens, a * SentencePiece model's control and unknown pieces) are never pooled; the sum is divided by the * count of pooled pieces, not the sum of weights. A text with no in-vocabulary pieces yields a * zero vector.</p> @@ -81,14 +81,6 @@ public final class StaticEmbeddingModel implements TextEmbedder { 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"; - private static final String SAFETENSORS_FILE_NAME = "model.safetensors"; - private static final String CONFIG_FILE_NAME = "config.json"; - private static final String TOKENIZER_CONFIG_FILE_NAME = "tokenizer_config.json"; - private static final String TOKENIZER_JSON_FILE_NAME = "tokenizer.json"; - // The file names SentencePiece models ship their trained .model under, by convention family. - private static final List<String> SENTENCEPIECE_MODEL_FILE_NAMES = - List.of("sentencepiece.bpe.model", "spiece.model", "tokenizer.model"); private static final int[] NO_EXCLUDED_ROWS = new int[0]; // Excluded from neighbor results, including [PAD] and [MASK] that a distilled table keeps. private static final Set<String> WORDPIECE_SPECIAL_TOKENS = @@ -102,7 +94,7 @@ public final class StaticEmbeddingModel implements TextEmbedder { private final int dimension; private final EmbeddingVocabulary vocabulary; private final SubwordTokenizer tokenizer; - // Tokenizer-id test for pieces that are never pooled (frame, control, and unknown pieces). + // Tokenizer-id test for pieces that are never pooled (delimiter, control, unknown pieces). private final IntPredicate skipPieceId; private final boolean normalize; // Per-row L2 norms and special-token mask, precomputed at load time for the neighbor scan. @@ -160,26 +152,27 @@ public final class StaticEmbeddingModel implements TextEmbedder { throw new IllegalArgumentException( "Model directory does not exist or is not a directory: " + modelDirectory); } - final Path vocabularyFile = modelDirectory.resolve(VOCABULARY_FILE_NAME); + final Path vocabularyFile = modelDirectory.resolve(ModelFileNames.VOCABULARY); if (Files.isRegularFile(vocabularyFile)) { return loadWordpieceDirectory(modelDirectory, vocabularyFile); } final Path sentencePieceModelFile = firstRegularFile(modelDirectory, - SENTENCEPIECE_MODEL_FILE_NAMES); - final Path tokenizerJsonFile = modelDirectory.resolve(TOKENIZER_JSON_FILE_NAME); + ModelFileNames.SENTENCEPIECE_MODELS); + final Path tokenizerJsonFile = modelDirectory.resolve(ModelFileNames.TOKENIZER_JSON); if (sentencePieceModelFile != null && Files.isRegularFile(tokenizerJsonFile)) { return loadSentencePiece(sentencePieceModelFile, tokenizerJsonFile, - requiredFile(modelDirectory, SAFETENSORS_FILE_NAME), - requiredNormalize(requiredFile(modelDirectory, CONFIG_FILE_NAME))); + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG))); } if (Files.isRegularFile(tokenizerJsonFile)) { throw new IllegalArgumentException("Model directory " + modelDirectory + " has a " - + TOKENIZER_JSON_FILE_NAME + " but no trained SentencePiece file (" - + String.join(", ", SENTENCEPIECE_MODEL_FILE_NAMES) + "); copy the .model file " + + ModelFileNames.TOKENIZER_JSON + " but no trained SentencePiece file (" + + String.join(", ", ModelFileNames.SENTENCEPIECE_MODELS) + "); copy the .model file " + "from the model's base tokenizer next to it"); } throw new IllegalArgumentException("Model directory " + modelDirectory + " has neither a " - + VOCABULARY_FILE_NAME + " (WordPiece layout) nor a " + TOKENIZER_JSON_FILE_NAME + + ModelFileNames.VOCABULARY + " (WordPiece layout) nor a " + + ModelFileNames.TOKENIZER_JSON + " with a trained SentencePiece file (SentencePiece layout)"); } @@ -195,10 +188,11 @@ public final class StaticEmbeddingModel implements TextEmbedder { private static StaticEmbeddingModel loadWordpieceDirectory(Path modelDirectory, Path vocabularyFile) throws IOException { - final Path safetensorsFile = requiredFile(modelDirectory, SAFETENSORS_FILE_NAME); - final Path tokenizerConfigFile = requiredFile(modelDirectory, TOKENIZER_CONFIG_FILE_NAME); + final Path safetensorsFile = requiredFile(modelDirectory, ModelFileNames.SAFETENSORS); + final Path tokenizerConfigFile = + requiredFile(modelDirectory, ModelFileNames.TOKENIZER_CONFIG); final Normalization normalization = - requiredNormalize(requiredFile(modelDirectory, CONFIG_FILE_NAME)); + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)); final Boolean lowerCase = FlatJsonFields.topLevelBoolean(tokenizerConfigFile, "do_lower_case"); if (lowerCase == null) { @@ -277,8 +271,8 @@ public final class StaticEmbeddingModel implements TextEmbedder { * @param vocabularyFile The {@code vocab.txt} file: one token per line, line number is the * token's row id. Must not be {@code null}, must exist, and must * contain the {@code [UNK]} token. The {@code [CLS]} and {@code [SEP]} - * frame tokens are optional: a distilled table that dropped them (as - * Model2Vec does) still loads, because the frame is never pooled. + * delimiter tokens are optional: a distilled table that dropped them + * (as Model2Vec does) still loads, because they are never pooled. * @param safetensorsFile The {@code model.safetensors} file. Must not be {@code null} and * must exist, and must contain exactly one 2-D float tensor * (the embedding matrix) whose row count matches the vocabulary size. @@ -319,8 +313,8 @@ public final class StaticEmbeddingModel implements TextEmbedder { } final WordpieceEncoder tokenizer = wordpieceEncoder(vocabulary, casing == Casing.UNCASED, unknownId); - // Pooling skips the [CLS]/[SEP] frame by id; an absent frame maps to the unknown id, which - // is skipped the same way. A negative id is the absent sentinel and matches no emitted piece. + // Pooling skips [CLS] and [SEP] by id; when absent they map to the unknown id, which is + // skipped the same way. A negative id is the absent sentinel and matches no emitted piece. final int classificationId = vocabulary.id(WordpieceTokenizer.BERT_CLS_TOKEN); final int separatorId = vocabulary.id(WordpieceTokenizer.BERT_SEP_TOKEN); final IntPredicate skipPieceId = @@ -332,17 +326,18 @@ public final class StaticEmbeddingModel implements TextEmbedder { } /** - * Builds the WordPiece encoder, caching {@code [CLS]} and {@code [SEP]} onto the unknown row + * Builds the WordPiece encoder, mapping {@code [CLS]} and {@code [SEP]} onto the unknown row * when the distilled vocabulary dropped them. A static embedding table mean-pools its content - * pieces and never frames, so distillers routinely remove {@code [CLS]}/{@code [SEP]} from the - * table; the encoder still frames every encoding and needs an id for the frame, and pooling - * skips the frame regardless of its id, so pointing the absent frame tokens at the unknown row - * makes the model loadable without changing which pieces are pooled. + * pieces and never pools the delimiters, so distillers routinely remove + * {@code [CLS]}/{@code [SEP]} from the table; the encoder still wraps every encoding in them + * and needs an id for each, and pooling skips them regardless of their ids, so pointing the + * absent delimiter tokens at the unknown row makes the model loadable without changing which + * pieces are pooled. * * @param vocabulary The matrix row vocabulary; must contain the unknown token. * @param lowerCase Whether the tokenizer lower-cases and strips accents. - * @param unknownId The unknown token's row, reused as the frame id when a frame token is - * absent. + * @param unknownId The unknown token's row, reused as the id of {@code [CLS]} or + * {@code [SEP]} when that token is absent. * @return The encoder. */ private static WordpieceEncoder wordpieceEncoder(EmbeddingVocabulary vocabulary, @@ -845,6 +840,8 @@ public final class StaticEmbeddingModel implements TextEmbedder { private int size; /** + * Creates an empty selection. + * * @param capacity The maximum number of rows to keep. */ TopK(int capacity) { 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 3234a2aa3..73b3993b1 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 @@ -56,8 +56,8 @@ public record TensorInfo(String name, String dtype, int[] shape, long dataOffset } /** - * @return The tensor's dimensions, outermost first, as a copy; mutating it does not affect - * this record. + * {@return the tensor's dimensions, outermost first, as a copy; mutating it does not affect + * this record} */ @Override public int[] shape() { diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java index c4e595920..6f5d96d5d 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java @@ -35,6 +35,7 @@ import java.util.List; */ final class TokenizerJsonVocab { + /** Not instantiable. */ private TokenizerJsonVocab() { } @@ -52,12 +53,18 @@ final class TokenizerJsonVocab { * * @param file The {@code tokenizer.json} file. Must not be {@code null} and must exist. * @return The pieces; the index is the matrix row. - * @throws IllegalArgumentException Thrown if the file is not a well-formed - * {@code tokenizer.json}, its model is not Unigram, or an added token's id neither matches - * an existing row nor appends as the next one. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null} or missing, the + * file is not a well-formed {@code tokenizer.json}, its model is not Unigram, or an added + * token's id neither matches an existing row nor appends as the next one. * @throws IOException Thrown if reading the file fails. */ static List<String> rows(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 String json = Files.readString(file); final JsonCursor cursor = new JsonCursor(json, file.getFileName().toString()); cursor.skipWhitespace(); diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java index c4a9bf6d7..5fd3a332d 100644 --- a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java @@ -56,16 +56,16 @@ public final class CLI { toolLookupMap = Collections.unmodifiableMap(toolLookupMap); } + /** Not instantiable. */ private CLI() { } - /** - * @return A set which contains all tool names. - */ + /** {@return the names of all tools this command line dispatcher can run} */ public static Set<String> getToolNames() { return toolLookupMap.keySet(); } + /** Logs the version banner and the list of available tools with their short descriptions. */ private static void usage() { logger.info("OpenNLP Static Embeddings {}.", Version.currentVersion()); logger.info("Usage: {} TOOL", CMD); @@ -93,6 +93,14 @@ public final class CLI { logger.info("Example: {} AssembleModel help", CMD); } + /** + * Runs the tool named by the first argument, passing it the remaining arguments. Without + * arguments it logs the usage overview instead, and a tool invoked with the {@code help} + * parameter logs that tool's help. Exits the JVM with the tool's error code when the tool + * terminates exceptionally. + * + * @param args The tool name followed by that tool's arguments; may be empty. + */ public static void main(String[] args) { if (args.length == 0) { 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 46b622645..07bcf3154 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 @@ -23,10 +23,13 @@ import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; +import java.util.List; import java.util.Set; 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.ValueSource; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -35,9 +38,19 @@ import static org.junit.jupiter.api.Assertions.assertTrue; class SafetensorsFileTest { - // Builds a well-formed safetensors file: an 8-byte little-endian header length, the header - // JSON verbatim, then the raw data bytes. The header's data_offsets are expected to already - // be correct for the given data layout; callers construct both together. + private static final String MODEL_FILE_NAME = "model.safetensors"; + + // Builds the header JSON of a file holding one tensor, for tests that hand-roll headers with + // deliberately odd dtypes, shapes, or data_offsets. + private static String singleTensorHeader(String name, String dtype, String shape, + long begin, long end) { + return "{\"" + name + "\":{\"dtype\":\"" + dtype + "\",\"shape\":" + shape + + ",\"data_offsets\":[" + begin + "," + end + "]}}"; + } + + // Builds a safetensors file byte for byte: an 8-byte little-endian header length, the header + // JSON verbatim, then the raw data bytes. Used by the negative tests whose headers + // SafetensorsTestFiles would refuse to write; well-formed fixtures use that helper instead. private static Path writeFile(Path dir, String name, String headerJson, byte[] data) throws IOException { final byte[] headerBytes = headerJson.getBytes(StandardCharsets.UTF_8); @@ -61,11 +74,9 @@ class SafetensorsFileTest { @Test void testRoundTripsAFloat32Matrix(@TempDir Path dir) throws IOException { - final float[] values = {1f, 2f, 3f, 4f, 5f, 6f}; - final byte[] data = floatsToLittleEndianBytes(values); - final String header = "{\"weight\":{\"dtype\":\"F32\",\"shape\":[2,3]," - + "\"data_offsets\":[0," + data.length + "]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.matrix("weight", new float[][] {{1f, 2f, 3f}, {4f, 5f, 6f}})); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -75,25 +86,19 @@ class SafetensorsFileTest { assertEquals("F32", info.dtype()); assertArrayEquals(new int[] {2, 3}, info.shape()); assertEquals(6, info.elementCount()); - assertArrayEquals(values, parsed.readFloat32("weight")); + assertArrayEquals(new float[] {1f, 2f, 3f, 4f, 5f, 6f}, parsed.readFloat32("weight")); } @Test void testMultipleTensorsPreserveHeaderOrder(@TempDir Path dir) throws IOException { - final byte[] a = floatsToLittleEndianBytes(1f, 2f); - final byte[] b = floatsToLittleEndianBytes(3f, 4f, 5f); - final String header = "{\"first\":{\"dtype\":\"F32\",\"shape\":[2]," - + "\"data_offsets\":[0," + a.length + "]}," - + "\"second\":{\"dtype\":\"F32\",\"shape\":[3]," - + "\"data_offsets\":[" + a.length + "," + (a.length + b.length) + "]}}"; - final ByteArrayOutputStream data = new ByteArrayOutputStream(); - data.write(a); - data.write(b); - final Path file = writeFile(dir, "model.safetensors", header, data.toByteArray()); + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.vector("first", new float[] {1f, 2f}), + SafetensorsTestFiles.vector("second", new float[] {3f, 4f, 5f})); final SafetensorsFile parsed = SafetensorsFile.read(file); - assertEquals(java.util.List.of("first", "second"), java.util.List.copyOf(parsed.tensorNames())); + assertEquals(List.of("first", "second"), List.copyOf(parsed.tensorNames())); assertArrayEquals(new float[] {1f, 2f}, parsed.readFloat32("first")); assertArrayEquals(new float[] {3f, 4f, 5f}, parsed.readFloat32("second")); } @@ -103,7 +108,7 @@ class SafetensorsFileTest { final byte[] data = floatsToLittleEndianBytes(1f); final String header = "{\"__metadata__\":{\"format\":\"pt\",\"note\":\"line\\nbreak\"}," + "\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0," + data.length + "]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -117,7 +122,7 @@ class SafetensorsFileTest { final byte[] data = floatsToLittleEndianBytes(1f, 2f); final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[2]," + "\"data_offsets\":[0," + data.length + "],\"future_field\":{\"nested\":[1,2,3]}}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -126,16 +131,10 @@ class SafetensorsFileTest { @Test void testSingleMatrixTensorNameFindsTheOnly2DFloat32Tensor(@TempDir Path dir) throws IOException { - final byte[] scalar = floatsToLittleEndianBytes(9f); - final byte[] matrix = floatsToLittleEndianBytes(1f, 2f, 3f, 4f); - final String header = "{\"bias\":{\"dtype\":\"F32\",\"shape\":[1]," - + "\"data_offsets\":[0," + scalar.length + "]}," - + "\"embeddings\":{\"dtype\":\"F32\",\"shape\":[2,2]," - + "\"data_offsets\":[" + scalar.length + "," + (scalar.length + matrix.length) + "]}}"; - final ByteArrayOutputStream data = new ByteArrayOutputStream(); - data.write(scalar); - data.write(matrix); - final Path file = writeFile(dir, "model.safetensors", header, data.toByteArray()); + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.vector("bias", new float[] {9f}), + SafetensorsTestFiles.matrix("embeddings", new float[][] {{1f, 2f}, {3f, 4f}})); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -144,16 +143,10 @@ class SafetensorsFileTest { @Test void testSingleMatrixTensorNameRejectsAmbiguity(@TempDir Path dir) throws IOException { - final byte[] a = floatsToLittleEndianBytes(1f, 2f, 3f, 4f); - final byte[] b = floatsToLittleEndianBytes(5f, 6f, 7f, 8f); - final String header = "{\"a\":{\"dtype\":\"F32\",\"shape\":[2,2]," - + "\"data_offsets\":[0," + a.length + "]}," - + "\"b\":{\"dtype\":\"F32\",\"shape\":[2,2]," - + "\"data_offsets\":[" + a.length + "," + (a.length + b.length) + "]}}"; - final ByteArrayOutputStream data = new ByteArrayOutputStream(); - data.write(a); - data.write(b); - final Path file = writeFile(dir, "model.safetensors", header, data.toByteArray()); + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.matrix("a", new float[][] {{1f, 2f}, {3f, 4f}}), + SafetensorsTestFiles.matrix("b", new float[][] {{5f, 6f}, {7f, 8f}})); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -162,10 +155,8 @@ class SafetensorsFileTest { @Test void testSingleMatrixTensorNameRejectsNoCandidate(@TempDir Path dir) throws IOException { - final byte[] data = floatsToLittleEndianBytes(1f); - final String header = "{\"bias\":{\"dtype\":\"F32\",\"shape\":[1]," - + "\"data_offsets\":[0," + data.length + "]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, SafetensorsTestFiles.vector("bias", new float[] {1f})); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -175,9 +166,8 @@ class SafetensorsFileTest { @Test void testReadFloat32RejectsWrongDtype(@TempDir Path dir) throws IOException { final byte[] data = new byte[] {1, 2}; - final String header = "{\"ids\":{\"dtype\":\"I64\",\"shape\":[1]," - + "\"data_offsets\":[0,2]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final String header = singleTensorHeader("ids", "I64", "[1]", 0, 2); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -189,8 +179,8 @@ class SafetensorsFileTest { @Test void testTensorInfoRejectsUnknownName(@TempDir Path dir) throws IOException { final byte[] data = floatsToLittleEndianBytes(1f); - final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final String header = singleTensorHeader("w", "F32", "[1]", 0, 4); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -228,7 +218,7 @@ class SafetensorsFileTest { // header parser itself does not reject it; SafetensorsFile's post-parse check does. final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}," + "\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}}"; - final Path file = writeFile(dir, "model.safetensors", header, new byte[] {1, 2, 3, 4}); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[] {1, 2, 3, 4}); final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> SafetensorsFile.read(file)); @@ -238,15 +228,15 @@ class SafetensorsFileTest { @Test void testRejectsTensorMissingRequiredField(@TempDir Path dir) throws IOException { final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]}}"; - final Path file = writeFile(dir, "model.safetensors", header, new byte[0]); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[0]); assertThrows(IllegalArgumentException.class, () -> SafetensorsFile.read(file)); } @Test void testRejectsDataOffsetsOutOfRange(@TempDir Path dir) throws IOException { - final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,999]}}"; - final Path file = writeFile(dir, "model.safetensors", header, new byte[] {1, 2, 3, 4}); + final String header = singleTensorHeader("w", "F32", "[1]", 0, 999); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[] {1, 2, 3, 4}); assertThrows(IllegalArgumentException.class, () -> SafetensorsFile.read(file)); } @@ -254,7 +244,7 @@ class SafetensorsFileTest { @Test void testRejectsUnterminatedString(@TempDir Path dir) throws IOException { final String header = "{\"w\":{\"dtype\":\"F32"; - final Path file = writeFile(dir, "model.safetensors", header, new byte[0]); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[0]); assertThrows(IllegalArgumentException.class, () -> SafetensorsFile.read(file)); } @@ -264,9 +254,8 @@ class SafetensorsFileTest { // 2_000_000 * 2_000 = 4 billion elements, over the float[] ceiling. The bogus small data // range keeps the file tiny; the array-ceiling check fires before the range-mismatch check // because it subsumes it for tensors this large. - final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[2000000,2000]," - + "\"data_offsets\":[0,4]}}"; - final Path file = writeFile(dir, "model.safetensors", header, new byte[] {1, 2, 3, 4}); + final String header = singleTensorHeader("w", "F32", "[2000000,2000]", 0, 4); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[] {1, 2, 3, 4}); final SafetensorsFile parsed = SafetensorsFile.read(file); @@ -280,12 +269,11 @@ class SafetensorsFileTest { // Tensor data is streamed on demand rather than held in memory, so a file that shrinks // between read() and readFloat32() must fail loud, not return partial data. final byte[] data = floatsToLittleEndianBytes(1f, 2f); - final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[2]," - + "\"data_offsets\":[0," + data.length + "]}}"; - final Path file = writeFile(dir, "model.safetensors", header, data); + final String header = singleTensorHeader("w", "F32", "[2]", 0, data.length); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); - writeFile(dir, "model.safetensors", header, floatsToLittleEndianBytes(1f)); + writeFile(dir, MODEL_FILE_NAME, header, floatsToLittleEndianBytes(1f)); final IllegalStateException e = assertThrows(IllegalStateException.class, () -> parsed.readFloat32("w")); @@ -296,9 +284,8 @@ class SafetensorsFileTest { 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 String header = singleTensorHeader("w", "F32", "[2]", 0, data.length); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); final SafetensorsFile parsed = SafetensorsFile.read(file); final IllegalArgumentException e = @@ -334,32 +321,23 @@ class SafetensorsFileTest { assertTrue(e.getMessage().contains("overflows"), e.getMessage()); } - @Test - void testReadsF16TensorWidenedToFloat(@TempDir Path dir) throws IOException { - // F16 is model2vec's default output dtype, so this is the common downloaded-model case. - final Path file = dir.resolve("f16.safetensors"); - final float[] expected = {1.0f, -2.0f, 0.5f, 3.5f}; // all exact in IEEE half - SafetensorsTestFiles.write(file, "F16", SafetensorsTestFiles.vector("w", expected)); - - final SafetensorsFile parsed = SafetensorsFile.read(file); - assertEquals("F16", parsed.tensorInfo("w").dtype()); - assertArrayEquals(expected, parsed.readFloats("w"), 1e-3f); - } - - @Test - void testReadsBf16TensorWidenedToFloat(@TempDir Path dir) throws IOException { - final Path file = dir.resolve("bf16.safetensors"); - final float[] expected = {1.0f, -2.0f, 0.5f, 100.0f}; // exact in bfloat16 - SafetensorsTestFiles.write(file, "BF16", SafetensorsTestFiles.vector("w", expected)); + // F16 is Model2Vec's default output dtype, so widening is the common downloaded-model case; + // BF16 takes the same path with a different bit layout. + @ParameterizedTest + @ValueSource(strings = {"F16", "BF16"}) + void testReads16BitTensorWidenedToFloat(String dtype, @TempDir Path dir) throws IOException { + final Path file = dir.resolve(MODEL_FILE_NAME); + final float[] expected = {1.0f, -2.0f, 0.5f, 3.5f}; // exact in both 16-bit formats + SafetensorsTestFiles.write(file, dtype, SafetensorsTestFiles.vector("w", expected)); final SafetensorsFile parsed = SafetensorsFile.read(file); - assertEquals("BF16", parsed.tensorInfo("w").dtype()); + assertEquals(dtype, parsed.tensorInfo("w").dtype()); assertArrayEquals(expected, parsed.readFloats("w"), 1e-3f); } @Test void testSingleMatrixTensorNameAcceptsF16(@TempDir Path dir) throws IOException { - final Path file = dir.resolve("f16-matrix.safetensors"); + final Path file = dir.resolve(MODEL_FILE_NAME); SafetensorsTestFiles.write(file, "F16", SafetensorsTestFiles.matrix("embeddings", new float[][] {{1f, 2f}, {3f, 4f}})); @@ -369,7 +347,7 @@ class SafetensorsFileTest { @Test void testReadFloat32StrictlyRejectsF16(@TempDir Path dir) throws IOException { - final Path file = dir.resolve("f16-strict.safetensors"); + final Path file = dir.resolve(MODEL_FILE_NAME); SafetensorsTestFiles.write(file, "F16", SafetensorsTestFiles.vector("w", new float[] {1f, 2f})); final SafetensorsFile parsed = SafetensorsFile.read(file); diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.java index 043f97a37..3e118f60d 100644 --- a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.java +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.java @@ -44,6 +44,13 @@ class TokenizerJsonVocabTest { return file; } + @Test + void testRejectsNullAndMissingFile() { + assertThrows(IllegalArgumentException.class, () -> TokenizerJsonVocab.rows(null)); + assertThrows(IllegalArgumentException.class, + () -> TokenizerJsonVocab.rows(dir.resolve("absent.json"))); + } + @Test void testVocabListOrderIsTheRowOrder() throws IOException { final Path file = write("{\"model\":{\"type\":\"Unigram\",\"unk_id\":1,"
