krickert commented on code in PR #1152: URL: https://github.com/apache/opennlp/pull/1152#discussion_r3554653843
########## opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java: ########## @@ -0,0 +1,511 @@ +/* + * 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.nio.file.Path; +import java.util.List; +import java.util.Set; +import java.util.SortedSet; +import java.util.TreeSet; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.tokenize.BertTokenizer; +import opennlp.tools.tokenize.WordpieceTokenizer; + +/** + * A static (non-contextual) sentence embedding model: a per-token vector table plus WordPiece + * tokenization, the pure-JVM word2vec/GloVe successor described in the design doc this module + * implements. Distilled tables in this shape (Model2Vec and compatible releases) carry a modern + * sentence-transformer's semantics in a flat lookup table, so embedding a sentence is tokenize, + * gather, (optionally) weight, mean-pool, and (optionally) normalize: no model forward pass, no + * GPU, no native runtime. + * + * <p>The pooling formula matches the reference Model2Vec implementations exactly (verified + * against MinishLab's Rust {@code model2vec-rs}, not assumed): {@code [CLS]}/{@code [SEP]} are + * never added to the pool (this class tokenizes for lookup, not for a transformer), unknown + * tokens are dropped rather than contributing a meaningless vector, each remaining token's + * vector is multiplied by its optional per-token weight, the sum is divided by the plain count + * of pooled tokens (not the sum of weights), and if the model calls for normalization the + * pooled vector is L2-normalized with an epsilon floor so a token-less input yields a zero + * vector rather than a division by zero.</p> + * + * <p><b>Thread safety.</b> Instances are immutable and safe for concurrent use after + * construction: every field is final, the loaded arrays are never exposed or mutated, and the + * tokenizer chain holds no per-call state. The one piece of global mutable state in that chain, + * the {@code keepNewLines} flag on the {@code WhitespaceTokenizer.INSTANCE} singleton that + * {@link WordpieceTokenizer} splits with, cannot affect results here: BERT basic tokenization + * has already replaced every whitespace character, line breaks included, with plain spaces + * before that split runs, so the flag's only behavioral branch never triggers on this input.</p> + */ +@ThreadSafe +public final class StaticEmbeddingModel { + + private static final float NORMALIZE_EPSILON = 1e-12f; + private static final String WEIGHTS_TENSOR_NAME = "weights"; + private static final int[] NO_EXCLUDED_ROWS = new int[0]; + // Never meaningful as a "similar word" result. + private static final Set<String> SPECIAL_TOKENS = Set.of(WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); + + private final float[] embeddings; + private final float[] weights; + private final int dimension; + private final WordPieceVocabulary vocabulary; + private final BertTokenizer tokenizer; + private final boolean normalize; + private final String unknownToken; + // Per-row L2 norms and the special-token mask are constants of the model, precomputed at + // load time so the nearest-neighbor scan does no per-row square-root or string hashing. + private final double[] rowNorms; + private final boolean[] specialRows; + + private StaticEmbeddingModel(float[] embeddings, float[] weights, int dimension, + WordPieceVocabulary vocabulary, BertTokenizer tokenizer, + boolean normalize, String unknownToken, double[] rowNorms, + boolean[] specialRows) { + this.embeddings = embeddings; + this.weights = weights; + this.dimension = dimension; + this.vocabulary = vocabulary; + this.tokenizer = tokenizer; + this.normalize = normalize; + this.unknownToken = unknownToken; + this.rowNorms = rowNorms; + this.specialRows = specialRows; + } + + /** + * Loads a static embedding model from a BERT-style {@code vocab.txt} and a safetensors weight + * file, the file pair a Model2Vec-family distillation publishes. No model is bundled with this + * module: the caller points at files they downloaded (see the module's design doc for the + * license posture). + * + * @param vocabularyFile The {@code vocab.txt} file: one token per line, line number is the + * token's row id. Must not be {@code null} and must exist. + * @param safetensorsFile The {@code model.safetensors} file. Must not be {@code null} and + * must exist, and must contain exactly one 2-D {@code F32} tensor + * (the embedding matrix) whose row count matches the vocabulary size. + * 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. + * @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. + */ + public static StaticEmbeddingModel load(Path vocabularyFile, Path safetensorsFile, Review Comment: Good idea, added in 2c0f0135. `load(Path modelDirectory)` now reads `normalize` from `config.json` and `do_lower_case` from `tokenizer_config.json`. I checked the field names against published model releases rather than assuming them. The JSON cost turned out to be small: the header parser's scanning primitives moved into a shared cursor class, and the config reader only pulls top-level booleans, skipping every other field structurally. One edge case I handled deliberately: a `tokenizer_config.json` that explicitly sets `strip_accents` against `do_lower_case` is rejected with a message pointing at the explicit overload, since the single lower-case flag follows the BERT convention (strip accents exactly when lower-casing) and loading it silently would mis-tokenize. ########## opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java: ########## @@ -0,0 +1,215 @@ +/* + * 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.io.UncheckedIOException; +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.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.commons.ThreadSafe; + +/** + * Reads a <a href="https://github.com/huggingface/safetensors";>safetensors</a> file: an 8-byte + * little-endian header length, a JSON header describing each tensor's dtype, shape, and byte + * range, followed by the raw tensor bytes. Deliberately not a general tensor-format library: + * only the {@code F32} decode path {@link #readFloat32(String)} needs is implemented, since + * that is what a distilled static-embedding table stores. + * + * <p><b>Security.</b> Unlike PyTorch's pickle-based checkpoint format, safetensors carries no + * executable content: the header is data-only JSON and the body is raw tensor bytes, so loading + * one cannot execute arbitrary code. No hardening beyond ordinary malformed-input handling is + * needed.</p> + * + * <p>The whole file is read into memory up front (matching the project's existing bundled-data + * readers), which is appropriate for the small (tens of megabytes) tables this module targets. + * Instances are immutable and safe for concurrent reads after construction; every + * {@link #readFloat32(String)} call decodes into a fresh array the caller owns.</p> + */ +@ThreadSafe +public final class SafetensorsFile { + + private static final int HEADER_LENGTH_PREFIX_BYTES = 8; + + private final byte[] bytes; + private final long dataStart; + private final Map<String, TensorInfo> tensorsByName; + private final Map<String, String> metadata; + + private SafetensorsFile(byte[] bytes, long dataStart, Map<String, TensorInfo> tensorsByName, + Map<String, String> metadata) { + this.bytes = bytes; + this.dataStart = dataStart; + this.tensorsByName = tensorsByName; + this.metadata = metadata; + } + + /** + * Reads a safetensors file. + * + * @param file The file to read. Must not be {@code null} and must exist. + * @return The parsed file, with every tensor's metadata resolved and validated against the + * 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. + */ + public static SafetensorsFile read(Path file) { + 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 byte[] bytes; + try { + bytes = Files.readAllBytes(file); Review Comment: You're right, and it was worth fixing properly rather than documenting away. `Files.readAllBytes` refuses files over ~2 GB with an `OutOfMemoryError`, so that was the real ceiling. Fixed in d2f388eb: the reader no longer holds the file in memory at all. The header is read eagerly through a `FileChannel`, and tensor data streams straight into the caller's `float[]` via positional reads at long offsets, so file size is now unbounded. The remaining limit is per decoded tensor (a `float[]` caps at roughly 2.1 billion elements, about 8.6 GB of F32), and that is checked explicitly with a clear message. For what it's worth, the cleaner long-term answer is a memory-mapped `MemorySegment` from `java.lang.foreign`, but that API only became final in JDK 22 and our baseline is 21, where it is preview-gated; a library can't ship preview-compiled classes. There's a comment on the read loop marking that upgrade path. Side benefit of the streaming approach: peak load memory roughly halves, since the file bytes and the decoded array no longer coexist. ########## opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java: ########## Review Comment: Agreed, added in ad188c26: 25 direct tests covering header order, all escape sequences, structural skipping of unknown fields, and every malformed-input branch, each asserting the offset-carrying error message. Writing them also caught a real gap: the parser stopped at the closing brace and silently accepted trailing garbage. It now rejects that, while still tolerating the trailing whitespace writers pad headers with for data alignment. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
