jzonthemtn commented on code in PR #1152: URL: https://github.com/apache/opennlp/pull/1152#discussion_r3552694154
########## 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: Is it possible to have a second `load()` that reads the `normalize` and `lowercase` params from the `config.json`/`tokenizer_config.json` files, something like `public static StaticEmbeddingModel load(Path modelDirectory);`? The extra code to parse the JSON might make it not worth it though. -- 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]
