diff --git a/dev/embeddings/README.md b/dev/embeddings/README.md new file mode 100644 index 0000000000..2a488743a1 --- /dev/null +++ b/dev/embeddings/README.md @@ -0,0 +1,38 @@ + + +# Embeddings scripts + +Developer scripts around the static embeddings module. None of them are part of the build; they +make the module's numbers and its worked example reproducible from a checkout. + +## `distill_bge_m3.py` + +The runnable form of the TRAINING.md worked example: distills the multilingual bge-m3 teacher +into a 256-dimension static table with Model2Vec. Needs a Python environment with +`model2vec[distill]` installed; the script's header shows the setup. After it finishes, copy the +teacher's `sentencepiece.bpe.model` next to the output and verify with the `AssembleModel` +command. + +## `parity/` + +The parity and single-thread speed comparison between this module and the model2vec Python +reference: the same model and the same multilingual sentences on both sides, the two vector sets +checked against each other, and both throughputs measured with the same fixed-duration +methodology. `sh run.sh` after building the project; see the script header for the environment +overrides. A run passes only when the vectors agree within float tolerance, so the two speeds it +prints are for implementations producing the same answer. diff --git a/dev/embeddings/distill_bge_m3.py b/dev/embeddings/distill_bge_m3.py new file mode 100644 index 0000000000..5a33620e03 --- /dev/null +++ b/dev/embeddings/distill_bge_m3.py @@ -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. + +"""Distills the multilingual bge-m3 teacher into a static embedding table. + +This is the worked example from opennlp-extensions/opennlp-embeddings/TRAINING.md as a +runnable script. It needs a Python environment with model2vec's distill extra installed: + + uv venv .venv-distill + uv pip install --python .venv-distill "model2vec[distill]" + .venv-distill/bin/python distill_bge_m3.py [output-dir] + +After it finishes, copy the teacher's trained SentencePiece file +(sentencepiece.bpe.model on the model hub) into the output directory and run the +AssembleModel command to verify the directory loads: + + opennlp-embeddings AssembleModel -modelDir + +256 dimensions is the deliberate default: distilling the same teacher at 512 gives the +same cross-lingual similarity within noise while doubling the matrix and halving embed +throughput, because PCA to 256 already captures the useful variance. +""" + +import os +import sys + +from model2vec.distill import distill + +out = sys.argv[1] if len(sys.argv) > 1 else "bge-m3-static" + +static = distill("BAAI/bge-m3", pca_dims=256) +static.save_pretrained(out) +print("SAVED:", out, "dim:", static.dim) + +print("=== output files ===") +for name in sorted(os.listdir(out)): + path = os.path.join(out, name) + print(f" {os.path.getsize(path):>12} {name}") +print("Now copy the teacher's sentencepiece.bpe.model into", out, + "and run: opennlp-embeddings AssembleModel -modelDir", out) diff --git a/dev/embeddings/parity/EmbedBenchM3.java b/dev/embeddings/parity/EmbedBenchM3.java new file mode 100644 index 0000000000..cfa366513c --- /dev/null +++ b/dev/embeddings/parity/EmbedBenchM3.java @@ -0,0 +1,92 @@ +/* + * 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. + */ + +import java.io.BufferedWriter; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import opennlp.embeddings.StaticEmbeddingModel; + +/** + * The JVM half of the parity and speed comparison (see run.sh). Loads the static table, writes + * one vector per input sentence for the parity check, then measures single-thread embed + * throughput with the same fixed-duration, warmup-discarded methodology the Python side uses. + * + *

Args: modelDir sentencesFile vectorsOut warmupSeconds measureSeconds

+ */ +public final class EmbedBenchM3 { + + /** Not instantiable. */ + private EmbedBenchM3() { + } + + /** + * Runs the parity dump and the single-thread throughput measurement. + * + * @param args modelDir, sentencesFile, vectorsOut, warmupSeconds, measureSeconds. + * @throws Exception Thrown if a file cannot be read or written. + */ + public static void main(String[] args) throws Exception { + final Path modelDir = Path.of(args[0]); + final List sentences = Files.readAllLines(Path.of(args[1]), StandardCharsets.UTF_8) + .stream().map(String::strip).filter(s -> !s.isEmpty()).toList(); + final Path vectorsOut = Path.of(args[2]); + final int warmupSeconds = Integer.parseInt(args[3]); + final int measureSeconds = Integer.parseInt(args[4]); + + final long loadStart = System.nanoTime(); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(modelDir); + final double loadMs = (System.nanoTime() - loadStart) / 1e6; + + // One vector per sentence, so the Python side can diff them for parity. + try (BufferedWriter writer = Files.newBufferedWriter(vectorsOut, StandardCharsets.UTF_8)) { + for (final String sentence : sentences) { + final float[] vector = model.embed(sentence); + final StringBuilder line = new StringBuilder(); + for (int i = 0; i < vector.length; i++) { + if (i > 0) { + line.append(' '); + } + line.append(Float.toString(vector[i])); + } + writer.write(line.toString()); + writer.newLine(); + } + } + + final long warmupEnd = System.nanoTime() + warmupSeconds * 1_000_000_000L; + int index = 0; + while (System.nanoTime() < warmupEnd) { + model.embed(sentences.get(index++ % sentences.size())); + } + + long embedded = 0; + final long measureStart = System.nanoTime(); + final long measureEnd = measureStart + measureSeconds * 1_000_000_000L; + index = 0; + while (System.nanoTime() < measureEnd) { + model.embed(sentences.get(index++ % sentences.size())); + embedded++; + } + final double seconds = (System.nanoTime() - measureStart) / 1e9; + + System.out.printf("JVM load %.0f ms | %,.0f texts/s single-thread (%d embeds in %.1fs)%n", + loadMs, embedded / seconds, embedded, seconds); + } +} diff --git a/dev/embeddings/parity/parity_speed.py b/dev/embeddings/parity/parity_speed.py new file mode 100644 index 0000000000..13a82f129c --- /dev/null +++ b/dev/embeddings/parity/parity_speed.py @@ -0,0 +1,93 @@ +# 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. + +"""The Python half of the parity and speed comparison, plus the final parity check. + +Loads the same static table with model2vec, writes one vector per sentence, measures +single-thread throughput with the same fixed-duration warmup-discarded loop the JVM side +uses, then loads the JVM's vectors (written first by run.sh) and reports the parity between +them. + +The point is not to declare a winner; it is to show that both implementations produce the +same vectors, and to let anyone reproduce both numbers on their own hardware. + +Usage: parity_speed.py [sentences-file] [jvm-vectors-file] +""" + +import sys +import time + +import numpy as np +from model2vec import StaticModel + +WARMUP_SECONDS = 3 +MEASURE_SECONDS = 5 + + +def read_sentences(path): + with open(path, encoding="utf-8") as handle: + return [line.strip() for line in handle if line.strip()] + + +def main(): + model_dir = sys.argv[1] + sentences_file = sys.argv[2] if len(sys.argv) > 2 else "sentences.txt" + jvm_vectors_file = sys.argv[3] if len(sys.argv) > 3 else "jvm_vectors.tsv" + sentences = read_sentences(sentences_file) + + load_start = time.time() + model = StaticModel.from_pretrained(model_dir) + load_ms = (time.time() - load_start) * 1000.0 + + python_vectors = np.array([model.encode(s) for s in sentences], dtype=np.float32) + + end = time.time() + WARMUP_SECONDS + i = 0 + while time.time() < end: + model.encode(sentences[i % len(sentences)]) + i += 1 + + embedded = 0 + i = 0 + start = time.time() + end = start + MEASURE_SECONDS + while time.time() < end: + model.encode(sentences[i % len(sentences)]) + embedded += 1 + i += 1 + seconds = time.time() - start + print(f"Python load {load_ms:.0f} ms | {embedded / seconds:,.0f} texts/s single-thread " + f"({embedded} embeds in {seconds:.1f}s)") + + jvm_vectors = np.loadtxt(jvm_vectors_file, dtype=np.float32) + if jvm_vectors.shape != python_vectors.shape: + print(f"PARITY FAIL: shape mismatch {jvm_vectors.shape} vs {python_vectors.shape}") + sys.exit(1) + + max_abs_diff = float(np.abs(python_vectors - jvm_vectors).max()) + cosines = [ + float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))) + for a, b in zip(python_vectors, jvm_vectors) + ] + print(f"Parity max abs diff {max_abs_diff:.2e} | min cosine {min(cosines):.6f} " + f"over {len(sentences)} sentences in {python_vectors.shape[1]} dims") + if min(cosines) < 0.9999: + print("PARITY FAIL: vectors diverge") + sys.exit(1) + print("Parity OK: the JVM and Python vectors are the same within float tolerance") + + +if __name__ == "__main__": + main() diff --git a/dev/embeddings/parity/run.sh b/dev/embeddings/parity/run.sh new file mode 100755 index 0000000000..6fe1e4a920 --- /dev/null +++ b/dev/embeddings/parity/run.sh @@ -0,0 +1,47 @@ +#!/bin/sh +# 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. + +# Reproduces the parity and single-thread speed comparison between opennlp-embeddings and the +# model2vec Python reference: the same model and the same sentences on both sides, with the +# vector sets checked against each other. Run from this directory after building the project +# (mvn install, or at least mvn compile from the repository root). +# +# Environment overrides: +# MODEL_DIR the static model directory (default: bge-m3-static in this directory; +# see ../distill_bge_m3.py and opennlp-extensions/opennlp-embeddings/TRAINING.md +# to produce one) +# PYTHON a Python interpreter with model2vec installed (default: python3) +set -e + +MODEL_DIR="${MODEL_DIR:-bge-m3-static}" +PYTHON="${PYTHON:-python3}" + +# The repository root is three levels above this script. +ROOT=$(cd "$(dirname "$0")/../../.." && pwd) +CP="$ROOT/opennlp-api/target/classes:$ROOT/opennlp-core/opennlp-runtime/target/classes:$ROOT/opennlp-extensions/opennlp-subword/target/classes:$ROOT/opennlp-extensions/opennlp-embeddings/target/classes" + +echo "Model: $MODEL_DIR" +echo "Sentences: $(grep -c . sentences.txt) lines, multilingual" +echo + +# JVM side first: it writes jvm_vectors.tsv, which the Python side then diffs. +javac -cp "$CP" -d . EmbedBenchM3.java +java -cp "$CP:." EmbedBenchM3 "$MODEL_DIR" sentences.txt jvm_vectors.tsv 3 5 + +# Python side: prints its own rate, then reports parity against the JVM's vectors. +"$PYTHON" parity_speed.py "$MODEL_DIR" diff --git a/dev/embeddings/parity/sentences.txt b/dev/embeddings/parity/sentences.txt new file mode 100644 index 0000000000..5e6b6a3370 --- /dev/null +++ b/dev/embeddings/parity/sentences.txt @@ -0,0 +1,20 @@ +The weather is beautiful today and the sky is clear. +Machine learning models turn text into vectors. +I would like a cup of coffee with milk please. +The quarterly financial results disappointed investors. +Das Wetter ist heute wunderschoen und der Himmel ist klar. +Maschinelles Lernen verwandelt Text in Vektoren. +Le temps est magnifique aujourd'hui et le ciel est degage. +Los modelos de aprendizaje automatico convierten texto en vectores. +今天天气很好，天空很晴朗。 +机器学习模型把文本转换成向量。 +今日はとても良い天気で空が澄んでいます。 +機械学習モデルはテキストをベクトルに変換します。 +Сегодня прекрасная погода и ясное небо. +Модели машинного обучения превращают текст в векторы. +La retrieval semantica trova documenti per significato non per parole. +Natural language processing is a field of artificial intelligence. +A quick brown fox jumps over the lazy dog near the river. +Embeddings place similar sentences close together in space. +Coffee, tea, and espresso are all popular hot drinks. +The library opens at nine in the morning on weekdays. diff --git a/opennlp-api/src/main/java/opennlp/tools/embeddings/TextEmbedder.java b/opennlp-api/src/main/java/opennlp/tools/embeddings/TextEmbedder.java new file mode 100644 index 0000000000..32b3396afa --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/embeddings/TextEmbedder.java @@ -0,0 +1,79 @@ +/* + * 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.tools.embeddings; + +import java.util.List; + +import opennlp.tools.util.java.Experimental; + +/** + * Encodes a piece of text into a single fixed-length vector. + * + *

A text embedder maps whole texts (a sentence, a paragraph, a document) into one dense + * vector whose geometry carries meaning: texts about the same thing land near each other. This + * is the text-level counterpart of {@link opennlp.tools.util.wordvector.WordVectorTable}, which + * looks up a stored vector for a single word; an embedder composes a vector for text it has + * never seen, handling tokenization and pooling internally.

+ * + *

Thread safety is implementation specific. Failures during encoding surface as unchecked + * exceptions carrying the underlying cause.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public interface TextEmbedder { + + /** + * Embeds a piece of text. + * + *

Behavior for empty text, or text with no tokens the embedder recognizes, is + * implementation-defined: an implementation may return a zero vector, the vector of a special + * or fallback token, or something else, and should document its choice. Callers that need a + * uniform response should handle it themselves.

+ * + * @param text The text to embed. Must not be {@code null}. + * @return The embedding vector, of length {@link #dimension()}. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + float[] embed(CharSequence text); + + /** + * Embeds several texts. + * + *

The default implementation embeds one text at a time. Implementations backed by a + * runtime that executes batches more efficiently than single inputs should override this + * method.

+ * + * @param texts The texts to embed. Must not be {@code null} and must not contain {@code null}. + * @return One embedding vector per input, in input order. + * @throws IllegalArgumentException Thrown if {@code texts} is {@code null} or contains + * {@code null}. + */ + default float[][] embedAll(List texts) { + if (texts == null) { + throw new IllegalArgumentException("Texts must not be null"); + } + final float[][] vectors = new float[texts.size()][]; + for (int i = 0; i < vectors.length; i++) { + vectors[i] = embed(texts.get(i)); + } + return vectors; + } + + /** {@return the dimension of every vector this embedder produces} */ + int dimension(); +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java index 455f67f110..1cbfffe312 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java @@ -19,7 +19,7 @@ /** * Character classifications and text transforms of the reference BERT - * {@code BasicTokenizer}, shared by {@link BertTokenizer} and + * {@code BasicTokenizer}, shared by {@link WordpieceEncoder} and * {@link WordpieceTokenizer}. */ final class BertNormalization { diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java index 5517548a42..cea1fd8f53 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java @@ -6,7 +6,7 @@ * (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 + * 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, @@ -14,69 +14,39 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package opennlp.tools.tokenize; -import java.text.Normalizer; -import java.util.Locale; -import java.util.Objects; +import java.util.ArrayList; import java.util.Set; import opennlp.tools.util.Span; /** * A {@link Tokenizer} implementation of the full BERT tokenization pipeline: - * basic tokenization (text normalization) followed by wordpiece tokenization. - *

- * The basic tokenization stage reproduces the reference BERT - * {@code BasicTokenizer}: - *

    - *
  1. Removal of control characters and normalization of all whitespace - * to single spaces.
    2. - *
  2. Whitespace isolation of CJK ideographs.
    3. - *
  3. For uncased models: lower casing and accent stripping - * (Unicode NFD decomposition with removal of combining marks).
    4. - *
  4. Isolation of every punctuation character as its own token.
    5. - *
- * The normalized text is then split into subwords by a - * {@link WordpieceTokenizer} sharing the same vocabulary and special tokens. - *

- * This pipeline is required for correct results with BERT-style models: - * feeding raw text directly to {@link WordpieceTokenizer} maps every token - * that does not literally appear in the vocabulary - for uncased models that - * includes every capitalized word - to the unknown token. - *

- * Whether to use the lower casing variant is a property of the model: uncased - * models (for example {@code bert-base-uncased} and the - * {@code sentence-transformers} models derived from it) require it, cased - * models must not use it. Accent stripping is coupled to lower casing, as in - * the reference implementation's default ({@code strip_accents} follows - * {@code do_lower_case} unless overridden). - *

- * For reference see: - *

+ * basic tokenization (text normalization) followed by wordpiece tokenization, + * with the classification and separator tokens framing every result. + * + * @deprecated Use {@link WordpieceEncoder} instead: + * {@link WordpieceEncoder#encodeToPieces(CharSequence)} returns the same + * {@code String[]} as {@link #tokenize(String)}, and {@code encode} additionally + * carries vocabulary ids and original-text spans. This class is scheduled for + * removal after one stable release. * - * @see WordpieceTokenizer + * @see WordpieceEncoder */ +@Deprecated(since = "3.0.0", forRemoval = true) public class BertTokenizer implements Tokenizer { - /** - * Maximum characters per word before the word is replaced with the unknown - * token, matching the reference BERT implementation. - */ - private static final int MAX_WORD_CHARACTERS = 100; - - private final WordpieceTokenizer wordpieceTokenizer; - private final boolean lowerCase; + private final WordpieceEncoder encoder; /** * Initializes a {@link BertTokenizer} for an uncased BERT model, * with lower casing and accent stripping enabled. * * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. + * + * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, + * contains {@code null}, or is missing a BERT special token. */ public BertTokenizer(Set vocabulary) { this(vocabulary, true); @@ -88,6 +58,9 @@ public BertTokenizer(Set vocabulary) { * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. * @param lowerCase {@code true} for uncased models (lower casing and accent * stripping), {@code false} for cased models. + * + * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, + * contains {@code null}, or is missing a BERT special token. */ public BertTokenizer(Set vocabulary, boolean lowerCase) { this(vocabulary, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, @@ -101,19 +74,23 @@ public BertTokenizer(Set vocabulary, boolean lowerCase) { * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. * @param lowerCase {@code true} for uncased models (lower casing and * accent stripping), {@code false} for cased models. - * @param classificationToken The CLS token. - * @param separatorToken The SEP token. - * @param unknownToken The UNK token. + * @param classificationToken The CLS token; must be in the vocabulary. + * @param separatorToken The SEP token; must be in the vocabulary. + * @param unknownToken The UNK token; must be in the vocabulary. + * + * @throws IllegalArgumentException Thrown if any argument is {@code null}, + * the vocabulary contains {@code null}, or a special token is missing + * from the vocabulary. */ public BertTokenizer(Set vocabulary, boolean lowerCase, String classificationToken, String separatorToken, String unknownToken) { - Objects.requireNonNull(vocabulary, "vocabulary must not be null"); - Objects.requireNonNull(classificationToken, "classificationToken must not be null"); - Objects.requireNonNull(separatorToken, "separatorToken must not be null"); - Objects.requireNonNull(unknownToken, "unknownToken must not be null"); - this.wordpieceTokenizer = new WordpieceTokenizer(vocabulary, - classificationToken, separatorToken, unknownToken, MAX_WORD_CHARACTERS); - this.lowerCase = lowerCase; + if (vocabulary == null) { + throw new IllegalArgumentException("vocabulary must not be null"); + } + // The encoder assigns each piece its list index as the id. Ids are unused on the + // tokenize() path, so synthesizing them from an arbitrary set order is fine. + this.encoder = new WordpieceEncoder(new ArrayList<>(vocabulary), lowerCase, + classificationToken, separatorToken, unknownToken); } /** @@ -123,15 +100,19 @@ public BertTokenizer(Set vocabulary, boolean lowerCase, * @param text The text to tokenize. Must not be {@code null}. * * @return The wordpiece tokens. + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. */ @Override public String[] tokenize(String text) { - return wordpieceTokenizer.tokenize(normalize(text)); + return encoder.encodeToPieces(text); } /** * Not supported: wordpiece tokens (subwords, {@code ##} continuations and * special tokens) have no faithful character spans in the original text. + * Use {@link WordpieceEncoder#encode(CharSequence)} for pieces with + * original-text spans. * * @throws UnsupportedOperationException Always. */ @@ -141,71 +122,4 @@ public Span[] tokenizePos(String text) { "Wordpiece tokens cannot be mapped to character spans of the original text"); } - /** - * Applies the BERT basic tokenization (normalization) stage. - * - * @param text The text to normalize. Must not be {@code null}. - * - * @return The normalized text, ready for wordpiece tokenization. - */ - String normalize(String text) { - Objects.requireNonNull(text, "text must not be null"); - String normalized = cleanText(text); - normalized = isolateCjkCharacters(normalized); - if (lowerCase) { - normalized = stripAccents(normalized.toLowerCase(Locale.ROOT)); - } - return BertNormalization.isolatePunctuation(normalized); - } - - /** - * Removes invalid and control characters and normalizes all whitespace - * characters to plain spaces. - */ - private static String cleanText(String text) { - final StringBuilder cleaned = new StringBuilder(text.length()); - text.codePoints().forEach(codePoint -> { - if (codePoint == 0 || codePoint == 0xFFFD || BertNormalization.isControl(codePoint)) { - return; - } - if (BertNormalization.isWhitespace(codePoint)) { - cleaned.append(' '); - } else { - cleaned.appendCodePoint(codePoint); - } - }); - return cleaned.toString(); - } - - /** - * Surrounds every CJK ideograph with spaces, so each ideograph becomes its - * own token, matching the reference BERT treatment of Chinese text. - */ - private static String isolateCjkCharacters(String text) { - final StringBuilder spaced = new StringBuilder(text.length()); - text.codePoints().forEach(codePoint -> { - if (BertNormalization.isCjk(codePoint)) { - spaced.append(' ').appendCodePoint(codePoint).append(' '); - } else { - spaced.appendCodePoint(codePoint); - } - }); - return spaced.toString(); - } - - /** - * Removes accents by Unicode NFD decomposition followed by removal of - * combining marks ({@code Mn}). - */ - private static String stripAccents(String text) { - final String decomposed = Normalizer.normalize(text, Normalizer.Form.NFD); - final StringBuilder stripped = new StringBuilder(decomposed.length()); - decomposed.codePoints().forEach(codePoint -> { - if (Character.getType(codePoint) != Character.NON_SPACING_MARK) { - stripped.appendCodePoint(codePoint); - } - }); - return stripped.toString(); - } - } diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java new file mode 100644 index 0000000000..62260b7929 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java @@ -0,0 +1,58 @@ +/* + * 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.tools.tokenize; + +import opennlp.tools.util.Span; + +/** + * One subword unit produced by a {@link SubwordTokenizer}, carrying both the vocabulary view + * (the piece string and its id) and the exact place in the caller's text it came from. + * + *

The piece string is in the tokenizer's normalized form, so it is generally not a substring of + * the input. {@code start} and {@code end} are UTF-16 offsets into the original text, so the + * surface that produced this piece is {@code text.subSequence(start, end)}. Pieces that carry no + * surface of their own, such as control symbols or the fill bytes of a byte-fallback expansion, + * report an empty span with {@code start == end}.

+ * + * @param piece The piece in the vocabulary's normalized form; never null or empty. + * @param id The vocabulary id of the piece. + * @param start The inclusive start offset in the original text. + * @param end The exclusive end offset in the original text; not less than {@code start}. + */ +public record SubwordPiece(String piece, int id, int start, int end) { + + /** + * Instantiates a {@link SubwordPiece}. + * + * @throws IllegalArgumentException Thrown if {@code piece} is null or empty, or the span is + * negative or inverted. + */ + public SubwordPiece { + if (piece == null || piece.isEmpty()) { + throw new IllegalArgumentException("piece must not be null or empty"); + } + if (start < 0 || end < start) { + throw new IllegalArgumentException( + "The span [" + start + ", " + end + ") must not be negative or inverted."); + } + } + + /** {@return the original-text span of this piece as a {@link Span}} */ + public Span span() { + return new Span(start, end); + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java new file mode 100644 index 0000000000..957ff62229 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java @@ -0,0 +1,73 @@ +/* + * 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.tools.tokenize; + +import java.util.List; + +/** + * Splits text into subword units against a fixed vocabulary, reporting for every unit its + * vocabulary id and the exact span of the original text it covers. + * + *

The segmentation is vocabulary-driven rather than linguistic, and each piece is in the + * model's normalized form, so a piece is generally not a substring of the input. The offsets + * carried by each {@link SubwordPiece} always refer to the caller's original text.

+ * + *

Thread safety is implementation specific.

+ */ +public interface SubwordTokenizer { + + /** + * Encodes text into subword pieces. + * + * @param text The text to encode; must not be null. + * @return The pieces in text order; empty when the text contains nothing encodable. + * @throws IllegalArgumentException Thrown if {@code text} is null. + */ + List encode(CharSequence text); + + /** + * Encodes text into vocabulary ids. + * + * @param text The text to encode; must not be null. + * @return The ids in text order; empty when the text contains nothing encodable. + * @throws IllegalArgumentException Thrown if {@code text} is null. + */ + default int[] encodeToIds(CharSequence text) { + final List pieces = encode(text); + final int[] ids = new int[pieces.size()]; + for (int i = 0; i < ids.length; i++) { + ids[i] = pieces.get(i).id(); + } + return ids; + } + + /** + * Encodes text into piece strings in the vocabulary's normalized form. + * + * @param text The text to encode; must not be null. + * @return The pieces in text order; empty when the text contains nothing encodable. + * @throws IllegalArgumentException Thrown if {@code text} is null. + */ + default String[] encodeToPieces(CharSequence text) { + final List pieces = encode(text); + final String[] out = new String[pieces.size()]; + for (int i = 0; i < out.length; i++) { + out[i] = pieces.get(i).piece(); + } + return out; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java new file mode 100644 index 0000000000..8a8ed7adbc --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java @@ -0,0 +1,533 @@ +/* + * 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.tools.tokenize; + +import java.text.Normalizer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.commons.ThreadSafe; + +/** + * A {@link SubwordTokenizer} running the full BERT tokenization pipeline of the reference + * implementation: basic tokenization (control removal, whitespace normalization, CJK + * isolation, optional lower casing with accent stripping, punctuation isolation) followed by + * greedy longest-match wordpiece segmentation. + * + *

Every piece carries its vocabulary id and the span of the original text it came + * from, surviving the normalization steps that change, insert, and remove characters. The + * classification and separator pieces frame every encoding, carrying empty spans at the + * text's boundaries, so {@link #encode(CharSequence)} is never empty.

+ * + *

Ids follow the line-number convention of BERT {@code vocab.txt} files: with the list + * constructors a piece's id is its index, and with the map constructor the ids are given + * explicitly. The classification, separator, and unknown tokens must all be present in the + * vocabulary, because every emitted piece must have an id. Vocabulary entries starting with + * {@code ##} are continuation pieces, matching a word's interior rather than its start.

+ * + *

Instances are immutable and safe for concurrent use by multiple threads.

+ * + * @see WordpieceTokenizer + */ +@ThreadSafe +public final class WordpieceEncoder implements SubwordTokenizer { + + // The wordpiece vocabulary convention: a piece with this prefix continues the current word, + // so it can only match after the word's first piece. + private static final String CONTINUATION_PREFIX = "##"; + + // The reference implementation's limit: longer words become the unknown piece. + private static final int MAX_WORD_CHARACTERS = 100; + + private final Set vocabulary; + private final Map ids; + private final boolean lowerCase; + private final String classificationToken; + private final String separatorToken; + private final String unknownToken; + private final int classificationId; + private final int separatorId; + private final int unknownId; + + /** + * Instantiates an encoder for an uncased BERT model with the BERT special tokens. + * + * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be null, + * must not contain nulls or duplicates. + * @throws IllegalArgumentException Thrown if the vocabulary is null, contains a null or + * duplicate entry, or a BERT special token is missing from it. + */ + public WordpieceEncoder(List vocabulary) { + this(vocabulary, true); + } + + /** + * Instantiates an encoder with the BERT special tokens. + * + * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be null, + * must not contain nulls or duplicates. + * @param lowerCase True for uncased models (lower casing and accent stripping), false for + * cased models. + * @throws IllegalArgumentException Thrown if the vocabulary is null, contains a null or + * duplicate entry, or a BERT special token is missing from it. + */ + public WordpieceEncoder(List vocabulary, boolean lowerCase) { + this(vocabulary, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); + } + + /** + * Instantiates an encoder with custom special tokens, for models that do not use the BERT + * defaults. + * + * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be + * null, must not contain nulls or duplicates. + * @param lowerCase True for uncased models (lower casing and accent stripping), + * false for cased models. + * @param classificationToken The CLS token; must be in the vocabulary. + * @param separatorToken The SEP token; must be in the vocabulary. + * @param unknownToken The UNK token; must be in the vocabulary. + * @throws IllegalArgumentException Thrown if any argument is null, the vocabulary contains + * a null or duplicate entry, or a special token is missing from the vocabulary. + */ + public WordpieceEncoder(List vocabulary, boolean lowerCase, + String classificationToken, String separatorToken, + String unknownToken) { + this(byPiece(vocabulary), lowerCase, classificationToken, separatorToken, unknownToken); + } + + /** + * Instantiates an encoder from an explicit piece-to-id mapping, for vocabularies whose ids + * are not contiguous line numbers. + * + * @param vocabularyIds The piece-to-id mapping. Must not be null, must not contain + * null keys or values. + * @param lowerCase True for uncased models (lower casing and accent stripping), + * false for cased models. + * @param classificationToken The CLS token; must be in the vocabulary. + * @param separatorToken The SEP token; must be in the vocabulary. + * @param unknownToken The UNK token; must be in the vocabulary. + * @throws IllegalArgumentException Thrown if any argument is null, the mapping contains a + * null key or value, or a special token is missing from the vocabulary. + */ + public WordpieceEncoder(Map vocabularyIds, boolean lowerCase, + String classificationToken, String separatorToken, + String unknownToken) { + if (vocabularyIds == null) { + throw new IllegalArgumentException("vocabularyIds must not be null"); + } + if (classificationToken == null) { + throw new IllegalArgumentException("classificationToken must not be null"); + } + if (separatorToken == null) { + throw new IllegalArgumentException("separatorToken must not be null"); + } + if (unknownToken == null) { + throw new IllegalArgumentException("unknownToken must not be null"); + } + final Map byPiece = new HashMap<>(vocabularyIds.size() * 2); + for (final Map.Entry entry : vocabularyIds.entrySet()) { + if (entry.getKey() == null || entry.getValue() == null) { + throw new IllegalArgumentException( + "The vocabulary must not contain null pieces or ids: " + entry); + } + byPiece.put(entry.getKey(), entry.getValue()); + } + this.vocabulary = new HashSet<>(byPiece.keySet()); + this.ids = byPiece; + this.lowerCase = lowerCase; + this.classificationToken = classificationToken; + this.separatorToken = separatorToken; + this.unknownToken = unknownToken; + this.classificationId = requiredId(byPiece, classificationToken); + this.separatorId = requiredId(byPiece, separatorToken); + this.unknownId = requiredId(byPiece, unknownToken); + } + + /** + * Converts an ordered vocabulary list into the piece-to-id mapping, assigning each piece its + * index as the id. + * + * @param vocabulary The ordered vocabulary. + * @return The piece-to-id mapping. + * @throws IllegalArgumentException Thrown if the list is null or contains a null or duplicate + * entry. + */ + private static Map byPiece(List vocabulary) { + if (vocabulary == null) { + throw new IllegalArgumentException("vocabulary must not be null"); + } + final Map byPiece = new HashMap<>(vocabulary.size() * 2); + for (int id = 0; id < vocabulary.size(); id++) { + final String piece = vocabulary.get(id); + if (piece == null) { + throw new IllegalArgumentException("The vocabulary contains null at index " + id + "."); + } + if (byPiece.putIfAbsent(piece, id) != null) { + throw new IllegalArgumentException("The vocabulary contains '" + piece + + "' more than once; ids would be ambiguous."); + } + } + return byPiece; + } + + /** + * Looks up the id of a special token that must be present in the vocabulary. + * + * @param ids The piece-to-id mapping. + * @param specialToken The token to look up. + * @return The token's id. + * @throws IllegalArgumentException Thrown if the token is not in the vocabulary. + */ + private static int requiredId(Map ids, String specialToken) { + final Integer id = ids.get(specialToken); + if (id == null) { + throw new IllegalArgumentException("The special token '" + specialToken + + "' is not in the vocabulary; every emitted piece must have an id."); + } + return id; + } + + /** {@inheritDoc} */ + @Override + public List encode(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final String original = text.toString(); + + // The normalized text, one original-text range per char, built through the reference + // pipeline's transformations in the reference order. + MappedText mapped = cleanAndIsolateCjk(original); + if (lowerCase) { + mapped = lowerCaseAndStripAccents(mapped); + } + mapped = isolatePunctuation(mapped); + + final List pieces = new ArrayList<>(); + pieces.add(new SubwordPiece(classificationToken, classificationId, 0, 0)); + int from = 0; + while (from < mapped.length) { + if (mapped.chars[from] == ' ') { + from++; + continue; + } + int to = from; + while (to < mapped.length && mapped.chars[to] != ' ') { + to++; + } + encodeWord(mapped, from, to, pieces); + from = to; + } + pieces.add(new SubwordPiece(separatorToken, separatorId, + original.length(), original.length())); + return pieces; + } + + /** + * Greedily longest-match segments one whitespace-delimited word; the pieces are emitted only if + * the whole word is representable, otherwise the word becomes a single unknown piece. + * + * @param mapped The normalized text with per-character original-text ranges. + * @param from The inclusive start of the word in {@code mapped}. + * @param to The exclusive end of the word in {@code mapped}. + * @param pieces The output list to append to. + */ + private void encodeWord(MappedText mapped, int from, int to, List pieces) { + final int wordStart = mapped.starts[from]; + final int wordEnd = mapped.ends[to - 1]; + if (to - from > MAX_WORD_CHARACTERS) { + pieces.add(new SubwordPiece(unknownToken, unknownId, wordStart, wordEnd)); + return; + } + final List wordPieces = new ArrayList<>(); + int start = from; + boolean found = true; + while (start < to) { + int end = to; + found = false; + while (start < end) { + String substring = new String(mapped.chars, start, end - start); + if (start > from) { + substring = CONTINUATION_PREFIX + substring; + } + if (vocabulary.contains(substring)) { + wordPieces.add(new SubwordPiece(substring, ids.get(substring), + mapped.starts[start], mapped.ends[end - 1])); + start = end; + found = true; + break; + } + end--; + } + if (!found) { + break; + } + } + if (found) { + pieces.addAll(wordPieces); + } else { + pieces.add(new SubwordPiece(unknownToken, unknownId, wordStart, wordEnd)); + } + } + + /** + * The normalized text with, for every char, the original-text range it came from. Characters + * inserted by the pipeline (isolation spaces) carry an empty range at the insertion point. + */ + private static final class MappedText { + private char[] chars; + private int[] starts; + private int[] ends; + private int length; + + /** + * Instantiates an empty mapped text. + * + * @param capacity The initial capacity hint in chars. + */ + private MappedText(int capacity) { + chars = new char[capacity]; + starts = new int[capacity]; + ends = new int[capacity]; + } + + /** + * Appends one char with the original-text range it came from. + * + * @param c The char to append. + * @param originalStart The inclusive original-text start of the char. + * @param originalEnd The exclusive original-text end of the char. + */ + private void add(char c, int originalStart, int originalEnd) { + if (length == chars.length) { + final int capacity = Math.max(16, length * 2); + chars = Arrays.copyOf(chars, capacity); + starts = Arrays.copyOf(starts, capacity); + ends = Arrays.copyOf(ends, capacity); + } + chars[length] = c; + starts[length] = originalStart; + ends[length] = originalEnd; + length++; + } + + /** + * Appends every char of a string, all sharing one original-text range. + * + * @param s The string to append. + * @param originalStart The inclusive original-text start shared by all chars. + * @param originalEnd The exclusive original-text end shared by all chars. + */ + private void add(String s, int originalStart, int originalEnd) { + for (int i = 0; i < s.length(); i++) { + add(s.charAt(i), originalStart, originalEnd); + } + } + } + + /** + * Cleans the text (control and whitespace normalization) and isolates CJK code points in one + * pass, recording the original-text range of every output character. + * + * @param original The original input text. + * @return The cleaned, CJK-isolated text with per-character ranges. + */ + private static MappedText cleanAndIsolateCjk(String original) { + final MappedText out = new MappedText(original.length() + 16); + int i = 0; + while (i < original.length()) { + final int codePoint = original.codePointAt(i); + final int width = Character.charCount(codePoint); + if (codePoint == 0 || codePoint == 0xFFFD || BertNormalization.isControl(codePoint)) { + i += width; + continue; + } + if (BertNormalization.isWhitespace(codePoint) || isLineOrParagraphSeparator(codePoint)) { + out.add(' ', i, i + width); + } else if (BertNormalization.isCjk(codePoint)) { + out.add(' ', i, i); + for (int c = 0; c < width; c++) { + out.add(original.charAt(i + c), i, i + width); + } + out.add(' ', i + width, i + width); + } else { + for (int c = 0; c < width; c++) { + out.add(original.charAt(i + c), i, i + width); + } + } + i += width; + } + return out; + } + + /** + * A Unicode line or paragraph separator ({@code Zl}, {@code Zp}). These are not whitespace in + * the BERT {@code _is_whitespace} sense, but the reference pipeline's + * {@code whitespace_tokenize} (Python's {@code str.split()}) still breaks words on them, so + * they must become word boundaries here as well. + * + * @param codePoint The code point to classify. + * @return True if the code point is a line or paragraph separator. + */ + private static boolean isLineOrParagraphSeparator(int codePoint) { + final int type = Character.getType(codePoint); + return type == Character.LINE_SEPARATOR || type == Character.PARAGRAPH_SEPARATOR; + } + + /** + * Isolates punctuation, surrounding each punctuation code point with spaces, preserving the + * original-text range of every character. + * + * @param in The input text with per-character ranges. + * @return The punctuation-isolated text with per-character ranges. + */ + private static MappedText isolatePunctuation(MappedText in) { + final MappedText out = new MappedText(in.length + 16); + int i = 0; + while (i < in.length) { + final int codePoint = codePointAt(in, i); + final int width = Character.charCount(codePoint); + if (BertNormalization.isPunctuation(codePoint)) { + out.add(' ', in.starts[i], in.starts[i]); + for (int c = 0; c < width; c++) { + out.add(in.chars[i + c], in.starts[i + c], in.ends[i + c]); + } + out.add(' ', in.ends[i + width - 1], in.ends[i + width - 1]); + } else { + for (int c = 0; c < width; c++) { + out.add(in.chars[i + c], in.starts[i + c], in.ends[i + c]); + } + } + i += width; + } + return out; + } + + /** + * Lower cases and strips accents, preserving the original-text range of every character. When a + * contextual case mapping prevents a per-character range from being recovered, the whole + * whitespace run falls back to its full range, which widens spans but never misplaces them. + * + * @param in The input text with per-character ranges. + * @return The lower-cased, accent-stripped text with per-character ranges. + */ + private static MappedText lowerCaseAndStripAccents(MappedText in) { + final MappedText out = new MappedText(in.length + 16); + int from = 0; + while (from < in.length) { + if (in.chars[from] == ' ') { + out.add(' ', in.starts[from], in.ends[from]); + from++; + continue; + } + int to = from; + while (to < in.length && in.chars[to] != ' ') { + to++; + } + transformRun(in, from, to, out); + from = to; + } + return out; + } + + /** + * Lower cases and accent-strips one non-space run, emitting per-character ranges when the + * transformation is reproducible per code point and the run's full range otherwise. + * + * @param in The input text with per-character ranges. + * @param from The inclusive start of the run in {@code in}. + * @param to The exclusive end of the run in {@code in}. + * @param out The output text to append to. + */ + private static void transformRun(MappedText in, int from, int to, MappedText out) { + final String run = new String(in.chars, from, to - from); + // Locale.ROOT lower casing is the reference behavior of BERT's do_lower_case: the reference + // pipeline applies the full locale-independent Unicode case mappings (including one-to-many + // ones like the dotted capital I), which a per-code-point mapping cannot reproduce. + final String content = stripAccents(run.toLowerCase(Locale.ROOT)); + + // Rerun per code point to learn how many output chars each input code point produces. + final StringBuilder rerun = new StringBuilder(content.length()); + final int[] produced = new int[to - from]; + int i = from; + while (i < to) { + final int codePoint = codePointAt(in, i); + final int width = Character.charCount(codePoint); + final String transformed = stripAccents( + new String(Character.toChars(codePoint)).toLowerCase(Locale.ROOT)); + rerun.append(transformed); + produced[i - from] = transformed.length(); + i += width; + } + + if (rerun.toString().equals(content)) { + int at = from; + int emitted = 0; + while (at < to) { + final int width = Character.charCount(codePointAt(in, at)); + for (int c = 0; c < produced[at - from]; c++) { + out.add(content.charAt(emitted++), in.starts[at], in.ends[at + width - 1]); + } + at += width; + } + } else { + // Contextual case mapping changed the content; the run's chars share the run's range. + out.add(content, in.starts[from], in.ends[to - 1]); + } + } + + /** + * Removes combining marks after NFD decomposition, the accent stripping of BERT's + * {@code do_lower_case} mode. + * + * @param text The text to strip. + * @return The text without non-spacing marks. + */ + private static String stripAccents(String text) { + final String decomposed = Normalizer.normalize(text, Normalizer.Form.NFD); + final StringBuilder stripped = new StringBuilder(decomposed.length()); + decomposed.codePoints().forEach(codePoint -> { + if (Character.getType(codePoint) != Character.NON_SPACING_MARK) { + stripped.appendCodePoint(codePoint); + } + }); + return stripped.toString(); + } + + /** + * Reads the code point at an index, joining a surrogate pair when one starts there. + * + * @param text The text to read from. + * @param index The char index to read at. + * @return The code point at {@code index}. + */ + private static int codePointAt(MappedText text, int index) { + final char c = text.chars[index]; + if (Character.isHighSurrogate(c) && index + 1 < text.length + && Character.isLowSurrogate(text.chars[index + 1])) { + return Character.toCodePoint(c, text.chars[index + 1]); + } + return c; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java index 97e240e87c..fa1014ab2c 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java @@ -34,7 +34,7 @@ * of BERT tokenization. It does not normalize the input text: no lower casing, * no accent stripping, no control character removal. Text that does not match * the vocabulary's casing - for uncased models that includes every capitalized - * word - is mapped to the unknown token. Use {@link BertTokenizer} for the + * word - is mapped to the unknown token. Use {@link WordpieceEncoder} for the * full BERT tokenization pipeline. *

* As of OpenNLP 3.0.0 the behavior matches the reference BERT wordpiece @@ -58,7 +58,7 @@ * * * - * @see BertTokenizer + * @see WordpieceEncoder */ public class WordpieceTokenizer implements Tokenizer { diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java index a598199270..8f59c6f844 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java @@ -105,7 +105,7 @@ protected AbstractDL(final File model, final File vocabulary, final OrtSession createdSession = env.createSession(model.getPath(), sessionOptions); try { this.vocab = Map.copyOf(loadVocabFile(vocabulary)); - this.tokenizer = createBertTokenizer(vocab, lowerCase); + this.tokenizer = createPipelineTokenizer(vocab, lowerCase); } catch (IOException | RuntimeException e) { // Vocabulary/tokenizer init failed after the native session was created; close it // so a partially constructed instance never leaks the ONNX session. @@ -136,7 +136,7 @@ protected AbstractDL(final OrtEnvironment env, final OrtSession session, this.env = env; this.session = session; this.vocab = vocab; - this.tokenizer = createBertTokenizer(vocab, lowerCase); + this.tokenizer = createPipelineTokenizer(vocab, lowerCase); } /** @@ -247,16 +247,31 @@ static WordpieceTokenizer createWordpieceTokenizer( * @param lowerCase {@code true} for uncased models (lower casing and accent * stripping), {@code false} for cased models. * @return A configured {@link BertTokenizer}. - * @throws IllegalArgumentException Thrown if a RoBERTa-style vocabulary - * contains no supported unknown token. + * @throws IllegalArgumentException Thrown if the selected special tokens + * are not all present in the vocabulary. */ + // The deprecated BertTokenizer stays the return type until its removal in 3.1, so that + // already-compiled subclasses overriding this method keep overriding it. + @SuppressWarnings("removal") protected BertTokenizer createTokenizer( final Map vocab, final boolean lowerCase) { - return createBertTokenizer(vocab, lowerCase); + return createPipelineTokenizer(vocab, lowerCase); } - static BertTokenizer createBertTokenizer( + /** + * Builds the pipeline tokenizer, selecting the RoBERTa special tokens when the vocabulary + * carries them and the BERT defaults otherwise. + * + * @param vocab The vocabulary map. + * @param lowerCase {@code true} for uncased models, {@code false} for cased models. + * @return A configured {@link BertTokenizer}. + * @throws IllegalArgumentException Thrown if the selected special tokens are not all present in + * the vocabulary. + */ + // BertTokenizer is deprecated for removal in 3.1; built here until then. + @SuppressWarnings("removal") + static BertTokenizer createPipelineTokenizer( final Map vocab, final boolean lowerCase) { if (vocab.containsKey( WordpieceTokenizer.ROBERTA_CLS_TOKEN) diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java index e5fbe2c6a4..7b17cb3028 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java @@ -51,7 +51,7 @@ * using ONNX models. * *

Tokenization performs BERT basic tokenization (text normalization) - * before wordpiece, see {@link opennlp.tools.tokenize.BertTokenizer}. Input + * before wordpiece, see {@link opennlp.tools.tokenize.WordpieceEncoder}. Input * text is lower cased and accent stripped by default, matching the uncased * models commonly used for classification. For cased models, set * {@link InferenceOptions#setLowerCase(boolean)} to {@code false}.

diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java index 95bb81d879..956f177d61 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java @@ -50,7 +50,7 @@ * An implementation of {@link opennlp.tools.namefind.TokenNameFinder} that uses ONNX models. * *

Tokenization performs BERT basic tokenization (text normalization) - * before wordpiece, see {@link opennlp.tools.tokenize.BertTokenizer}. Input + * before wordpiece, see {@link opennlp.tools.tokenize.WordpieceEncoder}. Input * text is not lower cased by default, because named entity recognition * models are commonly cased: capitalization is a strong signal for entity * boundaries. For uncased models, set diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java index 6bc76ce183..2bab92e0ef 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/vectors/SentenceVectorsDL.java @@ -20,17 +20,23 @@ import java.io.File; import java.io.IOException; import java.nio.LongBuffer; +import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; +import java.util.Iterator; +import java.util.List; import java.util.Map; +import ai.onnxruntime.NodeInfo; import ai.onnxruntime.OnnxTensor; import ai.onnxruntime.OrtException; import ai.onnxruntime.OrtSession; +import ai.onnxruntime.TensorInfo; import opennlp.dl.AbstractDL; import opennlp.dl.Tokens; import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.embeddings.TextEmbedder; import opennlp.tools.tokenize.Tokenizer; @@ -47,7 +53,7 @@ * so the encoder attended to nothing and the output vectors were * incorrect. Additionally, tokenization now performs BERT basic * tokenization (lower casing and accent stripping by default, see - * {@link opennlp.tools.tokenize.BertTokenizer}) before wordpiece. + * {@link opennlp.tools.tokenize.WordpieceEncoder}) before wordpiece. * Output vectors change with the corrected encoding and tokenization; * any embeddings persisted from the previous behavior are not * comparable with the corrected output and must be re-embedded.

@@ -56,9 +62,18 @@ * holds no per-call instance state and the underlying {@link OrtSession} supports * concurrent execution. This thread-safety guarantee applies until {@link #close()} * is called; callers must not race {@code close()} with inference methods.

+ * + *

{@link #getVectors(String)} is the primary entry point; {@link #embed(CharSequence)} + * adapts it to the {@link TextEmbedder} contract. {@link #embedAll(List)} runs one batched + * session per distinct tokenized length, so a batch of same-length inputs costs one + * inference instead of one per input.

*/ @ThreadSafe -public class SentenceVectorsDL extends AbstractDL { +public class SentenceVectorsDL extends AbstractDL implements TextEmbedder { + + // The hidden dimension declared by the model's output metadata, or a value <= 0 when the + // model declares it dynamically; dimension() then probes once and caches here. + private volatile int dimension; /** * Instantiates a {@link SentenceVectorsDL sentence vector generator} for an @@ -94,6 +109,7 @@ public SentenceVectorsDL(final File model, final File vocabulary, final boolean throws OrtException, IOException { super(model, vocabulary, new OrtSession.SessionOptions(), lowerCase); + this.dimension = declaredOutputDimension(session); } @@ -132,6 +148,157 @@ public float[] getVectors(final String sentence) throws OrtException { } + /** + * {@inheritDoc} + * + *

Adapts {@link #getVectors(String)} to the {@link TextEmbedder} contract. Empty or + * unrecognized input is still run through the model, which returns the vector for the + * wrapped {@code [CLS] ... [SEP]} sequence rather than a zero vector.

+ * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + * @throws IllegalStateException Thrown if inference fails; the cause carries the + * underlying {@link OrtException}. + */ + @Override + public float[] embed(final CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + try { + return getVectors(text instanceof String s ? s : text.toString()); + } catch (OrtException e) { + throw new IllegalStateException("Sentence vector inference failed.", e); + } + } + + /** + * {@inheritDoc} + * + *

Batched execution: the inputs are tokenized up front, grouped by tokenized length, + * and each group runs through the session once with shape {@code [group size, length]}. + * Grouping by length means a batch never pads, so every row is computed from exactly the + * tensors its single-input call would have used. A length group of one executes the + * same {@code [1, length]} shapes as {@link #getVectors(String)}.

+ * + * @throws IllegalArgumentException Thrown if {@code texts} is {@code null} or contains + * {@code null}. + * @throws IllegalStateException Thrown if inference fails; the cause carries the + * underlying {@link OrtException}. + */ + @Override + public float[][] embedAll(final List texts) { + if (texts == null) { + throw new IllegalArgumentException("Texts must not be null"); + } + final float[][] vectors = new float[texts.size()][]; + if (texts.isEmpty()) { + return vectors; + } + final Tokens[] encoded = new Tokens[texts.size()]; + final Map> byLength = new HashMap<>(); + for (int i = 0; i < texts.size(); i++) { + final CharSequence text = texts.get(i); + if (text == null) { + throw new IllegalArgumentException("Texts must not contain null"); + } + encoded[i] = tokenize(text instanceof String s ? s : text.toString(), tokenizer, vocab); + byLength.computeIfAbsent(encoded[i].ids().length, length -> new ArrayList<>()).add(i); + } + try { + for (final List group : byLength.values()) { + runBatch(encoded, group, vectors); + } + } catch (OrtException e) { + throw new IllegalStateException("Sentence vector inference failed.", e); + } + return vectors; + } + + /** + * Runs one inference over a group of same-length encodings and stores each row's + * {@code [CLS]}-position vector under its original input index. + * + * @param encoded The tokenized inputs, indexed by input position. + * @param group The input positions sharing one tokenized length, in input order. + * @param vectors The output array to fill, indexed by input position. + * @throws OrtException Thrown if an error occurs during inference. + */ + private void runBatch(final Tokens[] encoded, final List group, + final float[][] vectors) throws OrtException { + + final int batch = group.size(); + final int length = encoded[group.get(0)].ids().length; + final long[] ids = new long[batch * length]; + final long[] mask = new long[batch * length]; + final long[] types = new long[batch * length]; + for (int b = 0; b < batch; b++) { + final Tokens tokens = encoded[group.get(b)]; + System.arraycopy(tokens.ids(), 0, ids, b * length, length); + System.arraycopy(tokens.mask(), 0, mask, b * length, length); + System.arraycopy(tokens.types(), 0, types, b * length, length); + } + + final Map inputs = new HashMap<>(); + final long[] shape = {batch, length}; + + try { + inputs.put(INPUT_IDS, OnnxTensor.createTensor(env, LongBuffer.wrap(ids), shape)); + + inputs.put(ATTENTION_MASK, OnnxTensor.createTensor(env, LongBuffer.wrap(mask), shape)); + + inputs.put(TOKEN_TYPE_IDS, OnnxTensor.createTensor(env, LongBuffer.wrap(types), shape)); + + try (OrtSession.Result result = session.run(inputs)) { + // getValue() copies the tensor into Java arrays, so the result can be closed safely. + final float[][][] v = (float[][][]) result.get(0).getValue(); + for (int b = 0; b < batch; b++) { + vectors[group.get(b)] = v[b][0]; + } + } + } finally { + inputs.values().forEach(OnnxTensor::close); + } + + } + + /** + * {@inheritDoc} + * + *

Read from the model's declared output metadata when it is static; a model that declares + * the hidden dimension dynamically is probed with one inference on the first call and the + * result cached.

+ */ + @Override + public int dimension() { + final int declared = dimension; + if (declared > 0) { + return declared; + } + synchronized (this) { + if (dimension <= 0) { + dimension = embed("a").length; + } + return dimension; + } + } + + /** + * {@return the last dimension of the first output's declared shape, or {@code -1} when the + * model declares it dynamically} + * + * @param session The model's ONNX session. + * @throws OrtException Thrown if reading the output metadata fails. + */ + private static int declaredOutputDimension(final OrtSession session) throws OrtException { + final Iterator outputs = session.getOutputInfo().values().iterator(); + if (!outputs.hasNext() || !(outputs.next().getInfo() instanceof TensorInfo tensorInfo)) { + return -1; + } + final long[] shape = tensorInfo.getShape(); + final long last = shape.length > 0 ? shape[shape.length - 1] : -1; + return last > 0 && last <= Integer.MAX_VALUE ? (int) last : -1; + } + /** * Encodes text as model inputs: wordpiece token ids, an attention mask of ones, * and single-segment (all zero) token type ids. diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/CreateTokenizerTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/CreateTokenizerTest.java index 54c4600a8e..5ae7e3ebee 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/CreateTokenizerTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/CreateTokenizerTest.java @@ -22,7 +22,7 @@ import org.junit.jupiter.api.Test; -import opennlp.tools.tokenize.BertTokenizer; +import opennlp.tools.tokenize.Tokenizer; import opennlp.tools.tokenize.WordpieceTokenizer; import static org.junit.jupiter.api.Assertions.assertArrayEquals; @@ -52,8 +52,8 @@ private static Map robertaVocab() { } @Test - void testCreatesLowerCasingBertTokenizer() { - final BertTokenizer tokenizer = AbstractDL.createBertTokenizer(bertVocab(), true); + void testCreatesLowerCasingPipelineTokenizer() { + final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), true); // Capitalized input must be lower cased before the wordpiece lookup. assertArrayEquals(new String[] { @@ -62,8 +62,8 @@ void testCreatesLowerCasingBertTokenizer() { } @Test - void testCreatesCasePreservingBertTokenizer() { - final BertTokenizer tokenizer = AbstractDL.createBertTokenizer(bertVocab(), false); + void testCreatesCasePreservingPipelineTokenizer() { + final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), false); // Without lower casing, capitalized words miss the lowercase-only vocabulary. assertArrayEquals(new String[] { @@ -74,7 +74,7 @@ void testCreatesCasePreservingBertTokenizer() { @Test void testSelectsRobertaSpecialTokens() { - final BertTokenizer tokenizer = AbstractDL.createBertTokenizer(robertaVocab(), false); + final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(robertaVocab(), false); assertArrayEquals(new String[] { WordpieceTokenizer.ROBERTA_CLS_TOKEN, "hello", WordpieceTokenizer.ROBERTA_UNK_TOKEN, @@ -88,7 +88,7 @@ void testFallsBackToBertUnknownToken() { vocab.remove(WordpieceTokenizer.ROBERTA_UNK_TOKEN); vocab.put(WordpieceTokenizer.BERT_UNK_TOKEN, 2); - final BertTokenizer tokenizer = AbstractDL.createBertTokenizer(vocab, false); + final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(vocab, false); assertArrayEquals(new String[] { WordpieceTokenizer.ROBERTA_CLS_TOKEN, "hello", WordpieceTokenizer.BERT_UNK_TOKEN, @@ -101,10 +101,26 @@ void testRejectsRobertaVocabularyWithoutUnknownToken() { final Map vocab = robertaVocab(); vocab.remove(WordpieceTokenizer.ROBERTA_UNK_TOKEN); - assertThrows(IllegalArgumentException.class, () -> AbstractDL.createBertTokenizer(vocab, false)); + assertThrows(IllegalArgumentException.class, + () -> AbstractDL.createPipelineTokenizer(vocab, false)); assertThrows(IllegalArgumentException.class, () -> AbstractDL.createWordpieceTokenizer(vocab)); } + @Test + void testTokenizePosIsUnsupported() { + final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), true); + assertThrows(UnsupportedOperationException.class, () -> tokenizer.tokenizePos("the fox")); + } + + @Test + void testRejectsBertVocabularyMissingSpecialTokensAtCreation() { + final Map vocab = bertVocab(); + vocab.remove(WordpieceTokenizer.BERT_UNK_TOKEN); + + assertThrows(IllegalArgumentException.class, + () -> AbstractDL.createPipelineTokenizer(vocab, true)); + } + @Test void testResolveLowerCaseUsesComponentDefaultWhenUnset() { final InferenceOptions options = new InferenceOptions(); diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java new file mode 100644 index 0000000000..306f38595f --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java @@ -0,0 +1,131 @@ +/* + * 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.dl.vectors; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.embeddings.TextEmbedder; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * The {@link TextEmbedder} adapter driven through a real ONNX session. The bundled + * {@code tiny-vectors.onnx} (see {@code gen_tiny_vectors_model.py} next to it) computes + * {@code output[b][t] = float(input_ids[b][t]) * [0.5, -1, 2]}, so every expected vector is + * hand-computable from the vocabulary ids: {@code getVectors} returns the vector at the + * {@code [CLS]} position, and {@code [CLS]} sits at line 7 of the test vocabulary. + */ +class SentenceVectorsDLEmbedderTest { + + // 7 * [0.5, -1, 2] + private static final float[] CLS_VECTOR = {3.5f, -7f, 14f}; + + // Copy the model out of the classpath rather than resolving it in place: when this test runs + // from the opennlp-dl test-jar (as it does in opennlp-dl-gpu) the resource URI is inside a jar + // and is not hierarchical, so new File(uri) would fail. + private static File model(Path dir) throws IOException { + final Path file = dir.resolve("tiny-vectors.onnx"); + try (InputStream is = Objects.requireNonNull(SentenceVectorsDLEmbedderTest.class + .getResourceAsStream("/opennlp/dl/vectors/tiny-vectors.onnx"))) { + Files.copy(is, file, StandardCopyOption.REPLACE_EXISTING); + } + return file.toFile(); + } + + private static File vocab(Path dir) throws IOException { + final Path file = dir.resolve("vocab.txt"); + // Line number = id: [UNK]=2, [SEP]=3, hello=4, world=5, [CLS]=7. + Files.write(file, List.of("[PAD]", "unused1", "[UNK]", "[SEP]", "hello", "world", + "unused2", "[CLS]")); + return file.toFile(); + } + + @Test + void testEmbedderContractOverARealSession(@TempDir Path dir) throws Exception { + try (SentenceVectorsDL vectors = new SentenceVectorsDL(model(dir), vocab(dir))) { + + // The primary entry point, against which the adapter below is compared. + assertArrayEquals(CLS_VECTOR, vectors.getVectors("hello world"), 1e-5f); + + final TextEmbedder embedder = vectors; + + // The dimension comes from the model's declared output metadata, no inference needed. + assertEquals(3, embedder.dimension()); + + // The seam produces the same vector as the original entry point, for String and + // non-String inputs alike. + assertArrayEquals(CLS_VECTOR, embedder.embed("hello world"), 1e-5f); + assertArrayEquals(CLS_VECTOR, embedder.embed(new StringBuilder("hello world")), 1e-5f); + + // The batch method returns one vector per input, in input order; + // this model's [CLS]-position output is input-independent by construction. + final float[][] batch = embedder.embedAll(List.of("hello world", "hello")); + assertEquals(2, batch.length); + assertArrayEquals(CLS_VECTOR, batch[0], 1e-5f); + assertArrayEquals(CLS_VECTOR, batch[1], 1e-5f); + + assertThrows(IllegalArgumentException.class, () -> embedder.embed(null)); + assertThrows(IllegalArgumentException.class, () -> embedder.embedAll(null)); + } + } + + /** + * Drives the batched path over inputs of mixed tokenized lengths ("hello" encodes one + * token shorter than "hello world") and asserts every row reproduces its single-input + * vector exactly: the length-grouped batch never pads, so the computation per row is + * the computation the single call performs. + */ + @Test + void testEmbedAllMatchesSingleEmbedsExactly(@TempDir Path dir) throws Exception { + try (SentenceVectorsDL vectors = new SentenceVectorsDL(model(dir), vocab(dir))) { + final List texts = List.of("hello", "hello world", "world", "hello world", + "hello"); + final float[][] batch = vectors.embedAll(texts); + assertEquals(texts.size(), batch.length); + for (int i = 0; i < texts.size(); i++) { + assertArrayEquals(vectors.embed(texts.get(i)), batch[i]); + } + } + } + + /** + * Asserts the batch contract edges: an empty input yields an empty batch, and a + * {@code null} element is rejected rather than failing later inside the session. + */ + @Test + void testEmbedAllEdges(@TempDir Path dir) throws Exception { + try (SentenceVectorsDL vectors = new SentenceVectorsDL(model(dir), vocab(dir))) { + assertEquals(0, vectors.embedAll(List.of()).length); + assertThrows(IllegalArgumentException.class, + () -> vectors.embedAll(Arrays.asList("hello", null))); + } + } +} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py new file mode 100644 index 0000000000..05a67ccd5f --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py @@ -0,0 +1,57 @@ +# 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. + +# Generates tiny-vectors.onnx, the deterministic model behind SentenceVectorsDLEmbedderTest. +# +# The graph computes output[b][t][d] = float(input_ids[b][t]) * W[0][d] with +# W = [[0.5, -1.0, 2.0]], so the vector at any token position is that token's vocabulary id +# times W, hand-computable in the test. It declares the same three inputs a BERT-style +# encoder declares (input_ids, attention_mask, token_type_ids; the latter two are accepted +# and ignored) and one output of shape [batch, tokens, 3] so the hidden dimension is static +# in the model metadata. +# +# Regenerate with: python3 gen_tiny_vectors_model.py (requires the onnx package) + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +W = np.array([[0.5, -1.0, 2.0]], dtype=np.float32) + +cast = helper.make_node("Cast", ["input_ids"], ["ids_float"], to=TensorProto.FLOAT) +unsqueeze = helper.make_node("Unsqueeze", ["ids_float", "axes"], ["ids_3d"]) +matmul = helper.make_node("MatMul", ["ids_3d", "w"], ["last_hidden_state"]) + + +def encoder_input(name): + return helper.make_tensor_value_info(name, TensorProto.INT64, ["batch", "tokens"]) + + +graph = helper.make_graph( + [cast, unsqueeze, matmul], + "tiny-vectors", + [encoder_input("input_ids"), encoder_input("attention_mask"), + encoder_input("token_type_ids")], + [helper.make_tensor_value_info( + "last_hidden_state", TensorProto.FLOAT, ["batch", "tokens", 3])], + [numpy_helper.from_array(np.array([2], dtype=np.int64), name="axes"), + numpy_helper.from_array(W, name="w")], +) + +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) +model.ir_version = 8 +onnx.checker.check_model(model) +onnx.save(model, "tiny-vectors.onnx") +print("wrote tiny-vectors.onnx,", len(model.SerializeToString()), "bytes") diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx new file mode 100644 index 0000000000..7d63c91322 Binary files /dev/null and b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx differ diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java index d8f706f4ef..070834316b 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java @@ -14,149 +14,89 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package opennlp.tools.tokenize; +import java.util.HashSet; +import java.util.List; import java.util.Set; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +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; +import static org.junit.jupiter.api.Assertions.assertThrows; /** - * Tests {@link BertTokenizer}. - *

- * All expected token sequences in this test were generated with the HuggingFace - * {@code tokenizers} reference implementation ({@code BertWordPieceTokenizer}) - * using the same vocabulary, so they are verified to be identical to the - * reference BERT tokenization. + * Pins the deprecated {@link BertTokenizer} shim: constructor validation, delegation to + * {@link WordpieceEncoder#encodeToPieces(CharSequence)}, the default special tokens, and + * the unsupported {@link BertTokenizer#tokenizePos(String)}. */ -public class BertTokenizerTest { - - private static final Set VOCABULARY = Set.of( - "the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", - "em", "##bed", "##ding", "##s", - "wurttemberg", "strasse", "grosse", - "don", "t", "wait", "what", ".", ",", "?", "!", "'", - "\u6211", "\u7231", // CJK: 我 爱 - "natural", "language", "processing"); - - @Test - void testLowerCasesCapitalizedWords() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - final String[] tokens = tokenizer.tokenize("The quick brown fox jumps over the lazy dog."); - - final String[] expected = {"[CLS]", "the", "quick", "brown", "fox", "jumps", "over", - "the", "lazy", "dog", ".", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testLowerCasesBeforeWordpieceSplitting() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - final String[] tokens = tokenizer.tokenize("Embeddings"); - - final String[] expected = {"[CLS]", "em", "##bed", "##ding", "##s", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testStripsAccentsButKeepsNonCombiningCharacters() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - // ü decomposes to u + combining diaeresis and the mark is stripped; - // ß is not a combining mark and must survive, leaving an OOV token. - final String[] tokens = tokenizer.tokenize("W\u00fcrttemberg Stra\u00dfe"); - - final String[] expected = {"[CLS]", "wurttemberg", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } +@SuppressWarnings("removal") // Exercises BertTokenizer deliberately until its removal in 3.1. +class BertTokenizerTest { - @Test - void testSplitsPunctuationRunsIntoSingleCharacters() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - final String[] tokens = tokenizer.tokenize("Wait... what?!"); + private static final List VOCABULARY = List.of( + "[PAD]", "[UNK]", "[CLS]", "[SEP]", "hello", "world", "##s", "ca", "##fe", ",", "!"); - final String[] expected = {"[CLS]", "wait", ".", ".", ".", "what", "?", "!", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + private static Set vocabularySet() { + return new HashSet<>(VOCABULARY); } - @Test - void testSplitsApostrophesAsPunctuation() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - final String[] tokens = tokenizer.tokenize("don't"); - - final String[] expected = {"[CLS]", "don", "'", "t", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + @ParameterizedTest + @ValueSource(strings = {"Hello, world!", "Caf\u00E9 hellos", ""}) + void testTokenizeReturnsTheEncoderPieceSequence(String input) { + final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); + final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); + assertArrayEquals(encoder.encodeToPieces(input), tokenizer.tokenize(input), + "delegation broke on: " + input); } @Test - void testIsolatesCjkIdeographs() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - final String[] tokens = tokenizer.tokenize("\u6211\u7231natural language processing"); - - final String[] expected = {"[CLS]", "\u6211", "\u7231", "natural", "language", - "processing", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + void testDefaultConstructorsChainToBertSpecialTokensAndLowerCasing() { + final String[] uncasedDefault = new BertTokenizer(vocabularySet()).tokenize("Hello worldS"); + assertArrayEquals(new String[] {"[CLS]", "hello", "world", "##s", "[SEP]"}, uncasedDefault); + assertArrayEquals(uncasedDefault, + new BertTokenizer(vocabularySet(), true).tokenize("Hello worldS")); + assertArrayEquals(uncasedDefault, + new BertTokenizer(vocabularySet(), true, WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN) + .tokenize("Hello worldS")); } @Test - void testCleansControlCharactersAndNormalizesWhitespace() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - // Tab and no-break space are whitespace; the NUL character is removed, - // joining "brown" and "fox" into one out-of-vocabulary token. - final String[] tokens = tokenizer.tokenize("the\tquick\u00a0brown\u0000fox"); - - final String[] expected = {"[CLS]", "the", "quick", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + void testCasedTokenizerKeepsCase() { + // Without lower casing, the capitalized word misses the lowercase-only vocabulary. + assertArrayEquals(new String[] {"[CLS]", "[UNK]", "[SEP]"}, + new BertTokenizer(vocabularySet(), false).tokenize("Hello")); } @Test - void testRemovesPrivateUseAndUnassignedCharacters() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - // The reference implementation treats all C* categories as control - // characters: private use (U+E000, Co) and noncharacters (U+FDD0, Cn) - // are removed, joining the surrounding text into one OOV token. - final String[] tokens = tokenizer.tokenize("fox\ue000jumps and fox\ufdd0jumps"); - - final String[] expected = {"[CLS]", "[UNK]", "[UNK]", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + void testConstructorsRejectNullArguments() { + assertThrows(IllegalArgumentException.class, () -> new BertTokenizer(null)); + assertThrows(IllegalArgumentException.class, () -> new BertTokenizer(null, true)); + assertThrows(IllegalArgumentException.class, + () -> new BertTokenizer(null, true, "[CLS]", "[SEP]", "[UNK]")); + assertThrows(IllegalArgumentException.class, + () -> new BertTokenizer(vocabularySet(), true, null, "[SEP]", "[UNK]")); + assertThrows(IllegalArgumentException.class, + () -> new BertTokenizer(vocabularySet(), true, "[CLS]", null, "[UNK]")); + assertThrows(IllegalArgumentException.class, + () -> new BertTokenizer(vocabularySet(), true, "[CLS]", "[SEP]", null)); } @Test - void testRejectsNullSpecialTokens() { - Assertions.assertThrows(NullPointerException.class, - () -> new BertTokenizer(VOCABULARY, true, null, "[SEP]", "[UNK]")); - Assertions.assertThrows(NullPointerException.class, - () -> new BertTokenizer(VOCABULARY, true, "[CLS]", null, "[UNK]")); - Assertions.assertThrows(NullPointerException.class, - () -> new BertTokenizer(VOCABULARY, true, "[CLS]", "[SEP]", null)); + void testTokenizeRejectsNullText() { + final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); + assertThrows(IllegalArgumentException.class, () -> tokenizer.tokenize(null)); } @Test - void testCasedModeKeepsCaseAndAccents() { - final Tokenizer tokenizer = new BertTokenizer( - Set.of("The", "W\u00fcrttemberg", "fox"), false); - final String[] tokens = tokenizer.tokenize("The W\u00fcrttemberg fox"); - - final String[] expected = {"[CLS]", "The", "W\u00fcrttemberg", "fox", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + void testTokenizePosIsUnsupportedWithTheDocumentedMessage() { + final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); + final UnsupportedOperationException e = assertThrows(UnsupportedOperationException.class, + () -> tokenizer.tokenizePos("hello world")); + assertEquals("Wordpiece tokens cannot be mapped to character spans of the original text", + e.getMessage()); } - - @Test - void testCustomSpecialTokens() { - final Tokenizer tokenizer = new BertTokenizer(Set.of("the", "fox"), true, - WordpieceTokenizer.ROBERTA_CLS_TOKEN, WordpieceTokenizer.ROBERTA_SEP_TOKEN, - WordpieceTokenizer.ROBERTA_UNK_TOKEN); - final String[] tokens = tokenizer.tokenize("The unknown fox"); - - final String[] expected = {"", "the", "", "fox", ""}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testTokenizePosIsUnsupported() { - final Tokenizer tokenizer = new BertTokenizer(VOCABULARY); - Assertions.assertThrows(UnsupportedOperationException.class, - () -> tokenizer.tokenizePos("the fox")); - } - } diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java new file mode 100644 index 0000000000..bbaeadaca3 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java @@ -0,0 +1,129 @@ +/* + * 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.tools.tokenize; + +import java.util.List; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Reference token-sequence expectations for {@link WordpieceEncoder}, covering lower casing, + * accent stripping, punctuation and CJK isolation, and text cleaning. + *

+ * All expected token sequences in this test were generated with the HuggingFace + * {@code tokenizers} reference implementation ({@code BertWordPieceTokenizer}) + * using the same vocabulary, so they are verified to be identical to the + * reference BERT tokenization. The encoder requires its special tokens to be + * present in the vocabulary (every piece must have an id), so the vocabularies + * here include them; the token sequences are unchanged. + */ +class WordpieceEncoderReferenceSequencesTest { + + private static final List VOCABULARY = List.of( + "[CLS]", "[SEP]", "[UNK]", + "the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", + "em", "##bed", "##ding", "##s", + "wurttemberg", "strasse", "grosse", + "don", "t", "wait", "what", ".", ",", "?", "!", "'", + "\u6211", "\u7231", // CJK + "natural", "language", "processing"); + + /** + * The reference input and expected-sequence pairs, one argument set per pipeline behavior. + * + * @return The (input, expected pieces) pairs. + */ + static Stream referenceSequences() { + return Stream.of( + // Lower cases capitalized words. + Arguments.of("The quick brown fox jumps over the lazy dog.", + new String[] {"[CLS]", "the", "quick", "brown", "fox", "jumps", "over", + "the", "lazy", "dog", ".", "[SEP]"}), + // Lower cases before wordpiece splitting. + Arguments.of("Embeddings", + new String[] {"[CLS]", "em", "##bed", "##ding", "##s", "[SEP]"}), + // The u-umlaut decomposes to u plus a combining diaeresis and the mark is stripped; + // the sharp s is not a combining mark and must survive, leaving an OOV token. + Arguments.of("W\u00fcrttemberg Stra\u00dfe", + new String[] {"[CLS]", "wurttemberg", "[UNK]", "[SEP]"}), + // Splits punctuation runs into single characters. + Arguments.of("Wait... what?!", + new String[] {"[CLS]", "wait", ".", ".", ".", "what", "?", "!", "[SEP]"}), + // Splits apostrophes as punctuation. + Arguments.of("don't", + new String[] {"[CLS]", "don", "'", "t", "[SEP]"}), + // Isolates CJK ideographs into single-character pieces. + Arguments.of("\u6211\u7231natural language processing", + new String[] {"[CLS]", "\u6211", "\u7231", "natural", "language", + "processing", "[SEP]"}), + // Tab and no-break space are whitespace; the NUL character is removed, + // joining "brown" and "fox" into one out-of-vocabulary token. + Arguments.of("the\tquick\u00a0brown\u0000fox", + new String[] {"[CLS]", "the", "quick", "[UNK]", "[SEP]"}), + // The reference implementation treats all C* categories as control + // characters: private use (U+E000, Co) and noncharacters (U+FDD0, Cn) + // are removed, joining the surrounding text into one OOV token. + Arguments.of("fox\ue000jumps and fox\ufdd0jumps", + new String[] {"[CLS]", "[UNK]", "[UNK]", "[UNK]", "[SEP]"})); + } + + @ParameterizedTest + @MethodSource("referenceSequences") + void testEncodesTheReferenceSequence(String input, String[] expected) { + final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); + Assertions.assertArrayEquals(expected, encoder.encodeToPieces(input), + "sequence broke on: " + input); + } + + @Test + void testRejectsNullSpecialTokens() { + // The encoder's contract throws IllegalArgumentException for null special tokens. + Assertions.assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(VOCABULARY, true, null, "[SEP]", "[UNK]")); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(VOCABULARY, true, "[CLS]", null, "[UNK]")); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(VOCABULARY, true, "[CLS]", "[SEP]", null)); + } + + @Test + void testCasedModeKeepsCaseAndAccents() { + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[CLS]", "[SEP]", "[UNK]", "The", "W\u00fcrttemberg", "fox"), false); + final String[] tokens = encoder.encodeToPieces("The W\u00fcrttemberg fox"); + + final String[] expected = {"[CLS]", "The", "W\u00fcrttemberg", "fox", "[SEP]"}; + Assertions.assertArrayEquals(expected, tokens); + } + + @Test + void testCustomSpecialTokens() { + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("", "", "", "the", "fox"), true, + WordpieceTokenizer.ROBERTA_CLS_TOKEN, WordpieceTokenizer.ROBERTA_SEP_TOKEN, + WordpieceTokenizer.ROBERTA_UNK_TOKEN); + final String[] tokens = encoder.encodeToPieces("The unknown fox"); + + final String[] expected = {"", "the", "", "fox", ""}; + Assertions.assertArrayEquals(expected, tokens); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java new file mode 100644 index 0000000000..997036b092 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java @@ -0,0 +1,243 @@ +/* + * 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.tools.tokenize; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +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 encoder held against the deprecated {@link BertTokenizer} for piece-sequence parity (the + * shim's contract is "the encoder's piece sequence, nothing more"), plus exact hand-computed span + * assertions through every normalization step that changes, inserts, or removes characters. + */ +class WordpieceEncoderTest { + + // Ids are indices: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, hello=4, world=5, ##s=6, won=7, + // ##der=8, ##ful=9, ca=10, ##fe=11, istanbul=12, U+4E2D=13, U+56FD=14, .=15, ,=16, !=17, + // he=18, ##llo=19, Greek "sofos" with a final sigma=20. + private static final List VOCAB = List.of( + "[PAD]", "[UNK]", "[CLS]", "[SEP]", "hello", "world", "##s", "won", "##der", "##ful", + "ca", "##fe", "istanbul", "\u4E2D", "\u56FD", ".", ",", "!", "he", "##llo", + "\u03C3\u03BF\u03C6\u03BF\u03C2"); + + private static WordpieceEncoder uncased() { + return new WordpieceEncoder(VOCAB); + } + + private static void assertPiece(SubwordPiece piece, String expectedPiece, int expectedId, + int expectedStart, int expectedEnd) { + assertEquals(expectedPiece, piece.piece()); + assertEquals(expectedId, piece.id()); + assertEquals(expectedStart, piece.start(), "start of " + piece); + assertEquals(expectedEnd, piece.end(), "end of " + piece); + assertEquals(new Span(expectedStart, expectedEnd), piece.span(), "span of " + piece); + } + + /** + * The curated parity inputs, each exercising a normalization step of the pipeline. + * + * @return The inputs. + */ + static Stream curatedInputs() { + return Stream.of( + "", + " ", + "Hello, WORLD!", + "Wonderful", + "hellos", + // An accented e, stripped by NFD decomposition. + "Caf\u00E9", + // The Turkish dotted capital I: lower cases to two chars, then the dot strips away. + "\u0130stanbul", + // CJK ideographs are isolated into single-character tokens. + "\u4E2D\u56FD is CJK", + // Greek upper case: the trailing sigma takes the contextual final-sigma mapping. + "\u03A3\u039F\u03A6\u039F\u03A3", + // The NBSP is whitespace in the BERT sense. + "hello\u00A0world", + // NUL and the zero-width space are removed by the cleaning stage. + "a\u0000b\u200Bc", + // An emoji: unknown to the vocabulary, and a surrogate pair. + "\uD83D\uDE00", + "!!!", + "a".repeat(101), + "he said: \u00ABhello\u00BB."); + } + + @ParameterizedTest + @MethodSource("curatedInputs") + @SuppressWarnings("removal") // BertTokenizer is pinned to the encoder until its removal in 3.1. + void testPieceSequenceMatchesBertTokenizerOnCuratedInputs(String input) { + final BertTokenizer bertTokenizer = new BertTokenizer(new HashSet<>(VOCAB), true); + final WordpieceEncoder encoder = uncased(); + assertArrayEquals(bertTokenizer.tokenize(input), encoder.encodeToPieces(input), + "parity broke on: " + input); + } + + @Test + @SuppressWarnings("removal") // BertTokenizer is pinned to the encoder until its removal in 3.1. + void testPieceSequenceMatchesBertTokenizerOnRandomInputs() { + final int[] pool = {'a', 'b', 'A', 'B', 'z', ' ', ' ', '\t', 0x00A0, '.', '!', ',', + 0x0301, 0x00E9, 0x0130, 0x03A3, 0x03C3, 0x03BF, 0x4E2D, 0xFFFD, 0x200B, 0x1F600, 0}; + final Random random = new Random(42); + for (final boolean lowerCase : new boolean[] {true, false}) { + final BertTokenizer bertTokenizer = new BertTokenizer(new HashSet<>(VOCAB), lowerCase); + final WordpieceEncoder encoder = new WordpieceEncoder(VOCAB, lowerCase); + for (int round = 0; round < 400; round++) { + final StringBuilder text = new StringBuilder(); + final int length = random.nextInt(25); + for (int i = 0; i < length; i++) { + text.appendCodePoint(pool[random.nextInt(pool.length)]); + } + final String input = text.toString(); + assertArrayEquals(bertTokenizer.tokenize(input), encoder.encodeToPieces(input), + "parity broke on: " + input); + + // Span invariants: within bounds and never moving backwards. + int previousStart = 0; + for (final SubwordPiece piece : encoder.encode(input)) { + assertTrue(piece.start() >= previousStart && piece.end() <= input.length(), + "span out of order or bounds in " + input + ": " + piece); + previousStart = piece.start(); + } + } + } + } + + @Test + void testSpansSurvivePunctuationIsolationAndCaseFolding() { + final List pieces = uncased().encode("Hello, WORLD!"); + assertEquals(6, pieces.size()); + assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); + assertPiece(pieces.get(1), "hello", 4, 0, 5); + assertPiece(pieces.get(2), ",", 16, 5, 6); + assertPiece(pieces.get(3), "world", 5, 7, 12); + assertPiece(pieces.get(4), "!", 17, 12, 13); + assertPiece(pieces.get(5), "[SEP]", 3, 13, 13); + } + + @Test + void testSpansSurviveAccentStripping() { + // The accent is stripped by NFD, yet ##fe still covers the accented surface. + final List pieces = uncased().encode("Caf\u00E9"); + assertEquals(4, pieces.size()); + assertPiece(pieces.get(1), "ca", 10, 0, 2); + assertPiece(pieces.get(2), "##fe", 11, 2, 4); + } + + @Test + void testSpansSurviveLengthChangingLowerCasing() { + // The Turkish dotted capital I lower cases to two chars before the combining dot strips; + // the piece still covers the original eight chars. + final List pieces = uncased().encode("\u0130stanbul"); + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), "istanbul", 12, 0, 8); + } + + @Test + void testCjkIsolationYieldsOnePieceAndSpanPerIdeograph() { + final List pieces = uncased().encode("\u4E2D\u56FD"); + assertEquals(4, pieces.size()); + assertPiece(pieces.get(1), "\u4E2D", 13, 0, 1); + assertPiece(pieces.get(2), "\u56FD", 14, 1, 2); + } + + @Test + void testLineAndParagraphSeparatorsSplitWords() { + // Zl and Zp are not whitespace in the BERT _is_whitespace sense, but the reference + // pipeline's whitespace_tokenize (Python's str.split()) breaks words on them, as did the + // previous OpenNLP pipeline via WhitespaceTokenizer. + final List pieces = uncased().encode("hello\u2028world\u2029hello"); + assertEquals(5, pieces.size()); + assertPiece(pieces.get(1), "hello", 4, 0, 5); + assertPiece(pieces.get(2), "world", 5, 6, 11); + assertPiece(pieces.get(3), "hello", 4, 12, 17); + } + + @Test + void testUnknownWordCoversItsWholeSurfaceIncludingRemovedChars() { + // NUL and the zero-width space are removed by cleaning, so one word "abc" remains; it is + // not representable and becomes the unknown piece spanning the full original surface. + final List pieces = uncased().encode("a\u0000b\u200Bc"); + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), "[UNK]", 1, 0, 5); + } + + @Test + void testContextualCaseMappingFallsBackToWordWideSpans() { + // Greek final sigma is a contextual mapping the per-char rerun cannot reproduce, so the + // word's pieces fall back to spanning the whole word; the piece content is asserted + // exactly below. + final List pieces = + uncased().encode("\u03A3\u039F\u03A6\u039F\u03A3"); + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), + "\u03C3\u03BF\u03C6\u03BF\u03C2", 20, 0, 5); + } + + @Test + void testEncodeToIdsCarriesVocabularyLineNumbers() { + assertArrayEquals(new int[] {2, 4, 5, 6, 3}, uncased().encodeToIds("Hello worldS")); + } + + @Test + void testCasedEncoderKeepsCase() { + final List vocabulary = new ArrayList<>(VOCAB); + vocabulary.add("Hello"); + final WordpieceEncoder cased = new WordpieceEncoder(vocabulary, false); + final List pieces = cased.encode("Hello hello"); + assertPiece(pieces.get(1), "Hello", vocabulary.size() - 1, 0, 5); + assertPiece(pieces.get(2), "hello", 4, 6, 11); + } + + @ParameterizedTest + @ValueSource(strings = {"", " "}) + void testEmptyAndBlankTextEncodeToTheFramePiecesOnly(String input) { + final List pieces = uncased().encode(input); + assertEquals(2, pieces.size(), "frame pieces broke on <" + input + ">"); + assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); + assertPiece(pieces.get(1), "[SEP]", 3, input.length(), input.length()); + } + + @Test + void testValidationFailsLoudly() { + assertThrows(IllegalArgumentException.class, () -> new WordpieceEncoder(null)); + assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(List.of("[CLS]", "[SEP]"))); + assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(List.of("[CLS]", "[SEP]", "[UNK]", "dup", "dup"))); + final List withNull = new ArrayList<>(VOCAB); + withNull.add(null); + assertThrows(IllegalArgumentException.class, () -> new WordpieceEncoder(withNull)); + assertThrows(IllegalArgumentException.class, () -> uncased().encode(null)); + } +} diff --git a/opennlp-distr/pom.xml b/opennlp-distr/pom.xml index e9092d8821..c1cf8e841c 100644 --- a/opennlp-distr/pom.xml +++ b/opennlp-distr/pom.xml @@ -91,6 +91,15 @@ org.apache.opennlp opennlp-spellcheck + + org.apache.opennlp + opennlp-subword + + + + org.apache.opennlp + opennlp-embeddings + diff --git a/opennlp-distr/src/main/assembly/bin.xml b/opennlp-distr/src/main/assembly/bin.xml index 2db4eafc65..ad465a18d6 100644 --- a/opennlp-distr/src/main/assembly/bin.xml +++ b/opennlp-distr/src/main/assembly/bin.xml @@ -107,6 +107,13 @@ bin + + ../opennlp-extensions/opennlp-embeddings/src/main/bin + 755 + 755 + bin + + ../opennlp-tools/lang 644 @@ -232,6 +239,13 @@ docs/apidocs/opennlp-morfologik + + ../opennlp-extensions/opennlp-embeddings/target/reports/apidocs + 644 + 755 + docs/apidocs/opennlp-embeddings + + ../opennlp-extensions/opennlp-spellcheck/target/reports/apidocs 644 @@ -239,6 +253,13 @@ docs/apidocs/opennlp-spellcheck + + ../opennlp-extensions/opennlp-subword/target/reports/apidocs + 644 + 755 + docs/apidocs/opennlp-subword + + ../opennlp-extensions/opennlp-uima/target/reports/apidocs 644 diff --git a/opennlp-docs/src/docbkx/embeddings.xml b/opennlp-docs/src/docbkx/embeddings.xml new file mode 100644 index 0000000000..18dfe57696 --- /dev/null +++ b/opennlp-docs/src/docbkx/embeddings.xml @@ -0,0 +1,321 @@ + + + + + + + Static Embeddings + +

+ Introduction + + The opennlp-embeddings extension module produces sentence and word + embedding vectors from a static (non-contextual) embedding 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 of releases is the + primary target), and looking a sentence up in the table approximates the transformer's + semantics at a small fraction of the cost: embedding a text is tokenize, gather, + mean-pool, and normalize, with no model forward pass, no GPU, and no native runtime. + SentencePiece support covers multilingual tables distilled from encoders of the + XLM-RoBERTa family, whose vectors embed different languages into the same space. + + + Use it when embedding throughput and a pure-JVM deployment matter more than the last + few points of retrieval quality: semantic similarity and deduplication, candidate + retrieval for a heavier reranker, clustering, or classification features. A + contextual model remains the better choice when distinguishing word senses in context + is the point of the task. + + + OpenNLP supports contextual, ONNX-backed sentence vectors as well, in the + opennlp-dl module. Those are inherently more accurate; Model2Vec trades + some accuracy for a large speed gain. Both paths implement the same + TextEmbedder interface (in opennlp-api), so an application can + swap one for the other without changing its calling code. + + + No model is bundled with the module. Callers point it at a model directory they + downloaded; the table's own license applies to the table. + + + The public API of this module + (StaticEmbeddingModel, SafetensorsFile, + TensorInfo, Neighbor, ModelDistiller, and + ModelAssembler, VectorIndex, + FlatFloatIndex, and TurboQuantIndex), together with the + TextEmbedder interface + in opennlp-api, is experimental and may change in a later release. + +
+ +
+ Embedding Text with the API + + A model directory loads with a single call, and the tokenizer family is detected from + the files present. A WordPiece model carries vocab.txt, + model.safetensors, config.json, and + tokenizer_config.json; a Model2Vec Unigram model carries + tokenizer.json, model.safetensors, and + config.json. Its JSON file contains the Unigram vocabulary, scores, + normalizer, and pre-tokenizer. Legacy SentencePiece directories may also carry a + trained .model file. The tokenizer and pooling switches are read from the + model's own configuration files: + + + neighbors = model.mostSimilar("coffee", 5); +List analogy = model.analogy("man", "king", "woman", 1);]]> + + + StaticEmbeddingUsageExampleTest asserts the load-and-query workflow + shown here. + + + For a model laid out differently, the explicit overloads take the data files and the + switches directly. The WordPiece overload takes whether the tokenizer lower-cases + (and strips accents) and whether embeddings are L2-normalized; both are properties of + the model, published in its configuration. The SentencePiece overload has no casing + switch, because the trained .model file carries the model's own text + normalizer. + + + + + + The testExplicitOverloads case of + StaticEmbeddingUsageExampleTest asserts the two overloads shown here. + + + Matrix rows are resolved by piece string, never by tokenizer id, because the two + files of a legacy SentencePiece model may order and offset their ids differently. A + poolable piece with no matrix row fails loud at load time. Current Model2Vec Unigram + releases load directly from their self-contained tokenizer.json and do not + require a tokenizer file copied from the teacher. + + + 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 (matching the reference implementation), and + similarity reports 0 for them. This is a rare edge case for + text in the model's language, since subword tokenization backs off to smaller pieces; + it mostly happens for empty input or text outside the vocabulary's coverage. + +
+ +
+ Semantic Search + + The headline use case is ranking documents against a query by meaning rather than by + shared words. similarity embeds both texts and returns the cosine + similarity of their vectors, so a small document list ranks with one call per + document: + + + documents = List.of( + "How do I brew espresso at home?", + "The history of tea in East Asia", + "Best grinders for pour-over coffee"); + +record Scored(String document, double score) {} +List results = new ArrayList<>(); +for (String document : documents) { + results.add(new Scored(document, model.similarity(query, document))); +} +results.sort(Comparator.comparingDouble(Scored::score).reversed());]]> + + + StaticEmbeddingSearchExampleTest asserts the ranking shown here. For a + corpus too large to score per query, embed each document once with + embed, keep the vectors in any vector index, and embed only the query + at search time; the index does not care how the vectors were produced. + +
+ +
+ Bounded In-Memory Vector Search + + The module includes two build-once, read-many implementations of + VectorIndex. FlatFloatIndex keeps full-precision vectors and + scans every row, which makes it the exact baseline and the simplest choice for a small + collection. TurboQuantIndex also scans every row, but stores packed 2-bit, + 3-bit, or 4-bit vectors to reduce the working set. Neither implementation is an + approximate-nearest-neighbor graph or a distributed search engine. They are intended for + bounded collections held by one JVM, such as the passages of one document or a modest + working set. + + + Build the index on one thread, call freeze once, then safely publish the + frozen index. Frozen indexes accept concurrent topK calls. Added vectors are + copied, ids must be unique, and the query and every vector must match the configured + dimension: + + + hits = index.topK(model.embed("king"), 5);]]> + + + VectorIndexUsageExampleTest asserts this lifecycle. Use + FlatFloatIndex in the same listing when exact float scoring matters more + than storage. A frozen, non-empty TurboQuantIndex can be saved with + write(Path) and reopened with read(Path); the loaded index is + already frozen. The interface deliberately has no delete or update operation. Rebuild a + new index when the bounded collection changes. + +
+ +
+ Command Line Tools + + The module ships its own command line launcher, bin/embeddings, next to + bin/opennlp in the binary distribution. Invoked without arguments it lists + the available tools, and every tool prints its help when invoked with the + help parameter. + +
+ Distill Model Tool + + The DistillModel tool compresses a sentence-transformer teacher into a + static embedding table. The teacher is either a Hugging Face model id + (org/model, or org/model@revision to pin a branch, tag, or + commit; the required files download once into a local cache) or a local directory + holding the teacher's tokenizer.json and onnx/model.onnx. + The following command distills a teacher into the directory given by + -out: + + + + -pcaDims is the number of principal components to keep and defaults + to 256. The run ends by assembling the output directory and verifying it with + StaticEmbeddingModel.load, then prints a summary naming the tokenizer + family, the row count, the dimension reduction, and the variance the PCA kept, so a + run that prints a summary is a directory that works. + + + -terms names an optional term file: one term per line, with text after + a tab ignored, so a learned vocabulary TSV works unchanged. Each term, a whole word + or a multi-word phrase such as a domain vocabulary entry, is segmented by the + teacher's own tokenizer, encoded through the teacher as one sequence, and appended + to the table as an extra row, recorded in the model directory as + terms.txt. When such a model embeds text, it first matches the text + against its terms greedily longest-first, case-insensitively, and pools a matched + term's single row instead of the subword pieces of its words; text between matches + is tokenized as usual, and a model without a term file embeds exactly as before. + Terms should arrive sorted by descending corpus frequency, because the Zipf + weighting treats the subword rows and the term rows as one frequency ranking. A + term that equals a vocabulary token is dropped as a duplicate row, and terms are + also returned by the similarity search of mostSimilar like any + vocabulary token. + +
+
+ Assemble Model Tool + + The AssembleModel tool completes a downloaded distillation in place so + StaticEmbeddingModel.load can open it. A Model2Vec distillation writes + model.safetensors, tokenizer.json, and + config.json; for a WordPiece model the tool derives the missing + vocab.txt and tokenizer_config.json from + tokenizer.json, and for a SentencePiece model it checks that the + trained .model file copied from the teacher is present, naming the fix + if it is not. It never overwrites an existing file. + + + + The tool then verifies the directory by loading it and prints the loaded model's + family, row count, and dimension, along with a line for every file it wrote. + +
+
+ +
+ Quantized Models + + A static embedding table can be quantized to 2, 3, or 4 bits per dimension, shrinking + it by roughly 4 to 16 times against the 32-bit float matrix, depending on the bit + width and on how far the dimension is from a power of two: rows are padded to the + next power of two before coding, so a 500,000-row, 300-dimension table stores 512 + coded dimensions per row and drops from about 600 MB to about 132 MB at 4 bits, 4.5 + times smaller. Because embedding is + memory-bound row gathering, reading fewer bytes is also the throughput lever. The + method is the TurboQuant construction: each row is rotated so its coordinates become + near-independent and near-Gaussian, and each rotated coordinate is encoded against an + optimal scalar grid, with a least-squares scale fitted per row. + + + The QuantizeModel tool quantizes a model directory's + model.safetensors in place, writing model.quantized next to + it and reporting the sizes and the reconstruction quality it measured from the written + file: + + + + + + A directory presents exactly one matrix file. After quantizing, delete the + model.safetensors to deploy the quantized matrix; the loader then reads + model.quantized with no code change, and any per-token pooling weights + are carried inside the quantized file. A directory that still holds both files is + rejected at load time, so the deployment must choose which matrix is authoritative. + Embedding, similarity, and mostSimilar behave the same as on + the float model, up to the quantization error of the chosen bit width; fewer bits + trade fidelity for size. + + + StaticEmbeddingModelQuantizedTest asserts this workflow end to end: + quantizing a directory, loading it after the safetensors is removed, and matching the + float model's embeddings and rankings; QuantizedEmbeddingMatrixTest pins + the reconstruction quality at each bit width. + +
+ +
+ The safetensors Reader + + Weights are read with a purpose-built reader for the safetensors format. Unlike + pickle-based checkpoint formats, a safetensors file carries no executable content + (the header is data-only JSON and the body is raw tensor bytes), so loading one + cannot execute arbitrary code. Only the header is read eagerly; tensor data streams + directly into the decoded array, so the file size is not limited by Java's + int-indexed arrays. One decoded tensor is capped at the maximum Java array length + (about 2.1 billion float elements), checked explicitly. + +
+ diff --git a/opennlp-docs/src/docbkx/opennlp.xml b/opennlp-docs/src/docbkx/opennlp.xml index 36641c2c89..7d78e290bd 100644 --- a/opennlp-docs/src/docbkx/opennlp.xml +++ b/opennlp-docs/src/docbkx/opennlp.xml @@ -120,6 +120,7 @@ under the License. + diff --git a/opennlp-docs/src/docbkx/tokenizer.xml b/opennlp-docs/src/docbkx/tokenizer.xml index cd1d8a2ddf..6733728d3a 100644 --- a/opennlp-docs/src/docbkx/tokenizer.xml +++ b/opennlp-docs/src/docbkx/tokenizer.xml @@ -539,4 +539,94 @@ wordTokenizer.tokenize("The quick brown fox.", (start, end, type) -> { +
+ Subword Tokenization + + Neural models usually operate on subword units drawn from a fixed vocabulary rather than + on words. The SubwordTokenizer interface in + opennlp.tools.tokenize covers this case: encode splits text into + pieces, and every returned SubwordPiece carries the piece string in the + model's normalized form, its vocabulary id, and the exact span of the original text it + came from. Because the segmentation runs after the model's own normalization, a piece is + generally not a substring of the input; the spans always refer to the caller's original + text, so annotations computed over the pieces can be mapped back without guesswork. The + encodeToIds and encodeToPieces methods return just the ids or + the piece strings when the spans are not needed. Thread safety is implementation specific; + both implementations described below are immutable and safe for concurrent use by + multiple threads. + +
+ SentencePiece + + SentencePieceTokenizer in the opennlp-subword artifact + (package opennlp.subword.sentencepiece) runs a trained SentencePiece + .model file purely in Java, with no native library. The file itself + carries everything the class needs: the vocabulary with piece scores and types, the + segmentation algorithm (unigram language model or byte-pair encoding), and the text + normalizer the model was trained with. Loading is a one-time cost and the resulting + instance is immutable, so one tokenizer can be shared by any number of threads. + + id " + piece.id() + + ", original text [" + piece.start() + ", " + piece.end() + ")"); +} + +int[] ids = tokenizer.encodeToIds("Ready for the embedding layer.");]]> + + SentencePieceUsageExampleTest asserts the load-and-encode workflow + shown here. The vocabulary can be inspected through + vocabularySize, idToPiece, pieceToId, and + score, and the algorithm method reports whether the + model uses the unigram or the byte-pair encoding algorithm. + + + The model's own normalizer is also exposed directly: + SentencePieceTokenizer implements OffsetAwareNormalizer, so + normalize applies the model's normalization rules to arbitrary text and + normalizeAligned additionally returns the character alignment described + in . This is useful when other processing must see + text exactly as the subword model does. + + + Output parity with the reference implementation is pinned by test fixtures whose + expected pieces, ids, and spans are produced by the reference + sentencepiece + package rather than by the Java code under test. The fixtures README at + opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/README.md + documents how the bundled test models were created and walks through regenerating + them and validating the implementation, including against real published models. + +
+
+ WordPiece + + WordpieceEncoder in opennlp.tools.tokenize runs the full + BERT tokenization pipeline: control character removal, whitespace normalization, CJK + isolation, optional lower casing with accent stripping, punctuation isolation, and + greedy longest-match wordpiece segmentation. It is constructed from the model's + vocabulary, conventionally the lines of a vocab.txt file, where a piece's + id is its line number. + + vocabulary = Files.readAllLines(Path.of("vocab.txt")); + +WordpieceEncoder encoder = new WordpieceEncoder(vocabulary); // uncased model +WordpieceEncoder cased = new WordpieceEncoder(vocabulary, false); // cased model + +List pieces = encoder.encode("OpenNLP encodes text for BERT models."); +int[] ids = encoder.encodeToIds("OpenNLP encodes text for BERT models.");]]> + + Every encoding is framed by the classification and separator pieces, which carry empty + spans at the text's boundaries, and words the vocabulary cannot cover become the + unknown piece. Models with other special tokens or with non-contiguous ids are + supported through the constructors taking explicit special tokens or a piece-to-id + map. The older BertTokenizer is deprecated and scheduled for removal: + its tokenize method returns the same piece sequence that + encodeToPieces returns. Migrating means passing the vocabulary as a + List or a piece-to-id map instead of a Set; the output + is unchanged. + +
+
diff --git a/opennlp-extensions/opennlp-embeddings/README.md b/opennlp-extensions/opennlp-embeddings/README.md new file mode 100644 index 0000000000..cffd98ad3b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/README.md @@ -0,0 +1,195 @@ + + +# OpenNLP Static Embeddings + +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 sentence up in the table approximates the transformer's semantics at a fraction of the cost. Because SentencePiece models are supported, this includes multilingual tables distilled from encoders like the [XLM-RoBERTa](https://arxiv.org/abs/1911.02116) family. There is no model forward pass, no GPU, and no native runtime; it is pure JVM. + +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 + +Point `load` at a downloaded model directory, then embed: + +```java +StaticEmbeddingModel model = StaticEmbeddingModel.load(Path.of("/path/to/model-directory")); + +float[] vector = model.embed("The quick brown fox"); +double similarity = model.similarity("coffee", "espresso"); +List near = model.mostSimilar("coffee", 5); +``` + +The directory is the layout published releases use, and `load` detects the tokenizer family from the files present. A WordPiece model carries `vocab.txt`, `model.safetensors`, `config.json`, and `tokenizer_config.json`. A SentencePiece model carries a trained `.model` file (`sentencepiece.bpe.model`, `spiece.model`, or `tokenizer.model`) next to `tokenizer.json`, `model.safetensors`, and `config.json`. In both cases the tokenizer and pooling switches are read from the model's own config. One loaded model is immutable and thread-safe, so it can serve every thread of an application. + +A multilingual SentencePiece table embeds different languages into the same space, so similarity works across them: + +```java +model.similarity("The weather is beautiful today", "今天天气很好"); // same meaning, high score +``` + +## When to use it + +Reach for this when embedding throughput and deployment simplicity matter more than the last few points of retrieval quality: semantic similarity, deduplication, candidate retrieval in front of a heavier reranker, clustering, or features for a classifier. A contextual model is still the better choice when the task depends on distinguishing word senses in context. + +## How it works + +A static embedding model is a vocabulary and a matrix: one row per token, each row a vector of the model's dimension. Embedding runs entirely as table lookups and arithmetic: + +```mermaid +flowchart LR + A["text"] --> B["subword tokenize
(WordPiece or SentencePiece)"] + B --> C["gather piece rows by string
drop unknown, skip special"] + C --> D["weight + mean-pool"] + D --> E["L2 normalize"] + 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]`, 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-vocabulary pieces embeds to a zero vector rather than raising. +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. + +Per-row L2 norms and the special-token mask are precomputed at load time, so the neighbor scan and similarity calls do not recompute them on every query. + +### Loading + +The one-argument `load` reads the model's own configuration to resolve the tokenizer and pooling switches, so callers do not restate them: + +```mermaid +flowchart TD + L["StaticEmbeddingModel.load(dir)"] --> DET{"vocab.txt present?"} + DET -- "yes: WordPiece" --> WCFG["read config.json,
tokenizer_config.json"] + WCFG --> CAS["casing = do_lower_case"] + DET -- "no: SentencePiece" --> SPM["load the trained .model
(its own normalizer, no casing switch)"] + SPM --> TJ["tokenizer.json vocab
names the matrix rows"] + TJ --> COV["verify every poolable piece
has a matrix row"] + L --> NRM["normalization from config.json"] + L --> MAT["model.safetensors to matrix"] + CAS --> M["immutable, thread-safe model"] + COV --> M + NRM --> M + MAT --> M +``` + +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 + +```mermaid +flowchart TD + subgraph MODEL["StaticEmbeddingModel"] + EV["EmbeddingVocabulary
(piece string to matrix row)"] + ST["SubwordTokenizer"] + MX["embedding matrix"] + end + WE["WordpieceEncoder
(opennlp-api)"] -. one of .-> ST + SP["SentencePieceTokenizer
(opennlp-subword)"] -. one of .-> ST + SHP["SafetensorsHeaderParser"] --> SF["SafetensorsFile"] + SF --> MX + MODEL -. implements .-> TE["TextEmbedder
(opennlp-api)"] + DL["SentenceVectorsDL
(opennlp-dl, ONNX)"] -. implements .-> TE +``` + +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 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): + +| table | tokenizer, rows | `embed()` | `mostSimilar()` | +| --- | --- | --- | --- | +| 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, the repository's `dev/embeddings/parity/` directory 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. Treat all of these as a starting expectation and run the benchmark on the model you plan to use. + +## Usage + +### Loading a non-standard layout + +For a model laid out differently, the explicit overloads take the data files and the model properties directly. WordPiece: + +```java +StaticEmbeddingModel model = StaticEmbeddingModel.load( + Path.of("vocab.txt"), Path.of("model.safetensors"), + StaticEmbeddingModel.Casing.UNCASED, // from the model's do_lower_case + StaticEmbeddingModel.Normalization.L2); // from the model's config +``` + +SentencePiece (no casing switch, because the `.model` file carries the model's own text normalizer): + +```java +StaticEmbeddingModel model = StaticEmbeddingModel.loadSentencePiece( + Path.of("sentencepiece.bpe.model"), Path.of("tokenizer.json"), + Path.of("model.safetensors"), + StaticEmbeddingModel.Normalization.L2); +``` + +### Neighbors and analogies + +`Neighbor` is a small record of the token and its cosine similarity: + +```java +for (Neighbor n : model.mostSimilar("coffee", 5)) { + System.out.println(n.token() + " " + n.similarity()); +} + +List king = model.analogy("man", "king", "woman", 1); +``` + +### Retrieval + +Embed a small corpus once, then rank documents against a query by cosine similarity. Because the vectors are L2-normalized, cosine is a plain dot product: + +```java +StaticEmbeddingModel model = StaticEmbeddingModel.load(modelDir); + +List docs = List.of( + "How do I brew espresso at home?", + "The history of tea in East Asia", + "Best grinders for pour-over coffee"); + +float[][] docVectors = docs.stream().map(model::embed).toArray(float[][]::new); +float[] query = model.embed("home espresso machine"); + +IntStream.range(0, docs.size()) + .boxed() + .sorted(Comparator.comparingDouble(i -> -dot(query, docVectors[i]))) + .forEach(i -> System.out.println(docs.get(i))); +``` + +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 + +No model is bundled. Point the module at files you download, and the table's own license applies to the table. The Model2Vec distilled releases (for example potion-base-8M) publish the exact directory layout the one-argument `load` expects: download that release's `vocab.txt`, `model.safetensors`, `config.json`, and `tokenizer_config.json` into one directory and pass the directory to `load`. Or distill your own teacher with the module's `DistillModel` command (see `TRAINING.md`). + +For a multilingual SentencePiece table (for example one distilled from a bge-m3 or XLM-RoBERTa teacher), the distillation output ships `tokenizer.json`, `model.safetensors`, and `config.json` but usually not the trained SentencePiece `.model` file; copy that one file from the teacher model's own repository (it is named `sentencepiece.bpe.model` there) into the same directory. The loader tells you exactly this if the file is missing. + +## Notes and limits + +- Instances are immutable and safe for concurrent use, so one loaded model serves every thread. +- Static tables do not disambiguate word senses in context. If the task turns on context, use a contextual model. +- Out-of-vocabulary-only input embeds to a zero vector, matching the reference implementation. This is a rare edge case for text in the model's language (WordPiece backs off to subwords), and mostly happens for empty input or text outside the vocabulary's coverage. Decide in your code whether a zero vector means "no signal" for your use case. + +## See also + +- [`TRAINING.md`](TRAINING.md) for distilling your own table from a sentence-transformer teacher, including the multilingual SentencePiece worked example. +- The Dev Manual chapter (`opennlp-docs/src/docbkx/embeddings.xml`) for the same material in the manual. +- `opennlp-dl` for the contextual, ONNX-backed sentence vector path, which shares the `TextEmbedder` interface with this module. diff --git a/opennlp-extensions/opennlp-embeddings/TRAINING.md b/opennlp-extensions/opennlp-embeddings/TRAINING.md new file mode 100644 index 0000000000..680a6e40c8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/TRAINING.md @@ -0,0 +1,85 @@ + + +# Distilling a Model for OpenNLP Static Embeddings + +This module loads static embedding tables and, with the `DistillModel` command, also produces them: a table is distilled once from a sentence-transformer teacher and then loaded in the JVM as many times as you like. The distiller replicates [Model2Vec](https://github.com/MinishLab/model2vec) in Java — it runs the teacher's ONNX graph 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. Distill the teacher + +``` +opennlp-embeddings DistillModel -teacher BAAI/bge-m3 -out bge-m3-static -pcaDims 256 +``` + +`-teacher` is a Hugging Face model id (its `tokenizer.json`, `tokenizer_config.json`, `onnx/model.onnx`, and, for an export that splits its weights out, `onnx/model.onnx_data` download once into `~/.cache/opennlp-embeddings/-`) or a local directory holding those files. `-pcaDims` defaults to 256. For a SentencePiece teacher like bge-m3 the trained `sentencepiece.bpe.model` is fetched alongside, because the static table keeps the teacher's segmentation. The command ends by completing the directory (the `AssembleModel` step) and verifying it by loading it, so a run that prints a summary is a directory that works. + +Distil into a fresh directory. The command replaces the files it writes itself, but the assembly step never overwrites a `vocab.txt` or `tokenizer_config.json` an earlier run left behind, and a run that fails part way through leaves whatever it had written. + +bge-m3 is an [XLM-RoBERTa](https://arxiv.org/abs/1911.02116)/SentencePiece model with a 250k multilingual vocabulary, native dimension 1024. + +### On the dimension + +`pcaDims` is the one quality knob worth thinking about, and bigger is not better. Distilling bge-m3 at 256 and at 512 gives the same cross-lingual similarity within noise (English/Chinese paraphrase around 0.69 either way), while 512 doubles the matrix on disk and in memory and cuts embedding throughput. PCA to 256 already captures the useful variance of the teacher; the extra dimensions are mostly noise that dilutes the signal. 256 is a good default, and it is where the reference Model2Vec tables (the MinishLab "potion" series) sit too. + +## 2. Assemble the model directory + +The distiller writes `model.safetensors` (F32), the cleaned `tokenizer.json`, and `config.json`, and copies the teacher's SentencePiece `.model` file when there is one. `DistillModel` assembles and verifies its own output; `AssembleModel` completes a directory put together by hand — for a WordPiece model it derives `vocab.txt` and `tokenizer_config.json` from `tokenizer.json`, for a SentencePiece model it checks that the trained `.model` file is present: + +``` +opennlp-embeddings AssembleModel -modelDir bge-m3-static +``` + +A loadable SentencePiece directory then holds: + +``` +bge-m3-static/ + sentencepiece.bpe.model # copied from the teacher; segments the text + tokenizer.json # Unigram vocab; its row order maps to the matrix + model.safetensors # the embedding matrix + config.json # carries "normalize": true|false +``` + +`load` detects the SentencePiece layout from the `.model` file next to `tokenizer.json`; it does not need `tokenizer_config.json`, because the `.model` carries the model's own text normalizer. If you forget the `.model` file, the loader says so by name. + +## 3. Load and verify in the JVM + +```java +StaticEmbeddingModel model = StaticEmbeddingModel.load(Path.of("bge-m3-static")); + +// Multilingual: the same meaning across languages lands nearby. +double crossLingual = model.similarity( + "The weather is beautiful today", "今天天气很好"); // high +double unrelated = model.similarity( + "The weather is beautiful today", "quarterly earnings missed"); // low + +// Sanity: neighbors of a word are its translations and case variants. +model.mostSimilar("coffee", 5); // ▁coffee, ▁Coffee, ▁koffie, ▁kávé, ▁кофе +``` + +Confirm parity against the Python reference before trusting a fresh distillation: embed the same text on both sides and check the vectors match within floating-point tolerance. They should agree to a few parts in ten thousand, because the JVM path reproduces the reference tokenization and pooling exactly, not approximately. The reference Python flow lives in the repository as `dev/embeddings/distill_bge_m3.py`, and `dev/embeddings/parity/` holds a harness that reruns the parity check and the single-thread speed comparison against the Python reference on any machine. + +Two tables distilled independently from the same teacher (one with this command, one with Python Model2Vec) agree on their pairwise geometry to a few parts in a thousand — similarities, neighbors, and rankings match — but their raw vectors are not directly comparable axis by axis: PCA fixes only the subspace, and within the near-degenerate tail of the spectrum two independent decompositions choose different bases. + +## The WordPiece path + +A WordPiece teacher (a BERT-family model such as bge-large-en) distills the same way. Its directory layout is the BERT one instead: `vocab.txt` (one token per line, line number is the row), `model.safetensors`, `config.json`, and `tokenizer_config.json` (whose `do_lower_case` sets the casing). `load` detects WordPiece from the presence of `vocab.txt`. + +A distillation writes `tokenizer.json` rather than a `vocab.txt`, so the two BERT files are derived: `vocab.txt` from the `tokenizer.json` vocabulary in id order, `tokenizer_config.json` from the normalizer's lowercase flag (absent, it defaults to lower-casing). `DistillModel` does this itself as its final step; `AssembleModel` is the same step run on its own, for a directory assembled by hand. + +## Where a table's license comes from + +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/pom.xml b/opennlp-extensions/opennlp-embeddings/pom.xml new file mode 100644 index 0000000000..48a007ed5b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/pom.xml @@ -0,0 +1,165 @@ + + + + + + 4.0.0 + + org.apache.opennlp + opennlp-extensions + 3.0.0-SNAPSHOT + + + opennlp-embeddings + jar + Apache OpenNLP :: Ext :: Embeddings + + + + org.apache.opennlp + opennlp-api + + + + org.apache.opennlp + opennlp-runtime + + + + org.apache.opennlp + opennlp-subword + + + + org.apache.opennlp + opennlp-cli + + + + + com.microsoft.onnxruntime + onnxruntime + ${onnxruntime.version} + + + + org.junit.jupiter + junit-jupiter-api + test + + + + org.junit.jupiter + junit-jupiter-engine + test + + + + org.junit.jupiter + junit-jupiter-params + test + + + + + + + de.thetaphi + forbiddenapis + + + + opennlp/embeddings/HuggingFaceModelCacheTest*.class + + + + + + + + + jmh + + + org.openjdk.jmh + jmh-core + ${jmh.version} + test + + + + org.openjdk.jmh + jmh-generator-annprocess + ${jmh.version} + test + + + + + + org.codehaus.mojo + build-helper-maven-plugin + 3.6.1 + + + add-test-source + generate-test-sources + + add-test-source + + + + src/jmh/java + + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + jmh-compile + test-compile + + testCompile + + + + + org.openjdk.jmh + jmh-generator-annprocess + ${jmh.version} + + + + + + + + + + + + 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 new file mode 100644 index 0000000000..6593009a09 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java @@ -0,0 +1,194 @@ +/* + * 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.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.concurrent.TimeUnit; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; +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}: single-JVM embedding and neighbor-scan + * throughput. + * + *

The {@code modelDir} parameter selects the table to benchmark. Its default, + * {@code "synthetic"}, builds a fixture sized to {@code minishlab/potion-base-8M} (29,528 rows, + * 256 dimensions) in a temp directory, so the benchmark runs with no model download. Passing one + * or more real model directories instead, for example + * {@code -p modelDir=/models/potion-base-8M,/models/bge-large-en-v1.5-static}, benchmarks those + * tables directly and reports one row per directory, which is how the two-model comparison in the + * README is produced.

+ */ +@BenchmarkMode(Mode.Throughput) +@OutputTimeUnit(TimeUnit.SECONDS) +@Warmup(iterations = 5, time = 2) +@Measurement(iterations = 10, time = 2) +@Fork(2) +public class StaticEmbeddingModelBenchmark { + + /** The synthetic-fixture selector; any other value is treated as a model directory path. */ + private static final String SYNTHETIC = "synthetic"; + + // Matches minishlab/potion-base-8M: hidden_dim 256 and 7,559,168 / 256 = 29,528 rows. + private static final int VOCAB_SIZE = 29_528; + private static final int DIMENSION = 256; + + private static final String[] REAL_WORDS = { + "the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", "she", "told", "me", "he", + "lived", "in", "wrote", "letter", "right", "away", "opennlp", "provides", "tools", "for", + "language", "processing", "driver", "got", "badly", "injured", "by", "accident", + }; + + private static final String[] SENTENCES = { + "The quick brown fox jumps over the lazy dog.", + "She told me he lived in Edinburgh.", + "I wrote him a letter right away.", + "OpenNLP provides tools for natural language processing.", + "The driver got badly injured by the accident.", + }; + + @State(Scope.Benchmark) + public static class ModelState { + + /** + * The model to benchmark: {@code "synthetic"} for the built-in fixture, or a model directory + * path. Override with {@code -p modelDir=dir1,dir2} to benchmark real tables. + */ + @Param({SYNTHETIC}) + public String modelDir; + + StaticEmbeddingModel model; + private Path tempDir; + + @Setup(Level.Trial) + public void load() throws IOException { + if (SYNTHETIC.equals(modelDir)) { + tempDir = Files.createTempDirectory("opennlp-embeddings-jmh"); + final Path vocabFile = writeVocab(tempDir); + final Path safetensorsFile = writeSafetensors(tempDir); + model = StaticEmbeddingModel.load(vocabFile, safetensorsFile, + Casing.UNCASED, Normalization.L2); + } else { + model = StaticEmbeddingModel.load(Path.of(modelDir)); + } + } + + @TearDown(Level.Trial) + public void cleanup() throws IOException { + if (tempDir != null) { + Files.deleteIfExists(tempDir.resolve("vocab.txt")); + Files.deleteIfExists(tempDir.resolve("model.safetensors")); + Files.deleteIfExists(tempDir); + } + } + + private static Path writeVocab(Path dir) throws IOException { + final List tokens = new ArrayList<>(VOCAB_SIZE); + tokens.add("[CLS]"); + tokens.add("[SEP]"); + tokens.add("[UNK]"); + for (final String word : REAL_WORDS) { + tokens.add(word); + } + while (tokens.size() < VOCAB_SIZE) { + tokens.add("tok" + (tokens.size() - REAL_WORDS.length - 3)); + } + final Path file = dir.resolve("vocab.txt"); + Files.write(file, tokens); + return file; + } + + private static Path writeSafetensors(Path dir) throws IOException { + final Random random = new Random(42); + final ByteBuffer buffer = ByteBuffer.allocate(VOCAB_SIZE * DIMENSION * 4) + .order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < VOCAB_SIZE * DIMENSION; i++) { + buffer.putFloat((random.nextFloat() - 0.5f) * 2f); + } + final byte[] data = buffer.array(); + final String header = "{\"embeddings\":{\"dtype\":\"F32\",\"shape\":[" + VOCAB_SIZE + "," + + 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()); + return file; + } + } + + @Benchmark + @Threads(Threads.MAX) + public void embed(ModelState state, Blackhole blackhole) { + for (final String sentence : SENTENCES) { + blackhole.consume(state.model.embed(sentence)); + } + } + + @Benchmark + @Threads(Threads.MAX) + public void mostSimilarTop10(ModelState state, Blackhole blackhole) { + blackhole.consume(state.model.mostSimilar(SENTENCES[0], 10)); + } + + /** + * Quick local iteration only: {@code forks(0)} disables JVM fork isolation (unlike + * {@code mvn} with the {@code jmh} profile). Use the Maven-invoked configuration for + * publishable numbers. + */ + public static void main(String[] args) throws Exception { + final Options opt = new OptionsBuilder() + .include(StaticEmbeddingModelBenchmark.class.getSimpleName()) + .forks(0) + .warmupIterations(3) + .measurementIterations(5) + .build(); + new Runner(opt).run(); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings b/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings new file mode 100755 index 0000000000..3a105a7ce2 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings @@ -0,0 +1,56 @@ +#!/bin/sh + +# 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. + +# Note: Do not output anything in this script file, any output +# may be inadvertantly placed in any output files if +# output redirection is used. + +# determine OPENNLP_HOME - $0 may be a symlink to OpenNLP's home +PRG="$0" + +while [ -h "$PRG" ] ; do + ls=$(ls -ld "$PRG") + link=$(expr "$ls" : '.*-> \(.*\)$') + if expr "$link" : '/.*' > /dev/null; then + PRG="$link" + else + PRG="$(dirname "$PRG")/$link" + fi +done + +saveddir=$(pwd) + +OPENNLP_HOME=$(dirname "$PRG")/.. + +# make it fully qualified +OPENNLP_HOME=$(cd "$OPENNLP_HOME" && pwd) + +cd "$saveddir" || exit + +if [ -z "$JAVACMD" ] ; then + if [ -n "$JAVA_HOME" ] ; then + JAVACMD="$JAVA_HOME/bin/java" + else + JAVACMD="$(which java)" + fi +fi + +CLASSPATH=$(echo "$OPENNLP_HOME"/lib/*.jar | tr ' ' ':') + +$JAVACMD -Xmx1024m -Dlog4j.configurationFile="$OPENNLP_HOME/conf/log4j2.xml" -cp "$CLASSPATH" opennlp.embeddings.cmdline.CLI "$@" diff --git a/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings.bat b/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings.bat new file mode 100644 index 0000000000..199d5820ff --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/bin/embeddings.bat @@ -0,0 +1,51 @@ +@ECHO off + +REM # Licensed to the Apache Software Foundation (ASF) under one +REM # or more contributor license agreements. See the NOTICE file +REM # distributed with this work for additional information +REM # regarding copyright ownership. The ASF licenses this file +REM # to you under the Apache License, Version 2.0 (the +REM # "License"); you may not use this file except in compliance +REM # with the License. You may obtain a copy of the License at +REM # +REM # http://www.apache.org/licenses/LICENSE-2.0 +REM # +REM # Unless required by applicable law or agreed to in writing, +REM # software distributed under the License is distributed on an +REM # # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +REM # KIND, either express or implied. See the License for the +REM # specific language governing permissions and limitations +REM # under the License. + +REM # Note: Do not output anything in this script file, any output +REM # may be inadvertantly placed in any output files if +REM # output redirection is used. +SETLOCAL + +IF "%JAVA_CMD%" == "" ( + IF "%JAVA_HOME%" == "" ( + SET JAVA_CMD=java + ) ELSE ( + REM # Keep JAVA_HOME to short-name without spaces + FOR %%A IN ("%JAVA_HOME%") DO SET JAVA_CMD=%%~sfA\bin\java + ) +) + +REM # Should work with Windows XP and greater. If not, specify the path to where it is installed. +IF "%OPENNLP_HOME%" == "" ( + SET OPENNLP_HOME=%~sp0.. +) ELSE ( + REM # Keep OPENNLP_HOME to short-name without spaces + FOR %%A IN ("%OPENNLP_HOME%") DO SET OPENNLP_HOME=%%~sfA +) +setLocal EnableDelayedExpansion +set CLASSPATH=" + +FOR %%A IN ("%OPENNLP_HOME%\lib\*.jar") DO ( + set CLASSPATH=!CLASSPATH!;%%A +) +set CLASSPATH=!CLASSPATH!" + +%JAVA_CMD% -Xmx1024m "-Dlog4j.configurationFile=%OPENNLP_HOME%\conf\log4j2.xml" -cp %CLASSPATH% opennlp.embeddings.cmdline.CLI %* + +ENDLOCAL diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingTable.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingTable.java new file mode 100644 index 0000000000..efb4efacf8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingTable.java @@ -0,0 +1,87 @@ +/* + * 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; + +/** + * The row storage behind {@link StaticEmbeddingModel}: gathering rows into a pooled vector and + * scoring rows against a query. + * + *

A table may work in a space of its own choosing. {@link #addRow(int, float, float[])} + * accumulates into a vector of {@link #pooledLength()}, and {@link #finishPooling(float[])} + * maps the accumulated vector to original space once per pooled result. Scoring mirrors this: + * {@link #prepareQuery(float[])} maps a query into the working space once, and + * {@link #dot(int, float[])} scores every row against the prepared query there. The working + * space must preserve norms and dot products, so cosine math is space-independent.

+ * + *

Implementations must be safe for concurrent use; the accumulator and prepared query + * arrays belong to the caller.

+ */ +interface EmbeddingTable { + + /** {@return the number of rows} */ + int rowCount(); + + /** {@return the original row width, the length of a pooled result} */ + int dimension(); + + /** {@return the length of the pooling accumulator and of a prepared query} */ + int pooledLength(); + + /** + * Adds a row, times a weight, onto a pooling accumulator. + * + * @param row The row to add, between 0 and {@code rowCount() - 1}. + * @param weight The weight to multiply the row by. + * @param sum The accumulator, of length {@link #pooledLength()}. + */ + void addRow(int row, float weight, float[] sum); + + /** + * Maps an accumulated vector to original space. Called once per pooled result; the returned + * array may be {@code sum} itself. + * + * @param sum The accumulator, of length {@link #pooledLength()}. + * @return The pooled vector in original space, of length {@link #dimension()}. + */ + float[] finishPooling(float[] sum); + + /** + * Maps an original-space query into this table's working space, once per scan. The returned + * array may be {@code query} itself. + * + * @param query The query, of length {@link #dimension()}. Not modified. + * @return The prepared query, of length {@link #pooledLength()}. + */ + float[] prepareQuery(float[] query); + + /** + * The dot product of a row with a prepared query, equal to the original-space dot product up + * to float rounding. + * + * @param row The row to score, between 0 and {@code rowCount() - 1}. + * @param preparedQuery The query as returned by {@link #prepareQuery(float[])}. + * @return The dot product. + */ + double dot(int row, float[] preparedQuery); + + /** + * {@return the L2 norm of a row as this table stores it, for cosine scoring} + * + * @param row The row, between 0 and {@code rowCount() - 1}. + */ + double rowNorm(int row); +} 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 new file mode 100644 index 0000000000..205c0980a6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java @@ -0,0 +1,165 @@ +/* + * 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.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.InvalidFormatException; + +/** + * The row table of a static embedding matrix: piece string to row index and back. Row {@code id} + * of the matrix holds the vector of the piece at position {@code id} in this vocabulary. + * + *

Two layouts produce it: a BERT-style {@code vocab.txt} (one token per line, the line number + * is the row), and a {@code tokenizer.json} with a Unigram model (the {@code model.vocab} list + * order is the row order, with {@code added_tokens} overlaid).

+ * + *

Immutable and safe for concurrent reads after construction.

+ */ +@ThreadSafe +final class EmbeddingVocabulary { + + private final Map idByToken; + private final List tokenById; + + /** Holds the parsed piece-to-row and row-to-piece views; built by the {@code from*} factories. */ + private EmbeddingVocabulary(Map idByToken, List tokenById) { + this.idByToken = idByToken; + this.tokenById = tokenById; + } + + /** + * Reads a {@code vocab.txt} file: one token per line, the line number (0-based) is the row. + * + * @param file The vocabulary file. Must not be {@code null} and must exist. + * @return The parsed vocabulary. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null} or missing. + * @throws InvalidFormatException Thrown if the file contains a duplicate token. + * @throws IOException Thrown if reading the file fails. + */ + static EmbeddingVocabulary fromVocabTxt(Path file) throws IOException { + requireRegularFile(file); + return fromLines(Files.readAllLines(file), file.toString()); + } + + /** + * Reads the Unigram vocabulary of a {@code tokenizer.json} file: the {@code model.vocab} list + * order is the row order, with {@code added_tokens} overlaid. + * + * @param file The {@code tokenizer.json} file. Must not be {@code null} and must exist. + * @return The parsed vocabulary. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null} or missing. + * @throws InvalidFormatException Thrown if the file is not a well-formed Unigram + * {@code tokenizer.json} or a piece appears more than once. + * @throws IOException Thrown if reading the file fails. + */ + static EmbeddingVocabulary fromTokenizerJson(Path file) throws IOException { + requireRegularFile(file); + return fromLines(TokenizerJsonVocab.rows(file), file.toString()); + } + + /** + * Requires {@code file} to be an existing regular file. + * + * @param file The file to check. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null}, missing, or not a + * regular file. + */ + private static void requireRegularFile(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); + } + } + + /** + * Builds a vocabulary from in-memory lines, the token order. + * + * @param lines The tokens, one per element; the index is the token's row. + * @param sourceName The source's name, for error messages. + * @return The parsed vocabulary. + * @throws InvalidFormatException Thrown if a token appears more than once. + */ + static EmbeddingVocabulary fromLines(List lines, String sourceName) + throws InvalidFormatException { + final Map idByToken = new LinkedHashMap<>(lines.size() * 2); + for (int id = 0; id < lines.size(); id++) { + final String token = lines.get(id); + if (idByToken.putIfAbsent(token, id) != null) { + throw new InvalidFormatException( + "Vocabulary " + sourceName + " declares token '" + token + + "' more than once, at rows " + idByToken.get(token) + " and " + id); + } + } + return new EmbeddingVocabulary(Collections.unmodifiableMap(idByToken), List.copyOf(lines)); + } + + /** {@return every token in this vocabulary, without order} */ + Set tokens() { + return idByToken.keySet(); + } + + /** {@return every token in row order, suitable for an id-is-index tokenizer constructor} */ + List orderedTokens() { + return tokenById; + } + + /** + * Looks up a token's row id. + * + * @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) { + throw new IllegalArgumentException("Token must not be null"); + } + final Integer id = idByToken.get(token); + return id == null ? -1 : id; + } + + /** {@return the number of tokens in this vocabulary} */ + int size() { + return idByToken.size(); + } + + /** + * Looks up the token at a row id. + * + * @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()) { + throw new IllegalArgumentException( + "Id " + id + " is outside [0, " + tokenById.size() + ")"); + } + return tokenById.get(id); + } +} 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 new file mode 100644 index 0000000000..07511bcaac --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FlatJsonFields.java @@ -0,0 +1,174 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +/** + * Reads single top-level fields out of a small flat JSON configuration file (a model's + * {@code config.json} or {@code tokenizer_config.json}) without a JSON library dependency. Only + * top-level scalar look-ups are implemented; every other field is skipped structurally, and a + * nested occurrence of the looked-up name never matches. + */ +final class FlatJsonFields { + + /** The JSON null literal, accepted in place of any looked-up value. */ + private static final String NULL_LITERAL = "null"; + + /** Not instantiable. */ + private FlatJsonFields() { + } + + /** + * Reads one top-level boolean field from a JSON object file. + * + * @param file The JSON file, a single top-level object. Must not be {@code null} and must + * exist. + * @param field The top-level field name to read. Must not be {@code null}. + * @return The field's value, or {@code null} when the field is absent or explicitly JSON + * {@code null} (the formats treat those the same: fall back to the default). + * @throws IllegalArgumentException Thrown if an argument is {@code null}. + * @throws InvalidFormatException 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 IOException Thrown if reading the file fails. + */ + static Boolean topLevelBoolean(Path file, String field) throws IOException { + return topLevelField(file, field, cursor -> { + if (cursor.consumeLiteral("true")) { + return Boolean.TRUE; + } + if (cursor.consumeLiteral("false")) { + return Boolean.FALSE; + } + if (cursor.consumeLiteral(NULL_LITERAL)) { + return null; + } + throw cursor.malformed("Field '" + field + "' must be a boolean or null"); + }); + } + + /** + * Reads one top-level string field from a JSON object file. + * + * @param file The JSON file, a single top-level object. Must not be {@code null} and must + * exist. + * @param field The top-level field name to read. Must not be {@code null}. + * @return The field's value, or {@code null} when the field is absent or explicitly JSON + * {@code null} (the formats treat those the same: fall back to the default). + * @throws IllegalArgumentException Thrown if an argument is {@code null}. + * @throws InvalidFormatException Thrown if the file is not a well-formed JSON object, the + * field appears more than once, or its value is neither a string nor {@code null}. + * @throws IOException Thrown if reading the file fails. + */ + static String topLevelString(Path file, String field) throws IOException { + return topLevelField(file, field, cursor -> { + if (cursor.consumeLiteral(NULL_LITERAL)) { + return null; + } + if (cursor.peek() == '"') { + return cursor.parseString(); + } + throw cursor.malformed("Field '" + field + "' must be a string or null"); + }); + } + + /** + * Walks a JSON object file's top-level fields, skipping every field but {@code field} and + * handing that one's value to {@code valueReader}. + * + * @param file The JSON file, a single top-level object. Must not be {@code null} and + * must exist. + * @param field The top-level field name to read. Must not be {@code null}. + * @param valueReader Reads the matched field's value off the cursor. + * @param The value type the reader produces. + * @return The field's value, or {@code null} when the field is absent. + * @throws IllegalArgumentException Thrown if an argument is {@code null}. + * @throws InvalidFormatException Thrown if the file is not a well-formed JSON object or the + * field appears more than once. + * @throws IOException Thrown if reading the file fails. + */ + private static T topLevelField(Path file, String field, ValueReader valueReader) + throws IOException { + if (file == null) { + throw new IllegalArgumentException("File must not be null"); + } + if (field == null) { + throw new IllegalArgumentException("Field must not be null"); + } + final String json = Files.readString(file); + final JsonCursor cursor = new JsonCursor(json, file.getFileName().toString()); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + T value = null; + boolean seen = false; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if (field.equals(key)) { + if (seen) { + throw cursor.malformed("Field '" + field + "' appears more than once"); + } + seen = true; + value = valueReader.read(cursor); + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a field, got '" + next + "'"); + } + } + cursor.requireEnd("Trailing content after the top-level object"); + return value; + } + + /** + * Decodes the value of the looked-up field, positioned at its first character. + * + * @param The value type produced. + */ + @FunctionalInterface + private interface ValueReader { + + /** + * Reads one value off the cursor. + * + * @param cursor The cursor, positioned at the value's first character. + * @return The decoded value, or {@code null} for a JSON {@code null}. + * @throws InvalidFormatException Thrown if the value is malformed or not of the expected + * type. + */ + T read(JsonCursor cursor) throws InvalidFormatException; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FloatEmbeddingTable.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FloatEmbeddingTable.java new file mode 100644 index 0000000000..7b1f7f7658 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FloatEmbeddingTable.java @@ -0,0 +1,129 @@ +/* + * 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; + +/** + * The float {@link EmbeddingTable}: a flat row-major matrix, exactly the storage + * {@link StaticEmbeddingModel} always had. Its working space is original space, so query + * preparation and pooling finish are the identity, and per-row norms are precomputed once for + * the neighbor scan. + */ +final class FloatEmbeddingTable implements EmbeddingTable { + + private final float[] values; + private final int dimension; + private final int rowCount; + private final double[] rowNorms; + + /** + * Wraps a flat row-major matrix. The array is used as given; the loaders that construct this + * table own it exclusively. + * + * @param values The matrix, {@code rowCount * dimension} floats. + * @param dimension The row width. + * @param rowCount The number of rows. + */ + FloatEmbeddingTable(float[] values, int dimension, int rowCount) { + this.values = values; + this.dimension = dimension; + this.rowCount = rowCount; + this.rowNorms = new double[rowCount]; + for (int row = 0; row < rowCount; row++) { + final int base = row * dimension; + double sumOfSquares = 0; + for (int d = 0; d < dimension; d++) { + final float value = values[base + d]; + sumOfSquares += (double) value * value; + } + rowNorms[row] = Math.sqrt(sumOfSquares); + } + } + + /** {@inheritDoc} */ + @Override + public int rowCount() { + return rowCount; + } + + /** {@inheritDoc} */ + @Override + public int dimension() { + return dimension; + } + + /** {@inheritDoc} */ + @Override + public int pooledLength() { + return dimension; + } + + /** {@inheritDoc} */ + @Override + public void addRow(int row, float weight, float[] sum) { + final int base = row * dimension; + if (weight == 1f) { + for (int d = 0; d < dimension; d++) { + sum[d] += values[base + d]; + } + } else { + for (int d = 0; d < dimension; d++) { + sum[d] += values[base + d] * weight; + } + } + } + + /** {@inheritDoc} */ + @Override + public float[] finishPooling(float[] sum) { + return sum; + } + + /** {@inheritDoc} */ + @Override + public float[] prepareQuery(float[] query) { + return query; + } + + /** {@inheritDoc} */ + @Override + public double dot(int row, float[] preparedQuery) { + final int base = row * dimension; + // Four accumulators so the JIT can vectorize the dot product without reordering FP adds. + double dot0 = 0; + double dot1 = 0; + double dot2 = 0; + double dot3 = 0; + int d = 0; + for (final int limit = dimension - 3; d < limit; d += 4) { + dot0 += preparedQuery[d] * values[base + d]; + dot1 += preparedQuery[d + 1] * values[base + d + 1]; + dot2 += preparedQuery[d + 2] * values[base + d + 2]; + dot3 += preparedQuery[d + 3] * values[base + d + 3]; + } + double dot = dot0 + dot1 + dot2 + dot3; + for (; d < dimension; d++) { + dot += preparedQuery[d] * values[base + d]; + } + return dot; + } + + /** {@inheritDoc} */ + @Override + public double rowNorm(int row) { + return rowNorms[row]; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/GaussianQuantizer.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/GaussianQuantizer.java new file mode 100644 index 0000000000..1339d1376d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/GaussianQuantizer.java @@ -0,0 +1,232 @@ +/* + * 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.Arrays; + +/** + * An optimal scalar quantizer for standard-normal values: {@code 2^bits} representation levels + * minimizing the mean squared error over {@code N(0,1)}, with encoding by nearest level. The + * coordinates of a {@link HadamardRotation rotated} unit vector, scaled by the square root of the + * padded dimension, follow this distribution closely, which is what makes one fixed grid + * near-optimal for every coordinate of every vector (Zandieh et al., + * TurboQuant: Online Vector Quantization with + * Near-optimal Distortion Rate). + * + *

The levels are the classic Lloyd-Max quantizer of the Gaussian (Max, + * Quantizing for minimum + * distortion, IRE Transactions on Information Theory, 1960), computed here by Lloyd + * iteration over a fine discretization of the density rather than copied from published tables, + * so the derivation is in this file and reproducible. Computed grids are cached per bit width. + * Encoding compares against the midpoints between adjacent levels, which is exactly the + * nearest-level rule for a sorted grid.

+ * + *

A quantized file stores its grid, and reading rebuilds the quantizer from the stored levels + * through {@link #fromLevels(float[])}, so decoding never depends on this derivation matching the + * one that encoded the file.

+ * + *

Instances are immutable and safe for concurrent use.

+ */ +final class GaussianQuantizer { + + /** The smallest supported bit width. */ + static final int MIN_BITS = 2; + + /** The largest supported bit width. */ + static final int MAX_BITS = 4; + + // The density is discretized on [-RANGE, RANGE]; beyond eight standard deviations the + // remaining mass (~1e-15) is far below the iteration tolerance. + private static final double LLOYD_RANGE = 8.0; + private static final int LLOYD_SAMPLES = 200_001; + private static final double LLOYD_TOLERANCE = 1e-10; + private static final int LLOYD_MAX_ITERATIONS = 1_000; + + private static final GaussianQuantizer[] CACHE = new GaussianQuantizer[MAX_BITS + 1]; + + private final float[] levels; + // Midpoints between adjacent levels: level i is nearest exactly when the value lies in + // (thresholds[i-1], thresholds[i]], with the outermost intervals unbounded. + private final float[] thresholds; + + /** + * Holds a validated grid; callers reach this through {@link #forBits(int)} or + * {@link #fromLevels(float[])}. + * + * @param levels The representation levels, ascending. + */ + private GaussianQuantizer(float[] levels) { + this.levels = levels; + this.thresholds = new float[levels.length - 1]; + for (int i = 0; i < thresholds.length; i++) { + thresholds[i] = (levels[i] + levels[i + 1]) / 2f; + } + } + + /** + * {@return the quantizer for a bit width, computed once and cached} + * + * @param bits The bit width. Must be between {@link #MIN_BITS} and {@link #MAX_BITS}. + * @throws IllegalArgumentException Thrown if {@code bits} is outside the supported range. + */ + static GaussianQuantizer forBits(int bits) { + requireSupportedBits(bits); + synchronized (CACHE) { + if (CACHE[bits] == null) { + CACHE[bits] = new GaussianQuantizer(lloydMaxLevels(1 << bits)); + } + return CACHE[bits]; + } + } + + /** + * {@return a quantizer over a stored grid, as read back from a quantized file} + * + * @param levels The representation levels, strictly ascending and finite, of a power-of-two + * length between {@code 2^MIN_BITS} and {@code 2^MAX_BITS}. The array is copied. + * @throws IllegalArgumentException Thrown if {@code levels} is {@code null}, of an unsupported + * length, not strictly ascending, or not finite. + */ + static GaussianQuantizer fromLevels(float[] levels) { + if (levels == null) { + throw new IllegalArgumentException("Levels must not be null"); + } + if (levels.length != Integer.highestOneBit(levels.length) + || levels.length < 1 << MIN_BITS || levels.length > 1 << MAX_BITS) { + throw new IllegalArgumentException("Levels must have a power-of-two length between " + + (1 << MIN_BITS) + " and " + (1 << MAX_BITS) + ", got " + levels.length); + } + for (int i = 0; i < levels.length; i++) { + if (!Float.isFinite(levels[i])) { + throw new IllegalArgumentException("Level " + i + " is not finite: " + levels[i]); + } + if (i > 0 && levels[i] <= levels[i - 1]) { + throw new IllegalArgumentException("Levels must be strictly ascending, but level " + + i + " (" + levels[i] + ") is not above level " + (i - 1) + + " (" + levels[i - 1] + ")"); + } + } + return new GaussianQuantizer(Arrays.copyOf(levels, levels.length)); + } + + /** + * Requires a bit width within the supported range. + * + * @param bits The bit width to check. + * @throws IllegalArgumentException Thrown if {@code bits} is outside the supported range. + */ + static void requireSupportedBits(int bits) { + if (bits < MIN_BITS || bits > MAX_BITS) { + throw new IllegalArgumentException("Bits must be between " + MIN_BITS + " and " + + MAX_BITS + ", got " + bits); + } + } + + /** {@return the number of representation levels} */ + int levelCount() { + return levels.length; + } + + /** + * {@return the representation level of a code} + * + * @param code The code, between 0 and {@code levelCount() - 1}. + */ + float level(int code) { + return levels[code]; + } + + /** {@return a copy of the representation levels, ascending} */ + float[] levels() { + return Arrays.copyOf(levels, levels.length); + } + + /** + * {@return the code of the representation level nearest a value} Ties at a midpoint take the + * lower level, a fixed convention so encoding is deterministic. + * + * @param value The value to encode. + */ + int encode(float value) { + int low = 0; + int high = thresholds.length; + while (low < high) { + final int middle = (low + high) >>> 1; + if (value <= thresholds[middle]) { + high = middle; + } else { + low = middle + 1; + } + } + return low; + } + + /** + * {@return the Lloyd-Max representation levels for the standard normal} Lloyd iteration over a + * fine discretization of the density: assign each sample to its nearest level, move each level + * to the probability-weighted mean of its samples, repeat to a fixed point. The discretization, + * tolerance, and iteration cap are constants of this file, so the result is deterministic. + * + * @param levelCount The number of levels, a power of two. + */ + private static float[] lloydMaxLevels(int levelCount) { + final double step = 2 * LLOYD_RANGE / (LLOYD_SAMPLES - 1); + final double[] samples = new double[LLOYD_SAMPLES]; + final double[] weights = new double[LLOYD_SAMPLES]; + for (int i = 0; i < LLOYD_SAMPLES; i++) { + samples[i] = -LLOYD_RANGE + i * step; + weights[i] = Math.exp(-samples[i] * samples[i] / 2); + } + // Initial levels: evenly spaced over the central mass; Lloyd converges to the optimum for + // the log-concave Gaussian regardless of the starting spread. + final double[] levels = new double[levelCount]; + for (int i = 0; i < levelCount; i++) { + levels[i] = -3.0 + 6.0 * (i + 0.5) / levelCount; + } + final double[] weightSums = new double[levelCount]; + final double[] weightedValueSums = new double[levelCount]; + for (int iteration = 0; iteration < LLOYD_MAX_ITERATIONS; iteration++) { + Arrays.fill(weightSums, 0); + Arrays.fill(weightedValueSums, 0); + int level = 0; + for (int i = 0; i < LLOYD_SAMPLES; i++) { + while (level < levelCount - 1 + && samples[i] > (levels[level] + levels[level + 1]) / 2) { + level++; + } + weightSums[level] += weights[i]; + weightedValueSums[level] += weights[i] * samples[i]; + } + double largestMove = 0; + for (int i = 0; i < levelCount; i++) { + if (weightSums[i] > 0) { + final double moved = weightedValueSums[i] / weightSums[i]; + largestMove = Math.max(largestMove, Math.abs(moved - levels[i])); + levels[i] = moved; + } + } + if (largestMove < LLOYD_TOLERANCE) { + break; + } + } + final float[] result = new float[levelCount]; + for (int i = 0; i < levelCount; i++) { + result[i] = (float) levels[i]; + } + return result; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HadamardRotation.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HadamardRotation.java new file mode 100644 index 0000000000..9a4abd89a4 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HadamardRotation.java @@ -0,0 +1,188 @@ +/* + * 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; + +/** + * A seeded randomized Hadamard rotation: a deterministic random sign flip per coordinate followed + * by the normalized fast Walsh-Hadamard transform. Rotating a vector this way spreads its energy + * evenly across coordinates, so each coordinate of a rotated unit vector is approximately + * Gaussian with variance {@code 1/paddedDimension} and nearly independent of the others, which is + * the property {@link GaussianQuantizer}'s per-coordinate grids rely on. + * + *

The transform is orthonormal, so it preserves norms and inner products exactly (up to float + * rounding): two vectors rotated with the same instance have the same dot product as the + * originals, which lets similarity math stay in rotated space and never pay for the inverse. + * Writing {@code S} for the sign flip and {@code H} for the normalized Walsh-Hadamard matrix + * (which is its own inverse), the rotation is {@code H·S} and its inverse is {@code S·H}: the + * same two operations applied in the opposite order, so no second table is needed.

+ * + *

The Walsh-Hadamard transform needs a power-of-two length, so vectors are padded with zeros + * from their original dimension up to {@link #paddedDimension(int)}. The sign flips derive from + * the seed through an in-file + * splitmix64 step, not through a JDK + * generator, so the same seed produces the same rotation on every JVM and release; the seed is + * stored in the quantized file and the rotation is rebuilt from it on load.

+ * + *

Instances are immutable and safe for concurrent use.

+ */ +final class HadamardRotation { + + private static final long SPLITMIX64_GOLDEN_GAMMA = 0x9E3779B97F4A7C15L; + + private final int paddedDimension; + // True where the coordinate is negated before (rotate) or after (inverse) the transform. + private final boolean[] flip; + private final float inverseSquareRoot; + + /** + * Creates the rotation for vectors of the given original dimension. + * + * @param dimension The original vector dimension. Must be at least 1. + * @param seed The seed the sign flips derive from. + * @throws IllegalArgumentException Thrown if {@code dimension} is less than 1. + */ + HadamardRotation(int dimension, long seed) { + if (dimension < 1) { + throw new IllegalArgumentException("Dimension must be at least 1, got " + dimension); + } + this.paddedDimension = paddedDimension(dimension); + this.flip = new boolean[paddedDimension]; + long state = seed; + long bits = 0; + for (int i = 0; i < paddedDimension; i++) { + if ((i & 63) == 0) { + state += SPLITMIX64_GOLDEN_GAMMA; + bits = splitmix64(state); + } + flip[i] = (bits & 1L) != 0; + bits >>>= 1; + } + this.inverseSquareRoot = (float) (1.0 / Math.sqrt(paddedDimension)); + } + + /** + * {@return the power-of-two length vectors are padded to before the transform} + * + * @param dimension The original vector dimension. Must be at least 1. + * @throws IllegalArgumentException Thrown if {@code dimension} is less than 1. + */ + static int paddedDimension(int dimension) { + if (dimension < 1) { + throw new IllegalArgumentException("Dimension must be at least 1, got " + dimension); + } + if (dimension > 1 << 30) { + throw new IllegalArgumentException("Dimension must be at most " + (1 << 30) + + " so the padded length stays an int power of two, got " + dimension); + } + final int highestOneBit = Integer.highestOneBit(dimension); + return highestOneBit == dimension ? dimension : highestOneBit << 1; + } + + /** {@return the power-of-two length this instance transforms} */ + int paddedDimension() { + return paddedDimension; + } + + /** + * Rotates a vector in place: sign flips, then the normalized Walsh-Hadamard transform. + * + * @param vector The vector to rotate. Must not be {@code null} and must have length + * {@link #paddedDimension()}. + * @throws IllegalArgumentException Thrown if {@code vector} is {@code null} or has the wrong + * length. + */ + void rotate(float[] vector) { + requirePaddedLength(vector); + for (int i = 0; i < paddedDimension; i++) { + if (flip[i]) { + vector[i] = -vector[i]; + } + } + walshHadamard(vector); + } + + /** + * Applies the inverse rotation in place: the normalized Walsh-Hadamard transform, then the + * sign flips. + * + * @param vector The rotated vector. Must not be {@code null} and must have length + * {@link #paddedDimension()}. + * @throws IllegalArgumentException Thrown if {@code vector} is {@code null} or has the wrong + * length. + */ + void inverse(float[] vector) { + requirePaddedLength(vector); + walshHadamard(vector); + for (int i = 0; i < paddedDimension; i++) { + if (flip[i]) { + vector[i] = -vector[i]; + } + } + } + + /** + * Requires the vector to be non-null and of the padded length. + * + * @param vector The vector to check. + * @throws IllegalArgumentException Thrown if {@code vector} is {@code null} or has the wrong + * length. + */ + private void requirePaddedLength(float[] vector) { + if (vector == null) { + throw new IllegalArgumentException("Vector must not be null"); + } + if (vector.length != paddedDimension) { + throw new IllegalArgumentException("Vector has length " + vector.length + + " but this rotation transforms length " + paddedDimension); + } + } + + /** + * The in-place normalized fast Walsh-Hadamard transform, {@code O(n log n)} butterflies + * followed by a {@code 1/sqrt(n)} scale so the transform is orthonormal and self-inverse. + * + * @param vector The vector to transform, of the padded length. + */ + private void walshHadamard(float[] vector) { + for (int half = 1; half < paddedDimension; half <<= 1) { + for (int block = 0; block < paddedDimension; block += half << 1) { + for (int i = block; i < block + half; i++) { + final float a = vector[i]; + final float b = vector[i + half]; + vector[i] = a + b; + vector[i + half] = a - b; + } + } + } + for (int i = 0; i < paddedDimension; i++) { + vector[i] *= inverseSquareRoot; + } + } + + /** + * {@return the splitmix64 mix of a state word} The finalizer of the splitmix64 generator, + * reproduced here so the bit stream is fixed by this file rather than by a JDK class. + * + * @param state The state word to mix. + */ + private static long splitmix64(long state) { + long z = state; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + return z ^ (z >>> 31); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HuggingFaceModelCache.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HuggingFaceModelCache.java new file mode 100644 index 0000000000..65cfd6cfc3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HuggingFaceModelCache.java @@ -0,0 +1,688 @@ +/* + * 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.InputStream; +import java.net.ProxySelector; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.security.DigestInputStream; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HexFormat; +import java.util.List; +import java.util.Optional; + +/** + * Fetches the files a distillation needs from a Hugging Face model repository into a local cache + * directory, so a teacher can be named by its hub id ({@code org/model}) instead of a local path. + * + *

A download is pinned and verified. The teacher's ref is resolved to a commit sha once, every + * file is then requested at that sha, and every file is checked against the digest the hub + * publishes for it before it is published into the cache directory. A file the hub does not + * publish a digest for is refused rather than used, the same posture + * {@link opennlp.tools.util.DownloadUtil} takes for a model whose {@code .sha512} sidecar cannot + * be read. A file the repository does not have (a 404, e.g. a WordPiece teacher's + * {@code sentencepiece.bpe.model}) is reported as absent, not an error, but only when that file is + * optional.

+ * + *

The resolved commit sha is recorded in {@value #REVISION_FILE} in the cache directory, which + * also marks the directory as a complete snapshot of that one revision: while it is there, the + * files are reused without a request to the hub and without being digested again. The record is + * written only once every file has been verified, and re-reading a multi-gigabyte ONNX graph on + * every distillation would only guard against something that could rewrite the record just as + * easily. A directory without the record is not trusted: its files are checked against the + * revision now being downloaded before they are kept. A record that does not describe the + * directory, because a file it vouches for is gone or because it names a different commit than the + * reference asks for, is removed before anything is downloaded, so that a download failing halfway + * cannot leave a directory the next attempt would trust on the strength of it.

+ */ +final class HuggingFaceModelCache { + + /** The hub's base URL, the prefix of every download URL. */ + private static final String HUB_BASE = "https://huggingface.co/"; + + /** The hub's download path between the model id and the revision. */ + private static final String RESOLVE_PATH = "/resolve/"; + + /** The revision downloaded when a teacher does not name one: the repository's default branch. */ + private static final String DEFAULT_REVISION = "main"; + + /** + * A parsed teacher reference: an organization and a model name joined by {@code /}, optionally + * followed by {@code @} and the revision to pin; see {@link #parseTeacherReference(String)}. + * + * @param modelId The {@code org/model} id. + * @param revision The pinned revision, or {@code null} when none is given. + */ + private record TeacherReference(String modelId, String revision) { + } + + /** The directory the cache lives in, below the user's home directory. */ + private static final String CACHE_DIRECTORY = ".cache"; + + /** The cache's own directory, below {@link #CACHE_DIRECTORY}. */ + private static final String CACHE_NAME = "opennlp-embeddings"; + + /** The hex length of the digest suffix that makes a cache directory name injective. */ + private static final int CACHE_KEY_HEX_LENGTH = 16; + + /** + * The file recording the commit sha the cache directory holds. It is written only after every + * file of that revision has been downloaded and verified, so its presence means the directory is + * complete. The name starts with a dot so that it cannot collide with a repository file. + */ + static final String REVISION_FILE = ".opennlp-revision"; + + /** The suffix of the temporary file a download streams into before it is moved into place. */ + private static final String DOWNLOAD_SUFFIX = ".download"; + + /** The response header holding the commit sha a ref resolved to. */ + private static final String COMMIT_HEADER = "x-repo-commit"; + + /** The response header holding the digest of the file, quoted. */ + private static final String ETAG_HEADER = "x-linked-etag"; + + /** The length in hex of a SHA-1: the shape of a commit sha and of a git object name. */ + private static final int SHA1_HEX_LENGTH = 40; + + /** The length in hex of a SHA-256: the shape of the digest published for a Git LFS file. */ + private static final int SHA256_HEX_LENGTH = 64; + + + /** The header git hashes in front of a blob's bytes, completed by the length and a NUL byte. */ + private static final String GIT_BLOB_PREFIX = "blob "; + + /** The read size when digesting a downloaded file. */ + private static final int DIGEST_BUFFER_SIZE = 8192; + + /** The HTTP status a served file answers with. */ + private static final int HTTP_OK = 200; + + /** The HTTP status of a file the repository does not have at the requested revision. */ + private static final int HTTP_NOT_FOUND = 404; + + /** How long the client waits for a connection to the hub. */ + private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(30); + + /** How long a single file download may take; an ONNX graph can be gigabytes. */ + private static final Duration DOWNLOAD_TIMEOUT = Duration.ofHours(1); + + /** The files a distillation needs, relative to the repository root. */ + private static final List REQUIRED_FILES = + List.of(ModelFileNames.TOKENIZER_JSON, ModelFileNames.ONNX_MODEL); + + /** + * The files used when present: the pad-token config, the trained SentencePiece model under any + * of the names a repository may ship it as, and the external weights of an ONNX export that + * splits them out (as bge-m3 does). + */ + private static final List OPTIONAL_FILES = optionalFiles(); + + /** Not instantiable. */ + private HuggingFaceModelCache() { + } + + /** {@return the repository-relative names of the files downloaded when the repository has them} */ + private static List optionalFiles() { + final List files = new ArrayList<>(); + files.add(ModelFileNames.TOKENIZER_CONFIG); + files.addAll(ModelFileNames.SENTENCEPIECE_MODELS); + files.add(ModelFileNames.ONNX_MODEL_DATA); + return List.copyOf(files); + } + + /** + * Resolves a teacher reference to a local directory holding its files, downloading them from the + * Hugging Face hub when the reference is a model id. + * + * @param teacher A local directory, used as-is, or a Hugging Face model id ({@code org/model}, + * or {@code org/model@revision} to pin a branch, tag, or commit sha instead of + * the default branch), downloaded into + * {@code ~/.cache/opennlp-embeddings/org-model} on first use (the slash becomes + * a dash, dots and the revision separator become underscores). Must not be + * {@code null}. + * @param listener Receives one progress line per download; may be {@code null}. + * @return The local teacher directory. + * @throws IllegalArgumentException Thrown if {@code teacher} is {@code null}, or is neither a + * directory nor a well-formed model id. + * @throws IOException Thrown if a required file cannot be downloaded, or if a downloaded file + * cannot be verified against the digest the hub publishes for it. + */ + static Path resolve(String teacher, ModelDistiller.ProgressListener listener) throws IOException { + return resolve(teacher, HUB_BASE, defaultCacheRoot(), listener); + } + + /** + * Resolves a teacher reference against a given hub and cache location, the form used by the + * tests and by an installation that mirrors the hub. + * + * @param teacher The teacher reference, as in {@link #resolve(String, + * ModelDistiller.ProgressListener)}. Must not be {@code null}. + * @param hubBase The hub's base URL, ending in a slash. Must not be {@code null}. + * @param cacheRoot The directory the per-teacher cache directories live in. Must not be + * {@code null}. + * @param listener Receives one progress line per download; may be {@code null}. + * @return The local teacher directory. + * @throws IllegalArgumentException Thrown if {@code teacher}, {@code hubBase}, or + * {@code cacheRoot} is {@code null}, or if {@code teacher} is neither a directory nor a + * well-formed model id. + * @throws IOException Thrown if a required file cannot be downloaded, or if a downloaded file + * cannot be verified against the digest the hub publishes for it. + */ + static Path resolve(String teacher, String hubBase, Path cacheRoot, + ModelDistiller.ProgressListener listener) throws IOException { + if (teacher == null) { + throw new IllegalArgumentException("Teacher must not be null"); + } + if (hubBase == null) { + throw new IllegalArgumentException("HubBase must not be null"); + } + if (cacheRoot == null) { + throw new IllegalArgumentException("CacheRoot must not be null"); + } + final Path local = Path.of(teacher); + if (Files.isDirectory(local)) { + return local; + } + final TeacherReference reference = parseTeacherReference(teacher); + if (reference == null) { + throw new IllegalArgumentException("Teacher '" + teacher + "' is neither a local " + + "directory nor a Hugging Face model id (expected 'org/model' or 'org/model@revision')"); + } + final String modelId = reference.modelId(); + final String requestedRevision = reference.revision(); + final Path cache = cacheRoot.resolve(cacheDirectoryName(teacher)); + final String pinned = pinnedRevision(cache); + if (pinned != null && hasRequiredFiles(cache) + && (!isCommitSha(requestedRevision) || pinned.equalsIgnoreCase(requestedRevision))) { + return cache; + } + // The directory is not a complete snapshot of a revision this reference names, so the record + // it carries does not describe it either. The record goes before the first file is fetched: + // a download that then fails verification leaves a directory the next attempt distrusts and + // verifies file by file, rather than one it would hand out whole on the strength of a record + // written for an earlier revision. + Files.deleteIfExists(cache.resolve(REVISION_FILE)); + final String ref = requestedRevision == null ? DEFAULT_REVISION : requestedRevision; + // A client built through the builder has no proxy selector unless one is set, so the + // http.proxyHost / https.proxyHost system properties would otherwise be ignored. + final HttpClient client = HttpClient.newBuilder() + .followRedirects(HttpClient.Redirect.NORMAL) + .proxy(ProxySelector.getDefault()) + .connectTimeout(CONNECT_TIMEOUT) + .build(); + final String commit = resolveCommit(client, hubBase, modelId, ref, requestedRevision); + report(listener, "Teacher " + modelId + " at " + ref + " is commit " + commit); + for (final String file : REQUIRED_FILES) { + download(client, hubBase, modelId, commit, file, cache, true, listener); + } + for (final String file : OPTIONAL_FILES) { + download(client, hubBase, modelId, commit, file, cache, false, listener); + } + Files.writeString(cache.resolve(REVISION_FILE), commit + System.lineSeparator(), + StandardCharsets.UTF_8); + return cache; + } + + /** + * {@return the commit sha a cache directory was downloaded at, or {@code null} when the + * directory is not a complete cached snapshot of a hub revision} + * + * @param teacherDirectory The directory to read; need not exist. + */ + static String pinnedRevision(Path teacherDirectory) { + final Path file = teacherDirectory.resolve(REVISION_FILE); + if (!Files.isRegularFile(file)) { + return null; + } + try { + final String recorded = Files.readString(file, StandardCharsets.UTF_8).trim(); + return isCommitSha(recorded) ? recorded : null; + } catch (IOException e) { + return null; + } + } + + /** {@return the directory the per-teacher cache directories live in} */ + private static Path defaultCacheRoot() { + return Path.of(System.getProperty("user.home"), CACHE_DIRECTORY, CACHE_NAME); + } + + /** + * {@return the cache directory name for a teacher reference} + * + *

The readable part replaces the characters a path cannot carry, which alone is not + * injective: {@code acme/model@v1}, {@code acme/model.v1} and {@code acme/model_v1} would all + * name one directory, and the cached fast path answers from that directory without contacting + * the hub, so one teacher would be served another's files. The suffix is a digest of the exact + * reference, so distinct references never share a directory.

+ * + * @param teacher The teacher reference, as the caller wrote it. + */ + static String cacheDirectoryName(String teacher) { + final String readable = teacher.replace('/', '-').replace('.', '_').replace('@', '_'); + final MessageDigest digest; + try { + digest = MessageDigest.getInstance("SHA-256"); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 is required of every JVM", e); + } + final byte[] hash = digest.digest(teacher.getBytes(StandardCharsets.UTF_8)); + final String suffix = HexFormat.of().formatHex(hash, 0, CACHE_KEY_HEX_LENGTH / 2); + return readable + '-' + suffix; + } + + /** + * {@return whether every file a distillation needs is in the cache directory} + * + * @param cache The cache directory. + */ + private static boolean hasRequiredFiles(Path cache) { + for (final String file : REQUIRED_FILES) { + if (!Files.isRegularFile(cache.resolve(file))) { + return false; + } + } + return true; + } + + /** + * Resolves a ref to the commit sha it points at, so that the files of one download all come from + * one revision even if the ref moves while the download runs. The hub reports the sha on every + * resolve response, so the body of the probed file is not read. + * + * @param client The HTTP client. + * @param hubBase The hub's base URL. + * @param modelId The hub model id. + * @param ref The revision to resolve. + * @param requestedRevision The revision the teacher reference named, or {@code null} when it + * named none. + * @return The commit sha, 40 hex characters. + * @throws IOException Thrown if the ref cannot be resolved. + */ + private static String resolveCommit(HttpClient client, String hubBase, String modelId, String ref, + String requestedRevision) throws IOException { + final String probe = REQUIRED_FILES.get(0); + final HttpResponse response = send(client, hubBase, modelId, ref, probe); + // The headers carry everything this probe wants, so the body is closed unread. + final InputStream body = response.body(); + try (body) { + if (response.statusCode() != HTTP_OK) { + throw new IOException("Failed to resolve revision '" + ref + "' of " + modelId + ": HTTP " + + response.statusCode() + " for " + probe); + } + final String commit = originHeader(response, COMMIT_HEADER); + if (!isCommitSha(commit)) { + throw new IOException("Revision '" + ref + "' of " + modelId + " could not be pinned: the " + + "hub sent " + (commit == null ? "no " + COMMIT_HEADER + " header" + : COMMIT_HEADER + " '" + commit + "', which is not a commit sha") + + "; refusing to download files that cannot be attributed to one revision"); + } + if (isCommitSha(requestedRevision) && !commit.equalsIgnoreCase(requestedRevision)) { + throw new IOException("Revision '" + requestedRevision + "' of " + modelId + " resolved to " + + "commit " + commit + " instead"); + } + return commit; + } + } + + /** + * Downloads one repository file at a pinned revision into the cache, keeping a copy that is + * already there when it matches the revision's digest. + * + * @param client The HTTP client. + * @param hubBase The hub's base URL. + * @param modelId The hub model id. + * @param commit The commit sha every file of this download is requested at. + * @param file The repository-relative file name. + * @param cache The cache directory. + * @param required Whether a file the revision does not have is an error. + * @param listener The progress listener; may be {@code null}. + * @throws IOException Thrown if a required file cannot be downloaded, or if the download cannot + * be verified against the digest the hub publishes for it. + */ + private static void download(HttpClient client, String hubBase, String modelId, String commit, + String file, Path cache, boolean required, + ModelDistiller.ProgressListener listener) throws IOException { + final Path target = cache.resolve(file); + final HttpResponse response = send(client, hubBase, modelId, commit, file); + Path temporary = null; + try (InputStream body = response.body()) { + if (response.statusCode() == HTTP_NOT_FOUND && !required) { + // The cache directory holds one revision: a copy left by an earlier one has to go. + Files.deleteIfExists(target); + return; + } + if (response.statusCode() != HTTP_OK) { + throw new IOException("Failed to download " + file + " of " + modelId + " at commit " + + commit + ": HTTP " + response.statusCode() + + (required ? "; the distillation needs this file" : "")); + } + final Digest expected = expectedDigest(response, modelId, file); + if (Files.isRegularFile(target) && expected.matches(target)) { + return; + } + report(listener, "Downloading " + modelId + "/" + file + " ..."); + Files.createDirectories(target.getParent()); + // A temporary name unique per download: two processes sharing one cache directory must not + // stream two copies of the same file into one partial file and publish the interleaving. + temporary = Files.createTempFile(target.getParent(), target.getFileName().toString(), + DOWNLOAD_SUFFIX); + Files.copy(body, temporary, StandardCopyOption.REPLACE_EXISTING); + final String actual = expected.form().hexOf(temporary); + if (!expected.hex().equalsIgnoreCase(actual)) { + throw new IOException(expected.form().displayName() + " checksum validation failed for " + + file + " of " + modelId + " at commit " + commit + ". Expected: " + expected.hex() + + ", but got: " + actual); + } + Files.move(temporary, target, StandardCopyOption.REPLACE_EXISTING); + temporary = null; + } finally { + deleteIfPresent(temporary); + } + } + + /** + * Sends one GET to the hub. + * + * @param client The HTTP client. + * @param hubBase The hub's base URL. + * @param modelId The hub model id. + * @param revision The revision to request the file at. + * @param file The repository-relative file name. + * @return The response, whose body has not been read yet. + * @throws IOException Thrown if the request fails. + */ + private static HttpResponse send(HttpClient client, String hubBase, String modelId, + String revision, String file) throws IOException { + final HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create(hubBase + modelId + RESOLVE_PATH + revision + "/" + file)) + .timeout(DOWNLOAD_TIMEOUT) + .GET() + .build(); + try { + return client.send(request, HttpResponse.BodyHandlers.ofInputStream()); + } catch (IOException e) { + throw new IOException("Failed to download " + file + " of " + modelId + ": " + + e.getMessage(), e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while downloading " + file + " of " + modelId, e); + } + } + + /** + * Reads the digest the hub publishes for a file, refusing a file that carries none. + * + * @param response The response. + * @param modelId The hub model id, for the message. + * @param file The repository-relative file name, for the message. + * @return The expected digest, either a git blob SHA-1 or a SHA-256. + * @throws IOException Thrown if the header is absent or is not one of the two digest forms. + */ + private static Digest expectedDigest(HttpResponse response, String modelId, + String file) throws IOException { + final String header = originHeader(response, ETAG_HEADER); + if (header == null) { + throw new IOException("Expected checksum could not be retrieved for " + file + " of " + + modelId + ": the hub sent no " + ETAG_HEADER + " header; refusing to use a file that " + + "cannot be verified"); + } + final String hex = header.trim().replace("\"", ""); + final Checksum form = Checksum.of(hex); + if (form == null) { + throw new IOException("Expected checksum could not be retrieved for " + file + " of " + + modelId + ": " + ETAG_HEADER + " '" + header + "' is neither a git blob SHA-1 nor a " + + "SHA-256; refusing to use a file that cannot be verified"); + } + return new Digest(form, hex); + } + + /** + * {@return the value the hub sent for a header, or {@code null} when it sent none} + * + *

A resolve request answers with a redirect to a content delivery network, and the client + * does not copy the headers of that redirecting response onto the response it finally returns, + * so the redirect chain is walked back to its start. The earliest response wins: the digest to + * verify against is the one the hub itself stated, not one a later hop restated.

+ * + * @param response The response, at the end of its redirect chain. + * @param name The header name. + */ + private static String originHeader(HttpResponse response, String name) { + final List> chain = new ArrayList<>(); + for (HttpResponse hop = response; hop != null; + hop = hop.previousResponse().orElse(null)) { + chain.add(hop); + } + for (int i = chain.size() - 1; i >= 0; i--) { + final Optional value = chain.get(i).headers().firstValue(name); + if (value.isPresent()) { + return value.get(); + } + } + return null; + } + + /** + * {@return whether a value is a commit sha, 40 hex characters} + * + * @param value The value to check; may be {@code null}. + */ + private static boolean isCommitSha(String value) { + return value != null && value.length() == SHA1_HEX_LENGTH && isHex(value); + } + + /** + * {@return whether a value is one or more ASCII hex characters, the shape both the commit sha + * and the digests have} + * + * @param value The value to check. + */ + private static boolean isHex(String value) { + if (value.isEmpty()) { + return false; + } + for (int i = 0; i < value.length(); i++) { + final char c = value.charAt(i); + final boolean hex = (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') + || (c >= 'A' && c <= 'F'); + if (!hex) { + return false; + } + } + return true; + } + + /** + * Parses a teacher reference: an organization and a model name, both runs of ASCII word + * characters, dots, or dashes, joined by {@code /} and optionally followed by {@code @} and a + * revision of the same shape. + * + * @param teacher The reference to parse. + * @return The parsed reference, or {@code null} when the value does not have this form. + */ + private static TeacherReference parseTeacherReference(String teacher) { + final int slash = teacher.indexOf('/'); + if (slash < 0) { + return null; + } + final int at = teacher.indexOf('@', slash + 1); + final String organization = teacher.substring(0, slash); + final String model = at < 0 ? teacher.substring(slash + 1) : teacher.substring(slash + 1, at); + final String revision = at < 0 ? null : teacher.substring(at + 1); + if (!isReferencePart(organization) || !isReferencePart(model) + || (revision != null && !isReferencePart(revision))) { + return null; + } + return new TeacherReference(organization + "/" + model, revision); + } + + /** + * {@return whether a reference part is one or more ASCII word characters, dots, or dashes} + * + * @param part The part to check. + */ + private static boolean isReferencePart(String part) { + if (part.isEmpty()) { + return false; + } + for (int i = 0; i < part.length(); i++) { + final char c = part.charAt(i); + final boolean allowed = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') + || (c >= '0' && c <= '9') || c == '_' || c == '.' || c == '-'; + if (!allowed) { + return false; + } + } + return true; + } + + /** + * Reports one progress line, if anyone is listening. + * + * @param listener The listener; may be {@code null}. + * @param message The message. + */ + private static void report(ModelDistiller.ProgressListener listener, String message) { + if (listener != null) { + listener.progress(message); + } + } + + /** + * Deletes a partial download, if there is one, without reporting a failure to do so. + * + * @param file The file to delete; may be {@code null}. + */ + private static void deleteIfPresent(Path file) { + if (file == null) { + return; + } + try { + Files.deleteIfExists(file); + } catch (IOException e) { + // A leftover partial download costs disk space; the next attempt writes a fresh file. + } + } + + /** + * The digest the hub published for one file. + * + * @param form The digest form the hub stated it in. + * @param hex The digest value in hex, without the quotes the header carries. + */ + private record Digest(Checksum form, String hex) { + + /** + * {@return whether a file digests to the value the hub published} + * + * @param file The file to digest. + * @throws IOException Thrown if the file cannot be read. + */ + boolean matches(Path file) throws IOException { + return hex.equalsIgnoreCase(form.hexOf(file)); + } + } + + /** + * The two digest forms the hub publishes in its {@code x-linked-etag} header, told apart by the + * length of the hex value. + */ + private enum Checksum { + + /** The git object name of a file stored in git itself: its bytes behind a blob header. */ + GIT_BLOB_SHA1("git blob SHA-1", "SHA-1", SHA1_HEX_LENGTH), + + /** The digest of a file stored in Git LFS: its bytes alone. */ + LFS_SHA256("SHA-256", "SHA-256", SHA256_HEX_LENGTH); + + private final String displayName; + private final String algorithm; + private final int hexLength; + + /** + * Creates a digest form. + * + * @param displayName The name used in error messages. + * @param algorithm The {@link java.security.MessageDigest} algorithm name. + * @param hexLength The length of the digest's hex form. + */ + Checksum(String displayName, String algorithm, int hexLength) { + this.displayName = displayName; + this.algorithm = algorithm; + this.hexLength = hexLength; + } + + /** + * {@return the digest form a hex value of this length is, or {@code null} when the value is + * not a hex string of either length} + * + * @param value The digest value, without its quotes. Must not be {@code null}. + */ + static Checksum of(String value) { + for (final Checksum checksum : values()) { + if (value.length() == checksum.hexLength && isHex(value)) { + return checksum; + } + } + return null; + } + + /** {@return the name of this digest form, for a message} */ + String displayName() { + return displayName; + } + + /** + * {@return the digest of a file in this form, in lower case hex} + * + * @param file The file to digest. + * @throws IOException Thrown if the file cannot be read. + */ + String hexOf(Path file) throws IOException { + final MessageDigest digest; + try { + digest = MessageDigest.getInstance(algorithm); + } catch (NoSuchAlgorithmException e) { + throw new IOException(algorithm + " is not available", e); + } + if (this == GIT_BLOB_SHA1) { + digest.update((GIT_BLOB_PREFIX + Files.size(file) + '\0') + .getBytes(StandardCharsets.US_ASCII)); + } + try (InputStream in = Files.newInputStream(file); + DigestInputStream digesting = new DigestInputStream(in, digest)) { + final byte[] buffer = new byte[DIGEST_BUFFER_SIZE]; + while (digesting.read(buffer) != -1) { + // Reading the file is what updates the digest. + } + } + return HexFormat.of().formatHex(digest.digest()); + } + } +} 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 new file mode 100644 index 0000000000..1300fcd73d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java @@ -0,0 +1,357 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +/** + * Cursor primitives shared by this package's purpose-built JSON readers + * ({@link SafetensorsHeaderParser}, {@link FlatJsonFields}, {@link TokenizerJsonVocab}, + * {@link ModelAssembler}): string and integer scalars, literals, whitespace, and skipping one + * value of any type. Deliberately not a general JSON + * library: no floating-point decoding, no document model; each reader drives the cursor over + * its own known-shape input and fails loud on anything else, with the input's name and the + * offending offset in every message. Malformed input is a checked + * {@link InvalidFormatException}, the exception model content errors carry throughout this + * package. + */ +final class JsonCursor { + + private final String text; + private final String inputName; + 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. + */ + JsonCursor(String text, String inputName) { + this.text = text; + this.inputName = inputName; + } + + /** Advances the cursor past any run of whitespace. */ + void skipWhitespace() { + while (position < text.length() && Character.isWhitespace(text.charAt(position))) { + position++; + } + } + + /** {@return the cursor's current offset into the text, for readers that capture raw spans} */ + int position() { + return position; + } + + /** + * {@return the character at the cursor without advancing} + * + * @throws InvalidFormatException Thrown if the cursor is at the end of the input. + */ + char peek() throws InvalidFormatException { + if (position >= text.length()) { + throw malformed("Unexpected end of input"); + } + return text.charAt(position); + } + + /** + * {@return the character at the cursor, advancing past it} + * + * @throws InvalidFormatException Thrown if the cursor is at the end of the input. + */ + char consume() throws InvalidFormatException { + final char c = peek(); + position++; + return c; + } + + /** + * Consumes the next character, requiring it to be {@code c}. + * + * @param c The expected character. + * @throws InvalidFormatException Thrown if the next character is not {@code c}. + */ + void expect(char c) throws InvalidFormatException { + final char actual = consume(); + if (actual != c) { + throw malformed("Expected '" + c + "', got '" + actual + "'"); + } + } + + /** + * Consumes the given literal (for example {@code "true"}) when it starts at the cursor, + * leaving the cursor untouched when it does not. + * + * @param literal The literal to match. + * @return {@code true} when the literal was consumed. + */ + boolean consumeLiteral(String literal) { + if (text.startsWith(literal, position)) { + position += literal.length(); + return true; + } + return false; + } + + /** + * Requires the rest of the input to be whitespace only. + * + * @param message What to report when other content follows. + * @throws InvalidFormatException Thrown if non-whitespace content follows the cursor. + */ + void requireEnd(String message) throws InvalidFormatException { + skipWhitespace(); + if (position < text.length()) { + throw malformed(message); + } + } + + /** + * {@return the JSON string starting at the cursor, with escapes decoded} + * + * @throws InvalidFormatException Thrown if the string is unterminated or has a bad escape. + */ + String parseString() throws InvalidFormatException { + expect('"'); + final StringBuilder value = new StringBuilder(); + while (true) { + if (position >= text.length()) { + throw malformed("Unterminated string"); + } + final char c = text.charAt(position++); + if (c == '"') { + return value.toString(); + } + if (c == '\\') { + value.append(parseEscape()); + } else { + value.append(c); + } + } + } + + /** {@return the character named by the escape sequence following a backslash} */ + private char parseEscape() throws InvalidFormatException { + if (position >= text.length()) { + throw malformed("Unterminated escape sequence"); + } + final char escape = text.charAt(position++); + return switch (escape) { + case '"' -> '"'; + case '\\' -> '\\'; + case '/' -> '/'; + case 'b' -> '\b'; + case 'f' -> '\f'; + case 'n' -> '\n'; + case 'r' -> '\r'; + case 't' -> '\t'; + case 'u' -> parseUnicodeEscape(); + default -> throw malformed("Unknown escape sequence: \\" + escape); + }; + } + + /** {@return the character named by a {@code \\uXXXX} escape} */ + private char parseUnicodeEscape() throws InvalidFormatException { + if (position + 4 > text.length()) { + throw malformed("Truncated \\u escape sequence"); + } + final String hex = text.substring(position, position + 4); + position += 4; + // 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; + } + return (char) value; + } + + /** + * Skips one JSON number, holding it to the grammar (optional minus, digits, optional fraction, + * optional signed exponent) so malformed input fails loud even in a skipped field. + */ + private void skipNumber() throws InvalidFormatException { + 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++; + } + } + } + + /** + * {@return the integer starting at the cursor, parsed as a {@code long}} + * + * @throws InvalidFormatException Thrown if no integer is present or it overflows a long. + */ + long parseLong() throws InvalidFormatException { + final int start = position; + if (peek() == '-') { + position++; + } + if (position >= text.length() || !Character.isDigit(text.charAt(position))) { + throw malformed("Expected an integer"); + } + while (position < text.length() && Character.isDigit(text.charAt(position))) { + position++; + } + try { + return Long.parseLong(text.substring(start, position)); + } catch (NumberFormatException e) { + throw malformed("Malformed integer: " + text.substring(start, position)); + } + } + + /** + * {@return the finite JSON number starting at the cursor, parsed as a {@code double}} + * + * @throws InvalidFormatException Thrown if no JSON number is present or its value is not + * finite. + */ + double parseDouble() throws InvalidFormatException { + final int start = position; + skipNumber(); + final String number = text.substring(start, position); + try { + final double value = Double.parseDouble(number); + if (!Double.isFinite(value)) { + throw malformed("Number is not finite: " + number); + } + return value; + } catch (NumberFormatException e) { + throw malformed("Malformed number: " + number); + } + } + + /** + * {@return the JSON boolean starting at the cursor} + * + * @throws InvalidFormatException Thrown if the next value is not {@code true} or + * {@code false}. + */ + boolean parseBoolean() throws InvalidFormatException { + if (consumeLiteral("true")) { + return true; + } + if (consumeLiteral("false")) { + return false; + } + throw malformed("Expected a boolean"); + } + + /** + * Skips one JSON value of any type (string, number, array, object, true/false/null), so a + * reader tolerates fields it does not care about. + */ + void skipValue() throws InvalidFormatException { + skipWhitespace(); + final char c = peek(); + if (c == '"') { + parseString(); + } else if (c == '[') { + position++; + skipWhitespace(); + if (peek() != ']') { + while (true) { + skipValue(); + skipWhitespace(); + final char next = consume(); + if (next == ',') { + skipWhitespace(); + continue; + } + if (next == ']') { + return; + } + throw malformed("Expected ',' or ']' while skipping an array, got '" + next + "'"); + } + } + position++; + } else if (c == '{') { + position++; + skipWhitespace(); + if (peek() != '}') { + while (true) { + skipWhitespace(); + parseString(); + skipWhitespace(); + expect(':'); + skipValue(); + skipWhitespace(); + final char next = consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return; + } + throw malformed("Expected ',' or '}' while skipping an object, got '" + next + "'"); + } + } + position++; + } else if (c == '-' || Character.isDigit(c)) { + skipNumber(); + } else if (consumeLiteral("true") || consumeLiteral("false") || consumeLiteral("null")) { + // consumed, nothing to record + } else { + throw malformed("Unexpected character while skipping a value: '" + c + "'"); + } + } + + /** + * {@return an exception naming the input and the cursor offset} + * + * @param message What was wrong at the cursor. + */ + InvalidFormatException malformed(String message) { + return new InvalidFormatException( + "Malformed " + inputName + " at offset " + position + ": " + message); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Model2VecUnigramTokenizer.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Model2VecUnigramTokenizer.java new file mode 100644 index 0000000000..0db1734af5 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Model2VecUnigramTokenizer.java @@ -0,0 +1,680 @@ +/* + * 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.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import opennlp.subword.sentencepiece.SentencePieceTokenizer; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.tokenize.SubwordTokenizer; +import opennlp.tools.util.InvalidFormatException; + +/** + * Runs the Unigram tokenizer stored directly in a Model2Vec {@code tokenizer.json}. + * + *

Hugging Face stores the pieces and their scores in JSON, while OpenNLP's pure Java Unigram + * decoder reads the equivalent SentencePiece protobuf. This adapter creates that protobuf in + * memory. A tiny second tokenizer runs the JSON file's precompiled character map before the + * supported post-normalization steps are applied. No generated tokenizer file is written beside + * the user-supplied model.

+ */ +final class Model2VecUnigramTokenizer implements SubwordTokenizer { + + private static final int TYPE_NORMAL = 1; + private static final int TYPE_UNKNOWN = 2; + private static final int TYPE_CONTROL = 3; + private static final int TYPE_BYTE = 6; + private static final String UNIGRAM = "Unigram"; + private static final String METASPACE = "Metaspace"; + private static final String SEQUENCE = "Sequence"; + private static final String PRECOMPILED = "Precompiled"; + private static final String REPLACE = "Replace"; + private static final String STRIP = "Strip"; + private static final String MARKER = "▁"; + private static final char MARKER_CHAR = '▁'; + + private final SentencePieceTokenizer normalizer; + private final SentencePieceTokenizer segmenter; + private final List operations; + private final int unknownId; + private final Set controlIds; + + private Model2VecUnigramTokenizer(Parsed parsed) throws IOException { + final byte[] normalizerModel = modelBytes( + List.of(new Piece("", 0f, TYPE_UNKNOWN)), 0, false, + parsed.precompiledCharsMap(), true); + normalizer = SentencePieceTokenizer.load(new ByteArrayInputStream(normalizerModel)); + segmenter = SentencePieceTokenizer.load(new ByteArrayInputStream(modelBytes( + parsed.pieces(), parsed.unknownId(), parsed.byteFallback(), new byte[0], false))); + operations = List.copyOf(parsed.operations()); + unknownId = parsed.unknownId(); + controlIds = Set.copyOf(parsed.controlIds()); + } + + /** Reads and validates a supported Model2Vec Unigram tokenizer. */ + static Model2VecUnigramTokenizer load(Path tokenizerJson) throws IOException { + if (tokenizerJson == null) { + throw new IllegalArgumentException("tokenizerJson must not be null"); + } + if (!Files.isRegularFile(tokenizerJson)) { + throw new IllegalArgumentException( + "File does not exist or is not a regular file: " + tokenizerJson); + } + return new Model2VecUnigramTokenizer(parse(tokenizerJson)); + } + + /** {@inheritDoc} */ + @Override + public List encode(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + String normalized = normalizer.normalize(text).toString(); + for (NormalizationOperation operation : operations) { + normalized = operation.apply(normalized); + } + return segmenter.encode(normalized); + } + + /** {@return whether the row is the tokenizer's unknown piece} */ + boolean isUnknown(int id) { + return id == unknownId; + } + + /** {@return whether the row is a special control piece} */ + boolean isControl(int id) { + return controlIds.contains(id); + } + + /** {@return the number of tokenizer rows} */ + int vocabularySize() { + return segmenter.vocabularySize(); + } + + /** {@return the piece at the given tokenizer row} */ + String idToPiece(int id) { + return segmenter.idToPiece(id); + } + + private static Parsed parse(Path file) throws IOException { + final JsonCursor cursor = new JsonCursor(Files.readString(file), file.getFileName().toString()); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + ParsedModel model = null; + ParsedNormalizer normalizer = null; + boolean metaspace = false; + List addedTokens = List.of(); + if (cursor.peek() != '}') { + while (true) { + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "model" -> model = parseModel(cursor); + case "normalizer" -> normalizer = parseNormalizer(cursor); + case "pre_tokenizer" -> metaspace = parsePreTokenizer(cursor); + case "added_tokens" -> addedTokens = parseAddedTokens(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a top-level field"); + } + } else { + cursor.consume(); + } + cursor.requireEnd("Trailing content after the top-level object"); + if (model == null || !UNIGRAM.equals(model.type())) { + throw new InvalidFormatException(file + " does not define a Unigram tokenizer model"); + } + if (model.pieces() == null || model.pieces().isEmpty()) { + throw new InvalidFormatException(file + " has no model.vocab entries"); + } + if (normalizer == null || normalizer.precompiledCharsMap() == null) { + throw new InvalidFormatException(file + " has no supported Precompiled normalizer"); + } + if (!metaspace) { + throw new InvalidFormatException(file + " has no supported Metaspace pre-tokenizer"); + } + final List pieces = new ArrayList<>(model.pieces()); + final Set controls = new HashSet<>(); + final List sorted = new ArrayList<>(addedTokens); + sorted.sort(Comparator.comparingInt(AddedToken::id)); + for (AddedToken added : sorted) { + if (added.id() >= pieces.size()) { + throw new InvalidFormatException(file + " declares added token id " + added.id() + + " outside the model vocabulary of " + pieces.size() + " rows"); + } + if (!pieces.get(added.id()).text().equals(added.content())) { + throw new InvalidFormatException(file + " contradicts model.vocab at added token id " + + added.id()); + } + if (added.special() && added.id() != model.unknownId()) { + controls.add(added.id()); + } + } + if (model.unknownId() < 0 || model.unknownId() >= pieces.size()) { + throw new InvalidFormatException(file + " has an invalid model.unk_id"); + } + for (int id = 0; id < pieces.size(); id++) { + final Piece piece = pieces.get(id); + final int type = id == model.unknownId() ? TYPE_UNKNOWN + : controls.contains(id) ? TYPE_CONTROL + : model.byteFallback() && isBytePiece(piece.text()) ? TYPE_BYTE : TYPE_NORMAL; + pieces.set(id, new Piece(piece.text(), piece.score(), type)); + } + return new Parsed(pieces, model.unknownId(), model.byteFallback(), + normalizer.precompiledCharsMap(), normalizer.operations(), controls); + } + + private static ParsedModel parseModel(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + int unknownId = -1; + boolean byteFallback = false; + List pieces = null; + while (cursor.peek() != '}') { + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> type = cursor.parseString(); + case "unk_id" -> unknownId = checkedInt(cursor.parseLong(), cursor, "model.unk_id"); + case "byte_fallback" -> byteFallback = cursor.parseBoolean(); + case "vocab" -> pieces = parseVocabulary(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != '}') { + throw cursor.malformed("Expected ',' or '}' after a model field"); + } else { + return new ParsedModel(type, unknownId, byteFallback, pieces); + } + } + cursor.consume(); + return new ParsedModel(type, unknownId, byteFallback, pieces); + } + + private static List parseVocabulary(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List pieces = new ArrayList<>(); + while (cursor.peek() != ']') { + cursor.expect('['); + cursor.skipWhitespace(); + final String text = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(','); + cursor.skipWhitespace(); + final double score = cursor.parseDouble(); + if (score < -Float.MAX_VALUE || score > Float.MAX_VALUE) { + throw cursor.malformed("Unigram score is outside the float range"); + } + cursor.skipWhitespace(); + cursor.expect(']'); + pieces.add(new Piece(text, (float) score, TYPE_NORMAL)); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after a vocabulary entry"); + } else { + return pieces; + } + } + cursor.consume(); + return pieces; + } + + private static ParsedNormalizer parseNormalizer(JsonCursor cursor) + throws InvalidFormatException { + final NormalizerBuilder builder = new NormalizerBuilder(); + parseNormalizerObject(cursor, builder); + return new ParsedNormalizer(builder.precompiledCharsMap, builder.operations); + } + + private static void parseNormalizerObject(JsonCursor cursor, NormalizerBuilder builder) + throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + String precompiled = null; + List children = null; + PatternValue pattern = null; + String content = null; + boolean stripLeft = false; + boolean stripRight = false; + while (cursor.peek() != '}') { + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> type = cursor.parseString(); + case "precompiled_charsmap" -> precompiled = cursor.parseString(); + case "normalizers" -> children = parseNormalizerChildren(cursor); + case "pattern" -> pattern = parsePattern(cursor); + case "content" -> content = cursor.parseString(); + case "strip_left" -> stripLeft = cursor.parseBoolean(); + case "strip_right" -> stripRight = cursor.parseBoolean(); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != '}') { + throw cursor.malformed("Expected ',' or '}' after a normalizer field"); + } else { + break; + } + } + if (type == null) { + throw cursor.malformed("Normalizer has no type"); + } + switch (type) { + case SEQUENCE -> { + if (children == null) { + throw cursor.malformed("Sequence normalizer has no normalizers list"); + } + for (ParsedNormalizer child : children) { + if (child.precompiledCharsMap() != null) { + if (builder.precompiledCharsMap != null) { + throw cursor.malformed("More than one Precompiled normalizer is not supported"); + } + builder.precompiledCharsMap = child.precompiledCharsMap(); + } + builder.operations.addAll(child.operations()); + } + } + case PRECOMPILED -> { + if (precompiled == null) { + throw cursor.malformed("Precompiled normalizer has no character map"); + } + try { + builder.precompiledCharsMap = Base64.getDecoder().decode(precompiled); + } catch (IllegalArgumentException e) { + throw cursor.malformed("Precompiled normalizer has malformed base64"); + } + } + case REPLACE -> builder.operations.add(replacement(pattern, content, cursor)); + case STRIP -> builder.operations.add(new StripOperation(stripLeft, stripRight)); + default -> throw cursor.malformed("Unsupported normalizer type '" + type + "'"); + } + } + + private static List parseNormalizerChildren(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List children = new ArrayList<>(); + while (cursor.peek() != ']') { + final NormalizerBuilder child = new NormalizerBuilder(); + parseNormalizerObject(cursor, child); + children.add(new ParsedNormalizer(child.precompiledCharsMap, child.operations)); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after a normalizer"); + } else { + return children; + } + } + cursor.consume(); + return children; + } + + private static PatternValue parsePattern(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + final String kind = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + final String value = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect('}'); + return new PatternValue(kind, value); + } + + private static NormalizationOperation replacement( + PatternValue pattern, String content, JsonCursor cursor) throws InvalidFormatException { + if (pattern == null || content == null) { + throw cursor.malformed("Replace normalizer needs pattern and content"); + } + if ("String".equals(pattern.kind())) { + if (!content.equals(" " + pattern.value() + " ")) { + throw cursor.malformed("Only spacing literal replacements are supported"); + } + return new SurroundOperation(pattern.value()); + } + if ("Regex".equals(pattern.kind()) + && ("\\s+".equals(pattern.value()) || " {2,}".equals(pattern.value())) + && " ".equals(content)) { + return CollapseOperation.INSTANCE; + } + throw cursor.malformed("Unsupported Replace normalizer pattern"); + } + + private static boolean parsePreTokenizer(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + String replacement = null; + String prependScheme = null; + boolean split = true; + while (cursor.peek() != '}') { + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> type = cursor.parseString(); + case "replacement" -> replacement = cursor.parseString(); + case "prepend_scheme" -> prependScheme = cursor.parseString(); + case "split" -> split = cursor.parseBoolean(); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != '}') { + throw cursor.malformed("Expected ',' or '}' after a pre-tokenizer field"); + } else { + break; + } + } + return METASPACE.equals(type) && MARKER.equals(replacement) + && "always".equals(prependScheme) && !split; + } + + private static List parseAddedTokens(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List tokens = new ArrayList<>(); + while (cursor.peek() != ']') { + cursor.expect('{'); + cursor.skipWhitespace(); + int id = -1; + String content = null; + boolean special = false; + while (cursor.peek() != '}') { + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "id" -> id = checkedInt(cursor.parseLong(), cursor, "added token id"); + case "content" -> content = cursor.parseString(); + case "special" -> special = cursor.parseBoolean(); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != '}') { + throw cursor.malformed("Expected ',' or '}' after an added-token field"); + } else { + break; + } + } + if (id < 0 || content == null) { + throw cursor.malformed("Added token needs id and content"); + } + tokens.add(new AddedToken(id, content, special)); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after an added token"); + } else { + return tokens; + } + } + cursor.consume(); + return tokens; + } + + private static int checkedInt(long value, JsonCursor cursor, String field) + throws InvalidFormatException { + if (value < 0 || value > Integer.MAX_VALUE) { + throw cursor.malformed(field + " is outside the supported range"); + } + return (int) value; + } + + private static boolean isBytePiece(String piece) { + if (piece.length() != 6 || piece.charAt(0) != '<' || piece.charAt(1) != '0' + || piece.charAt(2) != 'x' || piece.charAt(5) != '>') { + return false; + } + return Character.digit(piece.charAt(3), 16) >= 0 + && Character.digit(piece.charAt(4), 16) >= 0; + } + + private static byte[] modelBytes(List pieces, int unknownId, boolean byteFallback, + byte[] precompiledCharsMap, boolean normalizing) { + final ProtoWriter model = new ProtoWriter(); + for (int id = 0; id < pieces.size(); id++) { + final Piece piece = pieces.get(id); + final ProtoWriter entry = new ProtoWriter(); + entry.string(1, piece.text()); + entry.float32(2, piece.score()); + entry.varintField(3, id == unknownId ? TYPE_UNKNOWN : piece.type()); + model.message(1, entry.bytes()); + } + final ProtoWriter trainer = new ProtoWriter(); + trainer.varintField(3, 1); + if (byteFallback) { + trainer.varintField(35, 1); + } + model.message(2, trainer.bytes()); + final ProtoWriter normalizer = new ProtoWriter(); + if (precompiledCharsMap.length > 0) { + normalizer.bytesField(2, precompiledCharsMap); + } + normalizer.varintField(3, normalizing ? 1 : 0); + normalizer.varintField(4, normalizing ? 1 : 0); + normalizer.varintField(5, normalizing ? 1 : 0); + model.message(3, normalizer.bytes()); + return model.bytes(); + } + + private interface NormalizationOperation { + String apply(String input); + } + + private record SurroundOperation(String literal) implements NormalizationOperation { + @Override + public String apply(String input) { + if (literal.isEmpty()) { + return input; + } + final StringBuilder out = new StringBuilder(input.length() + 8); + int cursor = 0; + while (cursor < input.length()) { + if (input.startsWith(literal, cursor)) { + appendMarker(out); + out.append(literal); + appendMarker(out); + cursor += literal.length(); + } else { + final int codePoint = input.codePointAt(cursor); + out.appendCodePoint(codePoint); + cursor += Character.charCount(codePoint); + } + } + return out.toString(); + } + } + + private enum CollapseOperation implements NormalizationOperation { + INSTANCE; + + @Override + public String apply(String input) { + final StringBuilder out = new StringBuilder(input.length()); + boolean marker = false; + for (int cursor = 0; cursor < input.length(); ) { + final int codePoint = input.codePointAt(cursor); + cursor += Character.charCount(codePoint); + if (codePoint == MARKER_CHAR) { + if (!marker) { + out.append(MARKER_CHAR); + } + marker = true; + } else { + out.appendCodePoint(codePoint); + marker = false; + } + } + return out.toString(); + } + } + + private record StripOperation(boolean left, boolean right) implements NormalizationOperation { + @Override + public String apply(String input) { + int start = 0; + int end = input.length(); + if (right) { + while (end > start && input.charAt(end - 1) == MARKER_CHAR) { + end--; + } + } + if (left) { + while (start < end && input.charAt(start) == MARKER_CHAR) { + start++; + } + if (start > 0 && start < end) { + start--; + } + } + return start == 0 && end == input.length() ? input : input.substring(start, end); + } + } + + private static void appendMarker(StringBuilder out) { + if (out.isEmpty() || out.charAt(out.length() - 1) != MARKER_CHAR) { + out.append(MARKER_CHAR); + } + } + + private record Piece(String text, float score, int type) { + } + + private record ParsedModel( + String type, int unknownId, boolean byteFallback, List pieces) { + } + + private record ParsedNormalizer( + byte[] precompiledCharsMap, List operations) { + } + + private record PatternValue(String kind, String value) { + } + + private record AddedToken(int id, String content, boolean special) { + } + + private record Parsed(List pieces, int unknownId, boolean byteFallback, + byte[] precompiledCharsMap, List operations, + Set controlIds) { + } + + private static final class NormalizerBuilder { + private byte[] precompiledCharsMap; + private final List operations = new ArrayList<>(); + } + + private static final class ProtoWriter { + private final ByteArrayOutputStream out = new ByteArrayOutputStream(); + + void message(int field, byte[] value) { + bytesField(field, value); + } + + void string(int field, String value) { + bytesField(field, value.getBytes(StandardCharsets.UTF_8)); + } + + void bytesField(int field, byte[] value) { + varint((long) field << 3 | 2); + varint(value.length); + out.writeBytes(value); + } + + void varintField(int field, long value) { + varint((long) field << 3); + varint(value); + } + + void float32(int field, float value) { + varint((long) field << 3 | 5); + final int bits = Float.floatToIntBits(value); + out.write(bits & 0xff); + out.write(bits >>> 8 & 0xff); + out.write(bits >>> 16 & 0xff); + out.write(bits >>> 24 & 0xff); + } + + void varint(long value) { + long remaining = value; + while ((remaining & ~0x7fL) != 0) { + out.write((int) (remaining & 0x7f) | 0x80); + remaining >>>= 7; + } + out.write((int) remaining); + } + + byte[] bytes() { + return out.toByteArray(); + } + } +} 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 new file mode 100644 index 0000000000..e4a436908d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java @@ -0,0 +1,361 @@ +/* + * 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.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.java.Experimental; + +/** + * Turns a distilled model directory (the layout the + * Model2Vec {@code save_pretrained} writes) + * into a directory {@link StaticEmbeddingModel#load(Path)} can open, then verifies it by loading + * it. + * + *

A distillation ships {@code model.safetensors}, {@code tokenizer.json}, and + * {@code config.json}, but not the two files the loader also needs for a WordPiece model + * ({@code vocab.txt} and {@code tokenizer_config.json}). This class fills the WordPiece gap from + * {@code tokenizer.json} itself: the matrix row order is the {@code model.vocab} dictionary in id + * order, and the casing is the {@code normalizer.lowercase} flag. A Model2Vec Unigram model is + * self-contained and loads directly from its JSON vocabulary, scores, and normalizer.

+ * + *

Assembly writes only the missing files and never overwrites an existing one, so a directory + * a caller already completed by hand is left intact.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public final class ModelAssembler { + + /** 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 legacy SentencePiece tokenizer family, when a separate model file is present. */ + private static final String FAMILY_SENTENCEPIECE = "SentencePiece"; + + /** Not instantiable. */ + private ModelAssembler() { + } + + /** + * The outcome of assembling a directory: what family it is, the files that were written, and the + * stats read back from the loaded model. + * + * @param family {@code "WordPiece"}, {@code "Unigram"}, or + * {@code "SentencePiece"}. + * @param dimension The embedding dimension of the loaded model. + * @param vocabularySize The number of subword rows in the loaded model's table. + * @param termCount The number of term rows after the subword rows; {@code 0} for a + * model without a term table. + * @param wroteVocabulary Whether a {@code vocab.txt} was written. + * @param wroteTokenizerConfig Whether a {@code tokenizer_config.json} was written. + */ + public record Result(String family, int dimension, int vocabularySize, int termCount, + boolean wroteVocabulary, boolean wroteTokenizerConfig) { + } + + /** + * Assembles and verifies a model directory in place. + * + * @param modelDirectory The distilled model directory. Must not be {@code null} and must be a + * directory holding at least {@code model.safetensors}, + * {@code tokenizer.json}, and {@code config.json}. + * @return The assembly result. + * @throws IllegalArgumentException Thrown if {@code modelDirectory} is {@code null}, is not a + * directory, or is missing a required distillation file. + * @throws InvalidFormatException Thrown if a model file is malformed, its tokenizer family is + * unsupported, or the directory does not load after assembly. + * @throws IOException Thrown if reading or writing a file fails. + */ + public static Result assemble(Path modelDirectory) throws IOException { + if (modelDirectory == null) { + throw new IllegalArgumentException("ModelDirectory must not be null"); + } + if (!Files.isDirectory(modelDirectory)) { + throw new IllegalArgumentException( + "Model directory does not exist or is not a directory: " + modelDirectory); + } + 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 FAMILY_WORDPIECE -> assembleWordpiece(modelDirectory, tokenizer); + case FAMILY_UNIGRAM -> assembleUnigram(modelDirectory); + default -> throw new InvalidFormatException(tokenizerJson + " has a '" + + tokenizer.modelType() + "' tokenizer model; only " + FAMILY_WORDPIECE + " and " + + FAMILY_UNIGRAM + " (" + FAMILY_SENTENCEPIECE + ") distillations are supported"); + }; + } + + /** + * Assembles a WordPiece directory, deriving {@code vocab.txt} and {@code tokenizer_config.json} + * from {@code tokenizer.json} when they are absent, then loading to verify. + * + * @param modelDirectory The model directory. + * @param tokenizer The parsed {@code tokenizer.json}. + * @return The assembly result. + * @throws IOException Thrown if reading or writing a file fails. + */ + private static Result assembleWordpiece(Path modelDirectory, TokenizerJson tokenizer) + throws IOException { + final Path vocabularyFile = modelDirectory.resolve(ModelFileNames.VOCABULARY); + boolean wroteVocabulary = false; + if (!Files.exists(vocabularyFile)) { + if (tokenizer.orderedVocabulary() == null) { + throw new InvalidFormatException("tokenizer.json in " + modelDirectory + + " has no model.vocab dictionary; cannot derive " + ModelFileNames.VOCABULARY); + } + Files.write(vocabularyFile, tokenizer.orderedVocabulary()); + wroteVocabulary = true; + } + 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 + // convention) when the tokenizer does not state it, which the load then reads back. + final boolean lowerCase = tokenizer.lowerCase() == null || tokenizer.lowerCase(); + Files.writeString(tokenizerConfigFile, + "{\n \"do_lower_case\": " + lowerCase + "\n}\n", StandardCharsets.UTF_8); + wroteTokenizerConfig = true; + } + final StaticEmbeddingModel model = load(modelDirectory); + return new Result(FAMILY_WORDPIECE, model.dimension(), model.vocabularySize(), + model.termCount(), wroteVocabulary, wroteTokenizerConfig); + } + + /** Loads and verifies a self-contained Model2Vec Unigram directory. */ + private static Result assembleUnigram(Path modelDirectory) throws IOException { + final StaticEmbeddingModel model = load(modelDirectory); + final boolean legacySentencePiece = ModelFileNames.firstRegularFile(modelDirectory, + ModelFileNames.SENTENCEPIECE_MODELS) != null; + return new Result(legacySentencePiece ? FAMILY_SENTENCEPIECE : FAMILY_UNIGRAM, + model.dimension(), model.vocabularySize(), model.termCount(), false, false); + } + + /** + * Loads the assembled directory to verify it, translating a load failure into an assembly + * failure with the same message and the same exception type. + * + * @param modelDirectory The assembled directory. + * @return The loaded model. + * @throws IOException Thrown if reading a file fails. + */ + private static StaticEmbeddingModel load(Path modelDirectory) throws IOException { + try { + return StaticEmbeddingModel.load(modelDirectory); + } catch (InvalidFormatException e) { + throw new InvalidFormatException("Assembled directory " + modelDirectory + + " does not load: " + e.getMessage(), e); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Assembled directory " + modelDirectory + + " does not load: " + e.getMessage(), e); + } + } + + /** + * {@return the required file in the directory} + * + * @param directory The model directory. + * @param name The required file name. + * @throws IllegalArgumentException Thrown if the file is absent. + */ + private static Path requireFile(Path directory, String name) { + final Path file = directory.resolve(name); + if (!Files.isRegularFile(file)) { + throw new IllegalArgumentException("Model directory " + directory + " has no " + name + + "; it does not look like a distilled model directory"); + } + return file; + } + + /** + * The fields read out of a {@code tokenizer.json} for assembly. + * + * @param modelType The {@code model.type}, e.g. {@code "WordPiece"} or {@code "Unigram"}. + * @param orderedVocabulary The matrix row order for a WordPiece dictionary vocabulary, or + * {@code null} when the model is not a WordPiece dictionary. + * @param lowerCase The {@code normalizer.lowercase} flag, or {@code null} when absent. + */ + private record TokenizerJson(String modelType, List orderedVocabulary, + Boolean lowerCase) { + } + + /** + * Reads the {@code model.type}, the WordPiece {@code model.vocab} dictionary in id order, and the + * {@code normalizer.lowercase} flag out of a {@code tokenizer.json}. + * + * @param file The {@code tokenizer.json} file. + * @return The parsed fields. + * @throws InvalidFormatException Thrown if the file is not a well-formed {@code tokenizer.json}. + * @throws IOException Thrown if reading the file fails. + */ + private static TokenizerJson readTokenizerJson(Path file) throws IOException { + final String json = Files.readString(file); + final JsonCursor cursor = new JsonCursor(json, file.getFileName().toString()); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + String modelType = null; + List orderedVocabulary = null; + Boolean lowerCase = null; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "model" -> { + final ModelSection model = parseModel(cursor); + modelType = model.type(); + orderedVocabulary = model.orderedVocabulary(); + } + case "normalizer" -> lowerCase = TeacherTokenizer.parseNormalizerLowercase(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a field, got '" + next + "'"); + } + } + cursor.requireEnd("Trailing content after the top-level object"); + if (modelType == null) { + throw new InvalidFormatException(file + " has no model.type"); + } + return new TokenizerJson(modelType, orderedVocabulary, lowerCase); + } + + /** The {@code model} object's type and, for a WordPiece dictionary, its rows in id order. */ + private record ModelSection(String type, List orderedVocabulary) { + } + + /** + * Parses the {@code model} object for its {@code type} and, when the vocabulary is a WordPiece + * dictionary, its rows in id order. + * + * @param cursor The cursor, positioned at the object's opening brace. + * @return The parsed type and, for a dictionary vocabulary, the ordered rows. + */ + private static ModelSection parseModel(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + List orderedVocabulary = null; + if (cursor.peek() == '}') { + cursor.consume(); + return new ModelSection(null, null); + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if ("type".equals(key)) { + type = cursor.parseString(); + } else if ("vocab".equals(key) && cursor.peek() == '{') { + orderedVocabulary = parseVocabularyDictionary(cursor); + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return new ModelSection(type, orderedVocabulary); + } + throw cursor.malformed("Expected ',' or '}' after a model field, got '" + next + "'"); + } + } + + /** + * Parses a WordPiece {@code vocab} dictionary of {@code "token": id} pairs into the token list in + * id order. + * + * @param cursor The cursor, positioned at the dictionary's opening brace. + * @return The tokens in id order. + * @throws InvalidFormatException Thrown if an id repeats or the ids are not a gapless range. + */ + private static List parseVocabularyDictionary(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + final Map tokenById = new LinkedHashMap<>(); + if (cursor.peek() == '}') { + cursor.consume(); + return List.of(); + } + while (true) { + cursor.skipWhitespace(); + final String token = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + final long id = cursor.parseLong(); + if (tokenById.putIfAbsent(id, token) != null) { + throw cursor.malformed("Vocabulary id " + id + " is assigned more than once"); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a vocab entry, got '" + next + "'"); + } + final List> entries = new ArrayList<>(tokenById.entrySet()); + entries.sort(Comparator.comparingLong(Map.Entry::getKey)); + final List ordered = new ArrayList<>(entries.size()); + for (int row = 0; row < entries.size(); row++) { + final Map.Entry entry = entries.get(row); + if (entry.getKey() != row) { + throw cursor.malformed("Vocabulary ids are not a gapless range: expected id " + row + + " but found " + entry.getKey()); + } + ordered.add(entry.getValue()); + } + return ordered; + } + +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelDistiller.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelDistiller.java new file mode 100644 index 0000000000..8aa207a8e5 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelDistiller.java @@ -0,0 +1,543 @@ +/* + * 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.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +import opennlp.tools.util.java.Experimental; + +/** + * Distills a sentence-transformer teacher into a static embedding table in the layout + * {@link StaticEmbeddingModel#load(Path)} opens, reproducing + * Model2Vec's distillation in Java so no + * Python environment is needed. The pipeline is Model2Vec's: + * + *
    + *
  1. The teacher's vocabulary is cleaned (unused tokens and special added tokens other than + * the unknown and pad tokens are dropped, the rest keeps its id order) and every surviving + * token is run through the teacher's ONNX graph as {@code [bos, token, eos]}; the token's + * embedding is the mean of the last hidden states.
    2. + *
  2. The matrix is projected onto its top principal components (a randomized SVD standing in + * for scikit-learn's dense one; see {@link RandomizedPca}).
    3. + *
  3. Each row is scaled by its Zipf weight {@code sif / (sif + p)}, where {@code p} is the + * row's share of a Zipf distribution over the vocabulary and {@code sif} is + * {@value #SIF_COEFFICIENT}, Model2Vec's default.
    4. + *
  4. The result is written as {@code model.safetensors} (F32), the cleaned + * {@code tokenizer.json}, and a {@code config.json} with {@code "normalize": true}; a + * SentencePiece teacher's {@code .model} file is copied alongside. The directory is then + * completed and verified by {@link ModelAssembler}.
    5. + *
+ * + *

The teacher directory must hold {@code tokenizer.json} and {@code onnx/model.onnx} (the + * ONNX export every sentence-transformer ships on the Hugging Face hub); a local + * {@code tokenizer_config.json} supplies the pad token when present.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public final class ModelDistiller { + + /** Model2Vec's default SIF coefficient for the Zipf weighting. */ + static final double SIF_COEFFICIENT = 1e-4; + + /** The number of id sequences per forward-pass batch, Model2Vec's batch size. */ + private static final int BATCH_SIZE = 256; + + /** The fixed seed of the PCA range finder, so a distillation is reproducible. */ + private static final long PCA_SEED = 42; + + /** Not instantiable. */ + private ModelDistiller() { + } + + /** Receives progress messages; the command-line tool prints them. */ + @FunctionalInterface + public interface ProgressListener { + + /** + * Reports a progress message. + * + * @param message The message. + */ + void progress(String message); + } + + /** + * The outcome of a distillation: the family, size, and dimension of the verified model, plus + * the variance the PCA kept. + * + * @param family {@code "WordPiece"} or {@code "SentencePiece"}. + * @param vocabularySize The number of subword rows in the distilled table. + * @param termCount The number of term rows appended after the subword rows. + * @param teacherDimension The teacher's hidden dimension. + * @param dimension The distilled table's dimension (after PCA). + * @param explainedVarianceRatio The share of the embedding variance the PCA kept. + */ + public record Result(String family, int vocabularySize, int termCount, int teacherDimension, + int dimension, double explainedVarianceRatio) { + } + + /** + * Distills a teacher into a model directory, resolving the teacher reference first: a local + * directory is used as-is, a Hugging Face model id ({@code org/model}, or + * {@code org/model@revision} to pin a revision) is downloaded into a local cache on first use. + * + * @param teacher The teacher: a local directory or a Hugging Face model id. Must not + * be {@code null}. + * @param outputDirectory The model directory to write. Must not be {@code null}. + * @param pcaDims The number of principal components to keep. + * @param listener Receives progress lines; may be {@code null}. + * @return The distillation result, read back from the verified directory. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or invalid, the + * teacher reference is malformed, or the teacher cannot be run. + * @throws IOException Thrown if reading or writing a file fails, or if a teacher cannot be + * downloaded and verified. + */ + public static Result distill(String teacher, Path outputDirectory, int pcaDims, + ProgressListener listener) throws IOException { + return distill(teacher, outputDirectory, pcaDims, List.of(), listener); + } + + /** + * Distills a teacher into a model directory with additional term rows, resolving the teacher + * reference the way {@link #distill(String, Path, int, ProgressListener)} does. + * + * @param teacher The teacher: a local directory or a Hugging Face model id. Must not + * be {@code null}. + * @param outputDirectory The model directory to write. Must not be {@code null}. + * @param pcaDims The number of principal components to keep. + * @param terms The terms to distill as extra rows; see + * {@link #distill(Path, Path, int, List, ProgressListener)}. Must not + * be {@code null}. + * @param listener Receives progress lines; may be {@code null}. + * @return The distillation result, read back from the verified directory. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or invalid, a term + * normalizes to nothing, the teacher reference is malformed, or the teacher cannot be run. + * @throws IOException Thrown if reading or writing a file fails, or if a teacher cannot be + * downloaded and verified. + */ + public static Result distill(String teacher, Path outputDirectory, int pcaDims, + List terms, ProgressListener listener) + throws IOException { + checkOutput(outputDirectory, pcaDims); + final List prepared = prepareTerms(terms); + return distill(HuggingFaceModelCache.resolve(teacher, listener), outputDirectory, pcaDims, + prepared, listener); + } + + /** + * Distills a teacher into a model directory. + * + * @param teacherDirectory The teacher's directory, holding {@code tokenizer.json} and + * {@code onnx/model.onnx}. Must not be {@code null} and must be a + * directory. + * @param outputDirectory The model directory to write. Created when missing. The four files + * a distillation produces are replaced, but files a previous run's + * assembly derived ({@code vocab.txt}, {@code tokenizer_config.json}) + * are not, so distil into a fresh or emptied directory when the + * vocabulary or the dimension changes. A failure part way through + * leaves whatever was written so far. Must not be {@code null}. + * @param pcaDims The number of principal components to keep; clamped to the teacher's + * hidden dimension, and skipped entirely when it would not reduce a + * tiny vocabulary. Model2Vec's default (and the recommended value) is + * 256. + * @param listener Receives one progress line per distillation phase and one per + * forward-pass batch; may be {@code null}. + * @return The distillation result, read back from the verified directory. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or invalid, the + * teacher directory lacks its files, or the teacher cannot be run. + * @throws IOException Thrown if reading or writing a file fails. + */ + public static Result distill(Path teacherDirectory, Path outputDirectory, int pcaDims, + ProgressListener listener) + throws IOException { + return distill(teacherDirectory, outputDirectory, pcaDims, List.of(), listener); + } + + /** + * Distills a teacher into a model directory with additional term rows: whole words and + * multi-word phrases (a learned corpus vocabulary) that are segmented by the teacher's own + * tokenizer, run through the teacher as full sequences, and appended to the table after the + * subword rows. The loaded model then matches text against these terms greedily + * longest-first before falling back to subword pieces. + * + *

Each term is normalized to lower-cased words joined by single spaces before use; terms + * that normalize to the same form are distilled once, and a term equal to a surviving + * vocabulary token is dropped, because its row would duplicate that token's. The terms are + * written to the model directory as {@code terms.txt}, one per line in row order, and should + * arrive sorted by descending corpus frequency: the Zipf weighting spans the subword rows and + * the term rows as one ranking.

+ * + * @param teacherDirectory The teacher's directory, as in + * {@link #distill(Path, Path, int, ProgressListener)}. A Unigram + * teacher must also hold its trained SentencePiece {@code .model} + * file. Must not be {@code null}. + * @param outputDirectory The model directory to write, as in + * {@link #distill(Path, Path, int, ProgressListener)}. Must not be + * {@code null}. + * @param pcaDims The number of principal components to keep. + * @param terms The terms to distill as extra rows; empty for none. Must not be + * {@code null} and must not contain {@code null}. + * @param listener Receives progress lines; may be {@code null}. + * @return The distillation result, read back from the verified directory. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or invalid, a term + * normalizes to nothing, the teacher directory lacks its files, or the teacher cannot be + * run. + * @throws IOException Thrown if reading or writing a file fails. + */ + public static Result distill(Path teacherDirectory, Path outputDirectory, int pcaDims, + List terms, ProgressListener listener) + throws IOException { + if (teacherDirectory == null) { + throw new IllegalArgumentException("TeacherDirectory must not be null"); + } + if (!Files.isDirectory(teacherDirectory)) { + throw new IllegalArgumentException("Teacher directory does not exist or is not a " + + "directory: " + teacherDirectory); + } + checkOutput(outputDirectory, pcaDims); + final Path onnxFile = teacherDirectory.resolve(ModelFileNames.ONNX_MODEL); + if (!Files.isRegularFile(onnxFile)) { + throw new IllegalArgumentException("Teacher directory " + teacherDirectory + " has no " + + ModelFileNames.ONNX_MODEL + "; the distillation runs the teacher's ONNX export, which " + + "sentence-transformers ship on the Hugging Face hub"); + } + final TeacherTokenizer tokenizer = TeacherTokenizer.read( + teacherDirectory.resolve(ModelFileNames.TOKENIZER_JSON), + teacherDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG)); + final int rows = tokenizer.vocabularySize(); + if (rows < 1) { + throw new IllegalArgumentException("Teacher directory " + teacherDirectory + " has no " + + "vocabulary token left after cleaning; there is nothing to distill"); + } + final List termList = new ArrayList<>(prepareTerms(terms)); + if (!termList.isEmpty()) { + // A term equal to a surviving vocabulary token would encode to the same teacher sequence + // and duplicate that token's row, so it is dropped; matching then reaches the token's row + // through the subword fallback instead. + final Set keptTokens = new HashSet<>(rows * 2); + for (int row = 0; row < rows; row++) { + keptTokens.add(tokenizer.rowToken(row)); + } + final int requestedTerms = termList.size(); + termList.removeIf(keptTokens::contains); + if (requestedTerms > termList.size()) { + report(listener, "Dropped " + (requestedTerms - termList.size()) + + " terms already present as vocabulary tokens"); + } + } + final int totalRows = rows + termList.size(); + + report(listener, "Encoding " + rows + " vocabulary tokens of " + teacherDirectory + + " through its ONNX graph"); + final float[] embeddings; + final int teacherDimension; + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(onnxFile)) { + float[][] first = encoder.encodeBatch(new long[][] {tokenizer.inputSequence(0)}); + teacherDimension = first[0].length; + embeddings = new float[totalRows * teacherDimension]; + System.arraycopy(first[0], 0, embeddings, 0, teacherDimension); + int row = 1; + while (row < rows) { + final int batchSize = Math.min(BATCH_SIZE, rows - row); + final long[][] batch = new long[batchSize][]; + for (int b = 0; b < batchSize; b++) { + batch[b] = tokenizer.inputSequence(row + b); + } + final float[][] pooled = encoder.encodeBatch(batch); + for (int b = 0; b < batchSize; b++) { + System.arraycopy(pooled[b], 0, embeddings, (row + b) * teacherDimension, + teacherDimension); + } + row += batchSize; + report(listener, "Encoded " + row + " / " + rows + " vocabulary tokens"); + } + encodeTerms(termList, tokenizer, teacherDirectory, encoder, embeddings, rows, + teacherDimension, listener); + } + nonFiniteToZero(embeddings); + + final int requested = Math.min(pcaDims, teacherDimension); + final float[] transformed; + final int components; + double explainedVarianceRatio = 1.0; + if (requested >= totalRows) { + // A PCA with more components than rows is not a reduction; Model2Vec skips it with a + // warning. Only reachable for toy vocabularies, which then keep the teacher's dimension. + transformed = embeddings; + components = teacherDimension; + } else { + report(listener, "Reducing " + totalRows + " x " + teacherDimension + " to " + requested + + " principal components"); + final RandomizedPca.Result pca = RandomizedPca.fitTransform(embeddings, totalRows, + teacherDimension, requested, PCA_SEED); + transformed = pca.transformed(); + components = requested; + explainedVarianceRatio = pca.explainedVarianceRatio(); + } + final float[] weights = zipfWeights(totalRows, SIF_COEFFICIENT); + for (int row = 0; row < totalRows; row++) { + final int base = row * components; + final float weight = weights[row]; + for (int d = 0; d < components; d++) { + transformed[base + d] *= weight; + } + } + + report(listener, "Writing and verifying the model directory " + outputDirectory); + Files.createDirectories(outputDirectory); + SafetensorsWriter.writeMatrix(outputDirectory.resolve(ModelFileNames.SAFETENSORS), totalRows, + components, transformed); + tokenizer.writeCleaned(outputDirectory.resolve(ModelFileNames.TOKENIZER_JSON)); + Files.writeString(outputDirectory.resolve(ModelFileNames.CONFIG), + configJson(teacherDirectory, pcaDims, components)); + copySentencePieceModel(teacherDirectory, outputDirectory); + final Path termsFile = outputDirectory.resolve(ModelFileNames.TERMS); + if (termList.isEmpty()) { + // A stale terms file from a previous run would no longer match the matrix's row count. + Files.deleteIfExists(termsFile); + } else { + Files.write(termsFile, termList); + } + final ModelAssembler.Result assembled = ModelAssembler.assemble(outputDirectory); + return new Result(assembled.family(), assembled.vocabularySize(), assembled.termCount(), + teacherDimension, assembled.dimension(), explainedVarianceRatio); + } + + /** + * Encodes the term rows: each term is segmented by the teacher's own tokenizer, wrapped as a + * full input sequence, and mean-pooled through the teacher, filling the matrix rows after the + * vocabulary rows. Sequences vary in length and a batch must not be ragged, so equal-length + * sequences are batched together. + * + * @param termList The normalized terms, in row order. + * @param tokenizer The teacher's parsed tokenizer. + * @param teacherDirectory The teacher's directory, for the segmenter. + * @param encoder The open teacher encoder. + * @param embeddings The matrix being filled, {@code totalRows * teacherDimension}. + * @param vocabularyRows The number of vocabulary rows preceding the term rows. + * @param teacherDimension The teacher's hidden dimension. + * @param listener Receives one progress line per batch; may be {@code null}. + * @throws IOException Thrown if reading the teacher's SentencePiece file fails. + */ + private static void encodeTerms(List termList, TeacherTokenizer tokenizer, + Path teacherDirectory, OnnxTeacherEncoder encoder, + float[] embeddings, int vocabularyRows, int teacherDimension, + ProgressListener listener) throws IOException { + if (termList.isEmpty()) { + return; + } + report(listener, "Encoding " + termList.size() + + " terms through the teacher's own segmentation"); + final TermSegmenter segmenter = TermSegmenter.forTeacher(tokenizer, teacherDirectory); + final long[][] sequences = new long[termList.size()][]; + for (int t = 0; t < sequences.length; t++) { + sequences[t] = tokenizer.inputSequence(segmenter.pieces(termList.get(t))); + } + final Integer[] byLength = new Integer[sequences.length]; + for (int t = 0; t < byLength.length; t++) { + byLength[t] = t; + } + Arrays.sort(byLength, Comparator.comparingInt(t -> sequences[t].length)); + int encoded = 0; + while (encoded < byLength.length) { + int end = encoded + 1; + while (end < byLength.length && end - encoded < BATCH_SIZE + && sequences[byLength[end]].length == sequences[byLength[encoded]].length) { + end++; + } + final long[][] batch = new long[end - encoded][]; + for (int b = 0; b < batch.length; b++) { + batch[b] = sequences[byLength[encoded + b]]; + } + final float[][] pooled = encoder.encodeBatch(batch); + for (int b = 0; b < batch.length; b++) { + System.arraycopy(pooled[b], 0, embeddings, + (vocabularyRows + byLength[encoded + b]) * teacherDimension, teacherDimension); + } + encoded = end; + report(listener, "Encoded " + encoded + " / " + termList.size() + " terms"); + } + } + + /** + * Normalizes and deduplicates the requested terms before any teacher work: each term becomes + * its lower-cased words joined by single spaces, and terms normalizing to the same form are + * kept once, in first-occurrence order. + * + * @param terms The requested terms. + * @return The normalized, duplicate-free terms. + * @throws IllegalArgumentException Thrown if {@code terms} is {@code null}, contains + * {@code null}, or contains a term with no letter or digit. + */ + private static List prepareTerms(List terms) { + if (terms == null) { + throw new IllegalArgumentException("Terms must not be null"); + } + final Set prepared = new LinkedHashSet<>(terms.size() * 2); + for (final String term : terms) { + if (term == null) { + throw new IllegalArgumentException("Terms must not contain null"); + } + final String normalized = TermTable.normalizeTerm(term); + if (normalized.isEmpty()) { + throw new IllegalArgumentException("Term '" + term + + "' has no letter or digit; it cannot be matched in text"); + } + prepared.add(normalized); + } + return List.copyOf(prepared); + } + + /** + * Validates the arguments that do not depend on the teacher, so that a distillation naming a + * hub teacher fails before it downloads anything. + * + * @param outputDirectory The model directory to write. + * @param pcaDims The number of principal components to keep. + * @throws IllegalArgumentException Thrown if the directory is {@code null} or {@code pcaDims} + * is below 1. + */ + private static void checkOutput(Path outputDirectory, int pcaDims) { + if (outputDirectory == null) { + throw new IllegalArgumentException("OutputDirectory must not be null"); + } + if (pcaDims < 1) { + throw new IllegalArgumentException("PcaDims must be at least 1, got " + pcaDims); + } + } + + /** + * Reports one progress line, if anyone is listening. + * + * @param listener The listener; may be {@code null}. + * @param message The message. + */ + private static void report(ProgressListener listener, String message) { + if (listener != null) { + listener.progress(message); + } + } + + /** + * {@return Model2Vec's Zipf weights: row {@code i} gets {@code sif / (sif + p_i)} with + * {@code p_i = (1 / (i + 2)) / sum_j (1 / (j + 2))}, a SIF weighting under the assumption that + * vocabulary order approximates frequency order (Zipf's law)} + * + * @param rows The number of rows. + * @param sifCoefficient The SIF coefficient. + */ + static float[] zipfWeights(int rows, double sifCoefficient) { + double harmonicSum = 0; + for (int j = 2; j <= rows + 1; j++) { + harmonicSum += 1.0 / j; + } + final float[] weights = new float[rows]; + for (int i = 0; i < rows; i++) { + final double probability = 1.0 / (i + 2) / harmonicSum; + weights[i] = (float) (sifCoefficient / (sifCoefficient + probability)); + } + return weights; + } + + /** + * Replaces non-finite values with zero, the guard against a teacher emitting a NaN or infinite + * hidden state. Model2Vec applies numpy's {@code nan_to_num} here, which maps an infinity to the + * largest finite float; zero is used instead because an infinity of that magnitude still leaves + * the principal component analysis with nothing but that one row. + * + * @param values The matrix, modified in place. + */ + private static void nonFiniteToZero(float[] values) { + for (int i = 0; i < values.length; i++) { + if (!Float.isFinite(values[i])) { + values[i] = 0; + } + } + } + + /** + * {@return the {@code config.json} of the distilled model, mirroring the fields Model2Vec + * writes; the loader reads only {@code normalize}} + * + * @param teacherDirectory The teacher's directory, for the name. + * @param pcaDims The requested PCA dimension. + * @param components The effective PCA dimension. + */ + private static String configJson(Path teacherDirectory, int pcaDims, int components) { + final Path name = teacherDirectory.getFileName(); + return "{\n" + + " \"model_type\": \"model2vec\",\n" + + " \"architectures\": [\"StaticModel\"],\n" + + " \"tokenizer_name\": \"" + (name == null ? teacherDirectory : name) + "\",\n" + + teacherRevisionField(teacherDirectory) + + " \"apply_pca\": " + pcaDims + ",\n" + + " \"sif_coefficient\": " + SIF_COEFFICIENT + ",\n" + + " \"hidden_dim\": " + components + ",\n" + + " \"seq_length\": 1000000,\n" + + " \"normalize\": true,\n" + + " \"pooling\": \"mean\",\n" + + " \"embedding_dtype\": \"float32\"\n" + + "}\n"; + } + + /** + * {@return the {@code config.json} field naming the commit the teacher's files came from, or an + * empty string when the teacher directory is not a cached hub download} + * + *

A distilled table is not reproducible without the exact revision of the teacher it was + * distilled from, and the teacher can move under its branch name, so the sha travels with the + * table rather than only staying in the cache directory it was downloaded into.

+ * + * @param teacherDirectory The teacher's directory. + */ + private static String teacherRevisionField(Path teacherDirectory) { + final String revision = HuggingFaceModelCache.pinnedRevision(teacherDirectory); + return revision == null ? "" : " \"teacher_revision\": \"" + revision + "\",\n"; + } + + /** + * Copies the teacher's trained SentencePiece {@code .model} file into the model directory when + * the teacher has one; the distillation cannot fabricate it and the loader needs it for the + * SentencePiece layout. + * + * @param teacherDirectory The teacher's directory. + * @param outputDirectory The model directory. + * @throws IOException Thrown if copying fails. + */ + private static void copySentencePieceModel(Path teacherDirectory, Path outputDirectory) + throws IOException { + for (final String name : ModelFileNames.SENTENCEPIECE_MODELS) { + final Path source = teacherDirectory.resolve(name); + if (Files.isRegularFile(source)) { + Files.copy(source, outputDirectory.resolve(name), + StandardCopyOption.REPLACE_EXISTING); + return; + } + } + } +} 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 0000000000..54c0eedd00 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelFileNames.java @@ -0,0 +1,92 @@ +/* + * 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.Files; +import java.nio.file.Path; +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 + * Unigram directory holds {@link #SAFETENSORS}, {@link #CONFIG}, and + * {@link #TOKENIZER_JSON}. Legacy SentencePiece directories may additionally hold one of + * {@link #SENTENCEPIECE_MODELS}. + * + *

{@link #ONNX_MODEL} and {@link #ONNX_MODEL_DATA} name files of a teacher directory + * rather than of a model directory; {@link ModelDistiller} and {@link HuggingFaceModelCache} share + * them.

+ */ +final class ModelFileNames { + + /** The safetensors file holding the embedding matrix and optional per-token weights. */ + static final String SAFETENSORS = "model.safetensors"; + + /** + * The quantized matrix file, written by the {@code QuantizeModel} tool. It carries the matrix + * and any per-token weights itself, and is the directory's matrix source in place of + * {@link #SAFETENSORS}, which a quantized deployment deletes. + */ + static final String QUANTIZED = "model.quantized"; + + /** 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 optional term rows of the matrix, one normalized term per line in row order. */ + static final String TERMS = "terms.txt"; + + /** The file names SentencePiece models ship their trained {@code .model} under, in try order. */ + static final List SENTENCEPIECE_MODELS = + List.of("sentencepiece.bpe.model", "spiece.model", "tokenizer.model"); + + /** The ONNX graph of a teacher, relative to the teacher directory's root. */ + static final String ONNX_MODEL = "onnx/model.onnx"; + + /** The external weights an ONNX export splits out of {@link #ONNX_MODEL}, if it splits them. */ + static final String ONNX_MODEL_DATA = "onnx/model.onnx_data"; + + /** Not instantiable. */ + private ModelFileNames() { + } + + /** + * {@return the first of the given file names that exists as a regular file in the directory, + * or {@code null} when none does} + * + * @param directory The directory to look in. + * @param names The file names to try, in order. + */ + static Path firstRegularFile(Path directory, List names) { + for (final String name : names) { + final Path file = directory.resolve(name); + if (Files.isRegularFile(file)) { + return file; + } + } + return null; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelQuantizer.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelQuantizer.java new file mode 100644 index 0000000000..81009b4e83 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelQuantizer.java @@ -0,0 +1,150 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +/** + * Quantizes a static embedding model directory in place: reads the matrix and optional + * per-token weights from the directory's {@code model.safetensors}, quantizes the matrix to the + * requested bit width (see {@link QuantizedEmbeddingMatrix}), and writes + * {@code model.quantized} next to it. {@code StaticEmbeddingModel.load} prefers the quantized + * file from then on; the safetensors may be deleted for a slim deployment. + * + *

The written file is verified by reading it back and measuring the mean cosine between the + * original and reconstructed rows over a deterministic sample, so a completed run reports the + * reconstruction quality actually on disk.

+ */ +public final class ModelQuantizer { + + // At most this many rows enter the verification sample, evenly strided so it is + // deterministic and covers the whole table. + private static final int VERIFICATION_SAMPLE_CAP = 1024; + + /** Not instantiable. */ + private ModelQuantizer() { + } + + /** + * What a quantization run produced and measured. + * + * @param rowCount The number of matrix rows. + * @param dimension The row width. + * @param bits The bit width per padded dimension. + * @param hasWeights Whether per-token pooling weights were carried over. + * @param safetensorsBytes The size of the source safetensors file. + * @param quantizedBytes The size of the written quantized file. + * @param sampledRows The number of rows in the verification sample. + * @param meanCosine The mean cosine between original and reconstructed sampled rows; + * {@code Double.NaN} when every sampled row was zero. + */ + public record Result(int rowCount, int dimension, int bits, boolean hasWeights, + long safetensorsBytes, long quantizedBytes, int sampledRows, + double meanCosine) { + } + + /** + * Quantizes the model directory's matrix and writes {@code model.quantized}. + * + * @param modelDirectory The model directory. Must not be {@code null}, must be a directory, + * and must hold a {@code model.safetensors}. + * @param bits The bit width, between {@link QuantizedEmbeddingMatrix#MIN_BITS} and + * {@link QuantizedEmbeddingMatrix#MAX_BITS}. + * @param seed The rotation seed; the same matrix, bits, and seed write the same + * file bytes. + * @return What was produced and measured. + * @throws IllegalArgumentException Thrown if an argument is invalid or the directory has no + * safetensors file. + * @throws IOException Thrown if reading or writing fails. + */ + public static Result quantize(Path modelDirectory, int bits, long seed) throws IOException { + if (modelDirectory == null) { + throw new IllegalArgumentException("ModelDirectory must not be null"); + } + if (!Files.isDirectory(modelDirectory)) { + throw new IllegalArgumentException( + "Model directory does not exist or is not a directory: " + modelDirectory); + } + final Path safetensorsFile = modelDirectory.resolve(ModelFileNames.SAFETENSORS); + if (!Files.isRegularFile(safetensorsFile)) { + throw new InvalidFormatException("Model directory " + modelDirectory + " has no " + + ModelFileNames.SAFETENSORS + " to quantize"); + } + final SafetensorsFile tensors = SafetensorsFile.read(safetensorsFile); + final String matrixName = tensors.singleMatrixTensorName(); + final TensorInfo matrixInfo = tensors.tensorInfo(matrixName); + final int rowCount = matrixInfo.shape()[0]; + final int dimension = matrixInfo.shape()[1]; + final float[] matrix = tensors.readFloats(matrixName); + float[] weights = null; + if (tensors.tensorNames().contains(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME)) { + weights = tensors.readFloats(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME); + if (weights.length != rowCount) { + throw new IllegalArgumentException("Tensor '" + + StaticEmbeddingModel.WEIGHTS_TENSOR_NAME + "' in " + safetensorsFile + " has " + + weights.length + " elements but the matrix has " + rowCount + " rows"); + } + } + final Path quantizedFile = modelDirectory.resolve(ModelFileNames.QUANTIZED); + QuantizedEmbeddingMatrix.quantize(matrix, rowCount, dimension, bits, seed) + .withPoolingWeights(weights) + .write(quantizedFile); + // Verify what is actually on disk, not the in-memory object. + final QuantizedEmbeddingMatrix written = QuantizedEmbeddingMatrix.read(quantizedFile); + final int stride = Math.max(1, rowCount / VERIFICATION_SAMPLE_CAP); + int sampled = 0; + int nonZero = 0; + double cosineSum = 0; + for (int row = 0; row < rowCount; row += stride) { + sampled++; + final double cosine = cosine(matrix, row * dimension, dimension, written.decodeRow(row)); + if (!Double.isNaN(cosine)) { + nonZero++; + cosineSum += cosine; + } + } + return new Result(rowCount, dimension, bits, weights != null, + Files.size(safetensorsFile), Files.size(quantizedFile), sampled, + nonZero == 0 ? Double.NaN : cosineSum / nonZero); + } + + /** + * {@return the cosine between a matrix row and its reconstruction, or {@code Double.NaN} when + * either has no direction} Also the shared fidelity measure of this package's tests. + * + * @param matrix The flat row-major matrix. + * @param base The row's first index. + * @param dimension The row width. + * @param decoded The reconstructed row. + */ + static double cosine(float[] matrix, int base, int dimension, float[] decoded) { + double dot = 0; + double normASquared = 0; + double normBSquared = 0; + for (int d = 0; d < dimension; d++) { + dot += (double) matrix[base + d] * decoded[d]; + normASquared += (double) matrix[base + d] * matrix[base + d]; + normBSquared += (double) decoded[d] * decoded[d]; + } + final double denominator = Math.sqrt(normASquared) * Math.sqrt(normBSquared); + return denominator == 0 ? Double.NaN : dot / denominator; + } +} 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 new file mode 100644 index 0000000000..e5f7a02cb1 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java @@ -0,0 +1,33 @@ +/* + * 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 opennlp.tools.util.java.Experimental; + +/** + * One vocabulary token found near a query vector by {@link StaticEmbeddingModel#mostSimilar} + * or {@link StaticEmbeddingModel#analogy}, most similar first. + * + *

Warning: Experimental new feature; the API might change in a later release.

+ * + * @param token The vocabulary token: one subword piece of the model's tokenizer, which is + * not necessarily a whole word. + * @param similarity Cosine similarity to the query vector, in {@code [-1, 1]}. + */ +@Experimental +public record Neighbor(String token, double similarity) { +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/OnnxTeacherEncoder.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/OnnxTeacherEncoder.java new file mode 100644 index 0000000000..7045702f56 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/OnnxTeacherEncoder.java @@ -0,0 +1,244 @@ +/* + * 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.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; + +import ai.onnxruntime.NodeInfo; +import ai.onnxruntime.OnnxJavaType; +import ai.onnxruntime.OnnxTensor; +import ai.onnxruntime.OnnxValue; +import ai.onnxruntime.OrtEnvironment; +import ai.onnxruntime.OrtException; +import ai.onnxruntime.OrtSession; +import ai.onnxruntime.TensorInfo; + +/** + * Runs a teacher transformer over id sequences through its ONNX graph and mean-pools the last + * hidden states, the forward pass + * Model2Vec's distillation performs per + * vocabulary token. The + * graph is fed exactly the inputs it declares: {@code input_ids} and {@code attention_mask} for + * every model, plus a zero {@code token_type_ids} for the BERT-family graphs that ask for one. + * The pooled output is the mean of the single rank-3 float output (the + * {@code last_hidden_state}) over all of a sequence's positions; the attention mask is all ones, + * because a batch is never padded (see {@link #encodeBatch(long[][])}). + * + *

Not thread-safe; a distillation drives one instance from a single thread. Close it to + * release the native session.

+ */ +final class OnnxTeacherEncoder implements AutoCloseable { + + /** The id-sequence input every transformer encoder graph declares. */ + private static final String INPUT_IDS = "input_ids"; + + /** The attention-mask input every transformer encoder graph declares. */ + private static final String ATTENTION_MASK = "attention_mask"; + + /** The segment input the BERT-family graphs declare; fed all zeros. */ + private static final String TOKEN_TYPE_IDS = "token_type_ids"; + + /** The rank of the last-hidden-state output: batch, position, hidden dimension. */ + private static final int HIDDEN_STATE_RANK = 3; + + private final OrtEnvironment environment; + private final OrtSession session; + private final boolean wantsTokenTypeIds; + private final String hiddenStateOutput; + private final AtomicBoolean closed = new AtomicBoolean(); + + /** Holds the open session; created by {@link #load(Path)}. */ + private OnnxTeacherEncoder(OrtEnvironment environment, OrtSession session, + boolean wantsTokenTypeIds, String hiddenStateOutput) { + this.environment = environment; + this.session = session; + this.wantsTokenTypeIds = wantsTokenTypeIds; + this.hiddenStateOutput = hiddenStateOutput; + } + + /** + * Loads a teacher's ONNX graph. + * + * @param onnxFile The ONNX file. Must not be {@code null} and must exist, must declare an + * {@code input_ids} input, and must produce a rank-3 float tensor output; the + * first such output is taken as the last hidden state and pooled. + * @return The encoder. + * @throws IllegalArgumentException Thrown if the file is missing, the graph has no + * {@code input_ids} input or no rank-3 float tensor output, or the runtime rejects the + * graph. + */ + static OnnxTeacherEncoder load(Path onnxFile) { + if (onnxFile == null) { + throw new IllegalArgumentException("OnnxFile must not be null"); + } + if (!Files.isRegularFile(onnxFile)) { + throw new IllegalArgumentException("File does not exist or is not a regular file: " + + onnxFile); + } + final OrtEnvironment environment = OrtEnvironment.getEnvironment(); + final OrtSession session; + try (OrtSession.SessionOptions options = new OrtSession.SessionOptions()) { + session = environment.createSession(onnxFile.toString(), options); + } catch (OrtException e) { + throw new IllegalArgumentException("Failed to load ONNX graph " + onnxFile + ": " + + e.getMessage(), e); + } + // The session is open from here on, so every exit below closes it: an inspection that + // rejects the graph, or that fails outright, must not leak the native handle. + try { + if (!session.getInputNames().contains(INPUT_IDS)) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " has no '" + INPUT_IDS + + "' input; it does not look like a transformer encoder (inputs: " + + session.getInputNames() + ")"); + } + final boolean wantsTokenTypeIds = session.getInputNames().contains(TOKEN_TYPE_IDS); + String hiddenStateOutput = null; + for (final Map.Entry output : session.getOutputInfo().entrySet()) { + if (output.getValue().getInfo() instanceof TensorInfo tensorInfo + && tensorInfo.type == OnnxJavaType.FLOAT + && tensorInfo.getShape().length == HIDDEN_STATE_RANK) { + hiddenStateOutput = output.getKey(); + break; + } + } + if (hiddenStateOutput == null) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " has no rank-3 float " + + "tensor output (a last hidden state) to pool (outputs: " + + session.getOutputInfo().keySet() + ")"); + } + return new OnnxTeacherEncoder(environment, session, wantsTokenTypeIds, hiddenStateOutput); + } catch (OrtException e) { + final IllegalArgumentException failure = new IllegalArgumentException( + "Failed to inspect ONNX graph " + onnxFile + ": " + e.getMessage(), e); + closeAfterFailure(session, failure); + throw failure; + } catch (RuntimeException e) { + closeAfterFailure(session, e); + throw e; + } + } + + /** + * Closes a session on a failing load path, reporting a close failure as a suppressed exception + * of the failure being thrown rather than in place of it. + * + * @param session The session to close. + * @param failure The exception the caller is about to throw. + */ + private static void closeAfterFailure(OrtSession session, RuntimeException failure) { + try { + session.close(); + } catch (OrtException e) { + failure.addSuppressed(e); + } + } + + /** + * Runs one batch of id sequences and mean-pools each sequence's hidden states. All sequences + * in a batch must have the same length (the distillation wraps every vocabulary token in the + * same bos/eos pair, so they do); the attention mask is all ones and no padding is needed. + * + * @param batch The id sequences, {@code [batchSize][sequenceLength]}. Must not be + * {@code null} or empty. + * @return The pooled vectors, {@code [batchSize][hiddenDimension]}. + * @throws IllegalArgumentException Thrown if the batch is empty or ragged, or the runtime + * rejects the input. + */ + float[][] encodeBatch(long[][] batch) { + if (batch == null || batch.length == 0) { + throw new IllegalArgumentException("Batch must not be null or empty"); + } + final int sequenceLength = batch[0].length; + for (final long[] sequence : batch) { + if (sequence.length != sequenceLength) { + throw new IllegalArgumentException("Batch is ragged: sequence lengths differ"); + } + } + final long[][] attentionMask = new long[batch.length][sequenceLength]; + for (final long[] mask : attentionMask) { + Arrays.fill(mask, 1L); + } + try { + final Map inputs = new HashMap<>(); + OnnxTensor tokenTypeIds = null; + try (OnnxTensor inputIds = OnnxTensor.createTensor(environment, batch); + OnnxTensor mask = OnnxTensor.createTensor(environment, attentionMask)) { + inputs.put(INPUT_IDS, inputIds); + inputs.put(ATTENTION_MASK, mask); + if (wantsTokenTypeIds) { + tokenTypeIds = OnnxTensor.createTensor(environment, + new long[batch.length][sequenceLength]); + inputs.put(TOKEN_TYPE_IDS, tokenTypeIds); + } + try (OrtSession.Result result = session.run(inputs)) { + final OnnxValue value = result.get(hiddenStateOutput) + .orElseThrow(() -> new IllegalStateException("Output '" + hiddenStateOutput + + "' missing from the graph's results")); + final float[][][] hidden = (float[][][]) value.getValue(); + final float[][] pooled = new float[batch.length][]; + for (int i = 0; i < batch.length; i++) { + final float[][] states = hidden[i]; + final float[] sum = new float[states[0].length]; + for (final float[] state : states) { + for (int d = 0; d < sum.length; d++) { + sum[d] += state[d]; + } + } + for (int d = 0; d < sum.length; d++) { + sum[d] /= states.length; + if (Float.isNaN(sum[d])) { + sum[d] = 0; + } + } + pooled[i] = sum; + } + return pooled; + } finally { + if (tokenTypeIds != null) { + tokenTypeIds.close(); + } + } + } + } catch (OrtException e) { + throw new IllegalArgumentException("ONNX forward pass failed: " + e.getMessage(), e); + } + } + + /** + * Closes the native session; calling this more than once is a no-op after the first call. + * + *

The {@link OrtEnvironment} is deliberately not closed: {@link + * OrtEnvironment#getEnvironment()} returns a process-wide singleton shared with every other + * ONNX component in the JVM, so closing it here would tear down an environment they still use. + * {@link OrtSession#close()} rejects a second call, hence the guard.

+ */ + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + try { + session.close(); + } catch (OrtException e) { + // Closing a native resource must not mask a distillation result. + } + } + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedEmbeddingMatrix.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedEmbeddingMatrix.java new file mode 100644 index 0000000000..5019a9ffc3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedEmbeddingMatrix.java @@ -0,0 +1,680 @@ +/* + * 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.BufferedInputStream; +import java.io.BufferedOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.InvalidFormatException; + +/** + * An embedding matrix quantized to {@code 2}-{@code 4} bits per dimension, following the + * TurboQuant construction (Zandieh, Daliri, Hadian, Mirrokni, + * TurboQuant: Online Vector Quantization with + * Near-optimal Distortion Rate): each row is rotated by + * a seeded {@link HadamardRotation}, so its coordinates become near-independent and + * near-Gaussian, and each rotated coordinate is encoded independently with the + * {@link GaussianQuantizer} grid of the chosen bit width. A row decodes to a per-row scale times + * its grid levels; the scale is least-squares fitted per row, which strictly reduces the squared + * error of the fixed grid. + * + *

The storage is {@code bits} per padded dimension plus two floats per row (the + * fitted scale and the decoded norm), against 32 bits per dimension for the float matrix. The + * rotation pads each row to the next power of two, so a 500,000-row, 300-dimension table stores + * 512 coded dimensions per row and shrinks from roughly 600 MB to about 132 MB at 4 bits, 4.5 + * times smaller. The workload this serves is memory-bound row gathering, so reading fewer bytes + * is also the throughput lever.

+ * + *

Rows live in rotated space, and the cheap operations stay there: the rotation is + * orthonormal, so dot products and norms of rotated vectors equal those of the originals, and + * pooling commutes with it because rotation is linear. A consumer embeds text by summing rows + * with {@link #addRowRotated(int, float, float[])} and applying {@link #toOriginal(float[])} + * once per text, not once per row; a similarity scan rotates the query once with + * {@link #rotate(float[])} and scores every row with {@link #dotRotated(int, float[])}, never + * leaving rotated space. {@link #decodeRow(int)} exists for callers that need one original-space + * row and for measuring reconstruction quality.

+ * + *

The file format is self-describing: it stores the grid levels and the rotation seed, so a + * reader reconstructs exactly the decoder the writer used and never depends on this class's grid + * derivation staying fixed. Quantizing is deterministic: the same matrix, bit width, and seed + * produce the same file bytes on every JVM.

+ * + *

Instances are immutable and safe for concurrent use after construction.

+ */ +@ThreadSafe +public final class QuantizedEmbeddingMatrix { + + /** The smallest supported bit width. */ + public static final int MIN_BITS = GaussianQuantizer.MIN_BITS; + + /** The largest supported bit width. */ + public static final int MAX_BITS = GaussianQuantizer.MAX_BITS; + + // "ONQ1": OpenNLP quantized matrix, format 1. + private static final int MAGIC = 0x4F4E5131; + + private final int rowCount; + private final int dimension; + private final int paddedDimension; + private final int bits; + private final long seed; + private final int rowBytes; + private final GaussianQuantizer quantizer; + private final HadamardRotation rotation; + // One scale per row: decoded rotated coordinate i of a row is scale * level(code_i). + private final float[] scales; + // Packed codes, row-major: row r's code i occupies bits [i*bits, (i+1)*bits) of the row's + // rowBytes region, little-endian within the region. + private final byte[] codes; + // The L2 norm of each decoded original-space row. Quantization noise leaves some energy in + // the padding coordinates, which truncation drops, so this is computed exactly at quantize + // time (one inverse rotation per row) and stored in the file rather than recomputed from the + // codes on load. + private final float[] decodedNorms; + // Optional per-row pooling weights carried alongside the matrix, so a quantized file can + // fully replace a safetensors file that bundled a "weights" tensor; null when absent. The + // weights are stored as they are, not quantized. + private final float[] poolingWeights; + + /** + * Holds validated state; callers reach this through {@link #quantize} or {@link #read}. + */ + private QuantizedEmbeddingMatrix(int rowCount, int dimension, int bits, long seed, + GaussianQuantizer quantizer, float[] scales, byte[] codes, + float[] decodedNorms, float[] poolingWeights) { + this.rowCount = rowCount; + this.dimension = dimension; + this.paddedDimension = HadamardRotation.paddedDimension(dimension); + this.bits = bits; + this.seed = seed; + this.rowBytes = rowByteCount(paddedDimension, bits); + this.quantizer = quantizer; + this.rotation = new HadamardRotation(dimension, seed); + this.scales = scales; + this.codes = codes; + this.decodedNorms = decodedNorms; + this.poolingWeights = poolingWeights; + } + + /** + * Quantizes a float matrix. + * + * @param rowMajor The matrix, row-major, {@code rowCount * dimension} floats. Must not be + * {@code null} and every value must be finite. + * @param rowCount The number of rows. Must be at least 1. + * @param dimension The row width. Must be at least 1. + * @param bits The bit width per (padded) dimension. Must be between {@link #MIN_BITS} and + * {@link #MAX_BITS}. + * @param seed The rotation seed. Any value; stored in the file so decoding rebuilds the + * same rotation. + * @return The quantized matrix. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or out of range, the + * array length does not match {@code rowCount * dimension}, or a value is not finite. + */ + public static QuantizedEmbeddingMatrix quantize(float[] rowMajor, int rowCount, int dimension, + int bits, long seed) { + if (rowMajor == null) { + throw new IllegalArgumentException("RowMajor must not be null"); + } + if (rowCount < 1) { + throw new IllegalArgumentException("RowCount must be at least 1, got " + rowCount); + } + if (dimension < 1) { + throw new IllegalArgumentException("Dimension must be at least 1, got " + dimension); + } + if (rowMajor.length != (long) rowCount * dimension) { + throw new IllegalArgumentException("RowMajor has " + rowMajor.length + " floats but " + + rowCount + " rows of dimension " + dimension + " need " + + ((long) rowCount * dimension)); + } + GaussianQuantizer.requireSupportedBits(bits); + final GaussianQuantizer quantizer = GaussianQuantizer.forBits(bits); + final HadamardRotation rotation = new HadamardRotation(dimension, seed); + final int paddedDimension = rotation.paddedDimension(); + final int rowBytes = rowByteCount(paddedDimension, bits); + requireStorableSize(rowCount, rowBytes); + final float[] scales = new float[rowCount]; + final byte[] codes = new byte[rowCount * rowBytes]; + final float[] decodedNorms = new float[rowCount]; + final float[] rotated = new float[paddedDimension]; + final float[] decoded = new float[paddedDimension]; + final double squareRootOfPadded = Math.sqrt(paddedDimension); + for (int row = 0; row < rowCount; row++) { + final int base = row * dimension; + double sumOfSquares = 0; + for (int d = 0; d < dimension; d++) { + final float value = rowMajor[base + d]; + if (!Float.isFinite(value)) { + throw new IllegalArgumentException("Row " + row + " has a non-finite value at " + + "dimension " + d + ": " + value + "; a quantized matrix cannot represent it"); + } + rotated[d] = value; + sumOfSquares += (double) value * value; + } + Arrays.fill(rotated, dimension, paddedDimension, 0f); + final double norm = Math.sqrt(sumOfSquares); + if (norm == 0) { + // A zero row has no direction; a zero scale decodes it to zero whatever the codes say, + // and encoding zeros keeps the bytes deterministic. + scales[row] = 0f; + final int zeroCode = quantizer.encode(0f); + for (int i = 0; i < paddedDimension; i++) { + writeCode(codes, row * rowBytes, bits, i, zeroCode); + } + continue; + } + rotation.rotate(rotated); + // Standardized coordinates are near N(0,1); encode each against the grid, then fit the + // one free scale to the row: the alpha minimizing ||z - alpha*g||^2 is (z.g)/(g.g). + final double standardize = squareRootOfPadded / norm; + double gridDot = 0; + double gridSquares = 0; + for (int i = 0; i < paddedDimension; i++) { + final float standardized = (float) (rotated[i] * standardize); + final int code = quantizer.encode(standardized); + writeCode(codes, row * rowBytes, bits, i, code); + final float level = quantizer.level(code); + gridDot += (double) standardized * level; + gridSquares += (double) level * level; + } + final double fitted = gridDot > 0 ? gridDot / gridSquares : 1.0; + scales[row] = (float) (norm / squareRootOfPadded * fitted); + // The decoded original-space norm: quantization noise leaves energy in the padding + // coordinates and truncation drops it, so the norm is measured on the truncated decode, + // not on the rotated codes. + for (int i = 0; i < paddedDimension; i++) { + decoded[i] = scales[row] * quantizer.level(readCode(codes, row * rowBytes, bits, i)); + } + rotation.inverse(decoded); + double decodedSumOfSquares = 0; + for (int d = 0; d < dimension; d++) { + decodedSumOfSquares += (double) decoded[d] * decoded[d]; + } + decodedNorms[row] = (float) Math.sqrt(decodedSumOfSquares); + } + return new QuantizedEmbeddingMatrix(rowCount, dimension, bits, seed, quantizer, scales, + codes, decodedNorms, null); + } + + /** + * {@return a copy of this matrix carrying per-row pooling weights} The weights ride along in + * the file unquantized, so a quantized file can fully replace a safetensors file that bundled + * a {@code weights} tensor. + * + * @param weights One weight per row, or {@code null} to carry none. Every weight must be + * finite. The array is copied. + * @return A matrix sharing this one's codes and scales, with the given weights. + * @throws IllegalArgumentException Thrown if {@code weights} has the wrong length or a + * non-finite value. + */ + public QuantizedEmbeddingMatrix withPoolingWeights(float[] weights) { + if (weights == null) { + return new QuantizedEmbeddingMatrix(rowCount, dimension, bits, seed, quantizer, scales, + codes, decodedNorms, null); + } + if (weights.length != rowCount) { + throw new IllegalArgumentException("Weights has " + weights.length + " values but this " + + "matrix has " + rowCount + " rows"); + } + for (int row = 0; row < rowCount; row++) { + if (!Float.isFinite(weights[row])) { + throw new IllegalArgumentException("Weight for row " + row + " is not finite: " + + weights[row]); + } + } + return new QuantizedEmbeddingMatrix(rowCount, dimension, bits, seed, quantizer, scales, + codes, decodedNorms, Arrays.copyOf(weights, weights.length)); + } + + /** + * {@return a copy of the per-row pooling weights, or {@code null} when this matrix carries + * none} + */ + public float[] poolingWeights() { + return poolingWeights == null ? null : Arrays.copyOf(poolingWeights, poolingWeights.length); + } + + /** + * Requires the packed code array to fit in one Java array. + * + * @param rowCount The number of rows. + * @param rowBytes The packed bytes per row. + * @throws IllegalArgumentException Thrown if the total exceeds what an array can hold. + */ + private static void requireStorableSize(int rowCount, int rowBytes) { + if ((long) rowCount * rowBytes > Integer.MAX_VALUE - 8) { + throw new IllegalArgumentException("The packed codes need " + ((long) rowCount * rowBytes) + + " bytes, more than one array can hold; split the matrix"); + } + } + + /** + * {@return the packed byte count of one row} Computed in long arithmetic and range-checked, so + * a padded dimension large enough to overflow {@code paddedDimension * bits} as a signed int + * (which would silently produce a negative or wrapped byte count) is rejected instead. + * + * @param paddedDimension The power-of-two padded dimension. + * @param bits The bit width per padded dimension. + * @throws IllegalArgumentException Thrown if the padded bit count exceeds what an {@code int} + * can address. + */ + private static int rowByteCount(int paddedDimension, int bits) { + final long paddedBits = (long) paddedDimension * bits; + if (paddedBits > Integer.MAX_VALUE - 7) { + throw new IllegalArgumentException("A padded dimension of " + paddedDimension + " at " + + bits + " bits needs " + paddedBits + " bits per row, more than a quantized matrix " + + "can address; use a smaller dimension"); + } + return (int) ((paddedBits + 7) / 8); + } + + /** {@return the number of rows} */ + public int rowCount() { + return rowCount; + } + + /** {@return the original row width} */ + public int dimension() { + return dimension; + } + + /** {@return the power-of-two width rows are padded to in rotated space} */ + public int paddedDimension() { + return paddedDimension; + } + + /** {@return the bit width per padded dimension} */ + public int bits() { + return bits; + } + + /** {@return the rotation seed} */ + public long seed() { + return seed; + } + + /** + * Rotates an original-space vector into this matrix's rotated space, padding it first. Rotate + * a query once, then score rows against it with {@link #dotRotated(int, float[])}. + * + * @param vector The original-space vector. Must not be {@code null} and must have length + * {@link #dimension()}. + * @return A new array of length {@link #paddedDimension()} holding the rotated vector. + * @throws IllegalArgumentException Thrown if {@code vector} is {@code null} or has the wrong + * length. + */ + public float[] rotate(float[] vector) { + if (vector == null) { + throw new IllegalArgumentException("Vector must not be null"); + } + if (vector.length != dimension) { + throw new IllegalArgumentException("Vector has length " + vector.length + + " but this matrix has dimension " + dimension); + } + final float[] padded = new float[paddedDimension]; + System.arraycopy(vector, 0, padded, 0, dimension); + rotation.rotate(padded); + return padded; + } + + /** + * Maps a rotated-space vector back to original space. Apply this once per pooled result, after + * accumulating rows with {@link #addRowRotated(int, float, float[])}; rotation is linear, so + * the sum of rotated rows is the rotation of the summed rows. + * + * @param rotated The rotated-space vector. Must not be {@code null} and must have length + * {@link #paddedDimension()}. Not modified. + * @return A new array of length {@link #dimension()} holding the original-space vector. + * @throws IllegalArgumentException Thrown if {@code rotated} is {@code null} or has the wrong + * length. + */ + public float[] toOriginal(float[] rotated) { + if (rotated == null) { + throw new IllegalArgumentException("Rotated must not be null"); + } + if (rotated.length != paddedDimension) { + throw new IllegalArgumentException("Rotated has length " + rotated.length + + " but this matrix's padded dimension is " + paddedDimension); + } + final float[] copy = rotated.clone(); + rotation.inverse(copy); + final float[] original = new float[dimension]; + System.arraycopy(copy, 0, original, 0, dimension); + return original; + } + + /** + * Adds a decoded row, times a weight, onto a rotated-space accumulator. This is the pooling + * primitive: decode stays in rotated space and costs one grid lookup per coordinate. + * + * @param row The row to add. Must be between 0 and {@code rowCount() - 1}. + * @param weight The weight to multiply the row by. + * @param sum The accumulator. Must not be {@code null} and must have length + * {@link #paddedDimension()}. + * @throws IllegalArgumentException Thrown if {@code row} is out of range or {@code sum} is + * {@code null} or has the wrong length. + */ + public void addRowRotated(int row, float weight, float[] sum) { + requireRow(row); + if (sum == null) { + throw new IllegalArgumentException("Sum must not be null"); + } + if (sum.length != paddedDimension) { + throw new IllegalArgumentException("Sum has length " + sum.length + + " but this matrix's padded dimension is " + paddedDimension); + } + final float scaledWeight = scales[row] * weight; + final int base = row * rowBytes; + for (int i = 0; i < paddedDimension; i++) { + sum[i] += scaledWeight * quantizer.level(readCode(codes, base, bits, i)); + } + } + + /** + * The dot product of a decoded row with a rotated-space query. Because the rotation is + * orthonormal, this equals the original-space dot product of the decoded row with the + * un-rotated query, up to float rounding. + * + * @param row The row to score. Must be between 0 and {@code rowCount() - 1}. + * @param rotatedQuery The query in rotated space, as returned by {@link #rotate(float[])}. + * Must not be {@code null} and must have length + * {@link #paddedDimension()}. + * @return The dot product. + * @throws IllegalArgumentException Thrown if {@code row} is out of range or + * {@code rotatedQuery} is {@code null} or has the wrong length. + */ + public double dotRotated(int row, float[] rotatedQuery) { + requireRow(row); + if (rotatedQuery == null) { + throw new IllegalArgumentException("RotatedQuery must not be null"); + } + if (rotatedQuery.length != paddedDimension) { + throw new IllegalArgumentException("RotatedQuery has length " + rotatedQuery.length + + " but this matrix's padded dimension is " + paddedDimension); + } + final int base = row * rowBytes; + double dot = 0; + for (int i = 0; i < paddedDimension; i++) { + dot += (double) rotatedQuery[i] * quantizer.level(readCode(codes, base, bits, i)); + } + return dot * scales[row]; + } + + /** + * {@return the L2 norm of the decoded original-space row, for cosine scoring} Computed + * exactly at quantize time and stored in the file: quantization noise leaves some energy in + * the padding coordinates, which decoding truncates away, so this norm matches + * {@link #decodeRow(int)}'s result rather than the rotated codes. + * + * @param row The row. Must be between 0 and {@code rowCount() - 1}. + * @throws IllegalArgumentException Thrown if {@code row} is out of range. + */ + public double rowNorm(int row) { + requireRow(row); + return decodedNorms[row]; + } + + /** + * Decodes one row back to original space. This pays the inverse rotation for a single row; + * pooling and scanning callers should stay in rotated space instead (see the class comment). + * + * @param row The row to decode. Must be between 0 and {@code rowCount() - 1}. + * @return A new array of length {@link #dimension()} holding the decoded row. + * @throws IllegalArgumentException Thrown if {@code row} is out of range. + */ + public float[] decodeRow(int row) { + requireRow(row); + final float[] rotated = new float[paddedDimension]; + addRowRotated(row, 1f, rotated); + return toOriginal(rotated); + } + + /** + * Requires a row index in range. + * + * @param row The row index to check. + * @throws IllegalArgumentException Thrown if {@code row} is out of range. + */ + private void requireRow(int row) { + if (row < 0 || row >= rowCount) { + throw new IllegalArgumentException("Row must be between 0 and " + (rowCount - 1) + + ", got " + row); + } + } + + /** + * Writes this matrix to a file, deterministically: the same matrix, bit width, and seed + * produce the same bytes. + * + * @param file The file to write. Must not be {@code null}; an existing file is replaced. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null}. + * @throws IOException Thrown if writing fails. + */ + public void write(Path file) throws IOException { + if (file == null) { + throw new IllegalArgumentException("File must not be null"); + } + try (OutputStream out = Files.newOutputStream(file); + DataOutputStream data = new DataOutputStream(new BufferedOutputStream(out))) { + data.writeInt(MAGIC); + data.writeInt(rowCount); + data.writeInt(dimension); + data.writeInt(bits); + data.writeLong(seed); + final float[] levels = quantizer.levels(); + data.writeInt(levels.length); + for (final float level : levels) { + data.writeFloat(level); + } + for (final float scale : scales) { + data.writeFloat(scale); + } + for (final float decodedNorm : decodedNorms) { + data.writeFloat(decodedNorm); + } + data.writeBoolean(poolingWeights != null); + if (poolingWeights != null) { + for (final float weight : poolingWeights) { + data.writeFloat(weight); + } + } + data.write(codes); + } + } + + /** + * Reads a matrix written by {@link #write(Path)}. The stored grid and seed rebuild exactly the + * decoder the writer used. + * + * @param file The file to read. Must not be {@code null}. + * @return The quantized matrix. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null}. + * @throws InvalidFormatException Thrown if the content is malformed: not a quantized matrix + * of a supported version, declaring sizes the file's bytes cannot back or this reader + * cannot store, storing an invalid grid or non-finite per-row values, or carrying + * trailing bytes after the declared content. + * @throws IOException Thrown if reading fails or the file is truncated. + */ + public static QuantizedEmbeddingMatrix read(Path file) throws IOException { + if (file == null) { + throw new IllegalArgumentException("File must not be null"); + } + try (InputStream in = Files.newInputStream(file); + DataInputStream data = new DataInputStream(new BufferedInputStream(in))) { + final int magic = data.readInt(); + if (magic != MAGIC) { + throw new InvalidFormatException(file + " is not a quantized embedding matrix " + + "(magic 0x" + Integer.toHexString(magic) + ", expected 0x" + + Integer.toHexString(MAGIC) + ")"); + } + final int rowCount = data.readInt(); + if (rowCount < 1) { + throw new InvalidFormatException(file + " declares " + rowCount + " rows; a " + + "quantized matrix has at least 1"); + } + final int dimension = data.readInt(); + if (dimension < 1) { + throw new InvalidFormatException(file + " declares dimension " + dimension + "; a " + + "quantized matrix's dimension is at least 1"); + } + final long fileSize = Files.size(file); + if (rowCount > fileSize || dimension > fileSize) { + throw new InvalidFormatException(file + " declares " + rowCount + " rows and dimension " + + dimension + " but holds only " + fileSize + " bytes"); + } + final int bits = data.readInt(); + try { + GaussianQuantizer.requireSupportedBits(bits); + } catch (IllegalArgumentException e) { + throw new InvalidFormatException(file + " declares an unsupported bit width: " + + e.getMessage()); + } + final long seed = data.readLong(); + final int levelCount = data.readInt(); + if (levelCount != 1 << bits) { + throw new InvalidFormatException(file + " declares " + levelCount + " grid levels " + + "for " + bits + " bits; expected " + (1 << bits)); + } + final int rowBytes; + try { + rowBytes = rowByteCount(HadamardRotation.paddedDimension(dimension), bits); + requireStorableSize(rowCount, rowBytes); + } catch (IllegalArgumentException e) { + throw new InvalidFormatException(file + " declares a matrix this reader cannot " + + "store: " + e.getMessage()); + } + // The header fully determines the file size: 28 fixed bytes (magic, row count, dimension, + // and bit width at 4 each, the 8-byte seed, the 4-byte level count), the grid levels, a + // scale and a decoded norm per row, the 1-byte pooling-weight flag, and the packed codes. + // Everything past this point allocates per-row storage, so the declared total is held + // against the bytes actually present first: a small hostile file must not force a giant + // allocation before its content is ever read. + final long declaredBytes = 28L + 4L * levelCount + (8L + rowBytes) * rowCount + 1L; + if (declaredBytes > fileSize) { + throw new InvalidFormatException(file + " declares " + rowCount + " rows and " + + "dimension " + dimension + " at " + bits + " bits, needing at least " + + declaredBytes + " bytes of scales, norms, and packed codes, but holds only " + + fileSize + " bytes"); + } + final float[] levels = new float[levelCount]; + for (int i = 0; i < levelCount; i++) { + levels[i] = data.readFloat(); + } + final GaussianQuantizer quantizer; + try { + quantizer = GaussianQuantizer.fromLevels(levels); + } catch (IllegalArgumentException e) { + throw new InvalidFormatException(file + " stores an invalid grid: " + e.getMessage()); + } + final float[] scales = new float[rowCount]; + for (int row = 0; row < rowCount; row++) { + scales[row] = data.readFloat(); + if (!Float.isFinite(scales[row])) { + throw new InvalidFormatException(file + " has a non-finite scale for row " + row + + ": " + scales[row]); + } + } + final float[] decodedNorms = new float[rowCount]; + for (int row = 0; row < rowCount; row++) { + decodedNorms[row] = data.readFloat(); + if (!Float.isFinite(decodedNorms[row]) || decodedNorms[row] < 0) { + throw new InvalidFormatException(file + " has an invalid decoded norm for row " + + row + ": " + decodedNorms[row]); + } + } + float[] poolingWeights = null; + if (data.readBoolean()) { + if (declaredBytes + 4L * rowCount > fileSize) { + throw new InvalidFormatException(file + " declares per-row pooling weights, " + + "needing at least " + (declaredBytes + 4L * rowCount) + " bytes in total, but " + + "holds only " + fileSize + " bytes"); + } + poolingWeights = new float[rowCount]; + for (int row = 0; row < rowCount; row++) { + poolingWeights[row] = data.readFloat(); + if (!Float.isFinite(poolingWeights[row])) { + throw new InvalidFormatException(file + " has a non-finite pooling weight for " + + "row " + row + ": " + poolingWeights[row]); + } + } + } + final byte[] codes = new byte[rowCount * rowBytes]; + try { + data.readFully(codes); + } catch (EOFException e) { + throw new IOException(file + " is truncated: the header declares " + rowCount + + " rows of " + rowBytes + " packed bytes, but the file ends early", e); + } + if (data.read() != -1) { + throw new InvalidFormatException(file + " has trailing bytes after the declared " + + "content; it is not a quantized matrix of this version"); + } + return new QuantizedEmbeddingMatrix(rowCount, dimension, bits, seed, quantizer, scales, + codes, decodedNorms, poolingWeights); + } + } + + /** + * {@return one packed code} Codes may straddle a byte boundary (3-bit widths do), so two + * adjacent bytes are read and the code's bits selected; the second byte is only touched when + * the code actually crosses into it, so the last code of a row never reads past its region. + * + * @param codes The packed code array. + * @param rowBase The row's first byte index. + * @param bits The code width. + * @param index The code's index within the row. + */ + private static int readCode(byte[] codes, int rowBase, int bits, int index) { + final int bitPosition = index * bits; + final int byteIndex = rowBase + (bitPosition >>> 3); + final int shift = bitPosition & 7; + int word = codes[byteIndex] & 0xFF; + if (shift + bits > 8) { + word |= (codes[byteIndex + 1] & 0xFF) << 8; + } + return (word >>> shift) & ((1 << bits) - 1); + } + + /** + * Writes one packed code, the mirror of {@link #readCode(byte[], int, int, int)}. + * + * @param codes The packed code array. + * @param rowBase The row's first byte index. + * @param bits The code width. + * @param index The code's index within the row. + * @param code The code value, within the bit width. + */ + private static void writeCode(byte[] codes, int rowBase, int bits, int index, int code) { + final int bitPosition = index * bits; + final int byteIndex = rowBase + (bitPosition >>> 3); + final int shift = bitPosition & 7; + codes[byteIndex] |= (byte) (code << shift); + if (shift + bits > 8) { + codes[byteIndex + 1] |= (byte) (code >>> (8 - shift)); + } + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedTableAdapter.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedTableAdapter.java new file mode 100644 index 0000000000..62dc180f78 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedTableAdapter.java @@ -0,0 +1,85 @@ +/* + * 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; + +/** + * The quantized {@link EmbeddingTable}: a {@link QuantizedEmbeddingMatrix} whose working space + * is rotated space, so pooling accumulates decoded rows there and pays one inverse rotation per + * pooled result, and a scan rotates the query once and scores every row without leaving rotated + * space (see the matrix's class comment for why both are safe). + */ +final class QuantizedTableAdapter implements EmbeddingTable { + + private final QuantizedEmbeddingMatrix matrix; + + /** + * Wraps a quantized matrix. + * + * @param matrix The matrix to serve rows from. + */ + QuantizedTableAdapter(QuantizedEmbeddingMatrix matrix) { + this.matrix = matrix; + } + + /** {@inheritDoc} */ + @Override + public int rowCount() { + return matrix.rowCount(); + } + + /** {@inheritDoc} */ + @Override + public int dimension() { + return matrix.dimension(); + } + + /** {@inheritDoc} */ + @Override + public int pooledLength() { + return matrix.paddedDimension(); + } + + /** {@inheritDoc} */ + @Override + public void addRow(int row, float weight, float[] sum) { + matrix.addRowRotated(row, weight, sum); + } + + /** {@inheritDoc} */ + @Override + public float[] finishPooling(float[] sum) { + return matrix.toOriginal(sum); + } + + /** {@inheritDoc} */ + @Override + public float[] prepareQuery(float[] query) { + return matrix.rotate(query); + } + + /** {@inheritDoc} */ + @Override + public double dot(int row, float[] preparedQuery) { + return matrix.dotRotated(row, preparedQuery); + } + + /** {@inheritDoc} */ + @Override + public double rowNorm(int row) { + return matrix.rowNorm(row); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/RandomizedPca.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/RandomizedPca.java new file mode 100644 index 0000000000..f64d9bde96 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/RandomizedPca.java @@ -0,0 +1,619 @@ +/* + * 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.Random; +import java.util.function.IntConsumer; +import java.util.stream.IntStream; + +/** + * Principal component analysis by randomized SVD + * (Halko, Martinsson, Tropp), the approximation + * Model2Vec's distillation performs with a dense LAPACK SVD through scikit-learn. A dense SVD of + * a vocabulary-size matrix (250k rows for a multilingual teacher) is not practical in pure Java, + * so the top components are found with a random range finder and {@value #POWER_ITERATIONS} power + * iterations, which for the fast-decaying spectrum of transformer token embeddings recovers the + * same subspace as the exact decomposition. + * + *

The column mean is subtracted before decomposition (the data matrix is modified in place), + * and the signs of the components are fixed the way scikit-learn's full solver fixes them + * ({@code svd_flip} with {@code u_based_decision=false}): each component's largest-magnitude + * coordinate is positive, so two distillations of the same teacher produce directly comparable + * vectors instead of mirror images.

+ * + *

The heavy loops are row-parallel over the common fork/join pool; all accumulation is in + * {@code double}.

+ */ +final class RandomizedPca { + + /** Extra dimensions the range finder samples beyond the requested components. */ + private static final int OVERSAMPLING = 10; + + /** Power iterations sharpening the range finder toward the dominant subspace. */ + private static final int POWER_ITERATIONS = 8; + + /** Number of row blocks the parallel loops split the matrix into. */ + private static final int BLOCKS = 32; + + /** Jacobi eigensolver convergence, relative to the largest diagonal element. */ + private static final double JACOBI_EPSILON = 1e-12; + + /** Jacobi eigensolver sweep cap; convergence arrives long before this. */ + private static final int JACOBI_MAX_SWEEPS = 100; + + /** Floor on a squared singular value, so a rank-deficient direction divides by a non-zero. */ + private static final double MIN_SQUARED_SINGULAR_VALUE = 1e-12; + + /** CholeskyQR diagonal jitter, relative to the Gram matrix's average diagonal element. */ + private static final double JITTER_RATIO = 1e-12; + + /** Factor the jitter grows by after a failed factorization. */ + private static final double JITTER_ESCALATION = 1000; + + /** Number of jitter values tried before the factorization is given up on. */ + private static final int JITTER_ATTEMPTS = 5; + + /** Not instantiable. */ + private RandomizedPca() { + } + + /** The outcome of a PCA: the projected data and how much variance the projection keeps. */ + record Result(float[] transformed, double explainedVarianceRatio) { + } + + /** + * Centers {@code data} and projects it onto its top {@code components} principal components. + * + * @param data The row-major {@code rows x cols} matrix; centered in place. + * @param rows The number of rows. + * @param cols The number of columns (the original dimension). + * @param components The number of principal components to keep; at most {@code cols} and less + * than {@code rows}. + * @param seed The random seed of the range finder; a fixed seed makes the projection + * deterministic. + * @return The projected row-major {@code rows x components} matrix and the ratio of total + * variance it explains. + * @throws IllegalArgumentException Thrown if the arguments are inconsistent, or if the data has + * no variance to decompose (every row is identical, or a value is not finite). + */ + static Result fitTransform(float[] data, int rows, int cols, int components, long seed) { + if (data == null) { + throw new IllegalArgumentException("Data must not be null"); + } + if (rows < 1 || cols < 1 || data.length != (long) rows * cols) { + throw new IllegalArgumentException("Data has " + data.length + " elements, not " + rows + + " x " + cols); + } + if (components < 1 || components > cols || components >= rows) { + throw new IllegalArgumentException("Components must be in [1, " + Math.min(cols, rows - 1) + + "], got " + components); + } + final double[] mean = columnMean(data, rows, cols); + subtractMean(data, rows, cols, mean); + final double totalVariance = totalVariance(data, rows, cols); + if (!Double.isFinite(totalVariance) || totalVariance <= 0) { + throw new IllegalArgumentException("Data has a total variance of " + totalVariance + + "; there is no subspace to find. Every row is identical, or a value is not finite."); + } + final int sampleDimensions = Math.min(components + OVERSAMPLING, cols); + final double[] omega = new double[cols * sampleDimensions]; + final Random random = new Random(seed); + for (int i = 0; i < omega.length; i++) { + omega[i] = random.nextGaussian(); + } + double[] sample = multiplyDataByDense(data, rows, cols, omega, sampleDimensions); + for (int iteration = 0; iteration < POWER_ITERATIONS; iteration++) { + orthonormalizeInPlace(sample, rows, sampleDimensions); + final double[] transposed = multiplyDataTransposedByDense(data, rows, cols, sample, + sampleDimensions); + sample = multiplyDataByDense(data, rows, cols, transposed, sampleDimensions); + } + orthonormalizeInPlace(sample, rows, sampleDimensions); + // The small matrix B = Q'X holds the data's action on the found subspace; its right singular + // vectors rotated back are the principal components. + final double[] small = multiplyBasisTransposedByData(sample, rows, sampleDimensions, + data, cols); + final double[] gram = new double[sampleDimensions * sampleDimensions]; + for (int a = 0; a < sampleDimensions; a++) { + for (int b = 0; b <= a; b++) { + double sum = 0; + for (int c = 0; c < cols; c++) { + sum += small[a * cols + c] * small[b * cols + c]; + } + gram[a * sampleDimensions + b] = sum; + gram[b * sampleDimensions + a] = sum; + } + } + final double[][] eigen = jacobiEigen(gram, sampleDimensions); + final double[] eigenvalues = eigen[0]; + final double[] eigenvectors = eigen[1]; // row-major, column j is eigenvector j + // Components in component-major layout: component j is eigenvector j of B's Gram matrix + // mapped back through B and normalized by its singular value. + final double[] componentsMajor = new double[components * cols]; + for (int j = 0; j < components; j++) { + final double singularValue = Math.sqrt(Math.max(eigenvalues[j], MIN_SQUARED_SINGULAR_VALUE)); + for (int c = 0; c < cols; c++) { + double sum = 0; + for (int a = 0; a < sampleDimensions; a++) { + sum += small[a * cols + c] * eigenvectors[a * sampleDimensions + j]; + } + componentsMajor[j * cols + c] = sum / singularValue; + } + fixSign(componentsMajor, j * cols, cols); + } + final float[] transformed = project(data, rows, cols, componentsMajor, components); + double keptVariance = 0; + for (int j = 0; j < components; j++) { + keptVariance += eigenvalues[j]; + } + return new Result(transformed, keptVariance / totalVariance); + } + + /** + * Fixes a component's sign the way scikit-learn's {@code svd_flip} with + * {@code u_based_decision=false} does: the largest-magnitude coordinate is made positive. + * + * @param componentMajor The component-major components array. + * @param offset The component's start offset. + * @param length The component's length. + */ + private static void fixSign(double[] componentMajor, int offset, int length) { + int maxIndex = 0; + double maxAbs = 0; + for (int c = 0; c < length; c++) { + final double abs = Math.abs(componentMajor[offset + c]); + if (abs > maxAbs) { + maxAbs = abs; + maxIndex = c; + } + } + if (componentMajor[offset + maxIndex] < 0) { + for (int c = 0; c < length; c++) { + componentMajor[offset + c] = -componentMajor[offset + c]; + } + } + } + + /** + * {@return the per-column means of the matrix, computed row-parallel} + * + * @param data The row-major matrix. + * @param rows The number of rows. + * @param cols The number of columns. + */ + private static double[] columnMean(float[] data, int rows, int cols) { + final double[][] partials = new double[BLOCKS][cols]; + forBlocks(rows, block -> { + final double[] partial = partials[block]; + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + for (int c = 0; c < cols; c++) { + partial[c] += data[i * cols + c]; + } + } + }); + final double[] mean = new double[cols]; + for (final double[] partial : partials) { + for (int c = 0; c < cols; c++) { + mean[c] += partial[c]; + } + } + for (int c = 0; c < cols; c++) { + mean[c] /= rows; + } + return mean; + } + + /** + * Subtracts the per-column means from the matrix in place, row-parallel. + * + * @param data The row-major matrix. + * @param rows The number of rows. + * @param cols The number of columns. + * @param mean The per-column means. + */ + private static void subtractMean(float[] data, int rows, int cols, double[] mean) { + forBlocks(rows, block -> { + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + for (int c = 0; c < cols; c++) { + final int index = i * cols + c; + data[index] = (float) (data[index] - mean[c]); + } + } + }); + } + + /** + * {@return the total variance of the centered matrix (the squared Frobenius norm), row-parallel} + * + * @param data The centered row-major matrix. + * @param rows The number of rows. + * @param cols The number of columns. + */ + private static double totalVariance(float[] data, int rows, int cols) { + final double[] partials = new double[BLOCKS]; + forBlocks(rows, block -> { + double sum = 0; + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start * cols; i < end * cols; i++) { + sum += (double) data[i] * data[i]; + } + partials[block] = sum; + }); + double total = 0; + for (final double partial : partials) { + total += partial; + } + return total; + } + + /** + * {@return the product {@code data * dense} of the float data matrix with a dense double + * matrix, row-parallel over the data} + * + * @param data The row-major {@code rows x cols} float matrix. + * @param rows The number of rows. + * @param cols The number of columns. + * @param dense The row-major {@code cols x width} double matrix. + * @param width The number of columns of {@code dense}. + */ + private static double[] multiplyDataByDense(float[] data, int rows, int cols, double[] dense, + int width) { + final double[] out = new double[rows * width]; + forBlocks(rows, block -> { + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int rowBase = i * cols; + final int outBase = i * width; + for (int c = 0; c < cols; c++) { + final float value = data[rowBase + c]; + if (value != 0) { + final int denseBase = c * width; + for (int j = 0; j < width; j++) { + out[outBase + j] += value * dense[denseBase + j]; + } + } + } + } + }); + return out; + } + + /** + * {@return the product {@code data' * dense} of the transposed float data matrix with a dense + * double matrix, row-parallel over the data with per-block partial results} + * + * @param data The row-major {@code rows x cols} float matrix. + * @param rows The number of rows. + * @param cols The number of columns. + * @param dense The row-major {@code rows x width} double matrix. + * @param width The number of columns of {@code dense}. + */ + private static double[] multiplyDataTransposedByDense(float[] data, int rows, int cols, + double[] dense, int width) { + final double[][] partials = new double[BLOCKS][cols * width]; + forBlocks(rows, block -> { + final double[] partial = partials[block]; + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int rowBase = i * cols; + final int denseBase = i * width; + for (int c = 0; c < cols; c++) { + final float value = data[rowBase + c]; + if (value != 0) { + final int outBase = c * width; + for (int j = 0; j < width; j++) { + partial[outBase + j] += value * dense[denseBase + j]; + } + } + } + } + }); + final double[] out = new double[cols * width]; + for (final double[] partial : partials) { + for (int i = 0; i < out.length; i++) { + out[i] += partial[i]; + } + } + return out; + } + + /** + * {@return the product {@code basis' * data} of the transposed orthonormal basis with the float + * data matrix, row-parallel over the data with per-block partial results} + * + * @param basis The row-major {@code rows x width} orthonormal basis. + * @param rows The number of rows. + * @param width The basis width. + * @param data The row-major {@code rows x cols} float matrix. + * @param cols The number of data columns. + */ + private static double[] multiplyBasisTransposedByData(double[] basis, int rows, int width, + float[] data, int cols) { + final double[][] partials = new double[BLOCKS][width * cols]; + forBlocks(rows, block -> { + final double[] partial = partials[block]; + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int basisBase = i * width; + final int rowBase = i * cols; + for (int a = 0; a < width; a++) { + final double value = basis[basisBase + a]; + final int outBase = a * cols; + for (int c = 0; c < cols; c++) { + partial[outBase + c] += value * data[rowBase + c]; + } + } + } + }); + final double[] out = new double[width * cols]; + for (final double[] partial : partials) { + for (int i = 0; i < out.length; i++) { + out[i] += partial[i]; + } + } + return out; + } + + /** + * {@return the projection {@code data * components'} onto the component-major components, + * row-parallel} + * + * @param data The centered row-major {@code rows x cols} float matrix. + * @param rows The number of rows. + * @param cols The number of columns. + * @param componentsMajor The row-major {@code components x cols} components. + * @param components The number of components. + */ + private static float[] project(float[] data, int rows, int cols, double[] componentsMajor, + int components) { + final float[] out = new float[rows * components]; + forBlocks(rows, block -> { + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int rowBase = i * cols; + final int outBase = i * components; + for (int j = 0; j < components; j++) { + final int componentBase = j * cols; + double sum = 0; + for (int c = 0; c < cols; c++) { + sum += data[rowBase + c] * componentsMajor[componentBase + c]; + } + out[outBase + j] = (float) sum; + } + } + }); + return out; + } + + /** + * Orthonormalizes the columns of the tall {@code rows x width} matrix in place by CholeskyQR: + * the Cholesky factor of the Gram matrix triangular-solves the basis. A diagonal jitter relative + * to the average pivot keeps the factorization alive when a power iteration has driven the + * columns toward linear dependence. + * + * @param matrix The row-major tall matrix, orthonormalized in place. + * @param rows The number of rows. + * @param width The number of columns. + */ + private static void orthonormalizeInPlace(double[] matrix, int rows, int width) { + final double[][] partials = new double[BLOCKS][width * width]; + forBlocks(rows, block -> { + final double[] partial = partials[block]; + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int base = i * width; + for (int a = 0; a < width; a++) { + final double value = matrix[base + a]; + final int gramBase = a * width; + for (int b = 0; b <= a; b++) { + partial[gramBase + b] += value * matrix[base + b]; + } + } + } + }); + final double[] gram = new double[width * width]; + for (final double[] partial : partials) { + for (int a = 0; a < width; a++) { + for (int b = 0; b <= a; b++) { + gram[a * width + b] += partial[a * width + b]; + } + } + } + for (int a = 0; a < width; a++) { + for (int b = 0; b < a; b++) { + gram[b * width + a] = gram[a * width + b]; + } + } + double trace = 0; + for (int a = 0; a < width; a++) { + trace += gram[a * width + a]; + } + // Relative to the average diagonal element, so the factorization is unchanged when the whole + // matrix is rescaled; an absolute jitter would swamp a Gram matrix of small magnitude. + double jitter = trace / width * JITTER_RATIO; + double[] lower = null; + for (int attempt = 0; attempt < JITTER_ATTEMPTS && lower == null; attempt++) { + lower = cholesky(gram, width, jitter); + jitter *= JITTER_ESCALATION; + } + if (lower == null) { + throw new IllegalStateException("Gram matrix is not positive definite even with jitter; " + + "the data columns are linearly dependent"); + } + // Solve Q L' = Y row-wise. Transposed, that is L q' = y': a forward substitution against + // the lower-triangular L. + final double[] factor = lower; + forBlocks(rows, block -> { + final int start = blockStart(rows, block); + final int end = blockStart(rows, block + 1); + for (int i = start; i < end; i++) { + final int base = i * width; + for (int j = 0; j < width; j++) { + double sum = matrix[base + j]; + for (int m = 0; m < j; m++) { + sum -= factor[j * width + m] * matrix[base + m]; + } + matrix[base + j] = sum / factor[j * width + j]; + } + } + }); + } + + /** + * {@return the lower-triangular Cholesky factor of the symmetric positive-definite matrix, or + * {@code null} when a pivot is not positive even after adding {@code jitter} to the diagonal} + * + * @param matrix The row-major symmetric matrix. + * @param width The matrix order. + * @param jitter The value added to the diagonal before factoring. + */ + private static double[] cholesky(double[] matrix, int width, double jitter) { + final double[] lower = new double[width * width]; + for (int a = 0; a < width; a++) { + for (int b = 0; b <= a; b++) { + double sum = matrix[a * width + b]; + if (a == b) { + sum += jitter; + } + for (int m = 0; m < b; m++) { + sum -= lower[a * width + m] * lower[b * width + m]; + } + if (a == b) { + if (sum <= 0) { + return null; + } + lower[a * width + a] = Math.sqrt(sum); + } else { + lower[a * width + b] = sum / lower[b * width + b]; + } + } + } + return lower; + } + + /** + * {@return the eigenpairs of a small symmetric matrix by the cyclic Jacobi method, eigenvalues + * descending; the returned array holds the eigenvalues at index 0 and the row-major eigenvector + * matrix (column j is eigenvector j) at index 1} + * + * @param matrix The row-major symmetric matrix; not modified. + * @param width The matrix order. + */ + private static double[][] jacobiEigen(double[] matrix, int width) { + final double[] a = matrix.clone(); + final double[] eigenvectors = new double[width * width]; + for (int i = 0; i < width; i++) { + eigenvectors[i * width + i] = 1; + } + for (int sweep = 0; sweep < JACOBI_MAX_SWEEPS; sweep++) { + double offDiagonal = 0; + double diagonal = 0; + for (int p = 0; p < width; p++) { + diagonal = Math.max(diagonal, Math.abs(a[p * width + p])); + for (int q = p + 1; q < width; q++) { + offDiagonal += a[p * width + q] * a[p * width + q]; + } + } + if (Math.sqrt(offDiagonal) <= JACOBI_EPSILON * Math.max(diagonal, JACOBI_EPSILON)) { + break; + } + for (int p = 0; p < width; p++) { + for (int q = p + 1; q < width; q++) { + final double apq = a[p * width + q]; + if (Math.abs(apq) <= JACOBI_EPSILON * Math.max(diagonal, JACOBI_EPSILON)) { + continue; + } + final double app = a[p * width + p]; + final double aqq = a[q * width + q]; + final double theta = (aqq - app) / (2 * apq); + final double sign = theta >= 0 ? 1 : -1; + final double t = sign / (Math.abs(theta) + Math.sqrt(theta * theta + 1)); + final double cosine = 1 / Math.sqrt(t * t + 1); + final double sine = t * cosine; + for (int k = 0; k < width; k++) { + final double akp = a[k * width + p]; + final double akq = a[k * width + q]; + a[k * width + p] = cosine * akp - sine * akq; + a[k * width + q] = sine * akp + cosine * akq; + } + for (int k = 0; k < width; k++) { + final double apk = a[p * width + k]; + final double aqk = a[q * width + k]; + a[p * width + k] = cosine * apk - sine * aqk; + a[q * width + k] = sine * apk + cosine * aqk; + } + for (int k = 0; k < width; k++) { + final double vkp = eigenvectors[k * width + p]; + final double vkq = eigenvectors[k * width + q]; + eigenvectors[k * width + p] = cosine * vkp - sine * vkq; + eigenvectors[k * width + q] = sine * vkp + cosine * vkq; + } + } + } + } + // Sort eigenpairs by eigenvalue, descending, with an insertion sort (the matrix is small). + final double[] eigenvalues = new double[width]; + for (int j = 0; j < width; j++) { + eigenvalues[j] = a[j * width + j]; + } + for (int j = 1; j < width; j++) { + int k = j; + while (k > 0 && eigenvalues[k - 1] < eigenvalues[k]) { + final double value = eigenvalues[k]; + eigenvalues[k] = eigenvalues[k - 1]; + eigenvalues[k - 1] = value; + for (int i = 0; i < width; i++) { + final double v = eigenvectors[i * width + k]; + eigenvectors[i * width + k] = eigenvectors[i * width + k - 1]; + eigenvectors[i * width + k - 1] = v; + } + k--; + } + } + return new double[][] {eigenvalues, eigenvectors}; + } + + /** + * Runs {@code action} for every row-block index in parallel over the common pool. + * + * @param rows The total number of rows. + * @param action Receives the block index, in {@code [0, BLOCKS)}. + */ + private static void forBlocks(int rows, IntConsumer action) { + IntStream.range(0, Math.min(BLOCKS, rows)).parallel().forEach(action); + } + + /** + * {@return the first row of a block} + * + * @param rows The total number of rows. + * @param block The block index; the effective block count yields the end sentinel. + */ + private static int blockStart(int rows, int block) { + return (int) ((long) rows * block / Math.min(BLOCKS, rows)); + } +} 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 new file mode 100644 index 0000000000..1d166e1002 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java @@ -0,0 +1,370 @@ +/* + * 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.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.ShortBuffer; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.java.Experimental; + +/** + * Reads a safetensors 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. The floating-point decode path + * {@link #readFloats(String)} supports the {@code F32}, {@code F16} (IEEE half) and {@code BF16} + * (bfloat16) dtypes, widening the two 16-bit types to {@code float}. + * + *

Only the header is read eagerly; tensor data is streamed into a fresh array with positional + * reads on request, so a decoded {@code float[]} is capped at {@link Integer#MAX_VALUE} - 8 + * elements. The file must stay in place and unchanged between {@link #read(Path)} and a later + * {@link #readFloats(String)} call; a file truncated in between fails loud rather than returning + * partial data.

+ * + *

Instances are immutable and safe for concurrent use: every {@link #readFloats(String)} + * call opens its own channel and decodes into a fresh array the caller owns.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +@ThreadSafe +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; + + // The JVM refuses array allocations slightly below Integer.MAX_VALUE; the exact headroom is + // implementation-specific, 8 is the commonly reserved amount. + private static final long MAX_ARRAY_LENGTH = Integer.MAX_VALUE - 8; + + private final Path file; + private final long dataStart; + private final Map tensorsByName; + private final Map metadata; + + /** Holds the parsed header; built by {@link #read(Path)}. */ + private SafetensorsFile(Path file, long dataStart, Map tensorsByName, + Map metadata) { + this.file = file; + this.dataStart = dataStart; + this.tensorsByName = tensorsByName; + this.metadata = metadata; + } + + /** + * Reads a safetensors file's header. + * + * @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. + * @throws InvalidFormatException Thrown if the file is malformed. + * @throws IOException Thrown if reading the file fails. + */ + public static SafetensorsFile 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); + } + try (FileChannel channel = FileChannel.open(file, StandardOpenOption.READ)) { + final long fileSize = channel.size(); + if (fileSize < HEADER_LENGTH_PREFIX_BYTES) { + throw new InvalidFormatException( + "File " + file + " is too short to be a safetensors file: " + fileSize + " bytes"); + } + final ByteBuffer prefix = ByteBuffer.allocate(HEADER_LENGTH_PREFIX_BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + readFully(channel, prefix, 0, file); + final long headerLength = prefix.flip().getLong(); + if (headerLength < 0 || headerLength > fileSize - HEADER_LENGTH_PREFIX_BYTES) { + throw new InvalidFormatException("File " + file + " declares a header length of " + + headerLength + ", which does not fit in a file of " + fileSize + " bytes"); + } + if (headerLength > MAX_ARRAY_LENGTH) { + throw new InvalidFormatException("File " + file + " declares a header length of " + + headerLength + " bytes, too large to decode as a single JSON string"); + } + final ByteBuffer headerBytes = ByteBuffer.allocate((int) headerLength); + readFully(channel, headerBytes, HEADER_LENGTH_PREFIX_BYTES, file); + final String headerJson = new String(headerBytes.array(), StandardCharsets.UTF_8); + final SafetensorsHeaderParser.Result parsed = SafetensorsHeaderParser.parse(headerJson); + final long dataStart = HEADER_LENGTH_PREFIX_BYTES + headerLength; + final long dataLength = fileSize - dataStart; + final Map tensorsByName = + new LinkedHashMap<>(parsed.tensors().size() * 2); + for (final TensorInfo tensor : parsed.tensors()) { + if (tensor.dataOffsetBegin() < 0 || tensor.dataOffsetEnd() < tensor.dataOffsetBegin() + || tensor.dataOffsetEnd() > dataLength) { + throw new InvalidFormatException("File " + file + " tensor '" + tensor.name() + + "' has a data range [" + tensor.dataOffsetBegin() + ", " + tensor.dataOffsetEnd() + + ") that does not fit in the file"); + } + if (tensorsByName.putIfAbsent(tensor.name(), tensor) != null) { + throw new InvalidFormatException( + "File " + file + " declares tensor '" + tensor.name() + "' more than once"); + } + } + return new SafetensorsFile(file, dataStart, Collections.unmodifiableMap(tensorsByName), + Collections.unmodifiableMap(parsed.metadata())); + } + } + + /** {@return the names of every tensor declared in the header, in header order} */ + public Set tensorNames() { + return tensorsByName.keySet(); + } + + /** + * Returns the header metadata for one tensor. + * + * @param name The tensor's name. Must not be {@code null}. + * @return The tensor's metadata. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null} or not a tensor in + * this file. + */ + public TensorInfo tensorInfo(String name) { + if (name == null) { + throw new IllegalArgumentException("Name must not be null"); + } + final TensorInfo info = tensorsByName.get(name); + if (info == null) { + throw new IllegalArgumentException( + "No tensor named '" + name + "' in this file; available: " + tensorsByName.keySet()); + } + return info; + } + + /** + * 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 + * default output dtype, so this is the common case for downloaded distilled tables. + * + * @param name The tensor's name. Must not be {@code null}. + * @return The tensor's elements in row-major (shape outermost-first) order. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null} or not a tensor in + * this file. + * @throws InvalidFormatException Thrown if the tensor is not a supported float dtype + * ({@code F32}, {@code F16}, {@code BF16}), its data range disagrees with its shape, or + * it is 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 IOException Thrown if reading the file fails. + */ + public float[] readFloats(String name) throws IOException { + final TensorInfo info = tensorInfo(name); + final int elementBytes = floatElementBytes(info.dtype(), name); + final long elementCount = info.elementCount(); + if (elementCount < 0 || elementCount > MAX_ARRAY_LENGTH) { + throw new InvalidFormatException("Tensor '" + name + "' declares " + elementCount + + " elements, more than a Java array can hold (" + MAX_ARRAY_LENGTH + + "); decoding to a float[] is capped there"); + } + final long byteLength = info.dataOffsetEnd() - info.dataOffsetBegin(); + if (byteLength != elementCount * elementBytes) { + throw new InvalidFormatException("Tensor '" + name + "' declares " + elementCount + " " + + info.dtype() + " elements but its data range is " + byteLength + " bytes"); + } + final float[] values = new float[(int) elementCount]; + final String dtype = info.dtype(); + try (FileChannel channel = FileChannel.open(file, StandardOpenOption.READ)) { + final ByteBuffer chunk = ByteBuffer.allocate((int) Math.min(READ_CHUNK_BYTES, byteLength)) + .order(ByteOrder.LITTLE_ENDIAN); + long position = dataStart + info.dataOffsetBegin(); + int decoded = 0; + while (decoded < values.length) { + chunk.clear(); + final long remainingBytes = byteLength - (long) decoded * elementBytes; + if (remainingBytes < chunk.capacity()) { + chunk.limit((int) remainingBytes); + } + readFully(channel, chunk, position, file); + chunk.flip(); + final int count = chunk.remaining() / elementBytes; + decodeInto(chunk, dtype, values, decoded, count); + decoded += count; + position += (long) count * elementBytes; + } + return values; + } + } + + /** + * Decodes an {@code F32} tensor, rejecting any other dtype. Use {@link #readFloats(String)} to + * also accept {@code F16} and {@code BF16}. + * + * @param name The tensor's name. Must not be {@code null}. + * @return The tensor's elements in row-major (shape outermost-first) order. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null} or not a tensor in + * this file. + * @throws InvalidFormatException Thrown if the tensor is not declared with dtype {@code F32}, + * its data range disagrees with its shape, or it is larger than a Java array can hold. + * @throws IllegalStateException Thrown if the file has been truncated since {@link #read(Path)}. + * @throws IOException Thrown if reading the file fails. + */ + float[] readFloat32(String name) throws IOException { + final TensorInfo info = tensorInfo(name); + if (!DTYPE_F32.equals(info.dtype())) { + throw new InvalidFormatException( + "Tensor '" + name + "' has dtype " + info.dtype() + ", not " + DTYPE_F32); + } + return readFloats(name); + } + + /** + * Widens one chunk of raw tensor bytes into the output array according to its dtype. + * + * @param chunk The raw little-endian bytes, positioned at the first element to decode. + * @param dtype The tensor's dtype ({@code F32}, {@code F16}, or {@code BF16}). + * @param out The destination array. + * @param offset The index in {@code out} to write the first decoded element to. + * @param count The number of elements to decode from {@code chunk}. + */ + private static void decodeInto(ByteBuffer chunk, String dtype, float[] out, int offset, + int count) { + switch (dtype) { + 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 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++) { + out[offset + i] = Float.intBitsToFloat((shorts.get() & 0xFFFF) << 16); + } + } + default -> throw new IllegalArgumentException("Unsupported float dtype: " + dtype); + } + } + + /** + * {@return the number of bytes one element of {@code dtype} occupies} + * + * @param dtype The tensor dtype. + * @param tensorName The tensor's name, for the error message. + * @throws InvalidFormatException Thrown if {@code dtype} is not a supported float type. + */ + private static int floatElementBytes(String dtype, String tensorName) + throws InvalidFormatException { + return switch (dtype) { + case DTYPE_F32 -> Float.BYTES; + case DTYPE_F16, DTYPE_BF16 -> Short.BYTES; + default -> throw new InvalidFormatException("Tensor '" + tensorName + "' has dtype " + + 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 DTYPE_F32.equals(dtype) || DTYPE_F16.equals(dtype) || DTYPE_BF16.equals(dtype); + } + + /** + * Fills the buffer with bytes starting at the given file position. + * + * @param channel The open channel to read from. + * @param buffer The buffer to fill. + * @param position The starting file position. + * @param file The file, for error messages. + * @throws IOException Thrown if reading fails. + * @throws IllegalStateException Thrown if the file ends before the buffer is full, which can + * only happen when the file shrank after {@link #read(Path)} validated its ranges. + */ + private static void readFully(FileChannel channel, ByteBuffer buffer, long position, Path file) + throws IOException { + while (buffer.hasRemaining()) { + final int read = channel.read(buffer, position + buffer.position()); + if (read < 0) { + throw new IllegalStateException("File " + file + " ended at byte " + + (position + buffer.position()) + + "; it has been truncated since its header was read"); + } + } + } + + /** + * Finds the single 2-dimensional floating-point tensor in this file (dtype {@code F32}, + * {@code F16}, or {@code BF16}), the shape a static embedding table's weight matrix takes + * (vocabulary size by hidden dimension). Strict rather than guessing a name convention, so a + * wrong guess cannot silently load the wrong tensor. + * + * @return The name of the single 2-D float tensor. + * @throws InvalidFormatException Thrown if the file has zero or more than one 2-D float + * tensor; the message lists every candidate so the caller can pick explicitly with + * {@link #readFloats(String)}. + */ + public String singleMatrixTensorName() throws InvalidFormatException { + String found = null; + for (final TensorInfo info : tensorsByName.values()) { + if (isFloatDtype(info.dtype()) && info.shape().length == 2) { + if (found != null) { + throw new InvalidFormatException( + "More than one 2-D float tensor in this file; specify the name explicitly. " + + "Candidates: " + tensorsByName.keySet()); + } + found = info.name(); + } + } + if (found == null) { + throw new InvalidFormatException( + "No 2-D float (F32/F16/BF16) tensor in this file. Available tensors: " + + tensorsByName.keySet()); + } + return found; + } + + /** {@return the file's {@code __metadata__} string map, empty when the header has none} */ + Map metadata() { + return metadata; + } + + /** {@return the total number of tensors declared in this file} */ + public int size() { + return tensorsByName.size(); + } +} 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 new file mode 100644 index 0000000000..8916008966 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java @@ -0,0 +1,240 @@ +/* + * 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.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import opennlp.tools.util.InvalidFormatException; + +/** + * A cursor parser for the JSON header of a safetensors file: a flat object of tensor name to a + * {@code dtype}/{@code shape}/{@code data_offsets} record, plus an optional {@code __metadata__} + * string map. Not a general-purpose JSON parser; it fails loud on anything outside that shape. + */ +final class SafetensorsHeaderParser { + + private static final String METADATA_KEY = "__metadata__"; + + private final JsonCursor cursor; + + /** Wraps the header text in a cursor; driven by {@link #parse(String)}. */ + private SafetensorsHeaderParser(String text) { + this.cursor = new JsonCursor(text, "safetensors header"); + } + + /** + * Parses a safetensors header. + * + * @param headerJson The header's JSON text, decoded from the file's header bytes. Must not be + * {@code null}. + * @return The parse result: the declared tensors, in header order, and the + * {@code __metadata__} string map (empty when the header has none). + * @throws IllegalArgumentException Thrown if {@code headerJson} is {@code null}. + * @throws InvalidFormatException Thrown if {@code headerJson} is malformed. + */ + static Result parse(String headerJson) throws InvalidFormatException { + if (headerJson == null) { + throw new IllegalArgumentException("HeaderJson must not be null"); + } + final SafetensorsHeaderParser parser = new SafetensorsHeaderParser(headerJson); + return parser.parseTop(); + } + + /** {@return the parsed header: its tensors in header order and the {@code __metadata__} map} */ + private Result parseTop() throws InvalidFormatException { + final List tensors = new ArrayList<>(); + Map metadata = Map.of(); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + if (cursor.peek() == '}') { + cursor.consume(); + requireEnd(); + return new Result(tensors, metadata); + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if (METADATA_KEY.equals(key)) { + metadata = parseStringMap(); + } else { + tensors.add(parseTensorInfo(key)); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a header entry, got '" + next + "'"); + } + requireEnd(); + return new Result(tensors, metadata); + } + + /** + * Requires the rest of the header to be whitespace only. Trailing whitespace is legal (writers + * space-pad the header to align the data section); other trailing content is a length mismatch. + */ + private void requireEnd() throws InvalidFormatException { + cursor.requireEnd("Trailing content after the header object"); + } + + /** + * {@return one tensor's metadata, parsed from its header record} + * + * @param name The tensor's name, the key it was declared under. + */ + private TensorInfo parseTensorInfo(String name) throws InvalidFormatException { + cursor.expect('{'); + String dtype = null; + int[] shape = null; + long dataOffsetBegin = -1; + long dataOffsetEnd = -1; + cursor.skipWhitespace(); + while (cursor.peek() != '}') { + cursor.skipWhitespace(); + final String field = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (field) { + case "dtype" -> dtype = cursor.parseString(); + case "shape" -> shape = parseIntArray(); + case "data_offsets" -> { + final long[] offsets = parseLongArray(); + if (offsets.length != 2) { + throw cursor.malformed("Tensor '" + name + "' data_offsets must have exactly 2 " + + "elements, got " + offsets.length); + } + dataOffsetBegin = offsets[0]; + dataOffsetEnd = offsets[1]; + } + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + continue; + } + if (next == '}') { + if (dtype == null || shape == null || dataOffsetBegin < 0) { + throw cursor.malformed("Tensor '" + name + + "' is missing dtype, shape, or data_offsets"); + } + return new TensorInfo(name, dtype, shape, dataOffsetBegin, dataOffsetEnd); + } + throw cursor.malformed("Expected ',' or '}' in tensor '" + name + "', got '" + next + "'"); + } + throw cursor.malformed("Tensor '" + name + "' has an empty object; missing dtype, shape, " + + "and data_offsets"); + } + + /** {@return a JSON object of string values, used for the {@code __metadata__} map} */ + private Map parseStringMap() throws InvalidFormatException { + final Map map = new LinkedHashMap<>(); + cursor.expect('{'); + cursor.skipWhitespace(); + if (cursor.peek() == '}') { + cursor.consume(); + return map; + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + map.put(key, cursor.parseString()); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return map; + } + throw cursor.malformed("Expected ',' or '}' in __metadata__, got '" + next + "'"); + } + } + + /** + * {@return a JSON array of non-negative integers as an {@code int[]}} + * + * @throws InvalidFormatException Thrown if any element is outside the {@code int} range. + */ + private int[] parseIntArray() throws InvalidFormatException { + final long[] longs = parseLongArray(); + final int[] ints = new int[longs.length]; + for (int i = 0; i < longs.length; i++) { + if (longs[i] < 0 || longs[i] > Integer.MAX_VALUE) { + throw cursor.malformed("Shape dimension out of int range: " + longs[i]); + } + ints[i] = (int) longs[i]; + } + return ints; + } + + /** {@return a JSON array of integers as a {@code long[]}} */ + private long[] parseLongArray() throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List values = new ArrayList<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return new long[0]; + } + while (true) { + cursor.skipWhitespace(); + values.add(cursor.parseLong()); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + break; + } + throw cursor.malformed("Expected ',' or ']' in a number array, got '" + next + "'"); + } + final long[] array = new long[values.size()]; + for (int i = 0; i < array.length; i++) { + array[i] = values.get(i); + } + return array; + } + + /** + * The parsed header: the declared tensors, in header order, and the {@code __metadata__} + * string map. + * + * @param tensors The declared tensors, in header order. Never {@code null}. + * @param metadata The {@code __metadata__} string map, empty when the header has none. Never + * {@code null}. + */ + record Result(List tensors, Map metadata) { + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsWriter.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsWriter.java new file mode 100644 index 0000000000..4475236932 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsWriter.java @@ -0,0 +1,126 @@ +/* + * 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.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; + +/** + * Writes a safetensors file holding a + * single 2-D {@code F32} tensor, the shape a distilled embedding table takes (vocabulary size by + * output dimension). This is the write side of the format {@link SafetensorsFile} reads; the data + * is streamed to the file in chunks so the writer's overhead beyond the caller's matrix is + * constant. + */ +final class SafetensorsWriter { + + /** The name of the embedding matrix tensor, the Model2Vec convention. */ + static final String EMBEDDINGS_TENSOR = "embeddings"; + + /** The size of the encoding buffer the matrix is streamed through; a multiple of Float.BYTES. */ + private static final int WRITE_CHUNK_BYTES = 1 << 20; + + /** + * The boundary the header is space-padded to, so the tensor data starts aligned. The reference + * safetensors writer pads the same way, and readers that memory-map the data section rely on it. + */ + private static final int HEADER_ALIGNMENT_BYTES = 8; + + /** The byte the header is padded with; JSON treats it as insignificant whitespace. */ + private static final byte HEADER_PADDING = ' '; + + /** Not instantiable. */ + private SafetensorsWriter() { + } + + /** + * Writes a row-major float matrix as a one-tensor safetensors file. + * + * @param file The file to write, replaced when it exists. Must not be {@code null}. + * @param rows The number of matrix rows. + * @param cols The number of matrix columns. + * @param values The matrix values in row-major order, {@code rows * cols} of them. Must not be + * {@code null}. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or the value count + * does not match the shape. + * @throws IOException Thrown if writing fails. + */ + static void writeMatrix(Path file, int rows, int cols, float[] values) throws IOException { + if (file == null) { + throw new IllegalArgumentException("File must not be null"); + } + if (values == null) { + throw new IllegalArgumentException("Values must not be null"); + } + if (rows < 1 || cols < 1 || values.length != (long) rows * cols) { + throw new IllegalArgumentException("Values has " + values.length + " elements, not " + rows + + " x " + cols); + } + final long dataBytes = (long) values.length * Float.BYTES; + final String header = "{\"" + EMBEDDINGS_TENSOR + "\":{\"dtype\":\"F32\",\"shape\":[" + rows + + "," + cols + "],\"data_offsets\":[0," + dataBytes + "]}}"; + final byte[] headerBytes = header.getBytes(StandardCharsets.UTF_8); + final int padding = (HEADER_ALIGNMENT_BYTES + - (Long.BYTES + headerBytes.length) % HEADER_ALIGNMENT_BYTES) % HEADER_ALIGNMENT_BYTES; + final Path parent = file.getParent(); + if (parent != null) { + Files.createDirectories(parent); + } + try (FileChannel channel = FileChannel.open(file, StandardOpenOption.CREATE, + StandardOpenOption.WRITE, StandardOpenOption.TRUNCATE_EXISTING)) { + final ByteBuffer prefix = ByteBuffer.allocate(Long.BYTES + headerBytes.length + padding) + .order(ByteOrder.LITTLE_ENDIAN); + prefix.putLong((long) headerBytes.length + padding); + prefix.put(headerBytes); + for (int i = 0; i < padding; i++) { + prefix.put(HEADER_PADDING); + } + prefix.flip(); + writeFully(channel, prefix); + final ByteBuffer chunk = ByteBuffer.allocate(WRITE_CHUNK_BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + int written = 0; + while (written < values.length) { + chunk.clear(); + final int count = Math.min(values.length - written, WRITE_CHUNK_BYTES / Float.BYTES); + chunk.asFloatBuffer().put(values, written, count); + chunk.limit(count * Float.BYTES); + writeFully(channel, chunk); + written += count; + } + } + } + + /** + * Writes the buffer's remaining bytes to the channel. + * + * @param channel The open channel. + * @param buffer The buffer to drain. + * @throws IOException Thrown if writing fails. + */ + private static void writeFully(FileChannel channel, ByteBuffer buffer) throws IOException { + while (buffer.hasRemaining()) { + channel.write(buffer); + } + } +} 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 new file mode 100644 index 0000000000..6fb031d96a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java @@ -0,0 +1,1254 @@ +/* + * 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.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.SortedSet; +import java.util.TreeSet; +import java.util.function.IntConsumer; +import java.util.function.IntPredicate; + +import opennlp.subword.sentencepiece.SentencePieceTokenizer; +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.embeddings.TextEmbedder; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.tokenize.SubwordTokenizer; +import opennlp.tools.tokenize.WordpieceEncoder; +import opennlp.tools.tokenize.WordpieceTokenizer; +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.java.Experimental; + +/** + * A static (non-contextual) sentence embedding model: a per-token vector table plus subword + * tokenization. Embedding a sentence is tokenize, gather each piece's row, optionally weight, + * mean-pool, and optionally L2-normalize; there is no model forward pass. + * + *

It loads distilled tables in the + * Model2Vec release layout for both + * tokenizer families: + * WordPiece models carry a {@code vocab.txt} whose line number is the matrix row, and + * Unigram models carry a {@code tokenizer.json} whose {@code model.vocab} list order is the row + * order and whose normalizer and scores drive segmentation. Legacy SentencePiece layouts with a + * separate trained {@code .model} file remain supported. In both cases the + * {@code model.safetensors} holds one 2-D float matrix, with + * an optional per-token {@code weights} tensor. Matrix rows are resolved by piece string, + * 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.

+ * + *

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.

+ * + *

Either layout may carry its matrix quantized in a {@code model.quantized} file (written by + * the {@code QuantizeModel} tool), which holds the matrix and any per-token pooling weights + * itself. A directory presents exactly one matrix file: the quantized file or the safetensors, + * not both. A directory carrying both is rejected, because the quantizer writes the quantized + * file next to the safetensors it read and so a directory holding both has not declared which + * is authoritative; delete one to choose. Embedding and similarity over a quantized matrix + * behave identically up to the quantization error of the chosen bit width; see + * {@link QuantizedEmbeddingMatrix} for the storage and its cost.

+ * + *

A model directory may additionally carry a {@code terms.txt}: whole words and multi-word + * phrases distilled through the teacher as units, owning the matrix rows after the subword rows + * (see {@link ModelDistiller}). Embedding then matches the text against these terms greedily + * longest-first, pools a matched term's single row in place of its words' subword pieces, and + * tokenizes only the text between matches; a model without the file embeds exactly as before. + * Term matching is case-insensitive regardless of the subword tokenizer's casing.

+ * + *

Instances are immutable and safe for concurrent use after construction.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +@ThreadSafe +public final class StaticEmbeddingModel implements TextEmbedder { + + /** 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; + // Shared with ModelQuantizer, which carries this tensor into the quantized file. + static final String WEIGHTS_TENSOR_NAME = "weights"; + // The only pooling this model implements; the value the distiller writes into config.json. + private static final String MEAN_POOLING = "mean"; + 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 WORDPIECE_SPECIAL_TOKENS = + Set.of(WordpieceTokenizer.BERT_CLS_TOKEN, WordpieceTokenizer.BERT_SEP_TOKEN, + WordpieceTokenizer.BERT_UNK_TOKEN, "[PAD]", "[MASK]"); + private static final Set SENTENCEPIECE_SPECIAL_TOKENS = + Set.of("", "", "", "", ""); + + private final EmbeddingTable table; + private final float[] weights; + private final int dimension; + private final EmbeddingVocabulary vocabulary; + private final SubwordTokenizer tokenizer; + // Tokenizer-id test for pieces that are never pooled (delimiter, control, unknown pieces). + private final IntPredicate skipPieceId; + private final boolean normalize; + // Special-token mask, precomputed at load time for the neighbor scan. + private final boolean[] specialRows; + // The term rows after the subword rows; empty for a model without a term table. + private final TermTable terms; + + /** Holds the loaded, validated state; callers reach this through the {@code load} factories. */ + private StaticEmbeddingModel(EmbeddingTable table, float[] weights, + EmbeddingVocabulary vocabulary, SubwordTokenizer tokenizer, + IntPredicate skipPieceId, boolean normalize, + boolean[] specialRows, TermTable terms) { + this.table = table; + this.weights = weights; + this.dimension = table.dimension(); + this.vocabulary = vocabulary; + this.tokenizer = tokenizer; + this.skipPieceId = skipPieceId; + this.normalize = normalize; + this.specialRows = specialRows; + this.terms = terms; + } + + /** An embedding table and the optional per-token pooling weights that came with it. */ + private record TableAndWeights(EmbeddingTable table, float[] weights) { + } + + /** + * Reads a quantized table, holding its row count to the vocabulary's size plus the term + * count. + * + * @param quantizedFile The quantized matrix file. + * @param vocabulary The matrix row vocabulary. + * @param termCount The number of term rows after the vocabulary rows. + * @param vocabularySourceName The vocabulary's source, for error messages. + * @return The table and the pooling weights the file carries, if any. + * @throws InvalidFormatException Thrown if the row count disagrees with the vocabulary. + * @throws IOException Thrown if reading the file fails. + */ + private static TableAndWeights readQuantizedTable(Path quantizedFile, + EmbeddingVocabulary vocabulary, + int termCount, + String vocabularySourceName) + throws IOException { + final QuantizedEmbeddingMatrix matrix = QuantizedEmbeddingMatrix.read(quantizedFile); + final int expectedRows = vocabulary.size() + termCount; + if (matrix.rowCount() != expectedRows) { + throw new InvalidFormatException("Vocabulary " + vocabularySourceName + " has " + + vocabulary.size() + " tokens" + + (termCount > 0 ? " plus " + termCount + " terms" : "") + + " but quantized matrix " + quantizedFile + " has " + + matrix.rowCount() + " rows; these files do not belong to the same model"); + } + return new TableAndWeights(new QuantizedTableAdapter(matrix), matrix.poolingWeights()); + } + + /** + * Loads a static embedding model from a model directory, detecting the tokenizer family from + * the files present and reading the pooling switch ({@code normalize}) from the model's + * {@code config.json}. + * + *

A directory with a {@code vocab.txt} is a WordPiece model; its casing is read from + * {@code do_lower_case} in {@code tokenizer_config.json}. A {@code strip_accents} that + * explicitly disagrees with {@code do_lower_case} cannot be represented by the single + * lower-case switch of {@link #load(Path, Path, Casing, Normalization)} and is rejected rather + * than silently mis-tokenized; when absent or {@code null} it follows the BERT convention of + * stripping accents exactly when lower-casing. When both layouts are present, the + * {@code vocab.txt} wins.

+ * + *

A directory with a Unigram {@code tokenizer.json} is a Model2Vec Unigram model. Its + * vocabulary, scores, precompiled normalizer, and supported post-normalization steps are read + * directly from JSON. A legacy directory may instead carry a trained SentencePiece file + * ({@code sentencepiece.bpe.model}, {@code spiece.model}, or {@code tokenizer.model}) next to + * the JSON vocabulary; that file remains supported.

+ * + *

In either layout, the matrix comes from a {@code model.quantized} file when the directory + * has one and no {@code model.safetensors}; a directory holding both is rejected (see the + * class comment). After quantizing, delete the safetensors to deploy the quantized matrix, or + * delete the quantized file to fall back to the float matrix.

+ * + * @param modelDirectory The model directory. Must not be {@code null} and must be a + * directory. + * @return The loaded model. + * @throws IllegalArgumentException Thrown if {@code modelDirectory} is {@code null} or not a + * directory. + * @throws InvalidFormatException Thrown if neither layout's files are present, a required + * file is missing, a configuration file is malformed or lacks its field, the accent + * handling is not representable, or the tokenizer and the embedding matrix disagree. + * @throws IOException Thrown if reading a file fails. + */ + public static StaticEmbeddingModel load(Path modelDirectory) throws IOException { + if (modelDirectory == null) { + throw new IllegalArgumentException("ModelDirectory must not be null"); + } + if (!Files.isDirectory(modelDirectory)) { + throw new IllegalArgumentException( + "Model directory does not exist or is not a directory: " + modelDirectory); + } + final Path termsFile = modelDirectory.resolve(ModelFileNames.TERMS); + final List termLines = Files.isRegularFile(termsFile) + ? Files.readAllLines(termsFile) : List.of(); + final Path vocabularyFile = modelDirectory.resolve(ModelFileNames.VOCABULARY); + if (Files.isRegularFile(vocabularyFile)) { + return loadWordpieceDirectory(modelDirectory, vocabularyFile, termLines, + termsFile.toString()); + } + final Path sentencePieceModelFile = ModelFileNames.firstRegularFile(modelDirectory, + ModelFileNames.SENTENCEPIECE_MODELS); + final Path tokenizerJsonFile = modelDirectory.resolve(ModelFileNames.TOKENIZER_JSON); + if (sentencePieceModelFile != null && Files.isRegularFile(tokenizerJsonFile)) { + final Normalization normalization = + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)); + final Path quantizedFile = quantizedMatrixFileOrNull(modelDirectory); + if (quantizedFile != null) { + return loadSentencePieceQuantized(sentencePieceModelFile, tokenizerJsonFile, + quantizedFile, normalization, termLines, termsFile.toString()); + } + return loadSentencePiece(sentencePieceModelFile, tokenizerJsonFile, + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), normalization, + termLines, termsFile.toString()); + } + if (Files.isRegularFile(tokenizerJsonFile)) { + return loadModel2VecUnigram(tokenizerJsonFile, + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)), + termLines, termsFile.toString()); + } + throw new InvalidFormatException("Model directory " + modelDirectory + " has neither a " + + ModelFileNames.VOCABULARY + " (WordPiece layout) nor a " + + ModelFileNames.TOKENIZER_JSON + " (Unigram layout)"); + } + + /** Loads a self-contained Model2Vec Unigram directory. */ + private static StaticEmbeddingModel loadModel2VecUnigram( + Path tokenizerJsonFile, Path safetensorsFile, Normalization normalization, + List termLines, String termsSourceName) throws IOException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.fromTokenizerJson(tokenizerJsonFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + final Model2VecUnigramTokenizer tokenizer; + try { + tokenizer = Model2VecUnigramTokenizer.load(tokenizerJsonFile); + } catch (InvalidFormatException e) { + throw new InvalidFormatException("Unigram model needs either a self-contained " + + "tokenizer.json or a trained SentencePiece .model file: " + e.getMessage(), e); + } + requireVocabularyCoverage(tokenizer, vocabulary, tokenizerJsonFile); + final Matrix matrix = readMatrix(vocabulary, terms.size(), safetensorsFile, + tokenizerJsonFile.toString()); + final TableAndWeights tableAndWeights = new TableAndWeights( + new FloatEmbeddingTable(matrix.embeddings(), matrix.dimension(), + vocabulary.size() + terms.size()), + matrix.weights()); + final IntPredicate skipPieceId = + id -> tokenizer.isUnknown(id) || tokenizer.isControl(id); + return new StaticEmbeddingModel(tableAndWeights.table(), tableAndWeights.weights(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + specialRows(vocabulary, SENTENCEPIECE_SPECIAL_TOKENS, + tableAndWeights.table().rowCount()), + terms); + } + + /** Verifies the self-contained tokenizer and matrix vocabulary agree. */ + private static void requireVocabularyCoverage( + Model2VecUnigramTokenizer tokenizer, EmbeddingVocabulary vocabulary, + Path tokenizerJsonFile) throws InvalidFormatException { + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (!tokenizer.isUnknown(id) && !tokenizer.isControl(id) + && vocabulary.id(tokenizer.idToPiece(id)) < 0) { + throw new InvalidFormatException(tokenizerJsonFile + " defines tokenizer piece '" + + tokenizer.idToPiece(id) + "' without a matrix row"); + } + } + } + + /** + * Loads the WordPiece directory layout, reading the tokenizer and pooling switches from the + * model's own configuration files. + * + * @param modelDirectory The model directory. + * @param vocabularyFile The directory's {@code vocab.txt}. + * @param termLines The directory's terms in row order; empty without a terms file. + * @param termsSourceName The terms' source, for error messages. + * @return The loaded model. + * @throws IOException Thrown if reading a file fails. + */ + private static StaticEmbeddingModel loadWordpieceDirectory(Path modelDirectory, + Path vocabularyFile, + List termLines, + String termsSourceName) + throws IOException { + final Path tokenizerConfigFile = + requiredFile(modelDirectory, ModelFileNames.TOKENIZER_CONFIG); + final Normalization normalization = + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)); + final Boolean lowerCase = + FlatJsonFields.topLevelBoolean(tokenizerConfigFile, "do_lower_case"); + if (lowerCase == null) { + throw new InvalidFormatException(tokenizerConfigFile + " has no boolean " + + "'do_lower_case' field; use load(vocabularyFile, safetensorsFile, casing, " + + "normalization) and choose explicitly"); + } + final Boolean stripAccents = + FlatJsonFields.topLevelBoolean(tokenizerConfigFile, "strip_accents"); + if (stripAccents != null && !stripAccents.equals(lowerCase)) { + throw new InvalidFormatException(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, casing, normalization) after choosing " + + "deliberately"); + } + final Casing casing = lowerCase ? Casing.UNCASED : Casing.CASED; + final Path quantizedFile = quantizedMatrixFileOrNull(modelDirectory); + if (quantizedFile != null) { + final EmbeddingVocabulary vocabulary = EmbeddingVocabulary.fromVocabTxt(vocabularyFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + return createWordpiece(vocabulary, + readQuantizedTable(quantizedFile, vocabulary, terms.size(), + vocabularyFile.toString()), + casing, normalization, vocabularyFile.toString(), terms); + } + return loadWordpiece(vocabularyFile, + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), + casing, normalization, termLines, termsSourceName); + } + + /** + * Resolves which matrix file a model directory presents, failing loud when the choice is + * ambiguous. The quantizer writes {@code model.quantized} next to the {@code model.safetensors} + * it read, so a directory holding both has not declared which is authoritative; deleting one + * makes the deployment's choice explicit rather than letting the loader guess. + * + * @param modelDirectory The model directory. + * @return the {@code model.quantized} file when it is the directory's only matrix file, or + * {@code null} when the directory presents only a {@code model.safetensors}. + * @throws InvalidFormatException Thrown if the directory holds both matrix files. + */ + private static Path quantizedMatrixFileOrNull(Path modelDirectory) throws InvalidFormatException { + final Path quantizedFile = modelDirectory.resolve(ModelFileNames.QUANTIZED); + if (!Files.isRegularFile(quantizedFile)) { + return null; + } + if (Files.isRegularFile(modelDirectory.resolve(ModelFileNames.SAFETENSORS))) { + throw new InvalidFormatException("Model directory " + modelDirectory + " has both " + + ModelFileNames.QUANTIZED + " and " + ModelFileNames.SAFETENSORS + "; delete one so " + + "the matrix source is unambiguous (keep " + ModelFileNames.QUANTIZED + " for a " + + "quantized deployment, or " + ModelFileNames.SAFETENSORS + " for the float matrix)"); + } + return quantizedFile; + } + + /** + * Reads the required {@code normalize} switch out of a model's {@code config.json}, rejecting + * a configuration whose {@code pooling} field declares anything but the mean pooling this + * model implements. Silently mean-pooling a table distilled for another pooling would produce + * plausible but wrong vectors, so such a model fails loud here. + * + * @param configFile The {@code config.json} file. + * @return The corresponding {@link Normalization}. + * @throws InvalidFormatException Thrown if the {@code normalize} field is missing or not a + * boolean, or the {@code pooling} field declares a pooling other than {@code "mean"}. + * @throws IOException Thrown if reading the file fails. + */ + private static Normalization requiredNormalize(Path configFile) throws IOException { + final String pooling = FlatJsonFields.topLevelString(configFile, "pooling"); + if (pooling != null && !MEAN_POOLING.equals(pooling)) { + throw new InvalidFormatException(configFile + " declares pooling '" + pooling + + "' but only '" + MEAN_POOLING + "' pooling is implemented; embedding this model " + + "would silently pool differently than its distiller intended"); + } + final Boolean normalize = FlatJsonFields.topLevelBoolean(configFile, "normalize"); + if (normalize == null) { + throw new InvalidFormatException(configFile + " has no boolean 'normalize' field; " + + "use the explicit load overloads and choose the normalization deliberately"); + } + return normalize ? Normalization.L2 : Normalization.NONE; + } + + /** + * {@return the named file in the directory, requiring it to exist as a regular file} + * + * @param modelDirectory The model directory. + * @param name The required file name. + * @throws InvalidFormatException Thrown if the file is absent. + */ + private static Path requiredFile(Path modelDirectory, String name) + throws InvalidFormatException { + final Path file = modelDirectory.resolve(name); + if (!Files.isRegularFile(file)) { + throw new InvalidFormatException("Model directory " + modelDirectory + " has no " + + name + "; for a different layout, use the explicit load overloads"); + } + return file; + } + + /** + * Loads a WordPiece static embedding model from a BERT-style {@code vocab.txt} and a + * safetensors weight file. No model is bundled with this module; the caller supplies the + * files. + * + * @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]} + * 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. + * 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 casing Whether the tokenizer lower-cases and strips accents + * ({@link Casing#UNCASED}) or preserves case ({@link Casing#CASED}). + * @param normalization Whether {@link #embed(String)} L2-normalizes its result + * ({@link Normalization#L2}) or not ({@link Normalization#NONE}). + * @return The loaded model. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or a file is + * missing. + * @throws InvalidFormatException Thrown if a file is malformed, the vocabulary lacks the + * {@code [UNK]} token, 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, + Casing casing, Normalization normalization) + throws IOException { + return loadWordpiece(vocabularyFile, safetensorsFile, casing, normalization, List.of(), + ModelFileNames.TERMS); + } + + /** + * Loads the WordPiece layout with an optional term table. + * + * @param vocabularyFile The {@code vocab.txt} file. + * @param safetensorsFile The {@code model.safetensors} file. + * @param casing The tokenizer's casing. + * @param normalization The pooling normalization. + * @param termLines The terms in row order; empty without a term table. + * @param termsSourceName The terms' source, for error messages. + * @return The loaded model. + * @throws IOException Thrown if reading a file fails. + */ + private static StaticEmbeddingModel loadWordpiece(Path vocabularyFile, Path safetensorsFile, + Casing casing, Normalization normalization, + List termLines, + String termsSourceName) + 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 EmbeddingVocabulary vocabulary = EmbeddingVocabulary.fromVocabTxt(vocabularyFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + final Matrix matrix = readMatrix(vocabulary, terms.size(), safetensorsFile, + vocabularyFile.toString()); + final TableAndWeights tableAndWeights = new TableAndWeights( + new FloatEmbeddingTable(matrix.embeddings(), matrix.dimension(), + vocabulary.size() + terms.size()), + matrix.weights()); + return createWordpiece(vocabulary, tableAndWeights, casing, normalization, + vocabularyFile.toString(), terms); + } + + /** + * Builds a WordPiece model over a loaded table, whatever its storage form. + * + * @param vocabulary The matrix row vocabulary. + * @param tableAndWeights The table and its optional pooling weights. + * @param casing The tokenizer casing. + * @param normalization The pooling normalization. + * @param vocabularySourceName The vocabulary's source, for error messages. + * @param terms The term rows after the subword rows; empty for none. + * @return The loaded model. + * @throws InvalidFormatException Thrown if the vocabulary has no unknown token. + */ + private static StaticEmbeddingModel createWordpiece(EmbeddingVocabulary vocabulary, + TableAndWeights tableAndWeights, + Casing casing, + Normalization normalization, + String vocabularySourceName, + TermTable terms) + throws InvalidFormatException { + final int unknownId = vocabulary.id(WordpieceTokenizer.BERT_UNK_TOKEN); + if (unknownId < 0) { + throw new InvalidFormatException("Vocabulary " + vocabularySourceName + " has no " + + WordpieceTokenizer.BERT_UNK_TOKEN + " token; a WordPiece embedding model needs an " + + "unknown token as the fallback for out-of-vocabulary text"); + } + final WordpieceEncoder tokenizer = + wordpieceEncoder(vocabulary, casing == Casing.UNCASED, unknownId); + // 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 = + id -> id == unknownId || id == classificationId || id == separatorId; + return new StaticEmbeddingModel(tableAndWeights.table(), tableAndWeights.weights(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + specialRows(vocabulary, WORDPIECE_SPECIAL_TOKENS, tableAndWeights.table().rowCount()), + terms); + } + + /** + * 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 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 id of {@code [CLS]} or + * {@code [SEP]} when that token is absent. + * @return The encoder. + */ + private static WordpieceEncoder wordpieceEncoder(EmbeddingVocabulary vocabulary, + boolean lowerCase, int unknownId) { + if (vocabulary.id(WordpieceTokenizer.BERT_CLS_TOKEN) >= 0 + && vocabulary.id(WordpieceTokenizer.BERT_SEP_TOKEN) >= 0) { + return new WordpieceEncoder(vocabulary.orderedTokens(), lowerCase); + } + final List tokens = vocabulary.orderedTokens(); + final Map ids = new HashMap<>(tokens.size() * 2); + for (int id = 0; id < tokens.size(); id++) { + ids.put(tokens.get(id), id); + } + ids.putIfAbsent(WordpieceTokenizer.BERT_CLS_TOKEN, unknownId); + ids.putIfAbsent(WordpieceTokenizer.BERT_SEP_TOKEN, unknownId); + return new WordpieceEncoder(ids, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); + } + + /** + * Loads a SentencePiece static embedding model from a trained SentencePiece {@code .model} + * file, the Unigram {@code tokenizer.json} naming the matrix rows, and a safetensors weight + * file. No model is bundled with this module; the caller supplies the files. + * + *

The {@code .model} file carries the model's own text normalizer and segmentation state, + * so there is no casing switch. The two vocabulary files may order or offset their ids + * differently: matrix rows are resolved by piece string, and every piece the tokenizer can + * emit (except its control and unknown pieces, which are never pooled) must be present in the + * {@code tokenizer.json} vocabulary, verified once at load time.

+ * + * @param sentencePieceModelFile The trained SentencePiece {@code .model} file. Must not be + * {@code null} and must exist. + * @param tokenizerJsonFile The Unigram {@code tokenizer.json} file; its + * {@code model.vocab} list order is the matrix row order, with + * {@code added_tokens} overlaid. 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 float 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 normalization Whether {@link #embed(String)} L2-normalizes its result + * ({@link Normalization#L2}) or not ({@link Normalization#NONE}). + * @return The loaded model. + * @throws IllegalArgumentException Thrown if an argument is {@code null} or a file is + * missing. + * @throws InvalidFormatException Thrown if a file is malformed, the vocabulary size and the + * embedding matrix's row count disagree, or the tokenizer emits pieces the vocabulary + * does not map. + * @throws IOException Thrown if reading a file fails. + */ + public static StaticEmbeddingModel loadSentencePiece(Path sentencePieceModelFile, + Path tokenizerJsonFile, + Path safetensorsFile, + Normalization normalization) + throws IOException { + return loadSentencePiece(sentencePieceModelFile, tokenizerJsonFile, safetensorsFile, + normalization, List.of(), ModelFileNames.TERMS); + } + + /** + * Loads the SentencePiece layout with an optional term table. + * + * @param sentencePieceModelFile The trained SentencePiece {@code .model} file. + * @param tokenizerJsonFile The Unigram {@code tokenizer.json} file. + * @param safetensorsFile The {@code model.safetensors} file. + * @param normalization The pooling normalization. + * @param termLines The terms in row order; empty without a term table. + * @param termsSourceName The terms' source, for error messages. + * @return The loaded model. + * @throws IOException Thrown if reading a file fails. + */ + private static StaticEmbeddingModel loadSentencePiece(Path sentencePieceModelFile, + Path tokenizerJsonFile, + Path safetensorsFile, + Normalization normalization, + List termLines, + String termsSourceName) + throws IOException { + if (sentencePieceModelFile == null) { + throw new IllegalArgumentException("SentencePieceModelFile must not be null"); + } + if (tokenizerJsonFile == null) { + throw new IllegalArgumentException("TokenizerJsonFile must not be null"); + } + if (safetensorsFile == null) { + throw new IllegalArgumentException("SafetensorsFile must not be null"); + } + if (normalization == null) { + throw new IllegalArgumentException("Normalization must not be null"); + } + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.fromTokenizerJson(tokenizerJsonFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + final SentencePieceTokenizer tokenizer = + SentencePieceTokenizer.load(sentencePieceModelFile); + requireVocabularyCoverage(tokenizer, vocabulary, sentencePieceModelFile, tokenizerJsonFile); + final Matrix matrix = readMatrix(vocabulary, terms.size(), safetensorsFile, + tokenizerJsonFile.toString()); + final TableAndWeights tableAndWeights = new TableAndWeights( + new FloatEmbeddingTable(matrix.embeddings(), matrix.dimension(), + vocabulary.size() + terms.size()), + matrix.weights()); + return createSentencePiece(tokenizer, vocabulary, tableAndWeights, normalization, terms); + } + + /** + * Loads the SentencePiece layout over a quantized matrix file. + * + * @param sentencePieceModelFile The trained SentencePiece {@code .model} file. + * @param tokenizerJsonFile The Unigram {@code tokenizer.json} naming the matrix rows. + * @param quantizedFile The quantized matrix file. + * @param normalization The pooling normalization. + * @param termLines The terms in row order; empty without a term table. + * @param termsSourceName The terms' source, for error messages. + * @return The loaded model. + * @throws IOException Thrown if reading a file fails. + */ + private static StaticEmbeddingModel loadSentencePieceQuantized(Path sentencePieceModelFile, + Path tokenizerJsonFile, + Path quantizedFile, + Normalization normalization, + List termLines, + String termsSourceName) + throws IOException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.fromTokenizerJson(tokenizerJsonFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + final SentencePieceTokenizer tokenizer = + SentencePieceTokenizer.load(sentencePieceModelFile); + requireVocabularyCoverage(tokenizer, vocabulary, sentencePieceModelFile, tokenizerJsonFile); + return createSentencePiece(tokenizer, vocabulary, + readQuantizedTable(quantizedFile, vocabulary, terms.size(), + tokenizerJsonFile.toString()), + normalization, terms); + } + + /** + * Builds a SentencePiece model over a loaded table, whatever its storage form. + * + * @param tokenizer The loaded SentencePiece tokenizer. + * @param vocabulary The matrix row vocabulary. + * @param tableAndWeights The table and its optional pooling weights. + * @param normalization The pooling normalization. + * @param terms The term rows after the subword rows; empty for none. + * @return The loaded model. + */ + private static StaticEmbeddingModel createSentencePiece(SentencePieceTokenizer tokenizer, + EmbeddingVocabulary vocabulary, + TableAndWeights tableAndWeights, + Normalization normalization, + TermTable terms) { + final IntPredicate skipPieceId = + id -> tokenizer.isUnknown(id) || tokenizer.isControl(id); + return new StaticEmbeddingModel(tableAndWeights.table(), tableAndWeights.weights(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + specialRows(vocabulary, SENTENCEPIECE_SPECIAL_TOKENS, + tableAndWeights.table().rowCount()), + terms); + } + + /** + * Verifies once at load time that every piece the tokenizer can emit maps to a matrix row, so + * embedding never meets an unmapped piece. Control and unknown pieces are exempt: they are + * never pooled, and a distillation legitimately drops them from the matrix. + * + * @param tokenizer The loaded SentencePiece tokenizer. + * @param vocabulary The matrix row vocabulary. + * @param sentencePieceModelFile The tokenizer's source file, for error messages. + * @param tokenizerJsonFile The vocabulary's source file, for error messages. + * @throws InvalidFormatException Thrown if a poolable piece has no matrix row. + */ + private static void requireVocabularyCoverage(SentencePieceTokenizer tokenizer, + EmbeddingVocabulary vocabulary, + Path sentencePieceModelFile, + Path tokenizerJsonFile) + throws InvalidFormatException { + int missing = 0; + final StringBuilder samples = new StringBuilder(); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (tokenizer.isUnknown(id) || tokenizer.isControl(id)) { + continue; + } + if (vocabulary.id(tokenizer.idToPiece(id)) < 0) { + if (missing < 5) { + if (missing > 0) { + samples.append(", "); + } + samples.append('\'').append(tokenizer.idToPiece(id)).append('\''); + } + missing++; + } + } + if (missing > 0) { + throw new InvalidFormatException(sentencePieceModelFile + " defines " + missing + + " pieces that " + tokenizerJsonFile + " does not map to a matrix row (first: " + + samples + "); these files do not belong to the same model"); + } + } + + /** The embedding matrix and its optional per-token weights, as read from a safetensors file. */ + private record Matrix(float[] embeddings, float[] weights, int dimension) { + } + + /** + * Reads the embedding matrix and the optional {@code weights} tensor, holding both to the + * model's row count: the vocabulary's size plus the term count. + * + * @param vocabulary The matrix row vocabulary. + * @param termCount The number of term rows after the vocabulary rows. + * @param safetensorsFile The safetensors file to read. + * @param vocabularySourceName The vocabulary's source, for error messages. + * @return The matrix, its optional weights, and its dimension. + * @throws InvalidFormatException Thrown if the matrix's row count or the weights tensor's + * length disagrees with the model's row count, or the matrix contains a non-finite value. + * @throws IOException Thrown if reading the file fails. + */ + private static Matrix readMatrix(EmbeddingVocabulary vocabulary, int termCount, + Path safetensorsFile, String vocabularySourceName) + throws IOException { + final int expectedRows = vocabulary.size() + termCount; + final SafetensorsFile tensors = SafetensorsFile.read(safetensorsFile); + final String matrixName = tensors.singleMatrixTensorName(); + final TensorInfo matrixInfo = tensors.tensorInfo(matrixName); + if (matrixInfo.shape()[0] != expectedRows) { + throw new InvalidFormatException("Vocabulary " + vocabularySourceName + " has " + + vocabulary.size() + " tokens" + + (termCount > 0 ? " plus " + termCount + " terms" : "") + + " but embedding matrix '" + matrixName + "' in " + + safetensorsFile + " has " + matrixInfo.shape()[0] + " rows; these files do not " + + "belong to the same model"); + } + final int dimension = matrixInfo.shape()[1]; + final float[] embeddings = tensors.readFloats(matrixName); + // Distillation replaces non-finite teacher values with zero before writing, so a NaN or + // infinity here marks a corrupt or foreign file; unrejected, a single NaN row silently + // defeats the norm guards and corrupts every similarity ranking it appears in. + for (int i = 0; i < embeddings.length; i++) { + if (!Float.isFinite(embeddings[i])) { + throw new InvalidFormatException("Embedding matrix '" + matrixName + "' in " + + safetensorsFile + " holds the non-finite value " + embeddings[i] + " in row " + + (i / dimension) + "; the matrix is corrupt"); + } + } + + float[] weights = null; + if (tensors.tensorNames().contains(WEIGHTS_TENSOR_NAME)) { + weights = tensors.readFloats(WEIGHTS_TENSOR_NAME); + if (weights.length != expectedRows) { + throw new InvalidFormatException("Tensor '" + WEIGHTS_TENSOR_NAME + "' in " + + safetensorsFile + " has " + weights.length + " elements but the model has " + + expectedRows + " rows"); + } + } + return new Matrix(embeddings, weights, dimension); + } + + /** + * {@return the mask of rows holding special tokens, excluded from neighbor results; term rows + * are never special} + * + * @param vocabulary The matrix row vocabulary. + * @param specialTokens The special-token strings of the model's convention; tokens absent + * from the vocabulary are simply not marked. + * @param totalRows The model's row count, the vocabulary's size plus the term count. + */ + private static boolean[] specialRows(EmbeddingVocabulary vocabulary, + Set specialTokens, int totalRows) { + final boolean[] specialRows = new boolean[totalRows]; + for (final String special : specialTokens) { + final int row = vocabulary.id(special); + if (row >= 0) { + specialRows[row] = true; + } + } + return specialRows; + } + + /** + * {@inheritDoc} + * + *

A text with no in-vocabulary tokens yields a zero vector.

+ * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + @Override + public float[] embed(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + return embed(text instanceof String s ? s : text.toString()); + } + + /** + * Embeds a piece of text. + * + * @param text The text to embed. Must not be {@code null}. + * @return The pooled embedding vector, of length {@link #dimension()}. A text with no + * in-vocabulary tokens yields a zero vector. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + public float[] embed(String text) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + // Pooling accumulates in the table's working space (original space for the float table, + // rotated space for the quantized one) and maps to original space once per text. + final float[] sum = new float[table.pooledLength()]; + // The count travels through the IntConsumer as a one-element array. + final int[] pooledCount = new int[1]; + forEachPooledRow(text, row -> { + table.addRow(row, weights == null ? 1f : weights[row], sum); + pooledCount[0]++; + }); + final int denominator = Math.max(pooledCount[0], 1); + for (int i = 0; i < sum.length; i++) { + sum[i] /= denominator; + } + final float[] pooled = table.finishPooling(sum); + if (normalize) { + double sumOfSquares = 0; + for (final float value : pooled) { + sumOfSquares += (double) value * value; + } + final float norm = (float) Math.max(Math.sqrt(sumOfSquares), NORMALIZE_EPSILON); + for (int d = 0; d < dimension; d++) { + pooled[d] /= norm; + } + } + return pooled; + } + + /** + * Feeds every matrix row a text pools to the action, in text order: matched terms' rows where + * the term table matches, and subword piece rows everywhere else. Without a term table this + * is exactly the piece walk over the whole text. + * + * @param text The text to fold into rows. + * @param action Receives each pooled row. + */ + private void forEachPooledRow(String text, IntConsumer action) { + if (terms.size() == 0) { + forEachPieceRow(text, action); + return; + } + int cursor = 0; + for (final TermTable.Match match : terms.matches(text)) { + if (match.start() > cursor) { + forEachPieceRow(text.substring(cursor, match.start()), action); + } + action.accept(match.row()); + cursor = match.end(); + } + if (cursor < text.length()) { + forEachPieceRow(text.substring(cursor), action); + } + } + + /** + * Feeds the matrix row of every poolable subword piece of a text to the action. + * + * @param text The text to tokenize. + * @param action Receives each piece's row. + */ + private void forEachPieceRow(String text, IntConsumer action) { + final List pieces = tokenizer.encode(text); + for (int i = 0; i < pieces.size(); i++) { + final SubwordPiece piece = pieces.get(i); + if (skipPieceId.test(piece.id())) { + continue; + } + final int row = vocabulary.id(piece.piece()); + if (row < 0) { + throw new IllegalStateException("Tokenizer produced piece '" + piece.piece() + + "' that has no matrix row; load-time validation admits no such piece, so this " + + "indicates a construction bug, not an input problem"); + } + action.accept(row); + } + } + + /** {@inheritDoc} */ + @Override + public int dimension() { + return dimension; + } + + /** {@return the number of subword tokens in this model's vocabulary, without term rows} */ + public int vocabularySize() { + return vocabulary.size(); + } + + /** + * {@return the number of term rows appended after the subword vocabulary, {@code 0} for a + * model without a term table} + */ + public int termCount() { + return terms.size(); + } + + /** + * Cosine similarity between two pieces of text's pooled embeddings. + * + * @param text1 The first text. Must not be {@code null}. + * @param text2 The second text. Must not be {@code null}. + * @return The cosine similarity, in {@code [-1, 1]}; {@code 0} when either text has no + * in-vocabulary tokens (an undefined direction, not an error). + * @throws IllegalArgumentException Thrown if {@code text1} or {@code text2} is {@code null}. + */ + public double similarity(String text1, String text2) { + if (text1 == null) { + throw new IllegalArgumentException("Text1 must not be null"); + } + if (text2 == null) { + throw new IllegalArgumentException("Text2 must not be null"); + } + return cosineSimilarity(embed(text1), embed(text2)); + } + + /** + * Finds the vocabulary tokens whose vectors are nearest a piece of text's pooled embedding, + * most similar first. This is a brute-force scan over the whole table; a model with a term + * table returns matching terms as neighbors like any token. + * + * @param text The query text. Must not be {@code null}. + * @param topK The maximum number of results. Must be at least 1. + * @return Up to {@code topK} neighbors, most similar first, excluding the model's special + * tokens; empty when {@code text} has no in-vocabulary tokens. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null} or {@code topK} is + * less than 1. + */ + public List mostSimilar(String text, int topK) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + requirePositive(topK); + return nearestNeighbors(embed(text), topK, NO_EXCLUDED_ROWS); + } + + /** + * The classic word2vec analogy: {@code b} is to {@code a} as the results are to {@code c} + * (computed as {@code embed(b) - embed(a) + embed(c)}), for example {@code analogy("man", + * "king", "woman", 1)} for "man is to king as woman is to ?". + * + * @param a The first term. Must not be {@code null}. + * @param b The second term. Must not be {@code null}. + * @param c The third term. Must not be {@code null}. + * @param topK The maximum number of results. Must be at least 1. + * @return Up to {@code topK} neighbors, most similar first, excluding the model's special + * tokens and every vocabulary token the three terms themselves tokenize to. The exclusion + * folds the terms exactly the way {@link #embed(String)} folds text, so on an uncased + * model a capitalized input excludes its lower-cased vocabulary row, and a multiword term + * excludes each of its word pieces. + * @throws IllegalArgumentException Thrown if {@code a}, {@code b}, or {@code c} is + * {@code null}, or {@code topK} is less than 1. + */ + public List analogy(String a, String b, String c, int topK) { + if (a == null) { + throw new IllegalArgumentException("A must not be null"); + } + if (b == null) { + throw new IllegalArgumentException("B must not be null"); + } + if (c == null) { + throw new IllegalArgumentException("C must not be null"); + } + requirePositive(topK); + final float[] va = embed(a); + final float[] vb = embed(b); + final float[] vc = embed(c); + final float[] target = new float[dimension]; + for (int d = 0; d < dimension; d++) { + target[d] = vb[d] - va[d] + vc[d]; + } + return nearestNeighbors(target, topK, excludedRows(a, b, c)); + } + + /** + * Requires {@code topK} to be at least 1. + * + * @param topK The requested result count. + * @throws IllegalArgumentException Thrown if {@code topK} is less than 1. + */ + private void requirePositive(int topK) { + if (topK < 1) { + throw new IllegalArgumentException("TopK must be at least 1, got " + topK); + } + } + + /** + * {@return the vocabulary rows the given terms tokenize to, ascending and duplicate-free} + * Folding the terms through the model's own tokenizer keeps the exclusion case- and + * accent-insensitive on models that normalize. + * + * @param terms The terms to fold and exclude. + */ + private int[] excludedRows(String... queryTerms) { + final SortedSet rows = new TreeSet<>(); + for (final String queryTerm : queryTerms) { + forEachPooledRow(queryTerm, rows::add); + } + final int[] sorted = new int[rows.size()]; + int i = 0; + for (final int row : rows) { + sorted[i++] = row; + } + return sorted; + } + + /** + * Scans the whole vocabulary for the rows nearest {@code query}, most similar first. + * + * @param query The query vector. + * @param topK The maximum number of neighbors to return. + * @param sortedExcludedRows Row ids to skip, in ascending order; the scan advances a single + * pointer through them as it visits rows in order. + * @return Up to {@code topK} neighbors, most similar first; empty when {@code query} has no + * direction. + */ + private List nearestNeighbors(float[] query, int topK, int[] sortedExcludedRows) { + final double queryNorm = norm(query); + if (queryNorm < NORMALIZE_EPSILON) { + return List.of(); + } + // The query maps into the table's working space once; every row is scored there. Norms are + // unchanged by the mapping, so the cosine denominator uses the original query norm. + final float[] preparedQuery = table.prepareQuery(query); + final int rowCount = table.rowCount(); + // The capacity sizes the candidate arrays; a topK beyond the vocabulary (the scan can never + // yield more than every row) would otherwise allocate topK-sized arrays or overflow. + final TopK best = new TopK(Math.min(topK, rowCount)); + int nextExcluded = 0; + for (int row = 0; row < rowCount; row++) { + if (nextExcluded < sortedExcludedRows.length && sortedExcludedRows[nextExcluded] == row) { + nextExcluded++; + continue; + } + if (specialRows[row]) { + continue; + } + final double rowNorm = table.rowNorm(row); + if (rowNorm < NORMALIZE_EPSILON) { + // A zero row has no direction; scored 0 rather than NaN from a 0/0 division. + best.offer(row, 0.0); + continue; + } + best.offer(row, table.dot(row, preparedQuery) / (queryNorm * rowNorm)); + } + final Neighbor[] ordered = new Neighbor[best.size()]; + for (int i = ordered.length - 1; i >= 0; i--) { + ordered[i] = new Neighbor(rowToken(best.minRow()), best.minSimilarity()); + best.removeMin(); + } + return List.of(ordered); + } + + /** + * {@return the string of a matrix row: the vocabulary token of a subword row, the term of a + * term row} + * + * @param row The matrix row. + */ + private String rowToken(int row) { + return row < vocabulary.size() ? vocabulary.token(row) : terms.term(row); + } + + /** + * {@return the cosine similarity of two vectors, or {@code 0} when either has no direction} + * + * @param a The first vector. + * @param b The second vector, of the same length as {@code a}. + */ + private double cosineSimilarity(float[] a, float[] b) { + double dot = 0; + double normASquared = 0; + double normBSquared = 0; + for (int d = 0; d < a.length; d++) { + dot += (double) a[d] * b[d]; + normASquared += (double) a[d] * a[d]; + normBSquared += (double) b[d] * b[d]; + } + final double denominator = Math.sqrt(normASquared) * Math.sqrt(normBSquared); + return denominator < NORMALIZE_EPSILON ? 0.0 : dot / denominator; + } + + /** + * {@return the L2 norm of a vector} + * + * @param vector The vector to measure. + */ + private double norm(float[] vector) { + double sumOfSquares = 0; + for (final float value : vector) { + sumOfSquares += (double) value * value; + } + return Math.sqrt(sumOfSquares); + } + + /** + * A bounded selection of the {@code k} highest-similarity rows, kept as a min-heap over + * primitive parallel arrays: the root is always the weakest kept candidate, so a full scan + * decides most rows with one comparison against it and allocates nothing per row. + */ + private static final class TopK { + + private final double[] similarities; + private final int[] rows; + private int size; + + /** + * Creates an empty selection. + * + * @param capacity The maximum number of rows to keep. + */ + TopK(int capacity) { + this.similarities = new double[capacity]; + this.rows = new int[capacity]; + } + + /** + * Offers a candidate row, keeping it only if it ranks among the top {@code capacity}. + * + * @param row The candidate row id. + * @param similarity The row's similarity to the query. + */ + void offer(int row, double similarity) { + if (size < similarities.length) { + int i = size++; + similarities[i] = similarity; + rows[i] = row; + while (i > 0) { + final int parent = (i - 1) >>> 1; + if (similarities[parent] <= similarities[i]) { + break; + } + swap(parent, i); + i = parent; + } + } else if (similarity > similarities[0]) { + similarities[0] = similarity; + rows[0] = row; + siftDown(); + } + } + + /** {@return the number of rows currently held} */ + int size() { + return size; + } + + /** {@return the row id of the weakest held candidate, the heap root} */ + int minRow() { + return rows[0]; + } + + /** {@return the similarity of the weakest held candidate, the heap root} */ + double minSimilarity() { + return similarities[0]; + } + + /** Removes the weakest held candidate, the heap root. */ + void removeMin() { + size--; + similarities[0] = similarities[size]; + rows[0] = rows[size]; + siftDown(); + } + + /** Restores the min-heap invariant from the root downward. */ + private void siftDown() { + int i = 0; + while (true) { + final int left = 2 * i + 1; + final int right = left + 1; + int smallest = i; + if (left < size && similarities[left] < similarities[smallest]) { + smallest = left; + } + if (right < size && similarities[right] < similarities[smallest]) { + smallest = right; + } + if (smallest == i) { + return; + } + swap(i, smallest); + i = smallest; + } + } + + /** + * Swaps two heap entries in both parallel arrays. + * + * @param i The first index. + * @param j The second index. + */ + private void swap(int i, int j) { + final double similarity = similarities[i]; + similarities[i] = similarities[j]; + similarities[j] = similarity; + final int row = rows[i]; + rows[i] = rows[j]; + rows[j] = row; + } + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TeacherTokenizer.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TeacherTokenizer.java new file mode 100644 index 0000000000..f10abce892 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TeacherTokenizer.java @@ -0,0 +1,1185 @@ +/* + * 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.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.util.InvalidFormatException; + +/** + * The tokenizer side of a teacher model, distilled the way + * Model2Vec distills it. The class reads + * the teacher's {@code tokenizer.json} (and, when present, its {@code tokenizer_config.json} for + * the pad token), decides which vocabulary rows survive into the static table, and rewrites the + * {@code tokenizer.json} so it describes the distilled table. + * + *

The cleaning mirrors Model2Vec: tokens matching {@code \[unused\d+\]} are removed, the + * added-token overlay is pruned to the unknown and pad tokens (the only special tokens a distilled + * table keeps), the post-processor is dropped (a static table is pooled from content pieces, never + * wrapped in {@code [CLS]}/{@code [SEP]}), and the surviving tokens keep their original id order + * but are renumbered to a gapless id space. That new order is the matrix row order.

+ * + *

For the forward pass the class reports, per surviving token, its id in the teacher's + * id space plus the teacher's begin/end-of-sequence wrapper ids: Model2Vec feeds each vocabulary + * token to the teacher as {@code [bos, token, eos]} and mean-pools the hidden states.

+ * + *

The rewrite copies every field it does not change byte for byte (the normalizer, the + * pre-tokenizer, the Unigram scores), so the cleaned {@code tokenizer.json} stays a faithful + * fast-tokenizer description of the distilled table.

+ */ +final class TeacherTokenizer { + + /** The prefix of the BERT-style placeholder tokens Model2Vec's cleaning drops. */ + private static final String UNUSED_TOKEN_PREFIX = "[unused"; + + /** Marks a template item as the sequence placeholder rather than a special token. */ + private static final String SEQUENCE_PLACEHOLDER_PREFIX = "$"; + + /** The WordPiece {@code model.type} of a BERT-family teacher. */ + static final String WORDPIECE = "WordPiece"; + + /** The Unigram {@code model.type} of a SentencePiece-family teacher. */ + static final String UNIGRAM = "Unigram"; + + private final String json; + private final String inputName; + private final String modelType; + private final List tokensByOriginalId; + private final int[] keptOriginalIds; + private final int originalUnkId; + private final String unkToken; + private final String padToken; + private final int padTokenId; + private final int[] bosIds; + private final int[] eosIds; + private final Map idByOriginalToken; + private final Boolean lowerCase; + + /** Holds the parsed state; built by {@link #read(Path, Path)}. */ + private TeacherTokenizer(String json, String inputName, String modelType, + List tokensByOriginalId, + Map idByOriginalToken, int[] keptOriginalIds, + int originalUnkId, String unkToken, String padToken, int padTokenId, + int[] bosIds, int[] eosIds, Boolean lowerCase) { + this.json = json; + this.inputName = inputName; + this.modelType = modelType; + this.tokensByOriginalId = tokensByOriginalId; + this.idByOriginalToken = idByOriginalToken; + this.keptOriginalIds = keptOriginalIds; + this.originalUnkId = originalUnkId; + this.unkToken = unkToken; + this.padToken = padToken; + this.padTokenId = padTokenId; + this.bosIds = bosIds; + this.eosIds = eosIds; + this.lowerCase = lowerCase; + } + + /** + * Reads a teacher's tokenizer configuration. + * + * @param tokenizerJsonFile The teacher's {@code tokenizer.json}. Must not be {@code null} + * and must exist. + * @param tokenizerConfigFile The teacher's {@code tokenizer_config.json}, consulted for the + * pad token only; may be {@code null} (no pad token then). + * @return The parsed teacher tokenizer. + * @throws IllegalArgumentException Thrown if {@code tokenizerJsonFile} is {@code null} or + * missing. + * @throws InvalidFormatException Thrown if a file is malformed, the tokenizer model is + * neither WordPiece nor Unigram, the vocabulary ids are not a gapless range, the unknown + * token is missing, or the post-processor is of an unsupported type. + * @throws IOException Thrown if reading a file fails. + */ + static TeacherTokenizer read(Path tokenizerJsonFile, Path tokenizerConfigFile) + throws IOException { + if (tokenizerJsonFile == null) { + throw new IllegalArgumentException("TokenizerJsonFile must not be null"); + } + if (!Files.isRegularFile(tokenizerJsonFile)) { + throw new IllegalArgumentException("File does not exist or is not a regular file: " + + tokenizerJsonFile); + } + final String padToken = tokenizerConfigFile != null && Files.isRegularFile(tokenizerConfigFile) + ? FlatJsonFields.topLevelString(tokenizerConfigFile, "pad_token") : null; + final String json = Files.readString(tokenizerJsonFile); + final String inputName = tokenizerJsonFile.getFileName().toString(); + final JsonCursor cursor = new JsonCursor(json, inputName); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + String modelType = null; + List tokensById = null; + String unkToken = null; + Long unkId = null; + Boolean lowerCase = null; + Set addedContents = Set.of(); + PostProcessor postProcessor = new PostProcessor(List.of(), List.of(), null, null, Map.of()); + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "model" -> { + final ModelSection model = parseModel(cursor); + modelType = model.type(); + tokensById = model.tokensById(); + unkToken = model.unkToken(); + unkId = model.unkId(); + } + case "added_tokens" -> addedContents = parseAddedTokenContents(cursor); + case "post_processor" -> postProcessor = parsePostProcessor(cursor); + case "normalizer" -> lowerCase = parseNormalizerLowercase(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a field, got '" + next + "'"); + } + } + cursor.requireEnd("Trailing content after the top-level object"); + if (modelType == null || tokensById == null) { + throw new InvalidFormatException(tokenizerJsonFile + " has no model with a vocabulary; " + + "it does not look like a teacher's tokenizer.json"); + } + if (!WORDPIECE.equals(modelType) && !UNIGRAM.equals(modelType)) { + throw new InvalidFormatException(tokenizerJsonFile + " has a '" + modelType + + "' tokenizer model; only " + WORDPIECE + " and " + UNIGRAM + + " teachers are supported"); + } + final Map idByToken = new HashMap<>(tokensById.size() * 2); + for (int id = 0; id < tokensById.size(); id++) { + idByToken.putIfAbsent(tokensById.get(id), id); + } + if (unkToken == null) { + if (unkId == null || unkId < 0 || unkId >= tokensById.size()) { + throw new InvalidFormatException(tokenizerJsonFile + " does not name an unknown token " + + "(no model.unk_token / model.unk_id); a distilled table needs one"); + } + unkToken = tokensById.get(unkId.intValue()); + } + final Integer originalUnkId = idByToken.get(unkToken); + if (originalUnkId == null) { + throw new InvalidFormatException(tokenizerJsonFile + " names the unknown token '" + + unkToken + "' but it is not in the vocabulary"); + } + // The wrapper ids come from the cls/sep pairs of a BertProcessing/RobertaProcessing + // post-processor, or from resolving a TemplateProcessing's special token names through its + // special_tokens table, falling back to the vocabulary. + final int[] bosIds = postProcessor.clsId() != null + ? new int[] {postProcessor.clsId().intValue()} + : resolveNames(postProcessor.bosNames(), postProcessor.specialTokenIds(), idByToken, + tokenizerJsonFile); + final int[] eosIds = postProcessor.sepId() != null + ? new int[] {postProcessor.sepId().intValue()} + : resolveNames(postProcessor.eosNames(), postProcessor.specialTokenIds(), idByToken, + tokenizerJsonFile); + final Integer padId = padToken == null ? null : idByToken.get(padToken); + final int padTokenId = padId == null ? 0 : padId; + final Set keepSpecial = new HashSet<>(); + keepSpecial.add(unkToken); + if (padToken != null) { + keepSpecial.add(padToken); + } + final List kept = new ArrayList<>(tokensById.size()); + for (int id = 0; id < tokensById.size(); id++) { + final String token = tokensById.get(id); + if (isUnusedToken(token)) { + continue; + } + if (addedContents.contains(token) && !keepSpecial.contains(token)) { + continue; + } + kept.add(id); + } + return new TeacherTokenizer(json, inputName, modelType, tokensById, idByToken, + kept.stream().mapToInt(Integer::intValue).toArray(), originalUnkId, unkToken, padToken, + padTokenId, bosIds, eosIds, lowerCase); + } + + /** + * Reads the flat {@code lowercase} boolean of a {@code normalizer} object, for the BERT + * normalizer a WordPiece tokenizer carries. Shared with {@link ModelAssembler}, which derives + * a distilled directory's {@code do_lower_case} from the same flag. + * + * @param cursor The cursor, positioned at the normalizer value. + * @return The {@code lowercase} flag, or {@code null} when the value is JSON null or the flag + * is absent (for example a nested normalizer with no flat flag). + * @throws InvalidFormatException Thrown if the normalizer object is malformed. + */ + static Boolean parseNormalizerLowercase(JsonCursor cursor) throws InvalidFormatException { + if (cursor.peek() != '{') { + cursor.skipValue(); + return null; + } + cursor.expect('{'); + cursor.skipWhitespace(); + Boolean lowerCase = null; + if (cursor.peek() == '}') { + cursor.consume(); + return null; + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if ("lowercase".equals(key)) { + if (cursor.consumeLiteral("true")) { + lowerCase = Boolean.TRUE; + } else if (cursor.consumeLiteral("false")) { + lowerCase = Boolean.FALSE; + } else { + cursor.skipValue(); + } + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return lowerCase; + } + throw cursor.malformed("Expected ',' or '}' after a normalizer field, got '" + next + "'"); + } + } + + /** + * {@return whether a token starts with a BERT-style unused placeholder, {@code [unused} + * followed by at least one ASCII digit and {@code ]}} + * + *

Model2Vec's cleaning drops these tokens by a prefix match, so a longer token starting + * with the placeholder form is dropped the same way.

+ * + * @param token The vocabulary token. + */ + private static boolean isUnusedToken(String token) { + if (!token.startsWith(UNUSED_TOKEN_PREFIX)) { + return false; + } + int i = UNUSED_TOKEN_PREFIX.length(); + final int digitsStart = i; + while (i < token.length() && token.charAt(i) >= '0' && token.charAt(i) <= '9') { + i++; + } + return i > digitsStart && i < token.length() && token.charAt(i) == ']'; + } + + /** + * {@return the ids the named special tokens resolve to, through the post-processor's + * special-token table first and the vocabulary second} + * + * @param names The special token names in order. + * @param specialTokenIds The post-processor's name-to-id table. + * @param idByToken The vocabulary, token to id. + * @param file The source file, for error messages. + * @throws InvalidFormatException Thrown if a name resolves nowhere. + */ + private static int[] resolveNames(List names, Map specialTokenIds, + Map idByToken, Path file) + throws InvalidFormatException { + final int[] ids = new int[names.size()]; + for (int i = 0; i < names.size(); i++) { + final Long specialId = specialTokenIds.get(names.get(i)); + final Integer vocabId = idByToken.get(names.get(i)); + if (specialId != null) { + ids[i] = specialId.intValue(); + } else if (vocabId != null) { + ids[i] = vocabId; + } else { + throw new InvalidFormatException(file + " wraps sequences in the special token '" + + names.get(i) + "' but neither the post-processor nor the vocabulary defines it"); + } + } + return ids; + } + + /** {@return the tokenizer family, {@code "WordPiece"} or {@code "Unigram"}} */ + String modelType() { + return modelType; + } + + /** {@return the number of surviving tokens, the matrix row count} */ + int vocabularySize() { + return keptOriginalIds.length; + } + + /** {@return the surviving tokens' ids in the teacher's id space, in matrix row order} */ + int[] keptOriginalIds() { + return keptOriginalIds.clone(); + } + + /** {@return the teacher's pad token id, used to pad batches; 0 when the teacher names none} */ + int padTokenId() { + return padTokenId; + } + + /** {@return the unknown token's string} */ + String unkToken() { + return unkToken; + } + + /** {@return the pad token's string, or {@code null} when the teacher names none} */ + String padToken() { + return padToken; + } + + /** + * The teacher input sequence for one matrix row: the begin-of-sequence ids, the token's + * original id, and the end-of-sequence ids. + * + * @param row The matrix row. + * @return The teacher input ids. + */ + long[] inputSequence(int row) { + final long[] sequence = new long[bosIds.length + 1 + eosIds.length]; + int i = 0; + for (final int id : bosIds) { + sequence[i++] = id; + } + sequence[i++] = keptOriginalIds[row]; + for (final int id : eosIds) { + sequence[i++] = id; + } + return sequence; + } + + /** + * The teacher input sequence of a segmented term: the begin-of-sequence ids, each piece's + * original id (the unknown token's id for a piece the vocabulary does not carry), and the + * end-of-sequence ids. + * + * @param pieces The term's piece strings, as the teacher's own segmenter produced them. Must + * not be {@code null}. + * @return The teacher input ids. + * @throws IllegalArgumentException Thrown if {@code pieces} is {@code null}. + */ + long[] inputSequence(List pieces) { + if (pieces == null) { + throw new IllegalArgumentException("Pieces must not be null"); + } + final long[] sequence = new long[bosIds.length + pieces.size() + eosIds.length]; + int i = 0; + for (final int id : bosIds) { + sequence[i++] = id; + } + for (final String piece : pieces) { + final Integer id = idByOriginalToken.get(piece); + sequence[i++] = id == null ? originalUnkId : id; + } + for (final int id : eosIds) { + sequence[i++] = id; + } + return sequence; + } + + /** + * Looks up the token string of a matrix row. + * + * @param row The matrix row, within {@code [0, vocabularySize())}. + * @return The surviving token at that row. + */ + String rowToken(int row) { + return tokensByOriginalId.get(keptOriginalIds[row]); + } + + /** {@return the whole vocabulary in the teacher's id order, for an id-is-index segmenter} */ + List tokensByOriginalId() { + return Collections.unmodifiableList(tokensByOriginalId); + } + + /** + * {@return the {@code normalizer.lowercase} flag of the teacher's {@code tokenizer.json}, or + * {@code null} when the tokenizer does not state it} + */ + Boolean lowerCase() { + return lowerCase; + } + + /** + * Writes the cleaned {@code tokenizer.json}: the surviving vocabulary renumbered, the + * added-token overlay pruned to the unknown and pad tokens, the post-processor nulled, and + * every other field copied byte for byte from the teacher's file. + * + * @param file The file to write. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null}. + * @throws IOException Thrown if writing fails. + */ + void writeCleaned(Path file) throws IOException { + if (file == null) { + throw new IllegalArgumentException("File must not be null"); + } + final Map newIdByOriginal = new HashMap<>(keptOriginalIds.length * 2); + for (int row = 0; row < keptOriginalIds.length; row++) { + newIdByOriginal.put(keptOriginalIds[row], row); + } + final JsonCursor cursor = new JsonCursor(json, inputName); + final StringBuilder out = new StringBuilder(json.length()); + cursor.skipWhitespace(); + cursor.expect('{'); + out.append('{'); + cursor.skipWhitespace(); + if (cursor.peek() == '}') { + cursor.consume(); + } else { + boolean first = true; + while (true) { + cursor.skipWhitespace(); + final int keyStart = cursor.position(); + final String key = cursor.parseString(); + final String rawKey = json.substring(keyStart, cursor.position()); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if (!first) { + out.append(','); + } + first = false; + out.append(rawKey).append(':'); + switch (key) { + case "model" -> rewriteModel(cursor, out, newIdByOriginal); + case "added_tokens" -> { + cursor.skipValue(); + out.append(rewrittenAddedTokens(newIdByOriginal)); + } + case "post_processor" -> { + cursor.skipValue(); + out.append("null"); + } + default -> out.append(copyRawValue(cursor)); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a field, got '" + next + "'"); + } + } + cursor.requireEnd("Trailing content after the top-level object"); + out.append('}'); + Files.writeString(file, out.toString()); + } + + /** + * Rewrites the {@code model} object: the vocabulary renumbered to the surviving rows, the + * Unigram {@code unk_id} remapped, every other field copied byte for byte. + * + * @param cursor The cursor, positioned at the object's opening brace. + * @param out The output accumulator. + * @param newIdByOriginal The original-to-new id map. + */ + private void rewriteModel(JsonCursor cursor, StringBuilder out, + Map newIdByOriginal) + throws InvalidFormatException { + cursor.expect('{'); + out.append('{'); + cursor.skipWhitespace(); + if (cursor.peek() == '}') { + cursor.consume(); + out.append('}'); + return; + } + boolean first = true; + while (true) { + cursor.skipWhitespace(); + final int keyStart = cursor.position(); + final String key = cursor.parseString(); + final String rawKey = json.substring(keyStart, cursor.position()); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if (!first) { + out.append(','); + } + first = false; + out.append(rawKey).append(':'); + switch (key) { + case "vocab" -> out.append(rewrittenVocab(cursor, newIdByOriginal)); + case "unk_id" -> { + cursor.skipValue(); + out.append(newIdByOriginal.getOrDefault(originalUnkId, 0)); + } + default -> out.append(copyRawValue(cursor)); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a model field, got '" + next + "'"); + } + out.append('}'); + } + + /** + * {@return the rewritten vocabulary value: for a WordPiece dictionary the kept entries with + * their new ids (raw key spans reused), for a Unigram list the kept {@code [piece, score]} + * entries byte for byte} + * + * @param cursor The cursor, positioned at the vocabulary's opening character. + * @param newIdByOriginal The original-to-new id map. + */ + private String rewrittenVocab(JsonCursor cursor, Map newIdByOriginal) + throws InvalidFormatException { + final StringBuilder out = new StringBuilder(); + if (cursor.peek() == '{') { + cursor.consume(); + out.append('{'); + cursor.skipWhitespace(); + if (cursor.peek() == '}') { + cursor.consume(); + } else { + boolean first = true; + while (true) { + cursor.skipWhitespace(); + final int keyStart = cursor.position(); + cursor.parseString(); + final String rawKey = json.substring(keyStart, cursor.position()); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + final long originalId = cursor.parseLong(); + final Integer row = newIdByOriginal.get((int) originalId); + if (row != null) { + if (!first) { + out.append(','); + } + first = false; + out.append(rawKey).append(':').append(row); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a vocab entry, got '" + next + "'"); + } + } + out.append('}'); + } else { + cursor.expect('['); + out.append('['); + cursor.skipWhitespace(); + if (cursor.peek() == ']') { + cursor.consume(); + } else { + boolean first = true; + int originalId = 0; + while (true) { + cursor.skipWhitespace(); + final int entryStart = cursor.position(); + cursor.expect('['); + cursor.skipWhitespace(); + cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(','); + cursor.skipWhitespace(); + cursor.skipValue(); + cursor.skipWhitespace(); + cursor.expect(']'); + if (newIdByOriginal.containsKey(originalId++)) { + if (!first) { + out.append(','); + } + first = false; + out.append(json, entryStart, cursor.position()); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + break; + } + throw cursor.malformed("Expected ',' or ']' after a vocab entry, got '" + next + "'"); + } + } + out.append(']'); + } + return out.toString(); + } + + /** + * {@return the pruned {@code added_tokens} value: the unknown and pad tokens at their new ids, + * with the flag convention Model2Vec writes (the pad token strips around itself, the unknown + * token does not)} + * + * @param newIdByOriginal The original-to-new id map. + */ + private String rewrittenAddedTokens(Map newIdByOriginal) { + record Added(int id, String content, boolean pad) { + } + final List kept = new ArrayList<>(2); + for (int id = 0; id < tokensByOriginalId.size(); id++) { + final String token = tokensByOriginalId.get(id); + final Integer row = newIdByOriginal.get(id); + if (row == null) { + continue; + } + if (token.equals(unkToken)) { + kept.add(new Added(row, token, false)); + } else if (token.equals(padToken)) { + kept.add(new Added(row, token, true)); + } + } + kept.sort(Comparator.comparingInt(Added::id)); + final StringBuilder out = new StringBuilder("["); + boolean first = true; + for (final Added added : kept) { + if (!first) { + out.append(','); + } + first = false; + out.append("{\"id\":").append(added.id()) + .append(",\"content\":").append(quoted(added.content())) + .append(",\"single_word\":").append(added.pad()) + .append(",\"lstrip\":").append(added.pad()) + .append(",\"rstrip\":").append(added.pad()) + .append(",\"normalized\":").append(added.pad()) + .append(",\"special\":true}"); + } + return out.append(']').toString(); + } + + /** + * {@return the JSON string literal for the given content, escaping the quote, the backslash, + * and control characters} + * + * @param content The string to quote. + */ + private static String quoted(String content) { + final StringBuilder out = new StringBuilder(content.length() + 2).append('"'); + for (int i = 0; i < content.length(); i++) { + final char c = content.charAt(i); + switch (c) { + case '"' -> out.append("\\\""); + case '\\' -> out.append("\\\\"); + default -> { + if (c < 0x20) { + out.append(String.format("\\u%04x", (int) c)); + } else { + out.append(c); + } + } + } + } + return out.append('"').toString(); + } + + /** + * {@return the raw text of the JSON value at the cursor, unchanged} + * + * @param cursor The cursor, positioned at the value. + */ + private String copyRawValue(JsonCursor cursor) throws InvalidFormatException { + final int start = cursor.position(); + cursor.skipValue(); + return json.substring(start, cursor.position()); + } + + /** The fields read out of the {@code model} object. */ + private record ModelSection(String type, List tokensById, String unkToken, Long unkId) { + } + + /** + * Parses the {@code model} object for its type, its vocabulary in id order, and its unknown + * token (by name for WordPiece, by id for Unigram). + * + * @param cursor The cursor, positioned at the object's opening brace. + * @return The parsed section. + */ + private static ModelSection parseModel(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + List tokensById = null; + String unkToken = null; + Long unkId = null; + if (cursor.peek() == '}') { + cursor.consume(); + return new ModelSection(null, null, null, null); + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> type = cursor.parseString(); + case "unk_token" -> unkToken = cursor.parseString(); + case "unk_id" -> { + if (!cursor.consumeLiteral("null")) { + unkId = cursor.parseLong(); + } + } + case "vocab" -> tokensById = parseVocab(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return new ModelSection(type, tokensById, unkToken, unkId); + } + throw cursor.malformed("Expected ',' or '}' after a model field, got '" + next + "'"); + } + } + + /** + * {@return the vocabulary in id order, either from a WordPiece {@code "token": id} dictionary + * or from a Unigram {@code [piece, score]} list; dictionary ids must form a gapless range} + * + * @param cursor The cursor, positioned at the vocabulary's opening character. + */ + private static List parseVocab(JsonCursor cursor) throws InvalidFormatException { + if (cursor.peek() == '{') { + cursor.consume(); + cursor.skipWhitespace(); + final Map tokenById = new HashMap<>(); + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String token = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + final long id = cursor.parseLong(); + if (tokenById.putIfAbsent(id, token) != null) { + throw cursor.malformed("Vocabulary id " + id + " is assigned more than once"); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a vocab entry, got '" + next + "'"); + } + } + final List> entries = new ArrayList<>(tokenById.entrySet()); + entries.sort(Comparator.comparingLong(Map.Entry::getKey)); + final List ordered = new ArrayList<>(entries.size()); + for (int row = 0; row < entries.size(); row++) { + if (entries.get(row).getKey() != row) { + throw cursor.malformed("Vocabulary ids are not a gapless range: expected id " + row + + " but found " + entries.get(row).getKey()); + } + ordered.add(entries.get(row).getValue()); + } + return ordered; + } + cursor.expect('['); + cursor.skipWhitespace(); + final List pieces = new ArrayList<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return pieces; + } + while (true) { + cursor.skipWhitespace(); + cursor.expect('['); + cursor.skipWhitespace(); + pieces.add(cursor.parseString()); + cursor.skipWhitespace(); + cursor.expect(','); + cursor.skipWhitespace(); + cursor.skipValue(); + cursor.skipWhitespace(); + cursor.expect(']'); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + return pieces; + } + throw cursor.malformed("Expected ',' or ']' after a vocab entry, got '" + next + "'"); + } + } + + /** + * {@return the contents of the {@code added_tokens} overlay} + * + * @param cursor The cursor, positioned at the list's opening bracket. + */ + private static Set parseAddedTokenContents(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final Set contents = new HashSet<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return contents; + } + while (true) { + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + String content = null; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if ("content".equals(key)) { + content = cursor.parseString(); + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after an added token field, got '" + next + + "'"); + } + } + if (content != null) { + contents.add(content); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + return contents; + } + throw cursor.malformed("Expected ',' or ']' after an added token, got '" + next + "'"); + } + } + + /** The wrapper names or ids of a post-processor, plus its special-token id table. */ + private record PostProcessor(List bosNames, List eosNames, Long clsId, + Long sepId, Map specialTokenIds) { + } + + /** + * Parses the {@code post_processor} for the wrapper a single-sequence encoding adds. Supports + * the {@code TemplateProcessing} form (string or structured template) and the + * {@code BertProcessing}/{@code RobertaProcessing} forms with their {@code cls}/{@code sep} + * pairs; a {@code null} post-processor means no wrapper. + * + * @param cursor The cursor, positioned at the value. + * @return The parsed post-processor. + * @throws InvalidFormatException Thrown if the type is not one of the supported forms. + */ + private static PostProcessor parsePostProcessor(JsonCursor cursor) + throws InvalidFormatException { + if (cursor.consumeLiteral("null")) { + return new PostProcessor(List.of(), List.of(), null, null, Map.of()); + } + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + List bosNames = List.of(); + List eosNames = List.of(); + Map specialTokenIds = Map.of(); + Long clsId = null; + Long sepId = null; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> type = cursor.parseString(); + case "single" -> { + final List> wrapper = parseTemplate(cursor); + bosNames = wrapper.get(0); + eosNames = wrapper.get(1); + } + case "special_tokens" -> specialTokenIds = parseSpecialTokenIds(cursor); + case "cls" -> clsId = parseTokenIdPair(cursor); + case "sep" -> sepId = parseTokenIdPair(cursor); + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a post-processor field, got '" + next + + "'"); + } + } + if (type == null) { + return new PostProcessor(List.of(), List.of(), null, null, Map.of()); + } + return switch (type) { + case "TemplateProcessing" -> + new PostProcessor(bosNames, eosNames, null, null, specialTokenIds); + case "BertProcessing", "RobertaProcessing" -> + new PostProcessor(List.of(), List.of(), clsId, sepId, specialTokenIds); + default -> throw new InvalidFormatException("The post_processor type '" + type + + "' is not supported; expected TemplateProcessing, BertProcessing, or " + + "RobertaProcessing"); + }; + } + + /** + * {@return a two-element list: the special token names before the sequence placeholder (the + * begin-of-sequence wrapper) and those after it (the end-of-sequence wrapper); the template is + * either a string like {@code "[CLS] $A [SEP]"} or a list of {@code SpecialToken}/{@code + * Sequence} items} + * + * @param cursor The cursor, positioned at the template value. + */ + private static List> parseTemplate(JsonCursor cursor) + throws InvalidFormatException { + final List bos = new ArrayList<>(1); + final List eos = new ArrayList<>(1); + if (cursor.peek() == '"') { + // The template is items separated by whitespace runs, such as "[CLS] $A [SEP]". + final String template = cursor.parseString(); + List current = bos; + final int length = template.length(); + int i = 0; + while (i < length) { + final int c = template.codePointAt(i); + if (Character.isWhitespace(c)) { + i += Character.charCount(c); + continue; + } + final int start = i; + while (i < length && !Character.isWhitespace(template.codePointAt(i))) { + i += Character.charCount(template.codePointAt(i)); + } + final String part = template.substring(start, i); + if (part.startsWith(SEQUENCE_PLACEHOLDER_PREFIX)) { + current = eos; + } else { + current.add(part); + } + } + return List.of(bos, eos); + } + cursor.expect('['); + cursor.skipWhitespace(); + List current = bos; + if (cursor.peek() == ']') { + cursor.consume(); + return List.of(bos, eos); + } + while (true) { + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + final String itemType = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + String id = null; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if ("id".equals(key)) { + id = cursor.parseString(); + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a template item field, got '" + next + + "'"); + } + } + cursor.skipWhitespace(); + cursor.expect('}'); + if ("SpecialToken".equals(itemType)) { + current.add(id); + } else if ("Sequence".equals(itemType)) { + current = eos; + } else { + throw cursor.malformed("Unknown template item type: '" + itemType + "'"); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + break; + } + throw cursor.malformed("Expected ',' or ']' after a template item, got '" + next + "'"); + } + return List.of(bos, eos); + } + + /** + * {@return the post-processor's special-token id table, name to the first of its ids} + * + * @param cursor The cursor, positioned at the table's opening brace. + */ + private static Map parseSpecialTokenIds(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + final Map ids = new HashMap<>(); + if (cursor.peek() == '}') { + cursor.consume(); + return ids; + } + while (true) { + cursor.skipWhitespace(); + final String name = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + cursor.expect('{'); + cursor.skipWhitespace(); + Long id = null; + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if ("ids".equals(key)) { + cursor.expect('['); + cursor.skipWhitespace(); + id = cursor.parseLong(); + cursor.skipWhitespace(); + while (cursor.consume() == ',') { + cursor.skipWhitespace(); + cursor.skipValue(); + cursor.skipWhitespace(); + } + } else { + cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a special token field, got '" + next + + "'"); + } + } + if (id != null) { + ids.put(name, id); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return ids; + } + throw cursor.malformed("Expected ',' or '}' after a special token, got '" + next + "'"); + } + } + + /** + * {@return the id of a {@code ["token", id]} pair, as {@code cls} and {@code sep} carry it} + * + * @param cursor The cursor, positioned at the pair's opening bracket. + */ + private static Long parseTokenIdPair(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(','); + cursor.skipWhitespace(); + final long id = cursor.parseLong(); + cursor.skipWhitespace(); + cursor.expect(']'); + return id; + } +} 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 new file mode 100644 index 0000000000..a083b4026a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.java @@ -0,0 +1,116 @@ +/* + * 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.Arrays; + +import opennlp.tools.util.java.Experimental; + +/** + * Header metadata for one tensor in a safetensors file, as declared by the file's own JSON + * header. Carries no data; {@link SafetensorsFile#readFloats(String)} resolves the bytes. + * + *

Warning: Experimental new feature; the API might change in a later release.

+ * + * @param name The tensor's name, the key it was declared under. Never {@code null}. + * @param dtype The declared element type (e.g. {@code "F32"}, {@code "F16"}, + * {@code "I64"}), exactly as written in the header. Never {@code null}. + * @param shape The tensor's dimensions, outermost first. Never {@code null}; empty + * for a scalar. + * @param dataOffsetBegin Start byte offset into the file's data section (relative to the end + * of the header, not the start of the file). + * @param dataOffsetEnd End byte offset (exclusive) into the data section. + */ +@Experimental +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. + * + * @throws IllegalArgumentException Thrown if {@code name}, {@code dtype}, or {@code shape} is + * {@code null}. + */ + public TensorInfo { + if (name == null) { + throw new IllegalArgumentException("Name must not be null"); + } + if (dtype == null) { + throw new IllegalArgumentException("Dtype must not be null"); + } + if (shape == null) { + throw new IllegalArgumentException("Shape must not be null"); + } + 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) { + 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; + } + + /** {@inheritDoc} */ + @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; + } + + /** {@inheritDoc} */ + @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; + } + + /** {@inheritDoc} */ + @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/TermSegmenter.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TermSegmenter.java new file mode 100644 index 0000000000..5b3979abaf --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TermSegmenter.java @@ -0,0 +1,129 @@ +/* + * 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.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.function.IntPredicate; + +import opennlp.subword.sentencepiece.SentencePieceTokenizer; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.tokenize.SubwordTokenizer; +import opennlp.tools.tokenize.WordpieceEncoder; +import opennlp.tools.tokenize.WordpieceTokenizer; + +/** + * Segments a term's text into the piece strings the teacher's own tokenizer would produce, so a + * distillation can run a whole word or phrase through the teacher the way the teacher would see + * it in running text. A WordPiece teacher segments through a {@link WordpieceEncoder} built over + * the teacher's full vocabulary; a Unigram teacher segments through its trained SentencePiece + * {@code .model} file. + * + *

The sequence-delimiter pieces the segmenter itself wraps around an encoding are removed; + * {@link TeacherTokenizer#inputSequence(List)} adds the teacher's own wrapping when the pieces + * are turned into an input sequence.

+ */ +final class TermSegmenter { + + private final SubwordTokenizer tokenizer; + private final Set dropPieces; + private final IntPredicate dropPieceId; + + /** Holds the segmenter and its piece filters; built by {@link #forTeacher}. */ + private TermSegmenter(SubwordTokenizer tokenizer, Set dropPieces, + IntPredicate dropPieceId) { + this.tokenizer = tokenizer; + this.dropPieces = dropPieces; + this.dropPieceId = dropPieceId; + } + + /** + * Builds the segmenter matching a teacher's tokenizer family. + * + * @param teacher The teacher's parsed tokenizer. Must not be {@code null}. + * @param teacherDirectory The teacher's directory, holding the trained SentencePiece + * {@code .model} file when the teacher is a Unigram model. Must not be + * {@code null}. + * @return The segmenter. + * @throws IllegalArgumentException Thrown if an argument is {@code null}, a Unigram teacher + * has no trained SentencePiece file, or a WordPiece teacher's vocabulary lacks the BERT + * special tokens the encoder wraps with. + * @throws IOException Thrown if reading the SentencePiece file fails. + */ + static TermSegmenter forTeacher(TeacherTokenizer teacher, Path teacherDirectory) + throws IOException { + if (teacher == null) { + throw new IllegalArgumentException("Teacher must not be null"); + } + if (teacherDirectory == null) { + throw new IllegalArgumentException("TeacherDirectory must not be null"); + } + if (TeacherTokenizer.WORDPIECE.equals(teacher.modelType())) { + // The lowercase default matches ModelAssembler's: absent means the uncased convention. + final boolean lowerCase = teacher.lowerCase() == null || teacher.lowerCase(); + final WordpieceEncoder encoder; + try { + encoder = new WordpieceEncoder(teacher.tokensByOriginalId(), lowerCase, + WordpieceTokenizer.BERT_CLS_TOKEN, WordpieceTokenizer.BERT_SEP_TOKEN, + teacher.unkToken()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("The teacher's WordPiece vocabulary cannot segment " + + "terms: " + e.getMessage(), e); + } + return new TermSegmenter(encoder, + Set.of(WordpieceTokenizer.BERT_CLS_TOKEN, WordpieceTokenizer.BERT_SEP_TOKEN), + id -> false); + } + final Path sentencePieceModelFile = ModelFileNames.firstRegularFile(teacherDirectory, + ModelFileNames.SENTENCEPIECE_MODELS); + if (sentencePieceModelFile == null) { + throw new IllegalArgumentException("Teacher directory " + teacherDirectory + " has no " + + "trained SentencePiece file (one of " + + String.join(", ", ModelFileNames.SENTENCEPIECE_MODELS) + "); distilling terms " + + "needs the teacher's own segmentation"); + } + final SentencePieceTokenizer sentencePiece = + SentencePieceTokenizer.load(sentencePieceModelFile); + return new TermSegmenter(sentencePiece, Set.of(), + id -> id >= 0 && sentencePiece.isControl(id)); + } + + /** + * Segments a term into the teacher's piece strings, without sequence delimiters. + * + * @param term The term text. Must not be {@code null}. + * @return The piece strings in order. + * @throws IllegalArgumentException Thrown if {@code term} is {@code null}. + */ + List pieces(String term) { + if (term == null) { + throw new IllegalArgumentException("Term must not be null"); + } + final List encoded = tokenizer.encode(term); + final List pieces = new ArrayList<>(encoded.size()); + for (final SubwordPiece piece : encoded) { + if (dropPieces.contains(piece.piece()) || dropPieceId.test(piece.id())) { + continue; + } + pieces.add(piece.piece()); + } + return pieces; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TermTable.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TermTable.java new file mode 100644 index 0000000000..80555a96db --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TermTable.java @@ -0,0 +1,254 @@ +/* + * 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.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.StringUtil; + +/** + * The term rows of a static embedding matrix: whole words and multi-word phrases that were + * distilled through the teacher as units and sit after the subword rows. Matching a text against + * the table finds the greedily longest term at each word position, so "writ of habeas corpus" + * wins over "habeas corpus" wins over the subword pieces of each word. + * + *

A term is stored in normalized form: the lower-cased letter-or-digit word runs of its text, + * joined by single spaces (see {@link #normalizeTerm(String)}). Matching folds each word run of + * the input text the same way, so "Habeas Corpus" and "habeas-corpus" both match the term + * "habeas corpus". The fold is {@link StringUtil#toLowerCase(CharSequence)}, locale-independent + * and one code point to one code point, so word-run boundaries are the same before and after + * folding.

+ * + *

Immutable and safe for concurrent reads after construction.

+ */ +@ThreadSafe +final class TermTable { + + private final List termsByOffset; + private final Map rowByTerm; + private final int firstRow; + private final int maxTermWords; + + /** Holds the validated term-to-row views; built by {@link #of(List, int, String)}. */ + private TermTable(List termsByOffset, Map rowByTerm, int firstRow, + int maxTermWords) { + this.termsByOffset = termsByOffset; + this.rowByTerm = rowByTerm; + this.firstRow = firstRow; + this.maxTermWords = maxTermWords; + } + + /** + * Builds a term table from terms in matrix row order. + * + * @param terms The terms; the term at index {@code i} owns matrix row + * {@code firstRow + i}. Every term must already be in its normalized form. + * Must not be {@code null}. + * @param firstRow The matrix row of the first term, the number of subword rows. + * @param sourceName The terms' source, for error messages. + * @return The table. + * @throws IllegalArgumentException Thrown if {@code terms} is {@code null}. + * @throws InvalidFormatException Thrown if a term is {@code null}, not in normalized form, or + * appears more than once. + */ + static TermTable of(List terms, int firstRow, String sourceName) + throws InvalidFormatException { + if (terms == null) { + throw new IllegalArgumentException("Terms must not be null"); + } + final Map rowByTerm = new HashMap<>(terms.size() * 2); + int maxTermWords = 0; + for (int i = 0; i < terms.size(); i++) { + final String term = terms.get(i); + if (term == null || !term.equals(normalizeTerm(term)) || term.isEmpty()) { + throw new InvalidFormatException("Term " + i + " in " + sourceName + " ('" + term + + "') is not in normalized form (lower-cased words joined by single spaces)"); + } + if (rowByTerm.putIfAbsent(term, firstRow + i) != null) { + throw new InvalidFormatException("Term '" + term + "' appears more than once in " + + sourceName); + } + maxTermWords = Math.max(maxTermWords, countWords(term)); + } + return new TermTable(List.copyOf(terms), Map.copyOf(rowByTerm), firstRow, maxTermWords); + } + + /** + * {@return a term's normalized form: its lower-cased letter-or-digit word runs joined by + * single spaces, or the empty string when the text contains no such run} + * + * @param text The term text. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + static String normalizeTerm(String text) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + final StringBuilder normalized = new StringBuilder(text.length()); + final String folded = StringUtil.toLowerCase(text); + final int length = folded.length(); + int i = 0; + while (i < length) { + final int c = folded.codePointAt(i); + if (Character.isLetterOrDigit(c)) { + if (normalized.length() > 0) { + normalized.append(' '); + } + while (i < length && Character.isLetterOrDigit(folded.codePointAt(i))) { + normalized.appendCodePoint(folded.codePointAt(i)); + i += Character.charCount(folded.codePointAt(i)); + } + } else { + i += Character.charCount(c); + } + } + return normalized.toString(); + } + + /** {@return the number of space-separated words of a normalized term} */ + private static int countWords(String term) { + int words = 1; + for (int i = 0; i < term.length(); i++) { + if (term.charAt(i) == ' ') { + words++; + } + } + return words; + } + + /** {@return the number of terms in this table} */ + int size() { + return termsByOffset.size(); + } + + /** + * Looks up the term owning a matrix row. + * + * @param row The matrix row. Must be within {@code [firstRow, firstRow + size())}. + * @return The term at that row. + * @throws IllegalArgumentException Thrown if {@code row} is outside the term rows. + */ + String term(int row) { + final int offset = row - firstRow; + if (offset < 0 || offset >= termsByOffset.size()) { + throw new IllegalArgumentException("Row " + row + " is outside the term rows [" + + firstRow + ", " + (firstRow + termsByOffset.size()) + ")"); + } + return termsByOffset.get(offset); + } + + /** + * A term match in a text: the term's matrix row and the character range it consumed, from the + * start of its first word to the end of its last. + * + * @param row The matched term's matrix row. + * @param start The inclusive start of the consumed range. + * @param end The exclusive end of the consumed range. + */ + record Match(int row, int start, int end) { + } + + /** + * Finds every term of this table in a text, greedily longest-first: at each word, the longest + * matching term consumes its words, and matching continues after them. Matched ranges never + * overlap and appear in text order. + * + * @param text The text to match. Must not be {@code null}. + * @return The matches in text order; empty when the table is empty or nothing matches. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + List matches(String text) { + if (text == null) { + throw new IllegalArgumentException("Text must not be null"); + } + if (termsByOffset.isEmpty()) { + return List.of(); + } + final List runs = wordRuns(text); + final List matches = new ArrayList<>(); + int i = 0; + while (i < runs.size()) { + int consumed = 0; + for (int n = Math.min(maxTermWords, runs.size() - i); n >= 1; n--) { + final Integer row = rowByTerm.get(joined(runs, i, n)); + if (row != null) { + matches.add(new Match(row, runs.get(i).start(), runs.get(i + n - 1).end())); + consumed = n; + break; + } + } + i += Math.max(consumed, 1); + } + return matches; + } + + /** A word run of the matched text: its character range and its case-folded form. */ + private record Run(int start, int end, String folded) { + } + + /** + * {@return the letter-or-digit word runs of a text, each with its character range and its + * case-folded form} + * + * @param text The text to scan. + */ + private static List wordRuns(String text) { + final List runs = new ArrayList<>(); + final int length = text.length(); + int i = 0; + while (i < length) { + final int c = text.codePointAt(i); + if (Character.isLetterOrDigit(c)) { + final int start = i; + while (i < length && Character.isLetterOrDigit(text.codePointAt(i))) { + i += Character.charCount(text.codePointAt(i)); + } + runs.add(new Run(start, i, StringUtil.toLowerCase(text.substring(start, i)))); + } else { + i += Character.charCount(c); + } + } + return runs; + } + + /** + * {@return the folded forms of {@code n} runs from {@code first}, joined by single spaces, the + * lookup key of a candidate term} + * + * @param runs The text's word runs. + * @param first The first run of the candidate. + * @param n The number of runs of the candidate. + */ + private static String joined(List runs, int first, int n) { + if (n == 1) { + return runs.get(first).folded(); + } + final StringBuilder key = new StringBuilder(); + for (int i = 0; i < n; i++) { + if (i > 0) { + key.append(' '); + } + key.append(runs.get(first + i).folded()); + } + return key.toString(); + } +} 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 new file mode 100644 index 0000000000..3d918ee371 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java @@ -0,0 +1,352 @@ +/* + * 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.ArrayList; +import java.util.Comparator; +import java.util.List; + +import opennlp.tools.util.InvalidFormatException; + +/** + * Reads the row order of a static embedding matrix out of a {@code tokenizer.json} file with a + * Unigram model: the {@code model.vocab} list holds {@code [piece, score]} pairs whose index is + * the piece's id, and the {@code added_tokens} list overlays extra pieces (appended when their id + * is the next row, checked for agreement when it is an existing row). Only the vocabulary is + * read; every other section, including the tokenizer's normalizer and segmentation state, is + * skipped structurally. Like the package's other readers this is purpose-built on + * {@link JsonCursor}, not a general JSON library, and fails loud on anything outside the known + * shape. + */ +final class TokenizerJsonVocab { + + /** Not instantiable. */ + private TokenizerJsonVocab() { + } + + /** + * One entry of the {@code added_tokens} list. + * + * @param id The token's id, the matrix row it claims. + * @param content The token's string. + */ + private record AddedToken(long id, String content) { + } + + /** + * Reads the pieces of a Unigram {@code tokenizer.json} in row order. + * + * @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 {@code file} is {@code null} or missing. + * @throws InvalidFormatException 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 IOException Thrown if reading the file fails. + */ + static List 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(); + cursor.expect('{'); + cursor.skipWhitespace(); + + List vocab = null; + String modelType = null; + List addedTokens = List.of(); + boolean modelSeen = false; + boolean addedTokensSeen = false; + + if (cursor.peek() == '}') { + cursor.consume(); + } else { + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "model" -> { + if (modelSeen) { + throw cursor.malformed("Field 'model' appears more than once"); + } + modelSeen = true; + final ParsedModel model = parseModel(cursor); + vocab = model.vocab; + modelType = model.type; + } + case "added_tokens" -> { + if (addedTokensSeen) { + throw cursor.malformed("Field 'added_tokens' appears more than once"); + } + addedTokensSeen = true; + addedTokens = parseAddedTokens(cursor); + } + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after a field, got '" + next + "'"); + } + } + cursor.requireEnd("Trailing content after the top-level object"); + + if (modelType != null && !"Unigram".equals(modelType)) { + throw new InvalidFormatException(file + " has a '" + modelType + "' tokenizer model; " + + "only the Unigram list layout maps pieces to matrix rows here. For a WordPiece " + + "model, load from its vocab.txt instead"); + } + if (vocab == null) { + throw new InvalidFormatException(file + " has no model.vocab list; it does not name " + + "the matrix rows"); + } + return overlayAddedTokens(vocab, addedTokens, file); + } + + /** The fields read out of the {@code model} object. */ + private record ParsedModel(String type, List vocab) { + } + + /** + * Parses the {@code model} object, collecting its {@code type} and its {@code vocab} pieces + * in list order. + * + * @param cursor The cursor, positioned at the object's opening brace. + * @return The parsed type and vocabulary; either may be absent ({@code null}). + */ + private static ParsedModel parseModel(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + List vocab = null; + if (cursor.peek() == '}') { + cursor.consume(); + return new ParsedModel(null, null); + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "type" -> { + if (type != null) { + throw cursor.malformed("Field 'model.type' appears more than once"); + } + type = cursor.parseString(); + } + case "vocab" -> { + if (vocab != null) { + throw cursor.malformed("Field 'model.vocab' appears more than once"); + } + if (cursor.peek() == '{') { + throw cursor.malformed("model.vocab is an object; only the Unigram list layout " + + "([piece, score] pairs) maps pieces to matrix rows here"); + } + vocab = parseVocabList(cursor); + } + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + return new ParsedModel(type, vocab); + } + throw cursor.malformed("Expected ',' or '}' after a model field, got '" + next + "'"); + } + } + + /** + * Parses the Unigram {@code vocab} list of {@code [piece, score]} pairs. + * + * @param cursor The cursor, positioned at the list's opening bracket. + * @return The pieces in list order. + */ + private static List parseVocabList(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List pieces = new ArrayList<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return pieces; + } + while (true) { + cursor.skipWhitespace(); + cursor.expect('['); + cursor.skipWhitespace(); + pieces.add(cursor.parseString()); + cursor.skipWhitespace(); + cursor.expect(','); + cursor.skipWhitespace(); + cursor.skipValue(); + cursor.skipWhitespace(); + cursor.expect(']'); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + return pieces; + } + throw cursor.malformed("Expected ',' or ']' after a vocab entry, got '" + next + "'"); + } + } + + /** + * Parses the {@code added_tokens} list of objects, keeping each entry's {@code id} and + * {@code content}. + * + * @param cursor The cursor, positioned at the list's opening bracket. + * @return The added tokens in list order. + */ + private static List parseAddedTokens(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List tokens = new ArrayList<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return tokens; + } + while (true) { + cursor.skipWhitespace(); + tokens.add(parseAddedToken(cursor)); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == ']') { + return tokens; + } + throw cursor.malformed("Expected ',' or ']' after an added token, got '" + next + "'"); + } + } + + /** + * Parses one {@code added_tokens} object, requiring its {@code id} and {@code content}. + * + * @param cursor The cursor, positioned at the object's opening brace. + * @return The parsed entry. + */ + private static AddedToken parseAddedToken(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + Long id = null; + String content = null; + if (cursor.peek() == '}') { + throw cursor.malformed("An added token must carry 'id' and 'content'"); + } + while (true) { + cursor.skipWhitespace(); + final String key = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (key) { + case "id" -> { + if (id != null) { + throw cursor.malformed("Field 'id' appears more than once in an added token"); + } + id = cursor.parseLong(); + } + case "content" -> { + if (content != null) { + throw cursor.malformed("Field 'content' appears more than once in an added token"); + } + content = cursor.parseString(); + } + default -> cursor.skipValue(); + } + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + continue; + } + if (next == '}') { + break; + } + throw cursor.malformed("Expected ',' or '}' after an added token field, got '" + + next + "'"); + } + if (id == null || content == null) { + throw cursor.malformed("An added token must carry 'id' and 'content'"); + } + if (id < 0) { + throw cursor.malformed("An added token's id must not be negative: " + id); + } + return new AddedToken(id, content); + } + + /** + * Overlays the added tokens onto the vocabulary in id order: an id equal to the current size + * appends, an id below it must agree with the piece already there, and a gap fails loud. + * + * @param vocab The {@code model.vocab} pieces in list order; extended in place. + * @param addedTokens The added tokens to overlay. + * @param file The source file, for error messages. + * @return The vocabulary with the added tokens applied. + * @throws InvalidFormatException Thrown if an added token contradicts the vocabulary or + * leaves a gap in the id space. + */ + private static List overlayAddedTokens(List vocab, + List addedTokens, Path file) + throws InvalidFormatException { + final List byId = new ArrayList<>(addedTokens); + byId.sort(Comparator.comparingLong(AddedToken::id)); + for (final AddedToken token : byId) { + if (token.id() == vocab.size()) { + vocab.add(token.content()); + } else if (token.id() < vocab.size()) { + final String existing = vocab.get((int) token.id()); + if (!existing.equals(token.content())) { + throw new InvalidFormatException(file + " declares added token '" + token.content() + + "' at id " + token.id() + " but model.vocab holds '" + existing + + "' there; the file contradicts itself"); + } + } else { + throw new InvalidFormatException(file + " declares added token '" + token.content() + + "' at id " + token.id() + " but the vocabulary only has " + vocab.size() + + " rows; the id space has a gap"); + } + } + return vocab; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelParams.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelParams.java new file mode 100644 index 0000000000..eb8419fd7d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelParams.java @@ -0,0 +1,34 @@ +/* + * 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.cmdline; + +import java.io.File; + +import opennlp.tools.cmdline.ArgumentParser.ParameterDescription; + +/** + * The command-line arguments of {@link AssembleModelTool}. + */ +interface AssembleModelParams { + + /** + * {@return the distilled model directory to assemble in place and verify} + */ + @ParameterDescription(valueName = "dir", + description = "The distilled model directory to complete in place and verify.") + File getModelDir(); +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelTool.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelTool.java new file mode 100644 index 0000000000..56b6d47486 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelTool.java @@ -0,0 +1,80 @@ +/* + * 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.cmdline; + +import java.io.File; +import java.io.IOException; + +import opennlp.embeddings.ModelAssembler; +import opennlp.tools.cmdline.BasicCmdLineTool; +import opennlp.tools.cmdline.TerminateToolException; +import opennlp.tools.util.InvalidFormatException; + +/** + * Completes a distilled embedding model directory so {@code StaticEmbeddingModel.load} can open it, + * then verifies it by loading it. + * + *

A Model2Vec distillation writes {@code model.safetensors}, {@code tokenizer.json}, and + * {@code config.json}. For a WordPiece model this tool derives the missing {@code vocab.txt} and + * {@code tokenizer_config.json} from {@code tokenizer.json}. A Model2Vec Unigram model is already + * self-contained. The tool loads the assembled directory and prints its family, dimension, and + * vocabulary size, so a run that prints a summary is a directory that works.

+ */ +public class AssembleModelTool extends BasicCmdLineTool { + + interface Params extends AssembleModelParams { + } + + @Override + public String getShortDescription() { + return "Completes and verifies a distilled static embedding model directory"; + } + + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + @Override + public void run(String[] args) { + final Params params = validateAndParseParams(args, Params.class); + final File modelDir = params.getModelDir(); + if (!modelDir.isDirectory()) { + throw new TerminateToolException(1, + "Model directory does not exist or is not a directory: " + modelDir); + } + final ModelAssembler.Result result; + try { + result = ModelAssembler.assemble(modelDir.toPath()); + } catch (IllegalArgumentException | InvalidFormatException e) { + throw new TerminateToolException(1, e.getMessage(), e); + } catch (IOException e) { + throw new TerminateToolException(-1, + "IO error while assembling " + modelDir + ": " + e.getMessage(), e); + } + if (result.wroteVocabulary()) { + System.out.println("Wrote vocab.txt derived from tokenizer.json"); + } + if (result.wroteTokenizerConfig()) { + System.out.println("Wrote tokenizer_config.json derived from tokenizer.json"); + } + System.out.println("Assembled and verified a " + result.family() + " model: " + + result.vocabularySize() + " rows" + + (result.termCount() > 0 ? " plus " + result.termCount() + " terms" : "") + + ", dimension " + result.dimension()); + } +} 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 new file mode 100644 index 0000000000..efd4d11717 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java @@ -0,0 +1,141 @@ +/* + * 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.cmdline; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import opennlp.tools.cmdline.BasicCmdLineTool; +import opennlp.tools.cmdline.CmdLineTool; +import opennlp.tools.cmdline.TerminateToolException; +import opennlp.tools.util.Version; + +/** + * The command line dispatcher for the OpenNLP static embeddings tools. + */ +public final class CLI { + + private static final Logger logger = LoggerFactory.getLogger(CLI.class); + static final String CMD = "opennlp-embeddings"; + + private static Map toolLookupMap; + + static { + toolLookupMap = new LinkedHashMap<>(); + + final List tools = new LinkedList<>(); + + tools.add(new AssembleModelTool()); + tools.add(new DistillModelTool()); + tools.add(new QuantizeModelTool()); + + for (CmdLineTool tool : tools) { + toolLookupMap.put(tool.getName(), tool); + } + + toolLookupMap = Collections.unmodifiableMap(toolLookupMap); + } + + /** Not instantiable. */ + private CLI() { + } + + /** {@return the names of all tools this command line dispatcher can run} */ + public static Set 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); + + // distance of tool name from line start + int numberOfSpaces = -1; + for (String toolName : toolLookupMap.keySet()) { + if (toolName.length() > numberOfSpaces) { + numberOfSpaces = toolName.length(); + } + } + numberOfSpaces = numberOfSpaces + 4; + + final StringBuilder sb = new StringBuilder("where TOOL is one of: \n\n"); + for (CmdLineTool tool : toolLookupMap.values()) { + + sb.append(" ").append(tool.getName()); + sb.append(" ".repeat(Math.max(0, StrictMath.abs( + tool.getName().length() - numberOfSpaces)))); + sb.append(tool.getShortDescription()).append("\n"); + } + logger.info(sb.toString()); + + logger.info("All tools print help when invoked with help parameter"); + 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) { + usage(); + System.exit(0); + } + + final String[] toolArguments = new String[args.length - 1]; + System.arraycopy(args, 1, toolArguments, 0, toolArguments.length); + + final String toolName = args[0]; + + final CmdLineTool tool = toolLookupMap.get(toolName); + + try { + if (null == tool) { + throw new TerminateToolException(1, "Tool " + toolName + " is not found."); + } + + if ((0 == toolArguments.length && tool.hasParams()) + || 0 < toolArguments.length && "help".equals(toolArguments[0])) { + logger.info(tool.getHelp()); + System.exit(0); + } + + if (tool instanceof BasicCmdLineTool basicTool) { + basicTool.run(toolArguments); + } else { + throw new TerminateToolException(1, "Tool " + toolName + " is not supported."); + } + } catch (TerminateToolException e) { + logger.error(e.getLocalizedMessage(), e); + System.exit(e.getCode()); + } + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelParams.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelParams.java new file mode 100644 index 0000000000..0f7100ca8d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelParams.java @@ -0,0 +1,60 @@ +/* + * 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.cmdline; + +import opennlp.tools.cmdline.ArgumentParser.OptionalParameter; +import opennlp.tools.cmdline.ArgumentParser.ParameterDescription; + +/** + * The command-line arguments of {@link DistillModelTool}. + */ +interface DistillModelParams { + + /** + * {@return the teacher to distill: a local directory or a Hugging Face model id} + */ + @ParameterDescription(valueName = "hf-id-or-path", + description = "The sentence-transformer teacher: a Hugging Face model id (org/model, or " + + "org/model@revision to pin a branch, tag, or commit) or a local directory holding " + + "tokenizer.json and onnx/model.onnx.") + String getTeacher(); + + /** + * {@return the model directory to write} + */ + @ParameterDescription(valueName = "dir", + description = "The output directory for the distilled static embedding model.") + String getOut(); + + /** + * {@return the number of PCA dimensions to keep} + */ + @OptionalParameter(defaultValue = "256") + @ParameterDescription(valueName = "num", + description = "The number of principal components to keep, default is 256.") + Integer getPcaDims(); + + /** + * {@return the term file to distill as extra rows, or {@code null} for none} + */ + @OptionalParameter + @ParameterDescription(valueName = "file", + description = "A term file: one term per line, text after a tab ignored, so a learned " + + "vocabulary TSV works as-is. Each term is encoded through the teacher as a unit and " + + "added as an extra row, matched greedily longest-first before subword tokenization.") + String getTerms(); +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelTool.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelTool.java new file mode 100644 index 0000000000..a30f21a498 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelTool.java @@ -0,0 +1,103 @@ +/* + * 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.cmdline; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import opennlp.embeddings.ModelDistiller; +import opennlp.tools.cmdline.BasicCmdLineTool; +import opennlp.tools.cmdline.TerminateToolException; +import opennlp.tools.util.InvalidFormatException; + +/** + * Distills a sentence-transformer teacher into a static embedding model directory, the + * {@code opennlp-embeddings DistillModel} command. This is Model2Vec's distillation + * (teacher forward pass over the vocabulary, PCA, Zipf weighting) in Java, so producing a table + * no longer needs a Python environment; see {@link ModelDistiller} for the pipeline. + * + *

The teacher is a Hugging Face model id (its files download once into a local cache, pinned to + * the commit its revision resolved to and verified against the digests the hub publishes for them) + * or a local directory holding {@code tokenizer.json} and {@code onnx/model.onnx}. A SentencePiece + * teacher also needs its trained {@code .model} file, downloaded or supplied alongside. The + * written directory is completed and verified by loading it, so a run that prints a summary is a + * directory that works.

+ */ +public class DistillModelTool extends BasicCmdLineTool { + + interface Params extends DistillModelParams { + } + + @Override + public String getShortDescription() { + return "Distills a sentence-transformer teacher into a static embedding model"; + } + + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + @Override + public void run(String[] args) { + // -teacher and -out are mandatory parameters, so validateAndParseParams has already + // rejected the invocation if either is absent. + final Params params = validateAndParseParams(args, Params.class); + final ModelDistiller.ProgressListener listener = System.out::println; + final ModelDistiller.Result result; + try { + final List terms = params.getTerms() == null + ? List.of() : readTerms(Path.of(params.getTerms())); + result = ModelDistiller.distill(params.getTeacher(), Path.of(params.getOut()), + params.getPcaDims(), terms, listener); + } catch (IllegalArgumentException | InvalidFormatException e) { + throw new TerminateToolException(1, e.getMessage(), e); + } catch (IOException e) { + throw new TerminateToolException(-1, + "IO error while distilling: " + e.getMessage(), e); + } + System.out.println("Distilled and verified a " + result.family() + " model: " + + result.vocabularySize() + " rows" + + (result.termCount() > 0 ? " plus " + result.termCount() + " terms" : "") + + ", " + result.teacherDimension() + "d -> " + + result.dimension() + "d, PCA kept " + + String.format("%.1f", result.explainedVarianceRatio() * 100) + "% of the variance"); + } + + /** + * Reads a term file: one term per line, text after the first tab ignored, blank lines + * skipped. A learned vocabulary TSV (term, count, source) therefore works unchanged. + * + * @param file The term file. + * @return The terms in file order. + * @throws IOException Thrown if reading the file fails. + */ + private List readTerms(Path file) throws IOException { + final List terms = new ArrayList<>(); + for (final String line : Files.readAllLines(file)) { + final int tab = line.indexOf('\t'); + final String term = (tab < 0 ? line : line.substring(0, tab)).strip(); + if (!term.isEmpty()) { + terms.add(term); + } + } + return terms; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelParams.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelParams.java new file mode 100644 index 0000000000..151035417f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelParams.java @@ -0,0 +1,51 @@ +/* + * 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.cmdline; + +import java.io.File; + +import opennlp.tools.cmdline.ArgumentParser.OptionalParameter; +import opennlp.tools.cmdline.ArgumentParser.ParameterDescription; + +/** + * The command-line arguments of {@link QuantizeModelTool}. + */ +interface QuantizeModelParams { + + /** + * {@return the model directory whose safetensors matrix is quantized in place} + */ + @ParameterDescription(valueName = "dir", + description = "the model directory whose model.safetensors is quantized in place") + File getModelDir(); + + /** + * {@return the bit width per dimension} + */ + @ParameterDescription(valueName = "bits", + description = "bits per dimension, 2 to 4; fewer bits, smaller file, lower fidelity") + @OptionalParameter(defaultValue = "4") + Integer getBits(); + + /** + * {@return the rotation seed; the same matrix, bits, and seed write the same file} + */ + @ParameterDescription(valueName = "seed", + description = "the rotation seed; the same matrix, bits, and seed write the same file") + @OptionalParameter(defaultValue = "0") + Long getSeed(); +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelTool.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelTool.java new file mode 100644 index 0000000000..8b83861e9d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelTool.java @@ -0,0 +1,71 @@ +/* + * 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.cmdline; + +import java.io.IOException; +import java.util.Locale; + +import opennlp.embeddings.ModelQuantizer; +import opennlp.tools.cmdline.BasicCmdLineTool; +import opennlp.tools.cmdline.TerminateToolException; + +/** + * Quantizes a static embedding model directory's matrix to 2-4 bits per dimension, writing + * {@code model.quantized} next to the {@code model.safetensors}, and prints the sizes and the + * measured reconstruction quality. {@code StaticEmbeddingModel.load} prefers the quantized file + * from then on; the safetensors may be deleted for a slim deployment. + */ +public class QuantizeModelTool extends BasicCmdLineTool { + + interface Params extends QuantizeModelParams { + } + + @Override + public String getShortDescription() { + return "Quantizes a static embedding model directory's matrix to 2-4 bits per dimension"; + } + + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + @Override + public void run(String[] args) { + final Params params = validateAndParseParams(args, Params.class); + final ModelQuantizer.Result result; + try { + result = ModelQuantizer.quantize(params.getModelDir().toPath(), params.getBits(), + params.getSeed()); + } catch (IllegalArgumentException e) { + throw new TerminateToolException(1, e.getMessage()); + } catch (IOException e) { + throw new TerminateToolException(-1, + "IO error while quantizing " + params.getModelDir() + ": " + e.getMessage(), e); + } + System.out.println("Quantized " + result.rowCount() + " rows of dimension " + + result.dimension() + " to " + result.bits() + " bits" + + (result.hasWeights() ? ", carrying the per-token weights" : "")); + System.out.println(String.format(Locale.ROOT, + "Size: %,d bytes safetensors, %,d bytes quantized (%.1fx smaller)", + result.safetensorsBytes(), result.quantizedBytes(), + result.safetensorsBytes() / (double) result.quantizedBytes())); + System.out.println(String.format(Locale.ROOT, + "Verified from disk: mean cosine %.4f between original and reconstructed rows " + + "(%d sampled)", result.meanCosine(), result.sampledRows())); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/FlatFloatIndex.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/FlatFloatIndex.java new file mode 100644 index 0000000000..8159c00b82 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/FlatFloatIndex.java @@ -0,0 +1,137 @@ +/* + * 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.index; + +import java.util.List; + +import opennlp.tools.util.java.Experimental; + +/** + * The exact index: full-precision float vectors scanned brute force with double-accumulated + * cosine similarity. Every query scores every vector, so this is the ground truth the + * quantized index is measured against, and the right choice outright when the collection is + * small. + * + *

Follows the {@link VectorIndex} lifecycle: single-threaded build, then + * {@link #freeze()}, then concurrent queries.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public final class FlatFloatIndex implements VectorIndex { + + private static final double NORM_EPSILON = 1e-12; + + private final int dimension; + private VectorBuffer buffer; + private List ids; + private float[] rowMajor; + private double[] norms; + + /** + * Creates an empty index. + * + * @param dimension The dimension every vector and query must have. Must be at least 1. + * @throws IllegalArgumentException Thrown if {@code dimension} is below 1. + */ + public FlatFloatIndex(int dimension) { + this.buffer = new VectorBuffer(dimension); + this.dimension = dimension; + } + + /** {@inheritDoc} */ + @Override + public void add(String id, float[] vector) { + if (buffer == null) { + throw new IllegalStateException("The index is frozen; vectors can no longer be added"); + } + buffer.add(id, vector); + } + + /** {@inheritDoc} */ + @Override + public void freeze() { + if (buffer == null) { + return; + } + ids = buffer.ids(); + rowMajor = buffer.rowMajor(); + norms = new double[ids.size()]; + for (int row = 0; row < norms.length; row++) { + final int base = row * dimension; + double sumOfSquares = 0; + for (int d = 0; d < dimension; d++) { + final float value = rowMajor[base + d]; + sumOfSquares += (double) value * value; + } + norms[row] = Math.sqrt(sumOfSquares); + } + buffer = null; + } + + /** {@inheritDoc} */ + @Override + public List topK(float[] query, int k) { + if (buffer != null) { + throw new IllegalStateException("The index is not frozen; freeze() ends the build phase"); + } + final double queryNorm = IndexQueries.checkedQueryNorm(query, k, dimension); + if (queryNorm < NORM_EPSILON || ids.isEmpty()) { + return List.of(); + } + final TopK best = new TopK(Math.min(k, ids.size())); + for (int row = 0; row < norms.length; row++) { + final double norm = norms[row]; + if (norm < NORM_EPSILON) { + // A zero vector has no direction; scored 0 rather than NaN from a 0/0 division. + best.offer(row, 0.0); + continue; + } + final int base = row * dimension; + // Four accumulators so the JIT can vectorize the dot product without reordering FP adds. + double dot0 = 0; + double dot1 = 0; + double dot2 = 0; + double dot3 = 0; + int d = 0; + for (final int limit = dimension - 3; d < limit; d += 4) { + dot0 += query[d] * rowMajor[base + d]; + dot1 += query[d + 1] * rowMajor[base + d + 1]; + dot2 += query[d + 2] * rowMajor[base + d + 2]; + dot3 += query[d + 3] * rowMajor[base + d + 3]; + } + double dot = dot0 + dot1 + dot2 + dot3; + for (; d < dimension; d++) { + dot += query[d] * rowMajor[base + d]; + } + best.offer(row, dot / (queryNorm * norm)); + } + return best.drain(ids); + } + + /** {@inheritDoc} */ + @Override + public int size() { + return buffer != null ? buffer.size() : ids.size(); + } + + /** {@inheritDoc} */ + @Override + public int dimension() { + return dimension; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/IndexQueries.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/IndexQueries.java new file mode 100644 index 0000000000..7cd5248436 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/IndexQueries.java @@ -0,0 +1,55 @@ +/* + * 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.index; + +/** + * The query-argument checks shared by the index implementations. + */ +final class IndexQueries { + + /** Not instantiable. */ + private IndexQueries() { + } + + /** + * Validates a query and returns its L2 norm. + * + * @param query The query vector. + * @param k The requested result count. + * @param dimension The index's dimension. + * @return The query's L2 norm. + * @throws IllegalArgumentException Thrown if {@code query} is {@code null} or has the wrong + * length, or {@code k} is less than 1. + */ + static double checkedQueryNorm(float[] query, int k, int dimension) { + if (query == null) { + throw new IllegalArgumentException("Query must not be null"); + } + if (query.length != dimension) { + throw new IllegalArgumentException("Query has length " + query.length + + " but this index has dimension " + dimension); + } + if (k < 1) { + throw new IllegalArgumentException("K must be at least 1, got " + k); + } + double sumOfSquares = 0; + for (final float value : query) { + sumOfSquares += (double) value * value; + } + return Math.sqrt(sumOfSquares); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TopK.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TopK.java new file mode 100644 index 0000000000..47058d9a25 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TopK.java @@ -0,0 +1,121 @@ +/* + * 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.index; + +import java.util.List; + +/** + * A bounded selection of the {@code k} highest-scoring rows, kept as a min-heap over primitive + * parallel arrays: the root is always the weakest kept candidate, so a full scan decides most + * rows with one comparison against it and allocates nothing per row. + */ +final class TopK { + + private final double[] scores; + private final int[] rows; + private int size; + + /** + * Creates an empty selection. + * + * @param capacity The maximum number of rows to keep. + */ + TopK(int capacity) { + this.scores = new double[capacity]; + this.rows = new int[capacity]; + } + + /** + * Offers a candidate row, keeping it only if it ranks among the top {@code capacity}. + * + * @param row The candidate row id. + * @param score The row's score against the query. + */ + void offer(int row, double score) { + if (size < scores.length) { + int i = size++; + scores[i] = score; + rows[i] = row; + while (i > 0) { + final int parent = (i - 1) >>> 1; + if (scores[parent] <= scores[i]) { + break; + } + swap(parent, i); + i = parent; + } + } else if (score > scores[0]) { + scores[0] = score; + rows[0] = row; + siftDown(); + } + } + + /** + * Drains the selection into hits, most similar first, mapping each kept row through the ids. + * + * @param ids The indexed ids; a kept row's id is {@code ids.get(row)}. + * @return The hits, most similar first. + */ + List drain(List ids) { + final VectorIndex.Hit[] ordered = new VectorIndex.Hit[size]; + for (int i = ordered.length - 1; i >= 0; i--) { + ordered[i] = new VectorIndex.Hit(ids.get(rows[0]), scores[0]); + size--; + scores[0] = scores[size]; + rows[0] = rows[size]; + siftDown(); + } + return List.of(ordered); + } + + /** Restores the min-heap invariant from the root downward. */ + private void siftDown() { + int i = 0; + while (true) { + final int left = 2 * i + 1; + final int right = left + 1; + int smallest = i; + if (left < size && scores[left] < scores[smallest]) { + smallest = left; + } + if (right < size && scores[right] < scores[smallest]) { + smallest = right; + } + if (smallest == i) { + return; + } + swap(i, smallest); + i = smallest; + } + } + + /** + * Swaps two heap entries in both parallel arrays. + * + * @param i The first index. + * @param j The second index. + */ + private void swap(int i, int j) { + final double score = scores[i]; + scores[i] = scores[j]; + scores[j] = score; + final int row = rows[i]; + rows[i] = rows[j]; + rows[j] = row; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TurboQuantIndex.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TurboQuantIndex.java new file mode 100644 index 0000000000..111af5cf94 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/TurboQuantIndex.java @@ -0,0 +1,250 @@ +/* + * 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.index; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import opennlp.embeddings.QuantizedEmbeddingMatrix; +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.java.Experimental; + +/** + * The quantized index: freezing quantizes the indexed vectors with the + * {@link QuantizedEmbeddingMatrix TurboQuant construction} (seeded rotation, per-coordinate + * grid codes, per-row fitted scale), and a query rotates once and scans every row's packed + * codes without ever decoding to original space. Against {@link FlatFloatIndex} this trades a + * little recall for {@code bits}-per-dimension storage instead of 32, and the scan reads + * proportionally fewer bytes. + * + *

Follows the {@link VectorIndex} lifecycle: single-threaded build, then + * {@link #freeze()}, then concurrent queries. A frozen, non-empty index persists as a + * directory of two files, the quantized matrix ({@value #VECTORS_FILE}, the self-describing + * TurboQuant format) and the ids in row order ({@value #IDS_FILE}); {@link #read(Path)} loads + * it back frozen.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public final class TurboQuantIndex implements VectorIndex { + + /** The quantized vectors of a persisted index, in the TurboQuant file format. */ + public static final String VECTORS_FILE = "vectors.onq"; + + /** The ids of a persisted index, one per line in matrix row order. */ + public static final String IDS_FILE = "ids.txt"; + + private static final double NORM_EPSILON = 1e-12; + + private final int dimension; + private final int bits; + private final long seed; + private VectorBuffer buffer; + private List ids; + // Null while building, and in the frozen empty index, which has no rows to quantize. + private QuantizedEmbeddingMatrix matrix; + + /** + * Creates an empty index. + * + * @param dimension The dimension every vector and query must have. Must be at least 1. + * @param bits The bit width per stored dimension. Must be between + * {@link QuantizedEmbeddingMatrix#MIN_BITS} and + * {@link QuantizedEmbeddingMatrix#MAX_BITS}. + * @param seed The rotation seed; the same vectors, bit width, and seed quantize + * identically on every JVM. + * @throws IllegalArgumentException Thrown if {@code dimension} or {@code bits} is out of + * range. + */ + public TurboQuantIndex(int dimension, int bits, long seed) { + this.buffer = new VectorBuffer(dimension); + if (bits < QuantizedEmbeddingMatrix.MIN_BITS || bits > QuantizedEmbeddingMatrix.MAX_BITS) { + throw new IllegalArgumentException("Bits must be between " + + QuantizedEmbeddingMatrix.MIN_BITS + " and " + QuantizedEmbeddingMatrix.MAX_BITS + + ", got " + bits); + } + this.dimension = dimension; + this.bits = bits; + this.seed = seed; + } + + /** Holds a loaded index; callers reach this through {@link #read(Path)}. */ + private TurboQuantIndex(List ids, QuantizedEmbeddingMatrix matrix) { + this.dimension = matrix.dimension(); + this.bits = matrix.bits(); + this.seed = matrix.seed(); + this.ids = ids; + this.matrix = matrix; + } + + /** {@inheritDoc} */ + @Override + public void add(String id, float[] vector) { + if (buffer == null) { + throw new IllegalStateException("The index is frozen; vectors can no longer be added"); + } + buffer.add(id, vector); + } + + /** + * {@inheritDoc} + * + *

Freezing quantizes every added vector; this is the index's one expensive step.

+ */ + @Override + public void freeze() { + if (buffer == null) { + return; + } + ids = buffer.ids(); + if (!ids.isEmpty()) { + matrix = QuantizedEmbeddingMatrix.quantize(buffer.rowMajor(), ids.size(), dimension, + bits, seed); + } + buffer = null; + } + + /** {@inheritDoc} */ + @Override + public List topK(float[] query, int k) { + if (buffer != null) { + throw new IllegalStateException("The index is not frozen; freeze() ends the build phase"); + } + final double queryNorm = IndexQueries.checkedQueryNorm(query, k, dimension); + if (queryNorm < NORM_EPSILON || ids.isEmpty()) { + return List.of(); + } + final float[] rotated = matrix.rotate(query); + final TopK best = new TopK(Math.min(k, ids.size())); + for (int row = 0; row < ids.size(); row++) { + final double rowNorm = matrix.rowNorm(row); + if (rowNorm < NORM_EPSILON) { + // A zero row has no direction; scored 0 rather than NaN from a 0/0 division. + best.offer(row, 0.0); + continue; + } + best.offer(row, matrix.dotRotated(row, rotated) / (queryNorm * rowNorm)); + } + return best.drain(ids); + } + + /** {@inheritDoc} */ + @Override + public int size() { + return buffer != null ? buffer.size() : ids.size(); + } + + /** {@inheritDoc} */ + @Override + public int dimension() { + return dimension; + } + + /** {@return the bit width per stored dimension} */ + public int bits() { + return bits; + } + + /** + * {@return the storage cost of one indexed vector: the packed codes over the padded + * dimension plus the per-row scale and norm floats} + * + * @throws IllegalStateException Thrown if the index is not frozen or is empty. + */ + public double bytesPerVector() { + if (buffer != null) { + throw new IllegalStateException("The index is not frozen; freeze() ends the build phase"); + } + if (matrix == null) { + throw new IllegalStateException("An empty index stores no vectors"); + } + return (matrix.paddedDimension() * bits + Byte.SIZE - 1) / Byte.SIZE + 2 * Float.BYTES; + } + + /** + * Writes this frozen, non-empty index as a directory of {@value #VECTORS_FILE} and + * {@value #IDS_FILE}. The directory is created when missing; the two files are replaced. + * + * @param directory The directory to write. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code directory} is {@code null}. + * @throws IllegalStateException Thrown if the index is not frozen or is empty. + * @throws IOException Thrown if writing fails. + */ + public void write(Path directory) throws IOException { + if (directory == null) { + throw new IllegalArgumentException("Directory must not be null"); + } + if (buffer != null) { + throw new IllegalStateException("The index is not frozen; freeze() ends the build phase"); + } + if (ids.isEmpty()) { + throw new IllegalStateException("An empty index has nothing to persist"); + } + Files.createDirectories(directory); + matrix.write(directory.resolve(VECTORS_FILE)); + Files.write(directory.resolve(IDS_FILE), ids); + } + + /** + * Reads an index a previous {@link #write(Path)} persisted. The loaded index is frozen. + * + * @param directory The index directory. Must not be {@code null} and must be a directory + * holding {@value #VECTORS_FILE} and {@value #IDS_FILE}. + * @return The loaded index. + * @throws IllegalArgumentException Thrown if {@code directory} is {@code null}, is not a + * directory, or lacks one of the two files. + * @throws InvalidFormatException Thrown if a file is malformed, an id repeats or is blank, + * or the id count and the matrix's row count disagree. + * @throws IOException Thrown if reading fails. + */ + public static TurboQuantIndex read(Path directory) throws IOException { + if (directory == null) { + throw new IllegalArgumentException("Directory must not be null"); + } + if (!Files.isDirectory(directory)) { + throw new IllegalArgumentException( + "Index directory does not exist or is not a directory: " + directory); + } + final Path vectorsFile = directory.resolve(VECTORS_FILE); + final Path idsFile = directory.resolve(IDS_FILE); + if (!Files.isRegularFile(vectorsFile) || !Files.isRegularFile(idsFile)) { + throw new IllegalArgumentException("Index directory " + directory + " does not hold " + + VECTORS_FILE + " and " + IDS_FILE); + } + final List ids = Files.readAllLines(idsFile); + final Set seen = new HashSet<>(ids.size() * 2); + for (final String id : ids) { + if (id.isBlank()) { + throw new InvalidFormatException(idsFile + " holds a blank id"); + } + if (!seen.add(id)) { + throw new InvalidFormatException(idsFile + " holds id '" + id + "' more than once"); + } + } + final QuantizedEmbeddingMatrix matrix = QuantizedEmbeddingMatrix.read(vectorsFile); + if (matrix.rowCount() != ids.size()) { + throw new InvalidFormatException(idsFile + " holds " + ids.size() + " ids but " + + vectorsFile + " holds " + matrix.rowCount() + " rows; these files do not belong " + + "to the same index"); + } + return new TurboQuantIndex(List.copyOf(ids), matrix); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorBuffer.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorBuffer.java new file mode 100644 index 0000000000..e9a7c44bee --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorBuffer.java @@ -0,0 +1,108 @@ +/* + * 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.index; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * The build-phase collector shared by the index implementations: validates and copies each + * added id and vector, and flattens the rows when the index freezes. Not safe for concurrent + * use; the build phase is single-threaded by contract. + */ +final class VectorBuffer { + + private final int dimension; + private final List ids = new ArrayList<>(); + private final Set seen = new HashSet<>(); + private final List rows = new ArrayList<>(); + + /** + * Creates an empty buffer. + * + * @param dimension The dimension every vector must have. Must be at least 1. + * @throws IllegalArgumentException Thrown if {@code dimension} is below 1. + */ + VectorBuffer(int dimension) { + if (dimension < 1) { + throw new IllegalArgumentException("Dimension must be at least 1, got " + dimension); + } + this.dimension = dimension; + } + + /** + * Validates and stores one id and vector. + * + * @param id The vector's id. Must not be {@code null} or blank, must not contain a line + * break, and must not already be present. + * @param vector The vector. Must not be {@code null}, must have the buffer's dimension, and + * every value must be finite. The array is copied. + * @throws IllegalArgumentException Thrown if {@code id} or {@code vector} is invalid. + */ + void add(String id, float[] vector) { + if (id == null || id.isBlank()) { + throw new IllegalArgumentException("Id must not be null or blank"); + } + if (id.indexOf('\n') >= 0 || id.indexOf('\r') >= 0) { + throw new IllegalArgumentException("Id must not contain a line break: '" + id + "'"); + } + if (vector == null) { + throw new IllegalArgumentException("Vector must not be null"); + } + if (vector.length != dimension) { + throw new IllegalArgumentException("Vector has length " + vector.length + + " but this index has dimension " + dimension); + } + for (int d = 0; d < dimension; d++) { + if (!Float.isFinite(vector[d])) { + throw new IllegalArgumentException("Vector '" + id + "' has a non-finite value at " + + "dimension " + d + ": " + vector[d]); + } + } + if (!seen.add(id)) { + throw new IllegalArgumentException("Id '" + id + "' is already indexed"); + } + ids.add(id); + rows.add(vector.clone()); + } + + /** {@return the number of stored vectors} */ + int size() { + return ids.size(); + } + + /** {@return the dimension every vector has} */ + int dimension() { + return dimension; + } + + /** {@return the ids in add order, as an immutable list} */ + List ids() { + return List.copyOf(ids); + } + + /** {@return the vectors flattened row-major, in add order} */ + float[] rowMajor() { + final float[] rowMajor = new float[rows.size() * dimension]; + for (int row = 0; row < rows.size(); row++) { + System.arraycopy(rows.get(row), 0, rowMajor, row * dimension, dimension); + } + return rowMajor; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorIndex.java b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorIndex.java new file mode 100644 index 0000000000..3b44eb9755 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/index/VectorIndex.java @@ -0,0 +1,84 @@ +/* + * 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.index; + +import java.util.List; + +import opennlp.tools.util.java.Experimental; + +/** + * A build-once, read-many index over embedding vectors: vectors are added under caller-chosen + * ids, the index is frozen, and queries return the nearest ids by cosine similarity. + * + *

The lifecycle is two-phase. During the build phase, {@link #add(String, float[])} collects + * vectors; neither adding nor freezing is safe from more than one thread. {@link #freeze()} ends + * the build phase, after which {@link #add(String, float[])} is rejected and + * {@link #topK(float[], int)} is available; a frozen index that is safely published is safe for + * concurrent queries from any number of threads. Implementations must support concurrent + * {@code topK} calls after safe publication. No other concurrent calls are permitted.

+ * + *

Warning: Experimental new feature; the API might change in a later release.

+ */ +@Experimental +public interface VectorIndex { + + /** + * A query result: an indexed id and its cosine similarity to the query. + * + * @param id The indexed id. + * @param score The cosine similarity, in {@code [-1, 1]}. + */ + record Hit(String id, double score) { + } + + /** + * Adds a vector under an id during the build phase. + * + * @param id The vector's id. Must not be {@code null} or blank, must not contain a line + * break, and must not already be indexed. + * @param vector The vector. Must not be {@code null}, must have the index's dimension, and + * every value must be finite. The array is copied. + * @throws IllegalArgumentException Thrown if {@code id} or {@code vector} is invalid. + * @throws IllegalStateException Thrown if the index is frozen. + */ + void add(String id, float[] vector); + + /** + * Ends the build phase. Calling this more than once has no further effect. + */ + void freeze(); + + /** + * Finds the indexed vectors nearest a query, most similar first. + * + * @param query The query vector. Must not be {@code null} and must have the index's + * dimension. + * @param k The maximum number of results. Must be at least 1. + * @return Up to {@code k} hits, most similar first; empty when the index is empty or the + * query has no direction. + * @throws IllegalArgumentException Thrown if {@code query} is {@code null} or has the wrong + * length, or {@code k} is less than 1. + * @throws IllegalStateException Thrown if the index is not frozen. + */ + List topK(float[] query, int k); + + /** {@return the number of indexed vectors} */ + int size(); + + /** {@return the dimension every indexed vector and query must have} */ + int dimension(); +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingTestFixtures.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingTestFixtures.java new file mode 100644 index 0000000000..741e4f30db --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingTestFixtures.java @@ -0,0 +1,197 @@ +/* + * 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.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; +import opennlp.subword.sentencepiece.SentencePieceTokenizer; + +/** + * Fixtures shared by more than one test in this module: the small WordPiece table the geometry + * tests load, and JSON string quoting for the hand-built {@code tokenizer.json} fixtures. + */ +final class EmbeddingTestFixtures { + + /** The analogy table's tokens; the list index is the matrix row. */ + static final List ANALOGY_VOCABULARY = + List.of("[CLS]", "[SEP]", "[UNK]", "king", "queen", "man", "woman", "apple"); + + /** + * The analogy table's rows, chosen so the classic word2vec analogy is exact: + * {@code king - man + woman = [3,3] - [2,1] + [1,2] = [2,4] = queen}. The directions genuinely + * differ, so pairwise cosine similarities are not trivially 1.0. + */ + static final float[][] ANALOGY_ROWS = { + {0f, 0f}, // [CLS] + {0f, 0f}, // [SEP] + {0f, 0f}, // [UNK] + {3f, 3f}, // king + {2f, 4f}, // queen + {2f, 1f}, // man + {1f, 2f}, // woman + {-3f, -1f}, // apple: unrelated, opposite-ish direction + }; + + /** Not instantiable. */ + private EmbeddingTestFixtures() { + } + + /** + * Writes {@link #ANALOGY_VOCABULARY} and {@link #ANALOGY_ROWS} into a directory and loads them + * through the explicit WordPiece overload. + * + * @param dir The directory to write the fixture files into. + * @param normalization Whether the loaded model L2-normalizes its pooled vectors. + * @return The loaded model. + * @throws IOException Thrown if writing or reading a fixture file fails. + */ + static StaticEmbeddingModel loadAnalogyModel(Path dir, Normalization normalization) + throws IOException { + writeVocabularyAndMatrix(dir); + return StaticEmbeddingModel.load(dir.resolve("vocab.txt"), dir.resolve("model.safetensors"), + Casing.UNCASED, normalization); + } + + /** + * Writes {@link #ANALOGY_VOCABULARY} and {@link #ANALOGY_ROWS} into a directory as a complete + * WordPiece model directory (with its two JSON configuration files), so a test can load it + * with {@code StaticEmbeddingModel.load(Path)} the way the manual's usage listing shows. + * + * @param dir The directory to write the model files into. + * @throws IOException Thrown if writing a fixture file fails. + */ + static void writeAnalogyDirectory(Path dir) throws IOException { + writeVocabularyAndMatrix(dir); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + Files.writeString(dir.resolve("tokenizer_config.json"), "{\"do_lower_case\":true}"); + } + + /** + * Writes the analogy table's {@code vocab.txt} and {@code model.safetensors} into a directory. + * + * @param dir The directory to write the fixture files into. + * @throws IOException Thrown if writing a fixture file fails. + */ + private static void writeVocabularyAndMatrix(Path dir) throws IOException { + Files.write(dir.resolve("vocab.txt"), ANALOGY_VOCABULARY); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ANALOGY_ROWS)); + } + + /** The classpath resource of the tiny trained SentencePiece model shared by the tests. */ + static final String TINY_UNIGRAM_RESOURCE = "/opennlp/embeddings/tiny-unigram.model"; + + /** The row width of the matrix {@link #writeSentencePieceDirectory(Path)} writes. */ + static final int SENTENCEPIECE_DIMENSION = 4; + + /** + * Writes a minimal loadable SentencePiece model into a directory: the trained + * {@code tiny-unigram.model} fixture copied as {@code sentencepiece.bpe.model}, a Unigram + * {@code tokenizer.json} whose vocabulary is the unknown piece followed by every poolable + * tokenizer piece, and a deterministic embedding matrix with one row per listed piece. A test + * can then load it through the explicit + * {@code StaticEmbeddingModel.loadSentencePiece(Path, Path, Path, Normalization)} overload + * the way the manual's listing shows. + * + * @param dir The directory to write the model files into. + * @throws IOException Thrown if reading the fixture resource or writing a file fails. + */ + static void writeSentencePieceDirectory(Path dir) throws IOException { + writeSentencePieceDirectory(dir, List.of()); + } + + /** + * Writes the SentencePiece model directory of {@link #writeSentencePieceDirectory(Path)} with + * additional term rows: the terms land in {@code terms.txt} and the matrix grows one row per + * term, keeping the deterministic {@code row + d * 0.25} cell formula, so a test can predict a + * term row's vector from the model's vocabulary size. + * + * @param dir The directory to write the model files into. + * @param terms The terms in row order; empty for none. + * @throws IOException Thrown if reading the fixture resource or writing a file fails. + */ + static void writeSentencePieceDirectory(Path dir, List terms) throws IOException { + final byte[] modelBytes; + try (InputStream in = + EmbeddingTestFixtures.class.getResourceAsStream(TINY_UNIGRAM_RESOURCE)) { + modelBytes = in.readAllBytes(); + } + Files.write(dir.resolve("sentencepiece.bpe.model"), modelBytes); + final SentencePieceTokenizer tokenizer = + SentencePieceTokenizer.load(new ByteArrayInputStream(modelBytes)); + final List rows = new ArrayList<>(); + rows.add(""); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (!tokenizer.isControl(id) && !tokenizer.isUnknown(id)) { + rows.add(tokenizer.idToPiece(id)); + } + } + final StringBuilder json = + new StringBuilder("{\"model\":{\"type\":\"Unigram\",\"unk_id\":0,\"vocab\":["); + for (int i = 0; i < rows.size(); i++) { + if (i > 0) { + json.append(','); + } + json.append('[').append(jsonString(rows.get(i))).append(",-1.5]"); + } + Files.writeString(dir.resolve("tokenizer.json"), json.append("]}}").toString()); + if (!terms.isEmpty()) { + Files.write(dir.resolve("terms.txt"), terms); + } + final float[][] matrix = new float[rows.size() + terms.size()][SENTENCEPIECE_DIMENSION]; + for (int row = 0; row < matrix.length; row++) { + for (int d = 0; d < SENTENCEPIECE_DIMENSION; d++) { + matrix[row][d] = row + d * 0.25f; + } + } + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + } + + /** + * {@return {@code value} as a JSON string literal, quoted and escaped} + * + * @param value The string to quote. + */ + static String jsonString(String value) { + final StringBuilder quoted = new StringBuilder("\""); + for (int i = 0; i < value.length(); i++) { + final char c = value.charAt(i); + switch (c) { + case '"' -> quoted.append("\\\""); + case '\\' -> quoted.append("\\\\"); + default -> { + if (c < 0x20) { + quoted.append(String.format("\\u%04x", (int) c)); + } else { + quoted.append(c); + } + } + } + } + return quoted.append('"').toString(); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingVocabularyTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingVocabularyTest.java new file mode 100644 index 0000000000..d11a6297c7 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingVocabularyTest.java @@ -0,0 +1,83 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +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 EmbeddingVocabularyTest { + + @Test + void testLineNumberIsTheTokenId() throws InvalidFormatException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.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() throws InvalidFormatException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.fromLines(List.of("hello"), "test"); + assertEquals(-1, vocabulary.id("missing")); + assertThrows(IllegalArgumentException.class, () -> vocabulary.id(null)); + } + + @Test + void testDuplicateTokenFailsLoudlyNamingBothLines() { + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> EmbeddingVocabulary.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() throws InvalidFormatException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.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 EmbeddingVocabulary read = EmbeddingVocabulary.fromVocabTxt(file); + assertEquals(2, read.size()); + assertEquals(1, read.id("token")); + } +} 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 new file mode 100644 index 0000000000..6ceb7a179d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/FlatJsonFieldsTest.java @@ -0,0 +1,175 @@ +/* + * 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 org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class FlatJsonFieldsTest { + + private static Path write(Path dir, String json) throws IOException { + final Path file = dir.resolve("config.json"); + Files.writeString(file, json); + return file; + } + + @Test + void testReadsTopLevelBooleans(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"normalize\":true,\"do_lower_case\":false}"); + + assertEquals(Boolean.TRUE, FlatJsonFields.topLevelBoolean(file, "normalize")); + assertEquals(Boolean.FALSE, FlatJsonFields.topLevelBoolean(file, "do_lower_case")); + } + + @Test + void testAbsentFieldAndExplicitNullBothReadAsNull(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"strip_accents\":null}"); + + assertNull(FlatJsonFields.topLevelBoolean(file, "strip_accents")); + assertNull(FlatJsonFields.topLevelBoolean(file, "missing")); + } + + @Test + void testSkipsFieldsOfEveryOtherType(@TempDir Path dir) throws IOException { + // The shapes real tokenizer_config.json files carry around the looked-up field: nested + // objects, arrays, floats, and strings must all be skipped structurally. + final Path file = write(dir, "{\"added_tokens_decoder\":{\"0\":{\"special\":true}}," + + "\"model_max_length\":1.0E9,\"architectures\":[\"StaticModel\"]," + + "\"cls_token\":\"[CLS]\",\"normalize\":true}"); + + assertEquals(Boolean.TRUE, FlatJsonFields.topLevelBoolean(file, "normalize")); + } + + @Test + void testNestedOccurrencesOfTheNameDoNotMatch(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"outer\":{\"normalize\":true}}"); + + assertNull(FlatJsonFields.topLevelBoolean(file, "normalize")); + } + + @Test + void testToleratesAnEmptyObjectAndTrailingWhitespace(@TempDir Path dir) throws IOException { + assertNull(FlatJsonFields.topLevelBoolean(write(dir, "{} \n"), "normalize")); + } + + @Test + void testRejectsANonBooleanValue(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"normalize\":\"yes\"}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelBoolean(file, "normalize")); + assertTrue(e.getMessage().contains("must be a boolean")); + } + + @Test + void testRejectsADuplicateField(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"normalize\":true,\"normalize\":false}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelBoolean(file, "normalize")); + assertTrue(e.getMessage().contains("more than once")); + } + + @Test + void testRejectsMalformedJsonWithTheFileNameInTheMessage(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"normalize\" true}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelBoolean(file, "normalize")); + assertTrue(e.getMessage().contains("config.json")); + } + + @Test + void testRejectsTrailingGarbage(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{} x"); + + assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelBoolean(file, "normalize")); + } + + @Test + void testMissingFileFailsAsAnIoProblem(@TempDir Path dir) { + assertThrows(IOException.class, + () -> FlatJsonFields.topLevelBoolean(dir.resolve("absent.json"), "normalize")); + } + + @Test + void testReadsTopLevelStrings(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"pad_token\":\"[PAD]\",\"unk_token\":\"esc\\\"aped\"}"); + + assertEquals("[PAD]", FlatJsonFields.topLevelString(file, "pad_token")); + assertEquals("esc\"aped", FlatJsonFields.topLevelString(file, "unk_token")); + } + + @Test + void testAbsentStringFieldAndExplicitNullBothReadAsNull(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"pad_token\":null}"); + + assertNull(FlatJsonFields.topLevelString(file, "pad_token")); + assertNull(FlatJsonFields.topLevelString(file, "missing")); + } + + @Test + void testNestedOccurrencesOfAStringNameDoNotMatch(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"outer\":{\"pad_token\":\"[PAD]\"}}"); + + assertNull(FlatJsonFields.topLevelString(file, "pad_token")); + } + + @Test + void testRejectsANonStringValue(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"pad_token\":true}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelString(file, "pad_token")); + assertTrue(e.getMessage().contains("must be a string")); + } + + @Test + void testRejectsADuplicateStringField(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"pad_token\":\"a\",\"pad_token\":\"b\"}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> FlatJsonFields.topLevelString(file, "pad_token")); + assertTrue(e.getMessage().contains("more than once")); + } + + @Test + void testRejectsNullFileAndFieldArguments(@TempDir Path dir) throws IOException { + final Path file = write(dir, "{\"normalize\":true}"); + + assertThrows(IllegalArgumentException.class, + () -> FlatJsonFields.topLevelBoolean(null, "normalize")); + assertThrows(IllegalArgumentException.class, + () -> FlatJsonFields.topLevelBoolean(file, null)); + assertThrows(IllegalArgumentException.class, + () -> FlatJsonFields.topLevelString(null, "pad_token")); + assertThrows(IllegalArgumentException.class, + () -> FlatJsonFields.topLevelString(file, null)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/GaussianQuantizerTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/GaussianQuantizerTest.java new file mode 100644 index 0000000000..0484718dc3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/GaussianQuantizerTest.java @@ -0,0 +1,123 @@ +/* + * 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.Random; + +import org.junit.jupiter.api.Test; + +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 quantizer contract: the derived grids match the published Lloyd-Max tables for the + * Gaussian, encoding picks the nearest level, grids are symmetric, and a grid read back from a + * file round-trips through {@code fromLevels}. + */ +class GaussianQuantizerTest { + + @Test + void testGridsMatchThePublishedLloydMaxTables() { + // Reference values from Max, "Quantizing for minimum distortion", IRE Transactions on + // Information Theory 6(1), 1960, table for the standard normal: the positive levels of the + // symmetric optimal quantizer. The derivation here discretizes the density, so agreement is + // to the published tables' precision, not bit-exact. + assertPositiveLevels(GaussianQuantizer.forBits(2), 0.4528, 1.5104); + assertPositiveLevels(GaussianQuantizer.forBits(3), 0.2451, 0.7560, 1.3439, 2.1520); + assertPositiveLevels(GaussianQuantizer.forBits(4), + 0.1284, 0.3881, 0.6568, 0.9424, 1.2562, 1.6181, 2.0690, 2.7326); + } + + /** + * Asserts the upper half of a symmetric grid, and by symmetry the lower half. + * + * @param quantizer The quantizer under test. + * @param expected The published positive levels, ascending. + */ + private static void assertPositiveLevels(GaussianQuantizer quantizer, double... expected) { + final int half = quantizer.levelCount() / 2; + assertEquals(expected.length, half); + for (int i = 0; i < half; i++) { + assertEquals(expected[i], quantizer.level(half + i), 2e-3, + "positive level " + i + " must match the published Lloyd-Max table"); + assertEquals(-expected[i], quantizer.level(half - 1 - i), 2e-3, + "the grid must be symmetric"); + } + } + + @Test + void testEncodePicksTheNearestLevel() { + final Random random = new Random(42); + for (int bits = GaussianQuantizer.MIN_BITS; bits <= GaussianQuantizer.MAX_BITS; bits++) { + final GaussianQuantizer quantizer = GaussianQuantizer.forBits(bits); + for (int trial = 0; trial < 10_000; trial++) { + final float value = (float) (random.nextGaussian() * 2); + final int code = quantizer.encode(value); + final double encodedDistance = Math.abs(value - quantizer.level(code)); + for (int other = 0; other < quantizer.levelCount(); other++) { + assertTrue(encodedDistance <= Math.abs(value - quantizer.level(other)) + 1e-6, + "encode(" + value + ") chose level " + code + " but level " + other + + " is nearer"); + } + } + } + } + + @Test + void testEncodeCoversTheFullCodeRange() { + final GaussianQuantizer quantizer = GaussianQuantizer.forBits(2); + assertEquals(0, quantizer.encode(-10f)); + assertEquals(quantizer.levelCount() - 1, quantizer.encode(10f)); + } + + @Test + void testFromLevelsRoundTripsAGrid() { + final GaussianQuantizer original = GaussianQuantizer.forBits(3); + final GaussianQuantizer restored = GaussianQuantizer.fromLevels(original.levels()); + assertEquals(original.levelCount(), restored.levelCount()); + for (int code = 0; code < original.levelCount(); code++) { + assertEquals(original.level(code), restored.level(code), 0f); + } + final Random random = new Random(7); + for (int trial = 0; trial < 1_000; trial++) { + final float value = (float) (random.nextGaussian() * 2); + assertEquals(original.encode(value), restored.encode(value), + "a restored grid must encode exactly like its source"); + } + } + + @Test + void testFromLevelsRejectsMalformedGrids() { + assertThrows(IllegalArgumentException.class, () -> GaussianQuantizer.fromLevels(null)); + assertThrows(IllegalArgumentException.class, + () -> GaussianQuantizer.fromLevels(new float[] {1f, 2f, 3f})); + assertThrows(IllegalArgumentException.class, + () -> GaussianQuantizer.fromLevels(new float[] {-1f, -1f, 1f, 2f})); + assertThrows(IllegalArgumentException.class, + () -> GaussianQuantizer.fromLevels(new float[] {-1f, Float.NaN, 1f, 2f})); + assertThrows(IllegalArgumentException.class, + () -> GaussianQuantizer.fromLevels(new float[] {1f, 2f})); + } + + @Test + void testUnsupportedBitWidthsFailLoud() { + assertThrows(IllegalArgumentException.class, () -> GaussianQuantizer.forBits(1)); + assertThrows(IllegalArgumentException.class, () -> GaussianQuantizer.forBits(5)); + assertThrows(IllegalArgumentException.class, () -> GaussianQuantizer.forBits(0)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HadamardRotationTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HadamardRotationTest.java new file mode 100644 index 0000000000..013f31086b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HadamardRotationTest.java @@ -0,0 +1,144 @@ +/* + * 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.Arrays; +import java.util.Random; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * The rotation contract: orthonormal (norms and dot products preserved), self-consistent + * (inverse restores the input), deterministic per seed and JVM-independent by construction, and + * padded to the next power of two. + */ +class HadamardRotationTest { + + @Test + void testPaddedDimensionIsTheNextPowerOfTwo() { + assertEquals(1, HadamardRotation.paddedDimension(1)); + assertEquals(2, HadamardRotation.paddedDimension(2)); + assertEquals(4, HadamardRotation.paddedDimension(3)); + assertEquals(256, HadamardRotation.paddedDimension(256)); + assertEquals(512, HadamardRotation.paddedDimension(300)); + assertEquals(1024, HadamardRotation.paddedDimension(1024)); + assertThrows(IllegalArgumentException.class, () -> HadamardRotation.paddedDimension(0)); + assertThrows(IllegalArgumentException.class, () -> HadamardRotation.paddedDimension(-5)); + } + + @Test + void testRotationPreservesNormsAndDotProducts() { + final Random random = new Random(42); + final HadamardRotation rotation = new HadamardRotation(300, 7L); + final float[] a = randomPadded(rotation, random); + final float[] b = randomPadded(rotation, random); + final double normBefore = norm(a); + final double dotBefore = dot(a, b); + rotation.rotate(a); + rotation.rotate(b); + assertEquals(normBefore, norm(a), 1e-3 * normBefore, + "an orthonormal transform preserves norms"); + assertEquals(dotBefore, dot(a, b), 1e-3 * (1 + Math.abs(dotBefore)), + "an orthonormal transform preserves dot products"); + } + + @Test + void testInverseRestoresTheInput() { + final Random random = new Random(43); + final HadamardRotation rotation = new HadamardRotation(256, 99L); + final float[] vector = randomPadded(rotation, random); + final float[] original = vector.clone(); + rotation.rotate(vector); + rotation.inverse(vector); + assertArrayEquals(original, vector, 1e-4f); + } + + @Test + void testSameSeedSameRotationDifferentSeedDifferentRotation() { + final Random random = new Random(44); + final float[] input = randomPadded(new HadamardRotation(64, 5L), random); + final float[] first = input.clone(); + final float[] second = input.clone(); + final float[] other = input.clone(); + new HadamardRotation(64, 5L).rotate(first); + new HadamardRotation(64, 5L).rotate(second); + new HadamardRotation(64, 6L).rotate(other); + assertArrayEquals(first, second, 0f, "the same seed must give bit-identical rotations"); + assertFalse(Arrays.equals(first, other), + "different seeds must give different rotations"); + } + + @Test + void testEnergySpreadsAcrossCoordinates() { + // A one-hot vector concentrates all its energy in one coordinate; after rotation every + // coordinate must hold a share, which is the property the per-coordinate quantizer needs. + final HadamardRotation rotation = new HadamardRotation(128, 11L); + final float[] oneHot = new float[rotation.paddedDimension()]; + oneHot[3] = 1f; + rotation.rotate(oneHot); + final double expectedMagnitude = 1.0 / Math.sqrt(rotation.paddedDimension()); + for (final float value : oneHot) { + assertEquals(expectedMagnitude, Math.abs(value), 1e-6, + "a rotated one-hot vector has equal magnitude everywhere"); + } + } + + @Test + void testRejectsWrongLengthAndNull() { + final HadamardRotation rotation = new HadamardRotation(300, 1L); + assertEquals(512, rotation.paddedDimension()); + assertThrows(IllegalArgumentException.class, () -> rotation.rotate(null)); + assertThrows(IllegalArgumentException.class, () -> rotation.rotate(new float[300])); + assertThrows(IllegalArgumentException.class, () -> rotation.inverse(new float[511])); + assertThrows(IllegalArgumentException.class, () -> new HadamardRotation(0, 1L)); + } + + @Test + void testDimensionOneIsTheIdentityUpToSign() { + final HadamardRotation rotation = new HadamardRotation(1, 123L); + final float[] vector = new float[] {2.5f}; + rotation.rotate(vector); + assertEquals(2.5f, Math.abs(vector[0]), 1e-6f); + rotation.inverse(vector); + assertEquals(2.5f, vector[0], 1e-6f); + } + + private static float[] randomPadded(HadamardRotation rotation, Random random) { + final float[] vector = new float[rotation.paddedDimension()]; + for (int i = 0; i < vector.length; i++) { + vector[i] = (float) random.nextGaussian(); + } + return vector; + } + + private static double norm(float[] vector) { + return Math.sqrt(dot(vector, vector)); + } + + private static double dot(float[] a, float[] b) { + double dot = 0; + for (int i = 0; i < a.length; i++) { + dot += (double) a[i] * b[i]; + } + return dot; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HuggingFaceModelCacheTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HuggingFaceModelCacheTest.java new file mode 100644 index 0000000000..d97e13e450 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HuggingFaceModelCacheTest.java @@ -0,0 +1,835 @@ +/* + * 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.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Stream; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +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; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The cache's teacher-reference contract and its download integrity, exercised against a hub + * served on the loopback interface: no test here reaches the network. A local directory is + * returned as-is, anything that is neither a directory nor an {@code org/model} hub id is rejected + * before a request is made, and a download is pinned to one commit and refused unless it matches + * the digest the hub published for it. + */ +class HuggingFaceModelCacheTest { + + /** The address the test hub binds to, so that a test cannot leave the machine. */ + private static final String LOOPBACK = "127.0.0.1"; + + /** The model id of the teacher the hub serves. */ + private static final String MODEL_ID = "acme/teacher"; + + /** The cache directory name {@link #MODEL_ID} maps to, derived rather than restated. */ + private static final String CACHE_NAME = HuggingFaceModelCache.cacheDirectoryName(MODEL_ID); + + /** + * {@return the cache directory name for {@link #MODEL_ID} pinned to a revision} + * + * @param revision The revision the reference names. + */ + private static String cacheNameAt(String revision) { + return HuggingFaceModelCache.cacheDirectoryName(MODEL_ID + "@" + revision); + } + + /** The ref a teacher reference without a revision resolves. */ + private static final String DEFAULT_REF = "main"; + + /** The commit {@link #DEFAULT_REF} resolves to, a sha of the shape the hub reports. */ + private static final String COMMIT = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"; + + /** A second commit, for the teacher that moved under its ref. */ + private static final String OTHER_COMMIT = "0f2b8b1d4c7e6a5938271605f4e3d2c1b0a99887"; + + /** The tokenizer the hub serves, a file small enough for git to store it as a blob. */ + private static final byte[] TOKENIZER = bytes("{\"model\":{\"type\":\"WordPiece\"}}\n"); + + /** + * The git blob SHA-1 of {@link #TOKENIZER}, the 40 character form of the etag: this value comes + * from {@code git hash-object} over the same bytes, not from the code under test. + */ + private static final String TOKENIZER_BLOB_SHA1 = "296101682cfaaf7c2d1e2394062858aea9dd3ea5"; + + /** + * The SHA-1 of {@link #TOKENIZER}'s bytes alone, which is not how git names a blob: git hashes + * the length and a NUL byte in front of the content. + */ + private static final String TOKENIZER_PLAIN_SHA1 = "4d02516eda32c9ae5c590766d9e055835e0bb2c7"; + + /** The ONNX graph the hub serves, large enough in reality to be stored in Git LFS. */ + private static final byte[] ONNX = bytes("ONNX GRAPH BYTES\n"); + + /** The SHA-256 of {@link #ONNX}, the 64 character form of the etag, from {@code sha256sum}. */ + private static final String ONNX_SHA256 = + "faffaa0a29c6cf303b7a0dfc59d54131b17b2658c22e02c5da3a66d7526360ef"; + + /** + * A tokenizer larger than the buffer a download is digested in, so that a digest taken from a + * single read instead of a loop over the whole file would not match. + */ + private static final byte[] BIG_TOKENIZER = repeated('x', 20000); + + /** The git blob SHA-1 of {@link #BIG_TOKENIZER}, from {@code git hash-object}. */ + private static final String BIG_TOKENIZER_BLOB_SHA1 = + "7eded2aa2b98c9f0d9d4bb82c277cbbd09dcd044"; + + /** An ONNX graph larger than that buffer, for the SHA-256 form. */ + private static final byte[] BIG_ONNX = repeated('y', 20000); + + /** The SHA-256 of {@link #BIG_ONNX}, from {@code sha256sum}. */ + private static final String BIG_ONNX_SHA256 = + "fdb7f88419c3dd0053ff7c3e9db63fda5bcedf3b8a7344fc1a955a17f4423b58"; + + /** The tokenizer configuration the hub serves, an optional file. */ + private static final byte[] TOKENIZER_CONFIG = bytes("{\"do_lower_case\":true}\n"); + + /** The git blob SHA-1 of {@link #TOKENIZER_CONFIG}, from {@code git hash-object}. */ + private static final String TOKENIZER_CONFIG_BLOB_SHA1 = + "67a56d358bc09865322d344d13922261a6277f26"; + + /** The SentencePiece model the hub serves, an optional file. */ + private static final byte[] SENTENCEPIECE = bytes("SPM\n"); + + /** The git blob SHA-1 of {@link #SENTENCEPIECE}, from {@code git hash-object}. */ + private static final String SENTENCEPIECE_BLOB_SHA1 = + "91a9c1344fe72a78cc937f3cc515050ab1b52f20"; + + /** The first of the SentencePiece file names the cache tries. */ + private static final String SENTENCEPIECE_MODEL = ModelFileNames.SENTENCEPIECE_MODELS.get(0); + + /** The HTTP status of a file a revision does not have. */ + private static final int NOT_FOUND = 404; + + private Hub hub; + + @BeforeEach + void startHub() throws IOException { + hub = new Hub(); + } + + @AfterEach + void stopHub() { + hub.close(); + } + + @Test + void testNullTeacherFailsLoudly() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(null, null)); + assertTrue(e.getMessage().contains("must not be null"), e.getMessage()); + } + + @Test + void testNullHubBaseFailsLoudly(@TempDir Path cacheRoot) { + assertEquals("HubBase must not be null", assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, null, cacheRoot, null)).getMessage()); + } + + @Test + void testNullCacheRootFailsLoudly() { + assertEquals("CacheRoot must not be null", assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), null, null)).getMessage()); + } + + @Test + void testLocalDirectoryIsUsedAsIs(@TempDir Path teacher) throws IOException { + assertEquals(teacher, HuggingFaceModelCache.resolve(teacher.toString(), null)); + } + + @ParameterizedTest + @ValueSource(strings = {"bge-m3", "BAAI/bge m3", "BAAI/bge-m3/onnx", "/BAAI/bge-m3", + "BAAI/bge-m3/", "BAAI//bge-m3", "BAAI/bge-m3@", "BAAI/bge-m3@a b", "BAAI/bge-m3@main@main"}) + void testMalformedTeacherReferenceIsRejectedBeforeAnyRequest(String teacher, + @TempDir Path cacheRoot) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(teacher, hub.base(), cacheRoot, null)); + assertTrue(e.getMessage().contains("org/model"), e.getMessage()); + assertTrue(hub.requests.isEmpty(), hub.requests.toString()); + } + + /** + * A local directory wins over the hub even when its path ends in something shaped like a model + * id, so an {@code org/model} directory on disk is never downloaded over instead. + */ + @Test + void testALocalDirectoryShapedLikeAModelIdIsUsedAsIs(@TempDir Path root) throws IOException { + final Path teacher = Files.createDirectories(root.resolve("BAAI").resolve("bge-m3")); + + assertEquals(teacher, HuggingFaceModelCache.resolve(teacher.toString(), null)); + } + + /** A path that exists but is a regular file is not a teacher directory. */ + @Test + void testAnExistingRegularFileIsRejected(@TempDir Path root) throws IOException { + final Path file = Files.writeString(root.resolve("teacher.txt"), "not a directory"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(file.toString(), null)); + assertTrue(e.getMessage().contains("org/model"), e.getMessage()); + } + + /** + * The two digest forms the hub uses, on the two files a distillation needs: a git blob SHA-1 for + * a file stored in git and a SHA-256 for one stored in Git LFS. + */ + @Test + void testDownloadsAndVerifiesBothEtagForms(@TempDir Path cacheRoot) throws IOException { + serveTeacher(); + final List progress = new ArrayList<>(); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, + progress::add); + + assertEquals(cacheRoot.resolve(CACHE_NAME), cache); + assertArrayEquals(TOKENIZER, Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + assertArrayEquals(ONNX, Files.readAllBytes(cache.resolve(ModelFileNames.ONNX_MODEL))); + assertEquals(COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + assertTrue(progress.stream().anyMatch(line -> line.contains(COMMIT)), progress.toString()); + } + + /** The recorded revision is what a reader of the cache directory finds, in plain text. */ + @Test + void testTheResolvedCommitIsRecordedInTheCacheDirectory(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertEquals(COMMIT, + Files.readString(cache.resolve(HuggingFaceModelCache.REVISION_FILE)).trim()); + } + + /** + * The ref is resolved once and every file is then asked for by commit sha, so that a ref moving + * mid-download cannot mix two revisions into one cache directory. + */ + @Test + void testEveryFileIsRequestedAtTheResolvedCommit(@TempDir Path cacheRoot) throws IOException { + serveTeacher(); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertEquals(1, hub.requests.stream().filter(p -> p.contains("/" + DEFAULT_REF + "/")).count(), + hub.requests.toString()); + assertTrue(hub.requests.stream().filter(p -> !p.contains("/" + DEFAULT_REF + "/")) + .allMatch(p -> p.startsWith("/" + MODEL_ID + "/resolve/" + COMMIT + "/")), + hub.requests.toString()); + } + + @Test + void testACorruptedBodyIsRejected(@TempDir Path cacheRoot) throws IOException { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, bytes("not the graph the hub promised\n"), + quoted(ONNX_SHA256)); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("SHA-256 checksum validation failed"), e.getMessage()); + assertTrue(e.getMessage().contains(ModelFileNames.ONNX_MODEL), e.getMessage()); + assertTrue(e.getMessage().contains(ONNX_SHA256), e.getMessage()); + assertTrue(e.getMessage().contains("but got:"), e.getMessage()); + assertNothingUsable(cacheRoot.resolve(CACHE_NAME), ModelFileNames.ONNX_MODEL); + } + + /** + * The 40 character etag is the git blob SHA-1, not the SHA-1 of the content, and a file that + * only matches the latter is a file whose length git would disagree about. + */ + @Test + void testThePlainSha1OfTheContentIsNotAcceptedAsTheGitBlobSha1(@TempDir Path cacheRoot) { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_PLAIN_SHA1)); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("git blob SHA-1 checksum validation failed"), + e.getMessage()); + assertTrue(e.getMessage().contains(TOKENIZER_BLOB_SHA1), e.getMessage()); + } + + /** + * A download is digested by reading it in a loop, so a file longer than one of those reads is + * digested whole, in both of the forms the hub publishes. The expected values come from + * {@code git hash-object} and {@code sha256sum} over the same bytes. + */ + @Test + void testABodyLongerThanTheDigestBufferIsDigestedWhole(@TempDir Path cacheRoot) + throws IOException { + hub.serve(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, BIG_TOKENIZER, + quoted(BIG_TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, BIG_TOKENIZER, + quoted(BIG_TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, BIG_ONNX, quoted(BIG_ONNX_SHA256)); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(BIG_TOKENIZER, + Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + assertArrayEquals(BIG_ONNX, Files.readAllBytes(cache.resolve(ModelFileNames.ONNX_MODEL))); + } + + /** Hex is hex: a hub that states its digests in upper case is verified against just the same. */ + @Test + void testAnEtagInUpperCaseIsAccepted(@TempDir Path cacheRoot) throws IOException { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, + quoted(TOKENIZER_BLOB_SHA1.toUpperCase(Locale.ROOT))); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, ONNX, + quoted(ONNX_SHA256.toUpperCase(Locale.ROOT))); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(TOKENIZER, Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + assertArrayEquals(ONNX, Files.readAllBytes(cache.resolve(ModelFileNames.ONNX_MODEL))); + assertEquals(COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + } + + @Test + void testAMissingEtagIsRefused(@TempDir Path cacheRoot) throws IOException { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, null); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("Expected checksum could not be retrieved"), e.getMessage()); + assertTrue(e.getMessage().contains(ModelFileNames.TOKENIZER_JSON), e.getMessage()); + assertNothingUsable(cacheRoot.resolve(CACHE_NAME), ModelFileNames.TOKENIZER_JSON); + } + + @ParameterizedTest + @ValueSource(strings = {"", "not-a-digest", "296101682cfaaf7c2d1e2394062858aea9dd3ea", + "296101682cfaaf7c2d1e2394062858aea9dd3ea55", "zzz101682cfaaf7c2d1e2394062858aea9dd3ea5", + "sha256:faffaa0a29c6cf303b7a0dfc59d54131b17b2658c22e02c5da3a66d7526360ef"}) + void testAMalformedEtagIsRefused(String etag, @TempDir Path cacheRoot) { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(etag)); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("Expected checksum could not be retrieved"), e.getMessage()); + assertTrue(e.getMessage().contains("neither a git blob SHA-1 nor a SHA-256"), e.getMessage()); + } + + /** A file the repository does not have is absent, and one it has is downloaded and verified. */ + @Test + void testOptionalFilesAreDownloadedWhenPresentAndAbsentOnA404(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_CONFIG, TOKENIZER_CONFIG, + quoted(TOKENIZER_CONFIG_BLOB_SHA1)); + hub.serve(COMMIT, SENTENCEPIECE_MODEL, SENTENCEPIECE, quoted(SENTENCEPIECE_BLOB_SHA1)); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(TOKENIZER_CONFIG, + Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_CONFIG))); + assertArrayEquals(SENTENCEPIECE, Files.readAllBytes(cache.resolve(SENTENCEPIECE_MODEL))); + // The hub was asked for the external ONNX weights and answered 404, which is not an error. + assertTrue(hub.requests.contains(resolvePath(COMMIT, ModelFileNames.ONNX_MODEL_DATA)), + hub.requests.toString()); + assertTrue(Files.notExists(cache.resolve(ModelFileNames.ONNX_MODEL_DATA))); + } + + @Test + void testAMissingRequiredFileFails(@TempDir Path cacheRoot) { + serveTeacher(); + hub.status(COMMIT, ModelFileNames.ONNX_MODEL, NOT_FOUND); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains(ModelFileNames.ONNX_MODEL), e.getMessage()); + assertTrue(e.getMessage().contains("the distillation needs this file"), e.getMessage()); + } + + /** Only a 404 means absent: a hub that is broken or refuses access must not look like one. */ + @Test + void testAnOptionalFileServedWithAnErrorStatusIsNotTreatedAsAbsent(@TempDir Path cacheRoot) { + serveTeacher(); + hub.status(COMMIT, ModelFileNames.TOKENIZER_CONFIG, 503); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains(ModelFileNames.TOKENIZER_CONFIG), e.getMessage()); + assertTrue(e.getMessage().contains("HTTP 503"), e.getMessage()); + } + + /** + * The hub answers a resolve request with a redirect to a content delivery network and states the + * commit and the digest on the redirecting response, which the client does not carry over to the + * response it finally returns. + */ + @Test + void testTheHeadersOfARedirectingResponseAreUsed(@TempDir Path cacheRoot) throws IOException { + hub.redirect(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, quoted(TOKENIZER_BLOB_SHA1), + "/cdn/tokenizer"); + hub.redirect(COMMIT, ModelFileNames.TOKENIZER_JSON, quoted(TOKENIZER_BLOB_SHA1), + "/cdn/tokenizer"); + hub.redirect(COMMIT, ModelFileNames.ONNX_MODEL, quoted(ONNX_SHA256), "/cdn/onnx"); + hub.reply("/cdn/tokenizer", new Reply(200, null, null, null, TOKENIZER)); + hub.reply("/cdn/onnx", new Reply(200, null, null, null, ONNX)); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(TOKENIZER, Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + assertArrayEquals(ONNX, Files.readAllBytes(cache.resolve(ModelFileNames.ONNX_MODEL))); + assertEquals(COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + } + + /** A complete cache directory is a usable teacher with the hub unreachable. */ + @Test + void testACompleteCacheIsReusedWithoutContactingTheHub(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + final Path first = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + hub.replies.clear(); + hub.requests.clear(); + + final Path second = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertEquals(first, second); + assertTrue(hub.requests.isEmpty(), hub.requests.toString()); + assertArrayEquals(TOKENIZER, Files.readAllBytes(second.resolve(ModelFileNames.TOKENIZER_JSON))); + } + + /** A recorded revision without the files it vouches for is not a cache directory. */ + @Test + void testAMarkedCacheMissingItsFilesIsDownloadedAgain(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + final Path cache = Files.createDirectories(cacheRoot.resolve(CACHE_NAME)); + Files.writeString(cache.resolve(HuggingFaceModelCache.REVISION_FILE), COMMIT); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(ONNX, Files.readAllBytes(cache.resolve(ModelFileNames.ONNX_MODEL))); + } + + /** + * A run that stops on a failed verification must leave nothing the next run would trust. The + * record of the revision the directory used to hold is dropped before the first file is fetched, + * so a retry checks what is on disk against the hub instead of handing out a directory half + * replaced by a revision it never finished downloading. + */ + @Test + void testAFailedVerificationLeavesNoTrustedCacheBehind(@TempDir Path cacheRoot) + throws IOException { + final Path cache = Files.createDirectories(cacheRoot.resolve(CACHE_NAME)); + Files.createDirectories(cache.resolve(ModelFileNames.ONNX_MODEL).getParent()); + // A directory marked complete whose tokenizer is gone: its graph is the earlier revision's. + Files.write(cache.resolve(ModelFileNames.ONNX_MODEL), bytes("an older revision\n")); + Files.writeString(cache.resolve(HuggingFaceModelCache.REVISION_FILE), COMMIT + "\n"); + serveTeacher(); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, bytes("not the graph the hub promised\n"), + quoted(ONNX_SHA256)); + + assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertNull(HuggingFaceModelCache.pinnedRevision(cache)); + hub.replies.clear(); + assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null), + "the half-replaced directory must not be handed out"); + } + + /** Only a commit sha names a teacher, so a stray file cannot make a directory look pinned. */ + @Test + void testAnUnusableRevisionFileIsNotAPin(@TempDir Path cache) throws IOException { + assertNull(HuggingFaceModelCache.pinnedRevision(cache)); + + Files.writeString(cache.resolve(HuggingFaceModelCache.REVISION_FILE), "not a commit sha"); + assertNull(HuggingFaceModelCache.pinnedRevision(cache)); + + Files.writeString(cache.resolve(HuggingFaceModelCache.REVISION_FILE), COMMIT + "\n"); + assertEquals(COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + } + + /** + * A cache directory left incomplete by an interrupted run does not carry the revision it was + * downloaded at, so its files are checked against the revision now being downloaded instead of + * being trusted, and the expensive ones are not fetched again when they already match. + */ + @Test + void testAnUnmarkedCachedFileThatMatchesTheRevisionIsKept(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + final Path cache = Files.createDirectories(cacheRoot.resolve(CACHE_NAME)); + Files.write(cache.resolve(ModelFileNames.TOKENIZER_JSON), TOKENIZER); + // A body that would fail verification: reaching it means the cached file was not reused. + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, bytes("re-downloaded\n"), + quoted(TOKENIZER_BLOB_SHA1)); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(TOKENIZER, Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + } + + @Test + void testAnUnmarkedCachedFileFromAnotherRevisionIsReplaced(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + final Path cache = Files.createDirectories(cacheRoot.resolve(CACHE_NAME)); + Files.write(cache.resolve(ModelFileNames.TOKENIZER_JSON), bytes("an older revision\n")); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(TOKENIZER, Files.readAllBytes(cache.resolve(ModelFileNames.TOKENIZER_JSON))); + } + + /** A cache directory holds one revision, so a file the new one does not have has to go. */ + @Test + void testAnOptionalFileTheRevisionDoesNotHaveIsRemovedFromTheCache(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + final Path cache = Files.createDirectories(cacheRoot.resolve(CACHE_NAME)); + Files.write(cache.resolve(SENTENCEPIECE_MODEL), SENTENCEPIECE); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertTrue(Files.notExists(cache.resolve(SENTENCEPIECE_MODEL))); + } + + /** An explicit revision is downloaded, and pinned into a cache directory of its own. */ + @Test + void testAnExplicitRevisionIsRequestedAndCachedApart(@TempDir Path cacheRoot) throws IOException { + hub.serve(OTHER_COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1), + OTHER_COMMIT); + hub.serve(OTHER_COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256), OTHER_COMMIT); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID + "@" + OTHER_COMMIT, hub.base(), + cacheRoot, null); + + assertEquals(cacheRoot.resolve(cacheNameAt(OTHER_COMMIT)), cache); + assertEquals(OTHER_COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + assertTrue(hub.requests.stream() + .allMatch(p -> p.startsWith("/" + MODEL_ID + "/resolve/" + OTHER_COMMIT + "/")), + hub.requests.toString()); + } + + /** A named branch or tag is a revision too, and resolves to the commit the hub reports. */ + @Test + void testAnExplicitBranchIsResolvedToItsCommit(@TempDir Path cacheRoot) throws IOException { + hub.serve("refs-pr-1", ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256)); + + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID + "@refs-pr-1", hub.base(), + cacheRoot, null); + + assertEquals(COMMIT, HuggingFaceModelCache.pinnedRevision(cache)); + } + + @Test + void testARequestedCommitTheHubResolvesElsewhereIsRefused(@TempDir Path cacheRoot) { + hub.serve(OTHER_COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + + final IOException e = assertThrows(IOException.class, () -> HuggingFaceModelCache.resolve( + MODEL_ID + "@" + OTHER_COMMIT, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("resolved to commit " + COMMIT), e.getMessage()); + } + + /** A directory recording one commit is not the answer to a reference naming another. */ + @Test + void testACacheRecordingAnotherCommitThanTheOneAskedForIsNotReused(@TempDir Path cacheRoot) + throws IOException { + final Path cache = Files.createDirectories(cacheRoot.resolve(cacheNameAt(OTHER_COMMIT))); + Files.createDirectories(cache.resolve(ModelFileNames.ONNX_MODEL).getParent()); + Files.write(cache.resolve(ModelFileNames.TOKENIZER_JSON), TOKENIZER); + Files.write(cache.resolve(ModelFileNames.ONNX_MODEL), ONNX); + Files.writeString(cache.resolve(HuggingFaceModelCache.REVISION_FILE), COMMIT + "\n"); + hub.serve(OTHER_COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1), + OTHER_COMMIT); + hub.serve(OTHER_COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256), OTHER_COMMIT); + + final Path resolved = HuggingFaceModelCache.resolve(MODEL_ID + "@" + OTHER_COMMIT, hub.base(), + cacheRoot, null); + + assertEquals(cache, resolved); + assertEquals(OTHER_COMMIT, HuggingFaceModelCache.pinnedRevision(resolved)); + assertFalse(hub.requests.isEmpty(), "the hub must be asked, not the stale record believed"); + } + + @Test + void testARevisionThatCannotBePinnedIsRefused(@TempDir Path cacheRoot) { + hub.reply(resolvePath(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON), + new Reply(200, null, quoted(TOKENIZER_BLOB_SHA1), null, TOKENIZER)); + + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("could not be pinned"), e.getMessage()); + assertTrue(Files.notExists(cacheRoot.resolve(CACHE_NAME))); + } + + @Test + void testAModelTheHubDoesNotHaveIsRefused(@TempDir Path cacheRoot) { + final IOException e = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(e.getMessage().contains("Failed to resolve revision 'main'"), e.getMessage()); + } + + /** + * Asserts that a failed download left nothing a distillation could pick up: neither the file it + * was verifying nor the temporary file it streamed into. + * + * @param cache The cache directory; need not exist. + * @param file The repository-relative name of the file that failed. + * @throws IOException Thrown if the directory cannot be walked. + */ + private void assertNothingUsable(Path cache, String file) throws IOException { + assertTrue(Files.notExists(cache.resolve(file)), file + " must not be published"); + if (Files.isDirectory(cache)) { + try (Stream entries = Files.walk(cache)) { + assertFalse(entries.anyMatch(p -> p.getFileName().toString().contains(".download")), + "a partial download must not be left behind"); + } + } + } + + /** Serves the ref and the two files a distillation needs, all at {@link #COMMIT}. */ + private void serveTeacher() { + hub.serve(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256)); + } + + /** + * {@return the request path of a file at a revision} + * + * @param revision The revision. + * @param file The repository-relative file name. + */ + private static String resolvePath(String revision, String file) { + return "/" + MODEL_ID + "/resolve/" + revision + "/" + file; + } + + /** + * {@return a header value in the quotes the hub puts around it} + * + * @param value The value. + */ + private static String quoted(String value) { + return "\"" + value + "\""; + } + + /** + * {@return the UTF-8 bytes of a fixture} + * + * @param content The content. + */ + private static byte[] bytes(String content) { + return content.getBytes(StandardCharsets.UTF_8); + } + + /** + * {@return a fixture of one character repeated, long enough to outrun a single read} + * + * @param content The character to repeat; must be an ASCII one, so that the fixture is as many + * bytes long as it is characters. + * @param length The number of characters. + */ + private static byte[] repeated(char content, int length) { + return bytes(String.valueOf(content).repeat(length)); + } + + /** + * One canned response. + * + * @param status The HTTP status. + * @param commit The {@code x-repo-commit} header value, or {@code null} to send none. + * @param etag The {@code x-linked-etag} header value, or {@code null} to send none. + * @param location The {@code Location} header value, or {@code null} to send none. + * @param body The response body, or {@code null} to send none. + */ + private record Reply(int status, String commit, String etag, String location, byte[] body) { + } + + /** + * A stand-in for the hub on the loopback interface, answering canned responses per request path + * and recording the paths it was asked for. + */ + private static final class Hub implements AutoCloseable { + + private final HttpServer server; + private final Map replies = new ConcurrentHashMap<>(); + private final List requests = Collections.synchronizedList(new ArrayList<>()); + + private Hub() throws IOException { + server = HttpServer.create(new InetSocketAddress(LOOPBACK, 0), 0); + server.createContext("/", this::answer); + server.start(); + } + + /** {@return the base URL of this hub, ending in a slash} */ + private String base() { + return "http://" + LOOPBACK + ":" + server.getAddress().getPort() + "/"; + } + + /** + * Serves a file at a revision, reporting {@link #COMMIT} as the commit the request resolved to. + * + * @param revision The revision to serve it at. + * @param file The repository-relative file name. + * @param body The response body. + * @param etag The {@code x-linked-etag} header value, or {@code null} to send none. + */ + private void serve(String revision, String file, byte[] body, String etag) { + serve(revision, file, body, etag, COMMIT); + } + + /** + * Serves a file at a revision. + * + * @param revision The revision to serve it at. + * @param file The repository-relative file name. + * @param body The response body. + * @param etag The {@code x-linked-etag} header value, or {@code null} to send none. + * @param commit The commit the request resolves to. + */ + private void serve(String revision, String file, byte[] body, String etag, String commit) { + reply(resolvePath(revision, file), new Reply(200, commit, etag, null, body)); + } + + /** + * Answers a file with a redirect carrying the headers, as the hub does for a file its content + * delivery network serves. + * + * @param revision The revision to serve it at. + * @param file The repository-relative file name. + * @param etag The {@code x-linked-etag} header value. + * @param target The path the redirect points at. + */ + private void redirect(String revision, String file, String etag, String target) { + reply(resolvePath(revision, file), new Reply(302, COMMIT, etag, target, null)); + } + + /** + * Answers a file with a status and nothing else. + * + * @param revision The revision to serve it at. + * @param file The repository-relative file name. + * @param status The HTTP status. + */ + private void status(String revision, String file, int status) { + reply(resolvePath(revision, file), new Reply(status, COMMIT, null, null, null)); + } + + /** + * Registers one canned response, replacing any response registered for the same path. + * + * @param path The request path. + * @param reply The response. + */ + private void reply(String path, Reply reply) { + replies.put(path, reply); + } + + /** + * Answers one request, with 404 when nothing is registered for its path. + * + * @param exchange The exchange. + * @throws IOException Thrown if the response headers cannot be sent. + */ + private void answer(HttpExchange exchange) throws IOException { + final String path = exchange.getRequestURI().getPath(); + requests.add(path); + final Reply reply = replies.get(path); + if (reply == null) { + exchange.sendResponseHeaders(NOT_FOUND, -1); + exchange.close(); + return; + } + if (reply.commit() != null) { + exchange.getResponseHeaders().add("x-repo-commit", reply.commit()); + } + if (reply.etag() != null) { + exchange.getResponseHeaders().add("x-linked-etag", reply.etag()); + } + if (reply.location() != null) { + exchange.getResponseHeaders().add("Location", reply.location()); + } + if (reply.body() == null) { + exchange.sendResponseHeaders(reply.status(), -1); + } else { + exchange.sendResponseHeaders(reply.status(), reply.body().length); + try (OutputStream out = exchange.getResponseBody()) { + out.write(reply.body()); + } catch (IOException e) { + // The client closes a body it does not need, which fails this write; that is the point + // of the header-only requests, so it is not a test failure. + } + } + exchange.close(); + } + + @Override + public void close() { + server.stop(0); + } + } + /** + * Verifies that references differing only in a character the readable part of the directory + * name flattens do not share a cache directory. They did: the name replaced '.', '@' and '/' + * without disambiguating them, so three distinct teachers mapped to one directory, and the + * cached fast path would answer from it without contacting the hub. + */ + @Test + void testDistinctTeachersDoNotShareACacheDirectory() throws Exception { + final java.util.Set names = new java.util.HashSet<>(); + for (final String teacher : java.util.List.of( + "acme/model_v1", "acme/model.v1", "acme/model@v1", "acme/model-v1")) { + names.add(HuggingFaceModelCache.cacheDirectoryName(teacher)); + } + assertEquals(4, names.size(), "each distinct teacher reference needs its own directory"); + } + +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/Model2VecUnigramTokenizerTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/Model2VecUnigramTokenizerTest.java new file mode 100644 index 0000000000..5fa5bfc470 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/Model2VecUnigramTokenizerTest.java @@ -0,0 +1,47 @@ +/* + * 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 org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class Model2VecUnigramTokenizerTest { + + @Test + void testReportsMissingVocabularyAsInvalidModelContent(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, + "{\"normalizer\":{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}," + + "\"pre_tokenizer\":{\"type\":\"Metaspace\",\"replacement\":\"▁\"," + + "\"prepend_scheme\":\"always\",\"split\":false}," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0}} "); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer)); + + assertTrue(error.getMessage().contains("model.vocab"), error.getMessage()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelAssemblerTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelAssemblerTest.java new file mode 100644 index 0000000000..be809557f9 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelAssemblerTest.java @@ -0,0 +1,205 @@ +/* + * 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.InputStream; +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 opennlp.embeddings.cmdline.AssembleModelTool; +import opennlp.subword.sentencepiece.SentencePieceTokenizer; +import opennlp.tools.cmdline.TerminateToolException; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The assembler completes a distilled directory into a loadable one: it derives the WordPiece + * {@code vocab.txt} and {@code tokenizer_config.json} from {@code tokenizer.json}, leaves existing + * files alone, and assembles a Model2Vec Unigram tokenizer directly from {@code tokenizer.json}. + * The CLI tool wraps it and turns failures into a {@link TerminateToolException}. + */ +class ModelAssemblerTest { + + // A WordPiece tokenizer.json with a five-entry vocab dictionary (no [CLS]/[SEP], as Model2Vec + // ships) and a BERT normalizer that lower-cases. The dictionary is written out of id order to + // prove the assembler sorts it. + private static final String WORDPIECE_TOKENIZER_JSON = + "{\"version\":\"1.0\"," + + "\"normalizer\":{\"type\":\"BertNormalizer\",\"strip_accents\":null," + + "\"lowercase\":true}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[PAD]\":0,\"hello\":2,\"[UNK]\":1,\"cat\":4,\"world\":3}}}"; + + private static final float[][] ROWS = { + {0f, 0f, 0f}, // [PAD] + {1f, 10f, 100f}, // [UNK] + {2f, 20f, 200f}, // hello + {3f, 30f, 300f}, // world + {4f, 40f, 400f}, // cat + }; + + private static Path writeWordpieceDistillation(Path dir) throws IOException { + Files.writeString(dir.resolve("tokenizer.json"), WORDPIECE_TOKENIZER_JSON); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + return dir; + } + + @Test + void testDerivesTheWordpieceVocabularyAndConfigInIdOrder(@TempDir Path dir) throws IOException { + writeWordpieceDistillation(dir); + + final ModelAssembler.Result result = ModelAssembler.assemble(dir); + + assertEquals("WordPiece", result.family()); + assertEquals(3, result.dimension()); + assertEquals(5, result.vocabularySize()); + assertTrue(result.wroteVocabulary()); + assertTrue(result.wroteTokenizerConfig()); + // The vocab.txt must be the dictionary in id order, not the order it was written. + assertEquals(List.of("[PAD]", "[UNK]", "hello", "world", "cat"), + Files.readAllLines(dir.resolve("vocab.txt"))); + // The casing comes from the BERT normalizer's lowercase flag. + assertTrue(Files.readString(dir.resolve("tokenizer_config.json")).contains("\"do_lower_case\": true")); + } + + @Test + void testAssembledDirectoryEmbeds(@TempDir Path dir) throws IOException { + writeWordpieceDistillation(dir); + ModelAssembler.assemble(dir); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + // (hello[row 2] + world[row 3]) / 2 = (2 + 3) / 2 in the first component; the model has no + // frame tokens, so only the two content pieces pool. + assertEquals(2.5f, model.embed("hello world")[0], 1e-5f); + } + + @Test + void testLeavesExistingFilesUntouched(@TempDir Path dir) throws IOException { + writeWordpieceDistillation(dir); + // A vocab.txt the caller already wrote must not be overwritten. + Files.write(dir.resolve("vocab.txt"), List.of("[PAD]", "[UNK]", "hello", "world", "cat")); + Files.writeString(dir.resolve("tokenizer_config.json"), "{\"do_lower_case\": false}"); + + final ModelAssembler.Result result = ModelAssembler.assemble(dir); + + assertFalse(result.wroteVocabulary()); + assertFalse(result.wroteTokenizerConfig()); + assertTrue(Files.readString(dir.resolve("tokenizer_config.json")).contains("false")); + } + + @Test + void testRejectsAMissingDistillationFile(@TempDir Path dir) throws IOException { + Files.writeString(dir.resolve("tokenizer.json"), WORDPIECE_TOKENIZER_JSON); + // no model.safetensors, no config.json + final IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> ModelAssembler.assemble(dir)); + assertTrue(e.getMessage().contains("model.safetensors"), e.getMessage()); + } + + @Test + void testLoadsAModel2VecUnigramTokenizerWithoutASeparateModelFile(@TempDir Path dir) + throws IOException { + Files.writeString(dir.resolve("tokenizer.json"), + "{\"normalizer\":{\"type\":\"Sequence\",\"normalizers\":[" + + "{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}," + + "{\"type\":\"Replace\",\"pattern\":{\"String\":\".\"}," + + "\"content\":\" . \"}," + + "{\"type\":\"Replace\",\"pattern\":{\"Regex\":\"\\\\s+\"}," + + "\"content\":\" \"}," + + "{\"type\":\"Strip\",\"strip_left\":true,\"strip_right\":true}]}," + + "\"pre_tokenizer\":{\"type\":\"Metaspace\",\"replacement\":\"▁\"," + + "\"prepend_scheme\":\"always\",\"split\":false}," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":1," + + "\"byte_fallback\":false,\"vocab\":[" + + "[\"[PAD]\",-10.0],[\"[UNK]\",-10.0],[\"▁hello\",-1.0]," + + "[\"▁world\",-1.0],[\"▁\",-2.0],[\".\",-1.0]]}}"); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {2f}, {4f}, {8f}, {16f} + })); + + final ModelAssembler.Result result = ModelAssembler.assemble(dir); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + assertEquals("Unigram", result.family()); + assertEquals(6, result.vocabularySize()); + assertEquals(7.5f, model.embed("hello world.")[0], 1e-6f); + } + + @Test + void testLoadsTheRealSentencePieceModelAfterItsFileIsPresent(@TempDir Path dir) + throws IOException { + // Assemble a SentencePiece directory around the bundled tiny model: once its .model file is + // present the assembler only has to verify it loads. + final byte[] modelBytes; + try (InputStream in = getClass().getResourceAsStream("/opennlp/embeddings/tiny-unigram.model")) { + modelBytes = in.readAllBytes(); + } + Files.write(dir.resolve("sentencepiece.bpe.model"), modelBytes); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + // A tokenizer.json whose vocab is the model's own poolable pieces, so the coverage check + // passes; the matrix carries one row per piece. + final SentencePieceTokenizer tokenizer = + SentencePieceTokenizer.load(dir.resolve("sentencepiece.bpe.model")); + final StringBuilder vocab = new StringBuilder("{\"model\":{\"type\":\"Unigram\",\"vocab\":["); + int rows = 0; + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (rows > 0) { + vocab.append(','); + } + vocab.append('[').append(EmbeddingTestFixtures.jsonString(tokenizer.idToPiece(id))) + .append(",-1.0]"); + rows++; + } + vocab.append("]}}"); + Files.writeString(dir.resolve("tokenizer.json"), vocab.toString()); + final float[][] matrix = new float[rows][2]; + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + + final ModelAssembler.Result result = ModelAssembler.assemble(dir); + assertEquals("SentencePiece", result.family()); + assertEquals(rows, result.vocabularySize()); + assertFalse(result.wroteVocabulary()); + } + + @Test + void testToolPrintsASummaryAndFailsLoudlyOnABadDirectory(@TempDir Path dir) throws IOException { + writeWordpieceDistillation(dir); + // The tool runs the assembly without throwing on a good directory. + new AssembleModelTool().run(new String[] {"-modelDir", dir.toString()}); + + // A directory that is not a model fails as a TerminateToolException, not a raw exception. + final Path empty = Files.createDirectory(dir.resolve("empty")); + final TerminateToolException e = assertThrows(TerminateToolException.class, + () -> new AssembleModelTool().run(new String[] {"-modelDir", empty.toString()})); + assertTrue(e.getMessage().contains("tokenizer.json") || e.getMessage().contains("distilled"), + e.getMessage()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerTest.java new file mode 100644 index 0000000000..0d4a82b233 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerTest.java @@ -0,0 +1,197 @@ +/* + * 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.Collections; +import java.util.List; + +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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The distiller's argument checking and its Zipf weighting. The forward pass itself needs a real + * ONNX teacher and is exercised by the distillation script, not here, so these tests pin the + * checks that must reject a bad call before any teacher is downloaded or run. + */ +class ModelDistillerTest { + + /** Model2Vec's SIF coefficient, the value the distiller uses. */ + private static final double SIF = 1e-4; + + @Test + void testZipfWeightsFollowTheModel2vecFormula() { + // Two rows: the Zipf distribution is over 1/2 and 1/3, normalized by their sum 5/6. + final float[] weights = ModelDistiller.zipfWeights(2, SIF); + + assertEquals(2, weights.length); + assertEquals(SIF / (SIF + 0.6), weights[0], 1e-10); + assertEquals(SIF / (SIF + 0.4), weights[1], 1e-10); + } + + @Test + void testZipfWeightsOfASingleRowUseTheWholeDistribution() { + // One row takes all the probability mass, so p is 1 regardless of the harmonic sum. + final float[] weights = ModelDistiller.zipfWeights(1, SIF); + + assertEquals(1, weights.length); + assertEquals(SIF / (SIF + 1.0), weights[0], 1e-10); + } + + @ParameterizedTest + @ValueSource(ints = {2, 3, 100, 1000}) + void testZipfWeightsDiscountEarlyRows(int rows) { + final float[] weights = ModelDistiller.zipfWeights(rows, SIF); + + assertEquals(rows, weights.length); + // Frequent (early) tokens are down-weighted relative to rare (late) ones, and every weight is + // a proper fraction: sif / (sif + p) with p in (0, 1]. + for (int i = 0; i < weights.length; i++) { + assertTrue(weights[i] > 0 && weights[i] < 1, "row " + i + " has weight " + weights[i]); + if (i > 0) { + assertTrue(weights[i] > weights[i - 1], + "row " + i + " (" + weights[i] + ") must outweigh row " + (i - 1) + " (" + + weights[i - 1] + ")"); + } + } + } + + @Test + void testZipfWeightsMatchTheHarmonicNormalizationOfTheLastRow() { + final int rows = 1000; + final float[] weights = ModelDistiller.zipfWeights(rows, SIF); + + double harmonicSum = 0; + for (int j = 2; j <= rows + 1; j++) { + harmonicSum += 1.0 / j; + } + assertEquals(SIF / (SIF + 1.0 / (rows + 1) / harmonicSum), weights[rows - 1], 1e-5); + } + + @Test + void testRejectsANullTeacherDirectory(@TempDir Path dir) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill((Path) null, dir, 256, null)); + assertEquals("TeacherDirectory must not be null", e.getMessage()); + } + + @Test + void testRejectsATeacherDirectoryThatIsNotADirectory(@TempDir Path dir) throws IOException { + final Path file = Files.writeString(dir.resolve("teacher"), "not a directory"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(file, dir.resolve("out"), 256, null)); + assertTrue(e.getMessage().contains("is not a directory"), e.getMessage()); + } + + @Test + void testRejectsANullOutputDirectory(@TempDir Path dir) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(dir, null, 256, null)); + assertEquals("OutputDirectory must not be null", e.getMessage()); + } + + @ParameterizedTest + @ValueSource(ints = {0, -1, Integer.MIN_VALUE}) + void testRejectsANonPositivePcaDimension(int pcaDims, @TempDir Path dir) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(dir, dir.resolve("out"), pcaDims, null)); + assertEquals("PcaDims must be at least 1, got " + pcaDims, e.getMessage()); + } + + @Test + void testRejectsATeacherDirectoryWithoutAnOnnxGraph(@TempDir Path dir) throws IOException { + final Path teacher = Files.createDirectory(dir.resolve("teacher")); + Files.writeString(teacher.resolve(ModelFileNames.TOKENIZER_JSON), "{}"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, dir.resolve("out"), 256, null)); + assertTrue(e.getMessage().contains(ModelFileNames.ONNX_MODEL), e.getMessage()); + // Nothing may be written before the teacher is known to be usable. + assertTrue(Files.notExists(dir.resolve("out")), "the output directory must not be created"); + } + + /** + * A teacher whose whole vocabulary is dropped by the cleaning has no rows to encode. That has to + * be caught before the ONNX graph is loaded, where it would otherwise surface as a zero-length + * matrix. The graph file here is a placeholder that would fail to load if it were reached. + */ + @Test + void testRejectsATeacherWhoseVocabularyIsFullyCleanedAway(@TempDir Path dir) throws IOException { + final Path teacher = Files.createDirectory(dir.resolve("teacher")); + Files.createDirectory(teacher.resolve("onnx")); + Files.writeString(teacher.resolve(ModelFileNames.ONNX_MODEL), "not a real graph"); + Files.writeString(teacher.resolve(ModelFileNames.TOKENIZER_JSON), + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[unused0]\"," + + "\"vocab\":{\"[unused0]\":0}}}"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, dir.resolve("out"), 256, null)); + + assertTrue(e.getMessage().contains("nothing to distill"), e.getMessage()); + assertTrue(Files.notExists(dir.resolve("out")), "the output directory must not be created"); + } + + /** + * A bad output argument must be rejected before the teacher reference is resolved, so that a + * mistyped command against a hub id does not download gigabytes first. The teacher here is a + * well-formed hub id that would otherwise be fetched. + */ + @ParameterizedTest + @ValueSource(strings = {"BAAI/bge-m3", "sentence-transformers/all-MiniLM-L6-v2"}) + void testRejectsABadOutputBeforeResolvingAHubTeacher(String teacher, @TempDir Path dir) { + assertEquals("OutputDirectory must not be null", + assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, null, 256, null)).getMessage()); + assertEquals("PcaDims must be at least 1, got 0", + assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, dir.resolve("out"), 0, null)).getMessage()); + } + + /** + * Term arguments are validated before the teacher reference is resolved, so a bad term list + * against a hub id fails before anything is downloaded. + */ + @Test + void testRejectsBadTermsBeforeResolvingAHubTeacher(@TempDir Path dir) { + assertEquals("Terms must not be null", + assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill("BAAI/bge-m3", dir.resolve("out"), 256, null, null)) + .getMessage()); + assertEquals("Terms must not contain null", + assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill("BAAI/bge-m3", dir.resolve("out"), 256, + Collections.singletonList(null), null)).getMessage()); + } + + @ParameterizedTest + @ValueSource(strings = {"&", "!!", " . "}) + void testRejectsATermWithoutALetterOrDigit(String term, @TempDir Path dir) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill("BAAI/bge-m3", dir.resolve("out"), 256, List.of(term), + null)); + assertTrue(e.getMessage().contains("no letter or digit"), e.getMessage()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelFileNamesTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelFileNamesTest.java new file mode 100644 index 0000000000..5250af9c90 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelFileNamesTest.java @@ -0,0 +1,86 @@ +/* + * 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.assertNull; + +/** + * The file lookup the loader and the assembler share to find a SentencePiece model under whichever + * of its several names a teacher shipped it as: the first name that is a regular file wins, in the + * order given. + */ +class ModelFileNamesTest { + + @Test + void testReturnsTheFirstNameThatExists(@TempDir Path dir) throws IOException { + Files.writeString(dir.resolve("spiece.model"), "second"); + Files.writeString(dir.resolve("tokenizer.model"), "third"); + + assertEquals(dir.resolve("spiece.model"), + ModelFileNames.firstRegularFile(dir, ModelFileNames.SENTENCEPIECE_MODELS)); + } + + @Test + void testPrefersTheEarlierNameWhenSeveralExist(@TempDir Path dir) throws IOException { + for (final String name : ModelFileNames.SENTENCEPIECE_MODELS) { + Files.writeString(dir.resolve(name), name); + } + + assertEquals(dir.resolve(ModelFileNames.SENTENCEPIECE_MODELS.get(0)), + ModelFileNames.firstRegularFile(dir, ModelFileNames.SENTENCEPIECE_MODELS)); + } + + /** + * A directory carrying one of the names is not the model file. Accepting it would hand the + * loader a path it cannot read, one step further from the cause. + */ + @Test + void testSkipsADirectoryWithAMatchingName(@TempDir Path dir) throws IOException { + Files.createDirectory(dir.resolve("sentencepiece.bpe.model")); + Files.writeString(dir.resolve("spiece.model"), "the real one"); + + assertEquals(dir.resolve("spiece.model"), + ModelFileNames.firstRegularFile(dir, ModelFileNames.SENTENCEPIECE_MODELS)); + } + + @Test + void testReturnsNullWhenNoNameExists(@TempDir Path dir) { + assertNull(ModelFileNames.firstRegularFile(dir, ModelFileNames.SENTENCEPIECE_MODELS)); + } + + @Test + void testReturnsNullForAnEmptyNameList(@TempDir Path dir) throws IOException { + Files.writeString(dir.resolve("spiece.model"), "not asked for"); + + assertNull(ModelFileNames.firstRegularFile(dir, List.of())); + } + + @Test + void testReturnsNullForADirectoryThatDoesNotExist(@TempDir Path dir) { + assertNull(ModelFileNames.firstRegularFile(dir.resolve("missing"), + ModelFileNames.SENTENCEPIECE_MODELS)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/OnnxTeacherEncoderTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/OnnxTeacherEncoderTest.java new file mode 100644 index 0000000000..88f71047bb --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/OnnxTeacherEncoderTest.java @@ -0,0 +1,53 @@ +/* + * 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 org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The encoder's argument contract, checked before the ONNX runtime is touched. + */ +class OnnxTeacherEncoderTest { + + @Test + void testNullFileFailsLoudly() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(null)); + assertTrue(e.getMessage().contains("must not be null"), e.getMessage()); + } + + @Test + void testMissingFileFailsLoudly(@TempDir Path directory) { + final Path missing = directory.resolve("model.onnx"); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(missing)); + assertTrue(e.getMessage().contains(missing.toString()), e.getMessage()); + } + + @Test + void testDirectoryIsNotARegularFile(@TempDir Path directory) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(directory)); + assertTrue(e.getMessage().contains("regular file"), e.getMessage()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixEdgeCaseTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixEdgeCaseTest.java new file mode 100644 index 0000000000..95c42ed339 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixEdgeCaseTest.java @@ -0,0 +1,133 @@ +/* + * 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.Random; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Edge cases for the quantized matrix: dimensions that are not powers of two (so the padding + * and truncation path runs), a one-dimensional matrix, adversarial rows the random rotation + * must still reconstruct (constant, one-hot, sign-alternating), and the dot/norm consistency + * over padded dimensions. + */ +class QuantizedEmbeddingMatrixEdgeCaseTest { + + private static final long SEED = 3L; + + @ParameterizedTest + @ValueSource(ints = {1, 2, 3, 5, 17, 100, 300, 513}) + void testDotRotatedMatchesOriginalDotAtEveryDimension(int dimension) { + final Random random = new Random(dimension); + final int rows = 8; + final float[] matrix = new float[rows * dimension]; + for (int i = 0; i < matrix.length; i++) { + matrix[i] = (float) random.nextGaussian(); + } + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, rows, dimension, 4, SEED); + final float[] query = new float[dimension]; + for (int d = 0; d < dimension; d++) { + query[d] = (float) random.nextGaussian(); + } + final float[] rotatedQuery = quantized.rotate(query); + for (int row = 0; row < rows; row++) { + final float[] decoded = quantized.decodeRow(row); + double originalDot = 0; + for (int d = 0; d < dimension; d++) { + originalDot += (double) decoded[d] * query[d]; + } + // The rotation is orthonormal, so scoring in rotated space over the padded coordinates + // must equal the original-space dot with the truncated decoded row. + assertEquals(originalDot, quantized.dotRotated(row, rotatedQuery), + 1e-3 * (1 + Math.abs(originalDot)), + "dot mismatch at dimension " + dimension + ", row " + row); + } + } + + @Test + void testConstantRowReconstructsDespiteBeingSpikyAfterRotation() { + // A constant vector is the worst case for the transform: its rotation concentrates all + // energy in one coordinate, which the grid clamps. The per-row least-squares scale must + // absorb that clamp, so the reconstruction still points the same way. + final int dimension = 300; + final float[] matrix = new float[dimension]; + java.util.Arrays.fill(matrix, 0.7f); + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, 1, dimension, 4, SEED); + assertTrue(cosine(matrix, 0, dimension, quantized.decodeRow(0)) > 0.98, + "a constant row must still reconstruct in direction"); + } + + @Test + void testOneHotAndAlternatingRowsReconstruct() { + final int dimension = 128; + final float[] oneHot = new float[dimension]; + oneHot[7] = 3.5f; + final float[] alternating = new float[dimension]; + for (int d = 0; d < dimension; d++) { + alternating[d] = (d % 2 == 0 ? 1f : -1f); + } + for (final float[] row : new float[][] {oneHot, alternating}) { + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(row, 1, dimension, 4, SEED); + assertTrue(cosine(row, 0, dimension, quantized.decodeRow(0)) > 0.95, + "an adversarial row must reconstruct in direction"); + } + } + + @Test + void testRowNormMatchesDecodedRowAtNonPowerOfTwoDimension() { + final int dimension = 17; + final Random random = new Random(17); + final int rows = 5; + final float[] matrix = new float[rows * dimension]; + for (int i = 0; i < matrix.length; i++) { + matrix[i] = 2f * (float) random.nextGaussian(); + } + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, rows, dimension, 3, SEED); + for (int row = 0; row < rows; row++) { + final float[] decoded = quantized.decodeRow(row); + double sumOfSquares = 0; + for (final float value : decoded) { + sumOfSquares += (double) value * value; + } + assertEquals(Math.sqrt(sumOfSquares), quantized.rowNorm(row), + 1e-4 * (1 + Math.sqrt(sumOfSquares)), + "rowNorm must equal the decoded row's norm at a padded dimension"); + } + } + + /** + * {@return the cosine between a matrix row and a decoded vector} + * + * @param matrix The flat row-major matrix. + * @param base The row's first index. + * @param dimension The row width. + * @param decoded The decoded row. + */ + private static double cosine(float[] matrix, int base, int dimension, float[] decoded) { + return ModelQuantizer.cosine(matrix, base, dimension, decoded); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixTest.java new file mode 100644 index 0000000000..3cbab13f46 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixTest.java @@ -0,0 +1,333 @@ +/* + * 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.Arrays; +import java.util.Random; + +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.CsvSource; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The quantized matrix contract: reconstruction quality per bit width, rotated-space math that + * agrees with original-space math, linear pooling, deterministic bytes, a self-describing file + * that round-trips exactly, and loud failure on malformed input. + */ +class QuantizedEmbeddingMatrixTest { + + private static final int ROWS = 50; + private static final int DIMENSION = 300; + private static final long SEED = 12345L; + + /** + * {@return a deterministic random test matrix with varied row norms} + */ + private static float[] testMatrix() { + final Random random = new Random(42); + final float[] matrix = new float[ROWS * DIMENSION]; + for (int row = 0; row < ROWS; row++) { + // Vary the norms so per-row scaling is actually exercised. + final float rowScale = 0.1f + 3f * random.nextFloat(); + for (int d = 0; d < DIMENSION; d++) { + matrix[row * DIMENSION + d] = rowScale * (float) random.nextGaussian(); + } + } + return matrix; + } + + // The mean squared error of the Gaussian Lloyd-Max grids translates to an expected cosine + // between a row and its reconstruction; these thresholds sit safely below the analytic + // expectation (about 0.945 at 2 bits, 0.983 at 3, 0.995 at 4) but far above what a broken + // rotation, grid, or scale would produce. + @ParameterizedTest + @CsvSource({"2, 0.92", "3, 0.97", "4, 0.99"}) + void testReconstructionQualityPerBitWidth(int bits, double threshold) { + final float[] matrix = testMatrix(); + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, bits, SEED); + double cosineSum = 0; + for (int row = 0; row < ROWS; row++) { + final float[] decoded = quantized.decodeRow(row); + cosineSum += ModelQuantizer.cosine(matrix, row * DIMENSION, DIMENSION, decoded); + } + final double meanCosine = cosineSum / ROWS; + assertTrue(meanCosine >= threshold, bits + " bits reconstructed a mean cosine of " + + meanCosine + ", below the acceptable " + threshold); + } + + @Test + void testRotatedDotEqualsOriginalSpaceDot() { + final float[] matrix = testMatrix(); + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 4, SEED); + final Random random = new Random(7); + final float[] query = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + query[d] = (float) random.nextGaussian(); + } + final float[] rotatedQuery = quantized.rotate(query); + for (int row = 0; row < ROWS; row++) { + final float[] decoded = quantized.decodeRow(row); + double originalDot = 0; + for (int d = 0; d < DIMENSION; d++) { + originalDot += (double) decoded[d] * query[d]; + } + assertEquals(originalDot, quantized.dotRotated(row, rotatedQuery), + 1e-3 * (1 + Math.abs(originalDot)), + "the rotation is orthonormal, so rotated-space and original-space dots must agree"); + } + } + + @Test + void testRowNormIsTheDecodedRowsNorm() { + final float[] matrix = testMatrix(); + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 3, SEED); + for (int row = 0; row < ROWS; row++) { + final float[] decoded = quantized.decodeRow(row); + double sumOfSquares = 0; + for (final float value : decoded) { + sumOfSquares += (double) value * value; + } + final double decodedNorm = Math.sqrt(sumOfSquares); + assertEquals(decodedNorm, quantized.rowNorm(row), 1e-3 * (1 + decodedNorm)); + } + } + + @Test + void testPoolingInRotatedSpaceEqualsPoolingDecodedRows() { + final float[] matrix = testMatrix(); + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 4, SEED); + final float[] rotatedSum = new float[quantized.paddedDimension()]; + quantized.addRowRotated(0, 1f, rotatedSum); + quantized.addRowRotated(1, 2.5f, rotatedSum); + quantized.addRowRotated(2, -0.5f, rotatedSum); + final float[] pooled = quantized.toOriginal(rotatedSum); + final float[] row0 = quantized.decodeRow(0); + final float[] row1 = quantized.decodeRow(1); + final float[] row2 = quantized.decodeRow(2); + final float[] expected = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + expected[d] = row0[d] + 2.5f * row1[d] - 0.5f * row2[d]; + } + assertArrayEquals(expected, pooled, 1e-3f, + "rotation is linear, so pooling commutes with it"); + } + + @Test + void testZeroRowDecodesToZero() { + final float[] matrix = new float[3 * 8]; + matrix[0] = 1f; + matrix[2 * 8 + 5] = -2f; + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, 3, 8, 2, SEED); + assertArrayEquals(new float[8], quantized.decodeRow(1), 0f); + assertEquals(0.0, quantized.rowNorm(1), 0.0); + } + + @Test + void testQuantizingIsDeterministic(@TempDir Path directory) throws IOException { + final float[] matrix = testMatrix(); + final Path first = directory.resolve("first.bin"); + final Path second = directory.resolve("second.bin"); + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 3, SEED).write(first); + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 3, SEED).write(second); + assertArrayEquals(Files.readAllBytes(first), Files.readAllBytes(second), + "the same matrix, bits, and seed must produce the same file bytes"); + } + + @Test + void testWriteReadRoundTrip(@TempDir Path directory) throws IOException { + final float[] matrix = testMatrix(); + final QuantizedEmbeddingMatrix written = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 4, SEED); + final Path file = directory.resolve("matrix.bin"); + written.write(file); + final QuantizedEmbeddingMatrix read = QuantizedEmbeddingMatrix.read(file); + assertEquals(written.rowCount(), read.rowCount()); + assertEquals(written.dimension(), read.dimension()); + assertEquals(written.paddedDimension(), read.paddedDimension()); + assertEquals(written.bits(), read.bits()); + assertEquals(written.seed(), read.seed()); + for (int row = 0; row < ROWS; row++) { + assertArrayEquals(written.decodeRow(row), read.decodeRow(row), 0f, + "a read matrix must decode exactly like the written one"); + } + // Writing the read matrix reproduces the file, so the format loses nothing. + final Path rewritten = directory.resolve("rewritten.bin"); + read.write(rewritten); + assertArrayEquals(Files.readAllBytes(file), Files.readAllBytes(rewritten)); + } + + @Test + void testQuantizedFileIsSmallerThanTheFloatMatrix(@TempDir Path directory) + throws IOException { + final float[] matrix = testMatrix(); + final Path file = directory.resolve("matrix.bin"); + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 4, SEED).write(file); + final long floatBytes = (long) ROWS * DIMENSION * Float.BYTES; + // 4 bits over the padded dimension (512 for 300) plus one scale per row and the header: + // still far under half the float size; at equal dimensions the ratio approaches 8x. + assertTrue(Files.size(file) < floatBytes / 2, + "the 4-bit file (" + Files.size(file) + " bytes) must be well under half the float " + + "matrix (" + floatBytes + " bytes)"); + } + + @Test + void testPoolingWeightsRoundTripThroughTheFile(@TempDir Path directory) throws IOException { + final float[] matrix = testMatrix(); + final float[] weights = new float[ROWS]; + for (int row = 0; row < ROWS; row++) { + weights[row] = 0.5f + row / 100f; + } + final QuantizedEmbeddingMatrix withWeights = + QuantizedEmbeddingMatrix.quantize(matrix, ROWS, DIMENSION, 4, SEED) + .withPoolingWeights(weights); + final Path file = directory.resolve("weighted.bin"); + withWeights.write(file); + final QuantizedEmbeddingMatrix read = QuantizedEmbeddingMatrix.read(file); + assertArrayEquals(weights, read.poolingWeights(), 0f); + // Rewriting reproduces the file, so the weights block loses nothing. + final Path rewritten = directory.resolve("rewritten.bin"); + read.write(rewritten); + assertArrayEquals(Files.readAllBytes(file), Files.readAllBytes(rewritten)); + // Without weights the accessor answers null and the file omits the block. + final QuantizedEmbeddingMatrix withoutWeights = withWeights.withPoolingWeights(null); + assertNull(withoutWeights.poolingWeights()); + assertTrue(Files.size(file) > sizeWithoutWeights(directory, withoutWeights), + "the weights block must add to the file size"); + } + + /** + * {@return the file size of a matrix written without weights} + * + * @param directory The directory to write into. + * @param matrix The matrix to write. + */ + private static long sizeWithoutWeights(Path directory, QuantizedEmbeddingMatrix matrix) + throws IOException { + final Path file = directory.resolve("unweighted.bin"); + matrix.write(file); + return Files.size(file); + } + + @Test + void testWithPoolingWeightsValidates() { + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(new float[4 * 8], 4, 8, 2, SEED); + assertThrows(IllegalArgumentException.class, + () -> quantized.withPoolingWeights(new float[3])); + assertThrows(IllegalArgumentException.class, + () -> quantized.withPoolingWeights(new float[] {1f, 2f, Float.NaN, 4f})); + } + + @Test + void testQuantizeValidatesItsArguments() { + final float[] matrix = new float[2 * 4]; + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(null, 2, 4, 2, SEED)); + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 0, 4, 2, SEED)); + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 2, 0, 2, SEED)); + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 2, 5, 2, SEED)); + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 2, 4, 1, SEED)); + assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 2, 4, 5, SEED)); + } + + @Test + void testQuantizeRejectsNonFiniteValuesNamingTheirPosition() { + final float[] matrix = new float[2 * 4]; + matrix[5] = Float.NaN; + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> QuantizedEmbeddingMatrix.quantize(matrix, 2, 4, 2, SEED)); + assertTrue(e.getMessage().contains("Row 1"), e.getMessage()); + assertTrue(e.getMessage().contains("dimension 1"), e.getMessage()); + } + + @Test + void testRotatedSpaceAccessorsValidate() { + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(new float[4 * 8], 4, 8, 2, SEED); + assertThrows(IllegalArgumentException.class, () -> quantized.decodeRow(-1)); + assertThrows(IllegalArgumentException.class, () -> quantized.decodeRow(4)); + assertThrows(IllegalArgumentException.class, () -> quantized.rotate(null)); + assertThrows(IllegalArgumentException.class, () -> quantized.rotate(new float[7])); + assertThrows(IllegalArgumentException.class, () -> quantized.toOriginal(new float[7])); + assertThrows(IllegalArgumentException.class, + () -> quantized.addRowRotated(0, 1f, new float[7])); + assertThrows(IllegalArgumentException.class, + () -> quantized.dotRotated(0, new float[7])); + } + + @Test + void testReadRejectsADimensionThatOverflowsTheRowByteCount(@TempDir Path directory) + throws IOException { + // A header declaring dimension 2^29 makes paddedDimension*bits overflow a signed int, so + // the per-row byte count goes negative. A corrupt or hostile file must be rejected cleanly, + // not crash the reader with an undocumented NegativeArraySizeException. + final Path file = directory.resolve("matrix.bin"); + QuantizedEmbeddingMatrix.quantize(new float[] {1f, 2f, 3f, 4f}, 1, 4, 4, SEED).write(file); + final byte[] bytes = Files.readAllBytes(file); + // dimension is the third big-endian int: after magic[4] and rowCount[4], at offset 8. + final int overflowingDimension = 1 << 29; + bytes[8] = (byte) (overflowingDimension >>> 24); + bytes[9] = (byte) (overflowingDimension >>> 16); + bytes[10] = (byte) (overflowingDimension >>> 8); + bytes[11] = (byte) overflowingDimension; + final Path patched = directory.resolve("overflow.bin"); + Files.write(patched, bytes); + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(patched)); + } + + @Test + void testReadRejectsForeignAndTruncatedFiles(@TempDir Path directory) throws IOException { + final Path foreign = directory.resolve("foreign.bin"); + Files.write(foreign, new byte[] {1, 2, 3, 4, 5, 6, 7, 8}); + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(foreign)); + + final Path file = directory.resolve("matrix.bin"); + QuantizedEmbeddingMatrix.quantize(testMatrix(), ROWS, DIMENSION, 2, SEED).write(file); + final byte[] full = Files.readAllBytes(file); + final Path truncated = directory.resolve("truncated.bin"); + Files.write(truncated, Arrays.copyOf(full, full.length - 10)); + assertThrows(IOException.class, () -> QuantizedEmbeddingMatrix.read(truncated)); + + final Path trailing = directory.resolve("trailing.bin"); + final byte[] extra = Arrays.copyOf(full, full.length + 1); + Files.write(trailing, extra); + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(trailing)); + } + +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixFormatTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixFormatTest.java new file mode 100644 index 0000000000..1b05604702 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixFormatTest.java @@ -0,0 +1,210 @@ +/* + * 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.DataOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Pins the loader contract of {@link QuantizedEmbeddingMatrix#read(Path)}: malformed + * file content fails with the checked {@link InvalidFormatException}, never with an + * unchecked exception or an allocation failure. + */ +public class QuantizedMatrixFormatTest { + + /** The on-disk magic of a quantized matrix, as written by the writer. */ + private static final int MAGIC = 0x4F4E5131; + + @Test + void testBadMagicIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("bad-magic.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(0xCAFEBABE); + } + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(file)); + assertTrue(e.getMessage().contains("magic"), e.getMessage()); + } + + @Test + void testNegativeRowCountIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("negative-rows.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(MAGIC); + out.writeInt(-5); + } + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(file)); + } + + @Test + void testImplausibleRowCountFailsBeforeAllocating(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("huge-rows.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(MAGIC); + out.writeInt(Integer.MAX_VALUE); + out.writeInt(8); + out.writeInt(4); + out.writeLong(17L); + out.writeInt(16); + } + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(file)); + } + + @Test + void testUnsupportedBitWidthIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("bad-bits.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(MAGIC); + out.writeInt(1); + out.writeInt(8); + out.writeInt(7); + } + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(file)); + } + + @Test + void testDeclaredPayloadBeyondFileSizeFailsBeforeAllocating(@TempDir Path dir) + throws IOException { + // A 1.1 MB file declaring 1,000,000 rows of 512 dimensions at 4 bits describes 256 MB of + // packed codes plus 8 MB of scales and norms. Both dimensions individually pass a + // "smaller than the file size" plausibility check, so the loader must hold the declared + // total against the bytes actually present, before allocating anything row-sized. + final Path file = dir.resolve("huge-payload.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(MAGIC); + out.writeInt(1_000_000); + out.writeInt(512); + out.writeInt(4); + out.writeLong(17L); + out.writeInt(16); + for (int i = 0; i < 16; i++) { + out.writeFloat(i - 7.5f); + } + out.write(new byte[1_100_000 - 92]); + } + final InvalidFormatException e = assertTimeoutPreemptively(Duration.ofSeconds(10), + () -> assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(file))); + assertTrue(e.getMessage().contains(file.toString()), e.getMessage()); + } + + @Test + void testInvalidStoredGridIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + // The first grid level is the big-endian float at offset 28, after the six header fields. + writeNaN(bytes, 28); + final Path patched = dir.resolve("bad-grid.quantized"); + Files.write(patched, bytes); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + assertTrue(e.getMessage().contains("grid"), e.getMessage()); + } + + @Test + void testNonFiniteDecodedNormIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + // Row 0's decoded norm is the big-endian float at offset 96: 28 header bytes, 64 bytes of + // grid levels, and 4 bytes for row 0's scale. + writeNaN(bytes, 96); + final Path patched = dir.resolve("bad-norm.quantized"); + Files.write(patched, bytes); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + assertTrue(e.getMessage().contains("norm"), e.getMessage()); + } + + @Test + void testNonFinitePoolingWeightIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("valid-weights.quantized"); + QuantizedEmbeddingMatrix.quantize(new float[] {1f, 2f, 3f, 4f}, 1, 4, 4, 17L) + .withPoolingWeights(new float[] {2f}) + .write(file); + final byte[] bytes = Files.readAllBytes(file); + // Row 0's pooling weight is the big-endian float at offset 101, right after the weight + // presence flag at offset 100. + writeNaN(bytes, 101); + final Path patched = dir.resolve("bad-weight.quantized"); + Files.write(patched, bytes); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + assertTrue(e.getMessage().contains("pooling"), e.getMessage()); + } + + @Test + void testPoolingWeightsDeclaredBeyondFileSizeAreRejectedAsFormatError(@TempDir Path dir) + throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + // Flipping the presence flag at offset 100 declares per-row pooling weights the file does + // not contain, so the declared total exceeds the file size. + bytes[100] = 1; + final Path patched = dir.resolve("flagged-weights.quantized"); + Files.write(patched, bytes); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + assertTrue(e.getMessage().contains("pooling"), e.getMessage()); + } + + /** + * {@return a valid one-row, four-dimension, 4-bit quantized file without pooling weights} + * Its layout is fixed: 28 header bytes, 64 bytes of grid levels, one scale at offset 92, one + * decoded norm at offset 96, the weight presence flag at offset 100, and two packed code + * bytes, 103 bytes in total. + * + * @param dir The directory to write into. + * @throws IOException Thrown if writing fails. + */ + private static Path validFile(Path dir) throws IOException { + final Path file = dir.resolve("valid.quantized"); + QuantizedEmbeddingMatrix.quantize(new float[] {1f, 2f, 3f, 4f}, 1, 4, 4, 17L).write(file); + return file; + } + + /** + * Overwrites four bytes with the big-endian bits of {@code Float.NaN}. + * + * @param bytes The file image to patch. + * @param offset The offset of the float to replace. + */ + private static void writeNaN(byte[] bytes, int offset) { + final int nan = Float.floatToIntBits(Float.NaN); + bytes[offset] = (byte) (nan >>> 24); + bytes[offset + 1] = (byte) (nan >>> 16); + bytes[offset + 2] = (byte) (nan >>> 8); + bytes[offset + 3] = (byte) nan; + } + + private static DataOutputStream out(Path file) throws IOException { + final OutputStream raw = Files.newOutputStream(file); + return new DataOutputStream(raw); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/RandomizedPcaTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/RandomizedPcaTest.java new file mode 100644 index 0000000000..aaebe5bfc6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/RandomizedPcaTest.java @@ -0,0 +1,261 @@ +/* + * 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.Random; +import java.util.concurrent.ForkJoinPool; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +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 randomized PCA recovers the dominant subspace of a low-rank matrix: for data that is + * exactly rank-k, projecting to k components preserves the pairwise geometry (dot products) of + * the centered rows almost exactly, and it reports nearly all variance kept. A fixed seed makes + * the projection deterministic. + */ +class RandomizedPcaTest { + + private static final int ROWS = 400; + private static final int COLS = 48; + private static final int RANK = 6; + + /** Small enough that a fixed absolute regularization would swamp the rescaled matrix. */ + private static final float SMALL_SCALE = 1e-6f; + + /** + * {@return an exactly rank-{@link #RANK} matrix: a random factor times a random loading + * matrix, plus a non-zero column mean so centering is exercised} + */ + private static float[] lowRankData() { + final Random random = new Random(7); + final float[][] factors = new float[ROWS][RANK]; + final float[][] loadings = new float[RANK][COLS]; + for (final float[] row : factors) { + for (int j = 0; j < RANK; j++) { + row[j] = (float) random.nextGaussian() * (RANK - j); + } + } + for (final float[] row : loadings) { + for (int c = 0; c < COLS; c++) { + row[c] = (float) random.nextGaussian(); + } + } + final float[] data = new float[ROWS * COLS]; + for (int i = 0; i < ROWS; i++) { + for (int c = 0; c < COLS; c++) { + float value = c; // a column mean the PCA must subtract + for (int j = 0; j < RANK; j++) { + value += factors[i][j] * loadings[j][c]; + } + data[i * COLS + c] = value; + } + } + return data; + } + + private static double dot(float[] data, int rowA, int rowB, int cols) { + double dot = 0; + for (int c = 0; c < cols; c++) { + dot += (double) data[rowA * cols + c] * data[rowB * cols + c]; + } + return dot; + } + + @Test + void testRecoversTheExactSubspaceOfLowRankData() { + final float[] original = lowRankData(); + // The centered reference, for the geometry comparison. + final float[] centered = original.clone(); + for (int c = 0; c < COLS; c++) { + float mean = 0; + for (int i = 0; i < ROWS; i++) { + mean += centered[i * COLS + c]; + } + mean /= ROWS; + for (int i = 0; i < ROWS; i++) { + centered[i * COLS + c] -= mean; + } + } + + final RandomizedPca.Result result = + RandomizedPca.fitTransform(original.clone(), ROWS, COLS, RANK, 42); + + assertEquals(ROWS * RANK, result.transformed().length); + assertTrue(result.explainedVarianceRatio() > 0.999, + "rank-6 data projected to 6 components keeps (almost) all variance, got " + + result.explainedVarianceRatio()); + // Projecting exactly rank-k data onto its k principal components preserves pairwise dot + // products up to numerical noise; check the diagonal and a few off-diagonal pairs. + for (final int[] pair : new int[][] {{0, 0}, {1, 2}, {17, 399}, {5, 5}, {123, 321}}) { + final double expected = dot(centered, pair[0], pair[1], COLS); + final double actual = dot(result.transformed(), pair[0], pair[1], RANK); + final double scale = Math.max(Math.abs(expected), 1); + assertEquals(expected, actual, 1e-3 * scale, + "pairwise dot product of rows " + pair[0] + " and " + pair[1]); + } + } + + @Test + void testDeterministicForAFixedSeed() { + final float[] first = + RandomizedPca.fitTransform(lowRankData(), ROWS, COLS, RANK, 42).transformed(); + final float[] second = + RandomizedPca.fitTransform(lowRankData(), ROWS, COLS, RANK, 42).transformed(); + assertArrayEquals(first, second); + } + + /** + * The decomposition is equivariant under a rescaling of the whole matrix: the projected + * coordinates scale with the input and the explained variance ratio, being a ratio, does not + * move. Any absolute (rather than relative) tolerance inside the pipeline breaks this. + */ + @Test + void testIsUnchangedByRescalingTheWholeMatrix() { + final RandomizedPca.Result unscaled = + RandomizedPca.fitTransform(lowRankData(), ROWS, COLS, RANK, 42); + final float[] scaledData = lowRankData(); + for (int i = 0; i < scaledData.length; i++) { + scaledData[i] *= SMALL_SCALE; + } + + final RandomizedPca.Result scaled = + RandomizedPca.fitTransform(scaledData, ROWS, COLS, RANK, 42); + + assertEquals(unscaled.explainedVarianceRatio(), scaled.explainedVarianceRatio(), 1e-6, + "the explained variance ratio must not depend on the magnitude of the input"); + double largest = 0; + for (final float value : unscaled.transformed()) { + largest = Math.max(largest, Math.abs(value)); + } + final double tolerance = 1e-4 * largest * SMALL_SCALE; + for (int i = 0; i < unscaled.transformed().length; i++) { + assertEquals(unscaled.transformed()[i] * (double) SMALL_SCALE, scaled.transformed()[i], + tolerance, "projected coordinate " + i); + } + } + + /** + * The parallel loops reduce per-block partial sums in a fixed block order, so the result does not + * depend on how many threads the fork/join pool runs the blocks on. + */ + @ParameterizedTest + @ValueSource(ints = {1, 2, 7}) + void testDeterministicAcrossThreadCounts(int parallelism) throws Exception { + final float[] expected = + RandomizedPca.fitTransform(lowRankData(), ROWS, COLS, RANK, 42).transformed(); + final ForkJoinPool pool = new ForkJoinPool(parallelism); + + try { + final float[] actual = pool.submit( + () -> RandomizedPca.fitTransform(lowRankData(), ROWS, COLS, RANK, 42).transformed()) + .get(); + assertArrayEquals(expected, actual); + } finally { + pool.shutdown(); + } + } + + @Test + void testCentersTheDataInPlace() { + final float[] data = lowRankData(); + RandomizedPca.fitTransform(data, ROWS, COLS, RANK, 42); + for (int c = 0; c < COLS; c++) { + double mean = 0; + for (int i = 0; i < ROWS; i++) { + mean += data[i * COLS + c]; + } + assertEquals(0, mean / ROWS, 1e-5, "column " + c + " is centered"); + } + } + + @Test + void testRejectsNullData() { + assertEquals("Data must not be null", assertThrows(IllegalArgumentException.class, + () -> RandomizedPca.fitTransform(null, 3, 4, 2, 42)).getMessage()); + } + + /** + * The shape must describe the array exactly: a wrong column count, a non-positive dimension, or + * a length that is not {@code rows * cols} is rejected rather than silently reinterpreted. + */ + @ParameterizedTest + @CsvSource({"3, 5, 2", "3, 3, 2", "0, 4, 2", "3, 0, 2", "-1, 4, 2", "4, 4, 2", "2, 4, 1"}) + void testRejectsAShapeThatDoesNotDescribeTheData(int rows, int cols, int components) { + final float[] data = new float[12]; + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> RandomizedPca.fitTransform(data, rows, cols, components, 42)); + assertTrue(e.getMessage().startsWith("Data has 12 elements, not " + rows + " x " + cols), + e.getMessage()); + } + + /** + * The component count must be a genuine reduction: at least one, no more than the column count, + * and strictly fewer than the row count (the randomized range finder has no subspace to find + * otherwise). + */ + @ParameterizedTest + @CsvSource({"0", "-1", "5", "3", "4"}) + void testRejectsAComponentCountThatIsNotAReduction(int components) { + final float[] data = new float[12]; + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> RandomizedPca.fitTransform(data, 3, 4, components, 42)); + assertTrue(e.getMessage().startsWith("Components must be in [1, 2], got " + components), + e.getMessage()); + } + + /** + * Data whose rows are all identical centers to exactly zero, so there is no subspace and the + * explained-variance ratio would be 0/0. That must fail loudly rather than return NaN. + */ + @Test + void testRejectsDataWithoutVariance() { + final float[] data = new float[ROWS * COLS]; + for (int i = 0; i < ROWS; i++) { + for (int c = 0; c < COLS; c++) { + data[i * COLS + c] = c; + } + } + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> RandomizedPca.fitTransform(data, ROWS, COLS, RANK, 42)); + assertTrue(e.getMessage().contains("total variance"), e.getMessage()); + } + + /** + * A non-finite value poisons the column mean, so every centered value becomes NaN and the total + * variance is NaN. The check must reject that too, not let NaN through into the table. + */ + @ParameterizedTest + @ValueSource(floats = {Float.NaN, Float.POSITIVE_INFINITY, Float.NEGATIVE_INFINITY}) + void testRejectsNonFiniteData(float value) { + final float[] data = lowRankData(); + data[0] = value; + + assertThrows(IllegalArgumentException.class, + () -> RandomizedPca.fitTransform(data, ROWS, COLS, RANK, 42)); + } +} 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 new file mode 100644 index 0000000000..6101d3ecca --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java @@ -0,0 +1,360 @@ +/* + * 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.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 opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class SafetensorsFileTest { + + 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); + 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(name); + Files.write(file, out.toByteArray()); + return file; + } + + private static byte[] floatsToLittleEndianBytes(float... values) { + final ByteBuffer buffer = ByteBuffer.allocate(values.length * 4).order(ByteOrder.LITTLE_ENDIAN); + for (float value : values) { + buffer.putFloat(value); + } + return buffer.array(); + } + + @Test + void testRoundTripsAFloat32Matrix(@TempDir Path dir) throws IOException { + 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); + + assertEquals(1, parsed.size()); + assertEquals(Set.of("weight"), parsed.tensorNames()); + final TensorInfo info = parsed.tensorInfo("weight"); + assertEquals("F32", info.dtype()); + assertArrayEquals(new int[] {2, 3}, info.shape()); + assertEquals(6, info.elementCount()); + assertArrayEquals(new float[] {1f, 2f, 3f, 4f, 5f, 6f}, parsed.readFloat32("weight")); + } + + @Test + void testMultipleTensorsPreserveHeaderOrder(@TempDir Path dir) throws IOException { + 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(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")); + } + + @Test + void testMetadataMapIsParsed(@TempDir Path dir) throws IOException { + 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_FILE_NAME, header, data); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + + assertEquals("pt", parsed.metadata().get("format")); + assertEquals("line\nbreak", parsed.metadata().get("note")); + assertEquals(1, parsed.size()); + } + + @Test + void testUnknownHeaderFieldsAreSkipped(@TempDir Path dir) throws IOException { + 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_FILE_NAME, header, data); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + + assertArrayEquals(new float[] {1f, 2f}, parsed.readFloat32("w")); + } + + @Test + void testSingleMatrixTensorNameFindsTheOnly2DFloat32Tensor(@TempDir Path dir) throws IOException { + 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); + + assertEquals("embeddings", parsed.singleMatrixTensorName()); + } + + @Test + void testSingleMatrixTensorNameRejectsAmbiguity(@TempDir Path dir) throws IOException { + 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); + + assertThrows(InvalidFormatException.class, parsed::singleMatrixTensorName); + } + + @Test + void testSingleMatrixTensorNameRejectsNoCandidate(@TempDir Path dir) throws IOException { + final Path file = dir.resolve(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, SafetensorsTestFiles.vector("bias", new float[] {1f})); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + + assertThrows(InvalidFormatException.class, parsed::singleMatrixTensorName); + } + + @Test + void testReadFloat32RejectsWrongDtype(@TempDir Path dir) throws IOException { + final byte[] data = new byte[] {1, 2}; + 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); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> parsed.readFloat32("ids")); + assertTrue(e.getMessage().contains("I64")); + } + + @Test + void testTensorInfoRejectsUnknownName(@TempDir Path dir) throws IOException { + final byte[] data = floatsToLittleEndianBytes(1f); + 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); + + assertThrows(IllegalArgumentException.class, () -> parsed.tensorInfo("missing")); + } + + @Test + void testRejectsNullAndMissingFile(@TempDir Path dir) { + assertThrows(IllegalArgumentException.class, () -> SafetensorsFile.read(null)); + assertThrows(IllegalArgumentException.class, + () -> SafetensorsFile.read(dir.resolve("absent.safetensors"))); + } + + @Test + void testRejectsFileShorterThanTheLengthPrefix(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("truncated.safetensors"); + Files.write(file, new byte[] {1, 2, 3}); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsHeaderLengthLargerThanTheFile(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("bad-length.safetensors"); + final byte[] prefix = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(1000L).array(); + Files.write(file, prefix); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsDuplicateTensorName(@TempDir Path dir) throws IOException { + // The same key twice is syntactically valid JSON (just semantically ambiguous), so the + // 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_FILE_NAME, header, new byte[] {1, 2, 3, 4}); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + assertTrue(e.getMessage().contains("more than once")); + } + + @Test + void testRejectsTensorMissingRequiredField(@TempDir Path dir) throws IOException { + final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]}}"; + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[0]); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsDataOffsetsOutOfRange(@TempDir Path dir) throws IOException { + 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(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsUnterminatedString(@TempDir Path dir) throws IOException { + final String header = "{\"w\":{\"dtype\":\"F32"; + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[0]); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsTensorLargerThanAJavaArray(@TempDir Path dir) throws IOException { + // 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 = 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); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> parsed.readFloat32("w")); + assertTrue(e.getMessage().contains("more than a Java array can hold")); + } + + @Test + void testFailsLoudWhenTheFileIsTruncatedAfterRead(@TempDir Path dir) throws IOException { + // 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 = 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_FILE_NAME, header, floatsToLittleEndianBytes(1f)); + + final IllegalStateException e = + 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 = singleTensorHeader("w", "F32", "[2]", 0, data.length); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, data); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + final InvalidFormatException e = + assertThrows(InvalidFormatException.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()); + } + + // 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(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(MODEL_FILE_NAME); + SafetensorsTestFiles.write(file, "F16", + SafetensorsTestFiles.matrix("embeddings", new float[][] {{1f, 2f}, {3f, 4f}})); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + assertEquals("embeddings", parsed.singleMatrixTensorName()); + } + + @Test + void testReadFloat32StrictlyRejectsF16(@TempDir Path dir) throws IOException { + 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); + // readFloats accepts it; the strict readFloat32 must not. + assertArrayEquals(new float[] {1f, 2f}, parsed.readFloats("w"), 1e-3f); + assertThrows(InvalidFormatException.class, () -> parsed.readFloat32("w")); + } +} 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 new file mode 100644 index 0000000000..95219488b3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java @@ -0,0 +1,191 @@ +/* + * 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; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Direct tests for {@link SafetensorsHeaderParser}, complementing the indirect coverage in + * {@link SafetensorsFileTest}: the file-level tests exercise headers as whole files, these pin + * the parser's own contract, its error offsets, and every malformed-input branch. + */ +class SafetensorsHeaderParserTest { + + @Test + void testParsesTensorsInHeaderOrder() throws InvalidFormatException { + final SafetensorsHeaderParser.Result result = SafetensorsHeaderParser.parse( + "{\"beta\":{\"dtype\":\"F32\",\"shape\":[2,3],\"data_offsets\":[0,24]}," + + "\"alpha\":{\"dtype\":\"I64\",\"shape\":[],\"data_offsets\":[24,32]}}"); + + assertEquals(2, result.tensors().size()); + final TensorInfo beta = result.tensors().get(0); + assertEquals("beta", beta.name()); + assertEquals("F32", beta.dtype()); + assertArrayEquals(new int[] {2, 3}, beta.shape()); + assertEquals(0, beta.dataOffsetBegin()); + assertEquals(24, beta.dataOffsetEnd()); + final TensorInfo alpha = result.tensors().get(1); + assertEquals("alpha", alpha.name()); + assertArrayEquals(new int[0], alpha.shape()); + assertEquals(1, alpha.elementCount()); + assertTrue(result.metadata().isEmpty()); + } + + @Test + void testParsesAnEmptyHeader() throws InvalidFormatException { + final SafetensorsHeaderParser.Result result = SafetensorsHeaderParser.parse("{}"); + + assertTrue(result.tensors().isEmpty()); + assertTrue(result.metadata().isEmpty()); + } + + @Test + void testParsesAMetadataOnlyHeader() throws InvalidFormatException { + final SafetensorsHeaderParser.Result result = + SafetensorsHeaderParser.parse("{\"__metadata__\":{\"format\":\"pt\"}}"); + + assertTrue(result.tensors().isEmpty()); + assertEquals("pt", result.metadata().get("format")); + } + + @Test + void testDecodesEveryEscapeSequence() throws InvalidFormatException { + final SafetensorsHeaderParser.Result result = SafetensorsHeaderParser.parse( + "{\"__metadata__\":{\"note\":\"\\\"\\\\\\/\\b\\f\\n\\r\\t\\u0041\"}}"); + + assertEquals("\"\\/\b\f\n\r\tA", result.metadata().get("note")); + } + + @Test + void testSkipsUnknownFieldsOfEveryValueType() throws InvalidFormatException { + // Fields safetensors may add over time must not break the reader: nested objects, arrays, + // floating-point numbers, booleans, null, and strings are all skipped structurally. + final SafetensorsHeaderParser.Result result = SafetensorsHeaderParser.parse( + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]," + + "\"future\":{\"nested\":[1,-2.5e3,true,false,null,\"s\",{\"deep\":[]}]}}}"); + + assertEquals(List.of("w"), result.tensors().stream().map(TensorInfo::name).toList()); + } + + @Test + void testToleratesTrailingWhitespacePadding() throws InvalidFormatException { + // Writers space-pad the header so the data section starts aligned; padding is part of the + // declared header length and must parse cleanly. + final SafetensorsHeaderParser.Result result = SafetensorsHeaderParser.parse( + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}} "); + + assertEquals(1, result.tensors().size()); + } + + @Test + void testRejectsTrailingGarbage() { + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> SafetensorsHeaderParser.parse("{} x")); + assertTrue(e.getMessage().contains("Trailing content")); + } + + @Test + void testRejectsNull() { + assertThrows(IllegalArgumentException.class, () -> SafetensorsHeaderParser.parse(null)); + } + + @ParameterizedTest + @ValueSource(strings = { + // unterminated string + "{\"w", + // unknown escape + "{\"a\\x\":{}}", + // truncated \_u escape (split so the Java lexer does not see a \_u sequence) + "{\"a\\" + "u00", + // malformed \_u escape + "{\"a\\" + "uZZZZ\":{}}", + // missing colon + "{\"w\" 1}", + // empty tensor object + "{\"w\":{}}", + // missing dtype + "{\"w\":{\"shape\":[1],\"data_offsets\":[0,4]}}", + // missing shape + "{\"w\":{\"dtype\":\"F32\",\"data_offsets\":[0,4]}}", + // missing data_offsets + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]}}", + // data_offsets arity 1 + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0]}}", + // data_offsets arity 3 + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4,8]}}", + // negative shape dimension + "{\"w\":{\"dtype\":\"F32\",\"shape\":[-1],\"data_offsets\":[0,4]}}", + // shape dimension over int range + "{\"w\":{\"dtype\":\"F32\",\"shape\":[4294967296],\"data_offsets\":[0,4]}}", + // non-numeric array element + "{\"w\":{\"dtype\":\"F32\",\"shape\":[\"x\"],\"data_offsets\":[0,4]}}", + // number too large for long + "{\"w\":{\"dtype\":\"F32\",\"shape\":[99999999999999999999],\"data_offsets\":[0,4]}}", + // bare value instead of an object + "42", + // truncated after a tensor entry + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]}" + }) + void testRejectsMalformedHeaders(String header) { + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> SafetensorsHeaderParser.parse(header)); + 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(InvalidFormatException.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(InvalidFormatException.class, () -> SafetensorsHeaderParser.parse(header)); + } + + @Test + void testLoneMinusInSkippedFieldFailsLoudly() { + // A bare "-" is not a JSON number; the skip path must reject it rather than treating it as one. + final String header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4],\"unknown\":-}}"; + assertThrows(InvalidFormatException.class, () -> SafetensorsHeaderParser.parse(header)); + } + + @Test + void testWellFormedNumbersInSkippedFieldsAreAccepted() throws InvalidFormatException { + 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 0000000000..877555c6a6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsTestFiles.java @@ -0,0 +1,128 @@ +/* + * 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. Negative tests that + * need deliberately malformed bytes still hand-roll them. + */ +final class SafetensorsTestFiles { + + /** Not instantiable. */ + 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) { + } + + /** + * {@return a tensor of the given 2-D matrix, row-major} + * + * @param name The tensor name. + * @param rows The matrix rows, each of the same length. + */ + 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); + } + + /** + * {@return a tensor of the given 1-D values} + * + * @param name The tensor name. + * @param values The values. + */ + static Tensor vector(String name, float[] values) { + return new Tensor(name, new int[] {values.length}, values); + } + + /** + * Writes a safetensors file holding the given tensors as {@code F32}, header first, data in + * declaration order. + * + * @param file The file to write. + * @param tensors The tensors, in the order they should appear in the header and data. + * @throws IOException Thrown if writing the file fails. + */ + static void write(Path file, Tensor... tensors) throws IOException { + write(file, "F32", tensors); + } + + /** + * Writes a safetensors file encoding each tensor value as {@code dtype}, one of {@code F32}, + * {@code F16} (IEEE half), or {@code BF16} (bfloat16). The {@link Tensor} values stay + * {@code float}; they are converted to the target dtype's bytes here. + * + * @param file The file to write. + * @param dtype The dtype to encode every value as. + * @param tensors The tensors, in the order they should appear in the header and data. + * @throws IllegalArgumentException Thrown if {@code dtype} is not one of the three supported. + * @throws IOException Thrown if writing the file fails. + */ + static void write(Path file, String dtype, Tensor... tensors) throws IOException { + final int elementBytes = switch (dtype) { + case "F32" -> Float.BYTES; + case "F16", "BF16" -> Short.BYTES; + default -> throw new IllegalArgumentException("unsupported test dtype: " + dtype); + }; + 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 * elementBytes).order(ByteOrder.LITTLE_ENDIAN); + for (final float value : tensor.values()) { + switch (dtype) { + case "F32" -> buffer.putFloat(value); + case "F16" -> buffer.putShort(Float.floatToFloat16(value)); + case "BF16" -> buffer.putShort((short) (Float.floatToIntBits(value) >>> 16)); + default -> throw new IllegalArgumentException("unsupported test dtype: " + dtype); + } + } + 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 * elementBytes; + header.add("\"" + tensor.name() + "\":{\"dtype\":\"" + dtype + "\",\"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/SafetensorsWriterTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsWriterTest.java new file mode 100644 index 0000000000..aea8fff148 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsWriterTest.java @@ -0,0 +1,173 @@ +/* + * 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.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; + +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.CsvSource; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +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 writer's output round-trips through the module's own reader, including a matrix larger than + * one encoding chunk; the bytes it lays down are the safetensors layout, header padded so the + * data starts aligned; and a shape that does not match the value count is rejected. + */ +class SafetensorsWriterTest { + + /** More floats than fit in one encoding chunk, so the streaming loop runs more than once. */ + private static final int MULTI_CHUNK_ROWS = 400; + + /** The column count of the multi-chunk fixture; rows times columns exceeds 1 MiB of floats. */ + private static final int MULTI_CHUNK_COLS = 1024; + + @Test + void testRoundTripsThroughTheReader(@TempDir Path dir) throws IOException { + final float[] values = {1.5f, -2.25f, 3e8f, 0, -0.5f, 42}; + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + + SafetensorsWriter.writeMatrix(file, 2, 3, values); + + final SafetensorsFile tensors = SafetensorsFile.read(file); + assertEquals(SafetensorsWriter.EMBEDDINGS_TENSOR, tensors.singleMatrixTensorName()); + assertArrayEquals(new int[] {2, 3}, tensors.tensorInfo(SafetensorsWriter.EMBEDDINGS_TENSOR) + .shape()); + assertArrayEquals(values, tensors.readFloats(SafetensorsWriter.EMBEDDINGS_TENSOR)); + } + + @Test + void testRoundTripsAMatrixSpanningSeveralWriteChunks(@TempDir Path dir) throws IOException { + final float[] values = new float[MULTI_CHUNK_ROWS * MULTI_CHUNK_COLS]; + for (int i = 0; i < values.length; i++) { + values[i] = i * 0.5f; + } + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + + SafetensorsWriter.writeMatrix(file, MULTI_CHUNK_ROWS, MULTI_CHUNK_COLS, values); + + final SafetensorsFile tensors = SafetensorsFile.read(file); + assertArrayEquals(new int[] {MULTI_CHUNK_ROWS, MULTI_CHUNK_COLS}, + tensors.tensorInfo(SafetensorsWriter.EMBEDDINGS_TENSOR).shape()); + assertArrayEquals(values, tensors.readFloats(SafetensorsWriter.EMBEDDINGS_TENSOR)); + } + + /** + * Pins the on-disk layout: an 8-byte little-endian header length, the JSON header, then the + * values as little-endian {@code F32}, with nothing after them. A change to any of the three + * fails here rather than in whatever tool reads the distilled model next. + */ + @Test + void testWritesTheSafetensorsByteLayout(@TempDir Path dir) throws IOException { + final float[] values = {1, -2, 0.5f, 0, 7, -0.25f}; + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + + SafetensorsWriter.writeMatrix(file, 3, 2, values); + + final byte[] bytes = Files.readAllBytes(file); + final ByteBuffer buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN); + final long headerLength = buffer.getLong(); + final byte[] headerBytes = new byte[(int) headerLength]; + buffer.get(headerBytes); + assertEquals("{\"embeddings\":{\"dtype\":\"F32\",\"shape\":[3,2],\"data_offsets\":[0,24]}}", + new String(headerBytes, StandardCharsets.UTF_8).stripTrailing()); + assertEquals(Long.BYTES + headerLength + (long) values.length * Float.BYTES, bytes.length, + "the file is the length prefix, the header, and the values, with nothing after"); + for (int i = 0; i < values.length; i++) { + assertEquals(values[i], buffer.getFloat(), "value " + i + " must be little-endian F32"); + } + } + + /** + * The header is space-padded so the tensor data starts on an 8-byte boundary, the way the + * reference safetensors writer emits it. The header text length varies with the digits of the + * shape and the byte count, so every width has to be checked. + */ + @ParameterizedTest + @CsvSource({"1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "99", "100", "1000"}) + void testPadsTheHeaderToAnEightByteBoundary(int cols, @TempDir Path dir) throws IOException { + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + + SafetensorsWriter.writeMatrix(file, 1, cols, new float[cols]); + + final long headerLength = ByteBuffer.wrap(Files.readAllBytes(file)) + .order(ByteOrder.LITTLE_ENDIAN).getLong(); + assertEquals(0, (Long.BYTES + headerLength) % 8, "shape [1," + cols + "] leaves the data " + + "unaligned at byte " + (Long.BYTES + headerLength)); + } + + @Test + void testCreatesTheMissingParentDirectory(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("nested").resolve("deeper") + .resolve(ModelFileNames.SAFETENSORS); + + SafetensorsWriter.writeMatrix(file, 1, 2, new float[] {1, 2}); + + assertTrue(Files.isRegularFile(file), file + " must exist"); + } + + @Test + void testReplacesAnExistingFile(@TempDir Path dir) throws IOException { + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + SafetensorsWriter.writeMatrix(file, 2, 3, new float[6]); + + SafetensorsWriter.writeMatrix(file, 1, 2, new float[] {7, 8}); + + final SafetensorsFile tensors = SafetensorsFile.read(file); + assertArrayEquals(new int[] {1, 2}, + tensors.tensorInfo(SafetensorsWriter.EMBEDDINGS_TENSOR).shape()); + assertArrayEquals(new float[] {7, 8}, + tensors.readFloats(SafetensorsWriter.EMBEDDINGS_TENSOR)); + } + + @Test + void testRejectsANullFile() { + assertEquals("File must not be null", assertThrows(IllegalArgumentException.class, + () -> SafetensorsWriter.writeMatrix(null, 1, 1, new float[1])).getMessage()); + } + + @Test + void testRejectsNullValues(@TempDir Path dir) { + assertEquals("Values must not be null", assertThrows(IllegalArgumentException.class, + () -> SafetensorsWriter.writeMatrix(dir.resolve(ModelFileNames.SAFETENSORS), 1, 1, null)) + .getMessage()); + } + + @ParameterizedTest + @CsvSource({"0, 2, 2", "2, 0, 2", "-1, 2, 2", "2, 3, 5", "2, 3, 7", "1, 1, 0"}) + void testRejectsAShapeThatDoesNotMatchTheValues(int rows, int cols, int valueCount, + @TempDir Path dir) { + final Path file = dir.resolve(ModelFileNames.SAFETENSORS); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> SafetensorsWriter.writeMatrix(file, rows, cols, new float[valueCount])); + + assertEquals("Values has " + valueCount + " elements, not " + rows + " x " + cols, + e.getMessage()); + assertTrue(Files.notExists(file), "a rejected write must not leave a file behind"); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SentencePieceModelFixture.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SentencePieceModelFixture.java new file mode 100644 index 0000000000..7c9571c980 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SentencePieceModelFixture.java @@ -0,0 +1,151 @@ +/* + * 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.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +import opennlp.subword.sentencepiece.SentencePieceTokenizer; + +/** + * Writes a SentencePiece-layout static embedding model directory around the bundled tiny + * Unigram test model, for tests that need one. The matrix vocabulary is written the way a + * distillation ships it: control pieces dropped, extra special rows in front, and one token + * appended through {@code added_tokens}, so lookups must go by piece string, not tokenizer id. + */ +final class SentencePieceModelFixture { + + static final String MODEL_RESOURCE = "/opennlp/embeddings/tiny-unigram.model"; + + private final byte[] modelBytes; + private final SentencePieceTokenizer tokenizer; + // The matrix rows: , , then every poolable tokenizer piece. + private final List rows; + + /** + * Loads the bundled model resource and derives the row order. + * + * @throws IOException Thrown if the resource cannot be read. + */ + SentencePieceModelFixture() throws IOException { + try (InputStream in = SentencePieceModelFixture.class.getResourceAsStream(MODEL_RESOURCE)) { + modelBytes = in.readAllBytes(); + } + tokenizer = SentencePieceTokenizer.load( + SentencePieceModelFixture.class.getResourceAsStream(MODEL_RESOURCE)); + rows = new ArrayList<>(); + rows.add(""); + rows.add(""); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (!tokenizer.isControl(id) && !tokenizer.isUnknown(id)) { + rows.add(tokenizer.idToPiece(id)); + } + } + } + + /** {@return the loaded tokenizer} */ + SentencePieceTokenizer tokenizer() { + return tokenizer; + } + + /** {@return the matrix row pieces, in row order, with the appended added token last} */ + List rowPieces() { + final List pieces = new ArrayList<>(rows); + pieces.add(""); + return pieces; + } + + /** + * Writes the {@code .model}, a synthesized Unigram {@code tokenizer.json}, a deterministic + * embedding matrix, and a {@code config.json} into a directory. + * + * @param directory The directory to write into. + * @param dimension The embedding dimension. + * @param normalize The {@code config.json} normalize value. + * @param seed The seed of the deterministic matrix values. + * @throws IOException Thrown if writing fails. + */ + void write(Path directory, int dimension, boolean normalize, long seed) throws IOException { + Files.write(directory.resolve("sentencepiece.bpe.model"), modelBytes); + Files.writeString(directory.resolve("tokenizer.json"), tokenizerJson(rows)); + final int rowCount = rows.size() + 1; + final float[][] matrix = new float[rowCount][dimension]; + final Random random = new Random(seed); + for (final float[] row : matrix) { + final float rowScale = 0.5f + 2f * random.nextFloat(); + for (int d = 0; d < dimension; d++) { + row[d] = rowScale * (float) random.nextGaussian(); + } + } + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", matrix)); + Files.writeString(directory.resolve(ModelFileNames.CONFIG), + "{\"model_type\":\"model2vec\",\"normalize\":" + normalize + "}"); + } + + /** + * {@return a Unigram {@code tokenizer.json} whose vocabulary is the given pieces plus an + * appended added token} The appended token is not {@code } because the fixture model + * defines {@code } as a user-defined piece that already owns a row. + * + * @param pieces The {@code model.vocab} pieces in row order. + */ + private static String tokenizerJson(List pieces) { + final StringBuilder json = new StringBuilder("{\"version\":\"1.0\",\"added_tokens\":["); + json.append("{\"id\":0,\"content\":\"\",\"special\":true},"); + json.append("{\"id\":").append(pieces.size()).append(",\"content\":\"\"," + + "\"special\":true}],"); + json.append("\"normalizer\":{\"type\":\"Precompiled\"},\"model\":{\"type\":\"Unigram\"," + + "\"unk_id\":1,\"vocab\":["); + for (int i = 0; i < pieces.size(); i++) { + if (i > 0) { + json.append(','); + } + json.append('[').append(quote(pieces.get(i))).append(",-").append(i % 7).append(".5]"); + } + return json.append("]}}").toString(); + } + + /** + * {@return {@code text} as a JSON string literal} + * + * @param text The text to quote. + */ + private static String quote(String text) { + final StringBuilder quoted = new StringBuilder("\""); + for (int i = 0; i < text.length(); i++) { + final char c = text.charAt(i); + switch (c) { + case '"' -> quoted.append("\\\""); + case '\\' -> quoted.append("\\\\"); + default -> { + if (c < 0x20) { + quoted.append(String.format("\\u%04x", (int) c)); + } else { + quoted.append(c); + } + } + } + } + return quoted.append('"').toString(); + } +} 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 new file mode 100644 index 0000000000..9cc19087d6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelConcurrencyTest.java @@ -0,0 +1,94 @@ +/* + * 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.Arrays; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +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. Every operation is deterministic, so any + * deviation under contention is a thread-safety defect. + */ +class StaticEmbeddingModelConcurrencyTest { + + private static final int THREADS = 8; + private static final int ITERATIONS_PER_THREAD = 200; + + @Test + void testConcurrentUseMatchesSingleThreadedReference(@TempDir Path dir) throws Exception { + final StaticEmbeddingModel model = + EmbeddingTestFixtures.loadAnalogyModel(dir, Normalization.L2); + final float[] referenceEmbedding = model.embed("The King and Queen"); + final double referenceSimilarity = model.similarity("king", "queen"); + final List referenceNeighbors = model.mostSimilar("king", 3); + final List referenceAnalogy = model.analogy("man", "king", "woman", 2); + + final Queue problems = new ConcurrentLinkedQueue<>(); + final CountDownLatch start = new CountDownLatch(1); + final ExecutorService executor = Executors.newFixedThreadPool(THREADS); + try { + for (int t = 0; t < THREADS; t++) { + executor.submit(() -> { + try { + start.await(); + for (int i = 0; i < ITERATIONS_PER_THREAD; i++) { + if (!Arrays.equals(referenceEmbedding, model.embed("The King and Queen"))) { + problems.add("embed deviated from the single-threaded reference"); + } + if (referenceSimilarity != model.similarity("king", "queen")) { + problems.add("similarity deviated from the single-threaded reference"); + } + if (!referenceNeighbors.equals(model.mostSimilar("king", 3))) { + problems.add("mostSimilar deviated from the single-threaded reference"); + } + if (!referenceAnalogy.equals(model.analogy("man", "king", "woman", 2))) { + problems.add("analogy deviated from the single-threaded reference"); + } + } + } + catch (Exception e) { + problems.add("Unexpected exception: " + e); + } + }); + } + start.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.MINUTES), + "Concurrent workers did not finish in time"); + } + finally { + executor.shutdownNow(); + } + assertTrue(problems.isEmpty(), () -> "Thread-safety violations: " + problems); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelQuantizedTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelQuantizedTest.java new file mode 100644 index 0000000000..250698c218 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelQuantizedTest.java @@ -0,0 +1,273 @@ +/* + * 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.ArrayList; +import java.util.List; +import java.util.Random; + +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.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The quantized WordPiece model directory contract: after {@link ModelQuantizer} runs and the + * safetensors is removed, the directory loads from {@code model.quantized}, embeds and ranks + * like the float model up to the quantization error, and carries per-token pooling weights + * through the quantized file. A directory holding both matrix files is rejected. + */ +class StaticEmbeddingModelQuantizedTest { + + private static final int DIMENSION = 32; + private static final long SEED = 7L; + private static final String[] WORDS = { + "hello", "world", "apple", "banana", "cherry", "river", "mountain", "guitar", + "piano", "silver", "copper", "window" + }; + private static final String[] SENTENCES = { + "hello world", "apple banana cherry", "a guitar by the river", + "no vocabulary hit here matches nothing" + }; + + /** + * Writes a small WordPiece model directory. + * + * @param directory The directory to write into. + * @param withWeights Whether to bundle a per-token {@code weights} tensor. + * @throws IOException Thrown if writing fails. + */ + private static void writeModelDirectory(Path directory, boolean withWeights) + throws IOException { + final List vocabulary = new ArrayList<>(List.of("[UNK]", "[CLS]", "[SEP]")); + vocabulary.addAll(List.of(WORDS)); + Files.write(directory.resolve("vocab.txt"), vocabulary); + final Random random = new Random(11); + final float[][] rows = new float[vocabulary.size()][DIMENSION]; + for (final float[] row : rows) { + final float rowScale = 0.5f + 2f * random.nextFloat(); + for (int d = 0; d < DIMENSION; d++) { + row[d] = rowScale * (float) random.nextGaussian(); + } + } + if (withWeights) { + final float[] weights = new float[vocabulary.size()]; + for (int row = 0; row < weights.length; row++) { + weights[row] = 0.5f + 1.5f * random.nextFloat(); + } + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", rows), + SafetensorsTestFiles.vector(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME, weights)); + } else { + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", rows)); + } + Files.writeString(directory.resolve(ModelFileNames.CONFIG), "{\"normalize\": true}"); + Files.writeString(directory.resolve(ModelFileNames.TOKENIZER_CONFIG), + "{\"do_lower_case\": true}"); + } + + /** + * Quantizes the directory and removes the safetensors, leaving the quantized deployment the + * loader accepts. + * + * @param directory The model directory to quantize in place. + * @param bits The bit width. + * @return What the quantizer measured. + * @throws IOException Thrown if quantizing or deleting fails. + */ + private static ModelQuantizer.Result deployQuantized(Path directory, int bits) + throws IOException { + final ModelQuantizer.Result result = ModelQuantizer.quantize(directory, bits, SEED); + Files.delete(directory.resolve(ModelFileNames.SAFETENSORS)); + return result; + } + + @ParameterizedTest + @ValueSource(ints = {2, 3, 4}) + void testQuantizedDirectoryEmbedsLikeTheFloatModel(int bits, @TempDir Path directory) + throws IOException { + writeModelDirectory(directory, false); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + deployQuantized(directory, bits); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + assertEquals(floatModel.dimension(), quantizedModel.dimension()); + // Pooling several rows accumulates independent quantization noise, so the pooled cosine + // sits a little below the single-row reconstruction; the floor loosens as bits shrink. + // Every value is far above what a broken rotation, grid, or scale would produce. + final double threshold = switch (bits) { + case 2 -> 0.88; + case 3 -> 0.95; + default -> 0.98; + }; + for (final String text : SENTENCES) { + final double cosine = cosine(floatModel.embed(text), quantizedModel.embed(text)); + if (Double.isNaN(cosine)) { + // Both pooled to the zero vector: an out-of-vocabulary text, identical behavior. + continue; + } + assertTrue(cosine >= threshold, + bits + "-bit embedding of '" + text + "' drifted to cosine " + cosine); + } + } + + @Test + void testQuantizedDirectoryLoadsAfterTheSafetensorsIsRemoved(@TempDir Path directory) + throws IOException { + writeModelDirectory(directory, false); + deployQuantized(directory, 4); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(directory); + assertEquals(DIMENSION, model.dimension()); + assertEquals(WORDS.length + 3, model.vocabularySize()); + // The model's own row is its nearest neighbor, so ranking works end to end. + assertEquals("hello", model.mostSimilar("hello", 1).get(0).token()); + } + + @Test + void testBothMatrixFilesPresentIsRejected(@TempDir Path directory) throws IOException { + writeModelDirectory(directory, false); + // ModelQuantizer writes model.quantized next to model.safetensors and leaves both. + ModelQuantizer.quantize(directory, 4, SEED); + assertTrue(Files.isRegularFile(directory.resolve(ModelFileNames.SAFETENSORS))); + assertTrue(Files.isRegularFile(directory.resolve(ModelFileNames.QUANTIZED))); + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(directory)); + assertTrue(e.getMessage().contains("has both"), e.getMessage()); + assertTrue(e.getMessage().contains(ModelFileNames.QUANTIZED), e.getMessage()); + assertTrue(e.getMessage().contains(ModelFileNames.SAFETENSORS), e.getMessage()); + } + + @Test + void testPoolingWeightsRideThroughTheQuantizedFile(@TempDir Path directory) + throws IOException { + writeModelDirectory(directory, true); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + final ModelQuantizer.Result result = deployQuantized(directory, 4); + assertTrue(result.hasWeights(), "the weights tensor must be carried over"); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + // Weighted pooling changes the pooled direction; matching the float model closely proves + // the weights arrived, since dropping them would score a visibly different vector. + final double cosine = + cosine(floatModel.embed("hello world apple"), quantizedModel.embed("hello world apple")); + assertTrue(cosine > 0.98, "weighted quantized embedding drifted to cosine " + cosine); + } + + @Test + void testMostSimilarAgreesWithTheFloatModel(@TempDir Path directory) throws IOException { + writeModelDirectory(directory, false); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + deployQuantized(directory, 4); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + for (final String word : new String[] {"hello", "river", "copper"}) { + assertEquals(floatModel.mostSimilar(word, 1).get(0).token(), + quantizedModel.mostSimilar(word, 1).get(0).token(), + "top neighbor of '" + word + "' must survive quantization"); + } + } + + @Test + void testAnalogyRunsOverTheQuantizedTable(@TempDir Path directory) throws IOException { + writeModelDirectory(directory, false); + deployQuantized(directory, 4); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(directory); + // The three query terms are excluded, so a fourth vocabulary word comes back; the point is + // that the analogy path (query build, exclusion, scan) runs end to end over rotated space. + final List neighbors = model.analogy("hello", "world", "apple", 1); + assertEquals(1, neighbors.size()); + // The analogy excludes its own terms, so none of them comes back. + assertFalse(neighbors.get(0).token().equals("apple")); + } + + @ParameterizedTest + @CsvSource({"2, 0.9", "4, 0.98"}) + void testQuantizedFileSmallerAndVerified(int bits, double minCosine, @TempDir Path directory) + throws IOException { + writeModelDirectory(directory, false); + final ModelQuantizer.Result result = ModelQuantizer.quantize(directory, bits, SEED); + assertTrue(result.quantizedBytes() < result.safetensorsBytes(), + result.quantizedBytes() + " must be smaller than " + result.safetensorsBytes()); + assertEquals(result.rowCount(), result.sampledRows(), + "a small table is verified row by row"); + assertTrue(result.meanCosine() > minCosine, + bits + "-bit reconstruction reported mean cosine " + result.meanCosine()); + } + + @Test + void testRowCountMismatchFailsLoud(@TempDir Path directory) throws IOException { + writeModelDirectory(directory, false); + deployQuantized(directory, 4); + final Path vocabularyFile = directory.resolve("vocab.txt"); + final List extended = new ArrayList<>(Files.readAllLines(vocabularyFile)); + extended.add("straggler"); + Files.write(vocabularyFile, extended); + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(directory)); + assertTrue(e.getMessage().contains("do not belong to the same model"), e.getMessage()); + } + + @Test + void testQuantizerRequiresTheSafetensors(@TempDir Path directory) { + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> ModelQuantizer.quantize(directory, 4, SEED)); + assertTrue(e.getMessage().contains(ModelFileNames.SAFETENSORS), e.getMessage()); + } + + @Test + void testQuantizerRejectsBadBitWidths(@TempDir Path directory) throws IOException { + writeModelDirectory(directory, false); + assertThrows(IllegalArgumentException.class, () -> ModelQuantizer.quantize(directory, 1, SEED)); + assertThrows(IllegalArgumentException.class, () -> ModelQuantizer.quantize(directory, 5, SEED)); + } + + @Test + void testQuantizedEmbeddingsAreDeterministic(@TempDir Path first, @TempDir Path second) + throws IOException { + writeModelDirectory(first, false); + writeModelDirectory(second, false); + deployQuantized(first, 4); + deployQuantized(second, 4); + final StaticEmbeddingModel modelA = StaticEmbeddingModel.load(first); + final StaticEmbeddingModel modelB = StaticEmbeddingModel.load(second); + for (final String text : SENTENCES) { + assertArrayEquals(modelA.embed(text), modelB.embed(text), 0f, + "the same table, bits, and seed must embed bit-identically"); + } + } + + /** + * {@return the cosine between two vectors, or {@code Double.NaN} when either has no + * direction} + * + * @param a The first vector. + * @param b The second vector, of the same length. + */ + private static double cosine(float[] a, float[] b) { + return ModelQuantizer.cosine(a, 0, a.length, b); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceQuantizedTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceQuantizedTest.java new file mode 100644 index 0000000000..5373bc986b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceQuantizedTest.java @@ -0,0 +1,104 @@ +/* + * 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 org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.util.InvalidFormatException; + +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 quantized SentencePiece loading path: a SentencePiece directory whose matrix has been + * quantized and whose safetensors removed loads through {@code loadSentencePieceQuantized}, + * resolves rows by piece string across the tokenizer id offset, and embeds like the float + * model up to the quantization error. Both matrix files present is rejected. + */ +class StaticEmbeddingModelSentencePieceQuantizedTest { + + private static final int DIMENSION = 16; + private static final long SEED = 7L; + + private static SentencePieceModelFixture fixture; + + @BeforeAll + static void loadFixture() throws IOException { + fixture = new SentencePieceModelFixture(); + } + + @Test + void testQuantizedSentencePieceEmbedsLikeTheFloatModel(@TempDir Path directory) + throws IOException { + fixture.write(directory, DIMENSION, true, SEED); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + ModelQuantizer.quantize(directory, 4, SEED); + Files.delete(directory.resolve(ModelFileNames.SAFETENSORS)); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + assertEquals(floatModel.dimension(), quantizedModel.dimension()); + assertEquals(floatModel.vocabularySize(), quantizedModel.vocabularySize()); + for (final String text : new String[] {"a", "the model", "hello there world"}) { + final double cosine = cosine(floatModel.embed(text), quantizedModel.embed(text)); + if (Double.isNaN(cosine)) { + continue; + } + assertTrue(cosine > 0.97, + "quantized SentencePiece embedding of '" + text + "' drifted to cosine " + cosine); + } + } + + @Test + void testQuantizedSentencePieceRanksLikeTheFloatModel(@TempDir Path directory) + throws IOException { + fixture.write(directory, DIMENSION, true, SEED); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + ModelQuantizer.quantize(directory, 4, SEED); + Files.delete(directory.resolve(ModelFileNames.SAFETENSORS)); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + // A piece is its own nearest neighbor under both storage forms. + final String piece = fixture.rowPieces().get(3); + assertEquals(floatModel.mostSimilar(piece, 1).get(0).token(), + quantizedModel.mostSimilar(piece, 1).get(0).token()); + } + + @Test + void testBothMatrixFilesPresentIsRejected(@TempDir Path directory) throws IOException { + fixture.write(directory, DIMENSION, true, SEED); + ModelQuantizer.quantize(directory, 4, SEED); + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(directory)); + assertTrue(e.getMessage().contains("has both"), e.getMessage()); + } + + /** + * {@return the cosine between two vectors, or {@code Double.NaN} when either has no + * direction} + * + * @param a The first vector. + * @param b The second vector, of the same length. + */ + private static double cosine(float[] a, float[] b) { + return ModelQuantizer.cosine(a, 0, a.length, b); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceTest.java new file mode 100644 index 0000000000..5aaec54d22 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceTest.java @@ -0,0 +1,312 @@ +/* + * 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.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.embeddings.StaticEmbeddingModel.Normalization; +import opennlp.subword.sentencepiece.SentencePieceTokenizer; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The SentencePiece loading path, exercised against a real trained tiny model (a copy of the + * opennlp-subword test fixture). The matrix vocabulary is written the way a distillation ships + * it: control pieces dropped, rows ordered differently from the tokenizer's ids, extra special + * rows in front, and an extra token appended through {@code added_tokens}; every lookup must + * therefore go by piece string, never by tokenizer id. + */ +class StaticEmbeddingModelSentencePieceTest { + + private static final String MODEL_RESOURCE = "/opennlp/embeddings/tiny-unigram.model"; + private static final int DIMENSION = 4; + + private static byte[] modelBytes; + private static SentencePieceTokenizer tokenizer; + // The matrix rows: , , then every poolable tokenizer piece, then . + private static List rows; + + @BeforeAll + static void loadFixture() throws IOException { + try (InputStream in = + StaticEmbeddingModelSentencePieceTest.class.getResourceAsStream(MODEL_RESOURCE)) { + modelBytes = in.readAllBytes(); + } + tokenizer = SentencePieceTokenizer.load( + StaticEmbeddingModelSentencePieceTest.class.getResourceAsStream(MODEL_RESOURCE)); + rows = new ArrayList<>(); + rows.add(""); + rows.add(""); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (!tokenizer.isControl(id) && !tokenizer.isUnknown(id)) { + rows.add(tokenizer.idToPiece(id)); + } + } + } + + /** + * {@return the value at {@code (row, d)} of the deterministic test matrix} + * + * @param row The matrix row. + * @param d The dimension index. + */ + private static float cell(int row, int d) { + return row + d * 0.25f; + } + + /** + * Writes the three SentencePiece-layout files (and optionally a {@code config.json}) into a + * directory: the copied {@code .model}, a synthesized Unigram {@code tokenizer.json} whose + * vocabulary is {@link #rows} with one token appended via {@code added_tokens}, and a + * deterministic embedding matrix with one extra row for it. + * + * @param dir The directory to write into. + * @param normalize The {@code config.json} normalize value, or {@code null} to omit the file. + * @return The directory. + * @throws IOException Thrown if writing fails. + */ + private static Path writeModelDirectory(Path dir, Boolean normalize) throws IOException { + Files.write(dir.resolve("sentencepiece.bpe.model"), modelBytes); + Files.writeString(dir.resolve("tokenizer.json"), tokenizerJson(rows)); + final float[][] matrix = new float[rows.size() + 1][DIMENSION]; + for (int row = 0; row < matrix.length; row++) { + for (int d = 0; d < DIMENSION; d++) { + matrix[row][d] = cell(row, d); + } + } + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + if (normalize != null) { + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":" + normalize + "}"); + } + return dir; + } + + /** + * {@return a Unigram {@code tokenizer.json} whose vocabulary is the given pieces plus an + * appended added token} + * + *

The appended token is not named {@code } because the fixture model itself defines + * {@code } as a user-defined piece, which already owns a row.

+ * + * @param pieces The {@code model.vocab} pieces in row order. + */ + private static String tokenizerJson(List pieces) { + final StringBuilder json = new StringBuilder("{\"version\":\"1.0\",\"added_tokens\":["); + json.append("{\"id\":0,\"content\":\"\",\"special\":true},"); + json.append("{\"id\":").append(pieces.size()).append(",\"content\":\"\"," + + "\"special\":true}],"); + json.append("\"normalizer\":{\"type\":\"Precompiled\"},\"model\":{\"type\":\"Unigram\"," + + "\"unk_id\":1,\"vocab\":["); + for (int i = 0; i < pieces.size(); i++) { + if (i > 0) { + json.append(','); + } + json.append('[').append(EmbeddingTestFixtures.jsonString(pieces.get(i))) + .append(",-").append(i % 7).append(".5]"); + } + return json.append("]}}").toString(); + } + + @Test + void testEmbedGathersRowsByPieceStringAcrossTheIdOffset(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = loadFromDirectory(writeModelDirectory(dir, null)); + + // "a" segments to the single piece U+2581 + "a"; the embedding must be exactly that piece's + // matrix row, found by string in the reordered vocabulary, not by the tokenizer's id. + final List pieces = tokenizer.encode("a"); + assertEquals(1, pieces.size()); + final int row = rows.indexOf(pieces.get(0).piece()); + assertTrue(row >= 2, "the fixture row must sit above the injected specials"); + final float[] expected = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + expected[d] = cell(row, d); + } + assertArrayEquals(expected, model.embed("a"), 1e-5f); + } + + @Test + void testEmbedMeanPoolsAllMappedPieces(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = loadFromDirectory(writeModelDirectory(dir, null)); + + // Expected: the mean over every non-control, non-unknown piece's row, resolved by string. + final List pieces = tokenizer.encode("Hello world"); + final float[] expected = new float[DIMENSION]; + int pooled = 0; + for (final SubwordPiece piece : pieces) { + if (tokenizer.isControl(piece.id()) || tokenizer.isUnknown(piece.id())) { + continue; + } + final int row = rows.indexOf(piece.piece()); + assertTrue(row >= 0, "fixture piece '" + piece.piece() + "' must have a row"); + for (int d = 0; d < DIMENSION; d++) { + expected[d] += cell(row, d); + } + pooled++; + } + assertTrue(pooled > 1, "the fixture text must pool more than one piece"); + for (int d = 0; d < DIMENSION; d++) { + expected[d] /= pooled; + } + assertArrayEquals(expected, model.embed("Hello world"), 1e-4f); + } + + @Test + void testUnknownPiecesAreSkippedInPooling(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = loadFromDirectory(writeModelDirectory(dir, null)); + + // The euro sign is outside the tiny training corpus, so it segments to the dummy-prefix + // piece plus an unknown piece carrying the surface text. The unknown piece's string is not + // a vocabulary entry, so pooling must skip it by its id, leaving only the mapped pieces. + final List pieces = tokenizer.encode("\u20AC"); + final float[] expected = new float[DIMENSION]; + int pooled = 0; + int unknown = 0; + for (final SubwordPiece piece : pieces) { + if (tokenizer.isUnknown(piece.id())) { + unknown++; + continue; + } + if (tokenizer.isControl(piece.id())) { + continue; + } + final int row = rows.indexOf(piece.piece()); + for (int d = 0; d < DIMENSION; d++) { + expected[d] += cell(row, d); + } + pooled++; + } + assertTrue(unknown > 0, "fixture assumption: the euro sign must produce an unknown piece"); + for (int d = 0; d < DIMENSION; d++) { + expected[d] /= Math.max(pooled, 1); + } + assertArrayEquals(expected, model.embed("\u20AC"), 1e-5f); + } + + @Test + void testDirectoryLoadDetectsTheSentencePieceLayout(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeModelDirectory(dir, true)); + + assertEquals(DIMENSION, model.dimension()); + assertEquals(rows.size() + 1, model.vocabularySize()); + // normalize=true from config.json: the pooled vector must have unit length. + final float[] vector = model.embed("a"); + double normSquared = 0; + for (final float v : vector) { + normSquared += (double) v * v; + } + assertEquals(1.0, Math.sqrt(normSquared), 1e-5); + } + + @Test + void testMostSimilarNeverReturnsSpecialRows(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = loadFromDirectory(writeModelDirectory(dir, null)); + + for (final Neighbor neighbor : model.mostSimilar("a", 5)) { + assertFalse(List.of("", "", "").contains(neighbor.token()), + "special row leaked into neighbors: " + neighbor.token()); + } + } + + @Test + void testLoadRejectsAVocabularyMissingAPoolablePiece(@TempDir Path dir) throws IOException { + writeModelDirectory(dir, null); + // Remove one poolable piece from the matrix vocabulary; the matrix shrinks with it, so only + // the coverage check can catch the mismatch. + final List truncated = new ArrayList<>(rows); + truncated.remove(truncated.size() - 1); + Files.writeString(dir.resolve("tokenizer.json"), tokenizerJson(truncated)); + final float[][] matrix = new float[truncated.size() + 1][DIMENSION]; + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> loadFromDirectory(dir)); + assertTrue(e.getMessage().contains("do not belong"), e.getMessage()); + } + + @Test + void testLoadRejectsARowCountMismatch(@TempDir Path dir) throws IOException { + writeModelDirectory(dir, null); + final float[][] matrix = new float[rows.size()][DIMENSION]; + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> loadFromDirectory(dir)); + assertTrue(e.getMessage().contains("rows"), e.getMessage()); + } + + @Test + void testDirectoryLoadExplainsAnIncompleteLegacyTokenizer(@TempDir Path dir) + throws IOException { + writeModelDirectory(dir, true); + Files.delete(dir.resolve("sentencepiece.bpe.model")); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("self-contained tokenizer.json"), e.getMessage()); + assertTrue(e.getMessage().contains("trained SentencePiece .model"), e.getMessage()); + } + + @Test + void testLoadSentencePieceRejectsNullArguments(@TempDir Path dir) throws IOException { + writeModelDirectory(dir, null); + final Path model = dir.resolve("sentencepiece.bpe.model"); + final Path json = dir.resolve("tokenizer.json"); + final Path tensors = dir.resolve("model.safetensors"); + + assertThrows(IllegalArgumentException.class, () -> + StaticEmbeddingModel.loadSentencePiece(null, json, tensors, Normalization.NONE)); + assertThrows(IllegalArgumentException.class, () -> + StaticEmbeddingModel.loadSentencePiece(model, null, tensors, Normalization.NONE)); + assertThrows(IllegalArgumentException.class, () -> + StaticEmbeddingModel.loadSentencePiece(model, json, null, Normalization.NONE)); + assertThrows(IllegalArgumentException.class, () -> + StaticEmbeddingModel.loadSentencePiece(model, json, tensors, null)); + } + + /** + * Loads through the explicit SentencePiece overload from a directory written by + * {@link #writeModelDirectory(Path, Boolean)}. + * + * @param dir The model directory. + * @return The loaded model. + * @throws IOException Thrown if reading fails. + */ + private static StaticEmbeddingModel loadFromDirectory(Path dir) throws IOException { + return StaticEmbeddingModel.loadSentencePiece(dir.resolve("sentencepiece.bpe.model"), + dir.resolve("tokenizer.json"), dir.resolve("model.safetensors"), Normalization.NONE); + } +} 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 new file mode 100644 index 0000000000..00d872e238 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java @@ -0,0 +1,232 @@ +/* + * 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.time.Duration; +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; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Exercises {@link StaticEmbeddingModel#similarity}, {@link StaticEmbeddingModel#mostSimilar}, + * and {@link StaticEmbeddingModel#analogy} against {@link EmbeddingTestFixtures}' analogy table, + * whose vectors point in genuinely different directions (unlike {@link StaticEmbeddingModelTest}'s + * collinear rows, which are ideal for pooling-math assertions but would make every pairwise cosine + * similarity trivially 1.0). + */ +class StaticEmbeddingModelSimilarityTest { + + private static StaticEmbeddingModel load(Path dir) throws IOException { + return EmbeddingTestFixtures.loadAnalogyModel(dir, Normalization.NONE); + } + + @Test + void testSimilarityOfIdenticalTextIsOne(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertEquals(1.0, model.similarity("king", "king"), 1e-5); + } + + @Test + void testSimilarityIsSymmetric(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertEquals(model.similarity("king", "queen"), model.similarity("queen", "king"), 1e-9); + } + + @Test + void testSimilarityOfUnrelatedTermsIsLow(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertTrue(model.similarity("king", "apple") < model.similarity("king", "queen")); + } + + @Test + void testSimilarityOfOutOfVocabularyTextIsZero(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertEquals(0.0, model.similarity("xyzzy", "king"), 1e-9); + } + + @Test + void testMostSimilarFindsSelfAsTopMatch(@TempDir Path dir) throws IOException { + // Documented, deliberate behavior: unlike gensim's convention of excluding the query word, + // mostSimilar only excludes special tokens, so a single-word query's own vocabulary row is + // (correctly) its own nearest neighbor. + final StaticEmbeddingModel model = load(dir); + + final List result = model.mostSimilar("king", 1); + + assertEquals(1, result.size()); + assertEquals("king", result.get(0).token()); + assertEquals(1.0, result.get(0).similarity(), 1e-5); + } + + @Test + void testMostSimilarExcludesSpecialTokensAndOrdersByDescendingSimilarity(@TempDir Path dir) + throws IOException { + final StaticEmbeddingModel model = load(dir); + + final List result = model.mostSimilar("king", 5); + + assertEquals(5, result.size()); + for (final Neighbor neighbor : result) { + assertFalse(List.of("[CLS]", "[SEP]", "[UNK]").contains(neighbor.token())); + } + // Descending order. + for (int i = 1; i < result.size(); i++) { + assertTrue(result.get(i - 1).similarity() >= result.get(i).similarity()); + } + // apple is the clear outlier (opposite-ish direction) and must rank last. + assertEquals("apple", result.get(result.size() - 1).token()); + } + + @Test + void testMostSimilarClampsTopKToTheVocabularySize(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + // topK sizes the candidate arrays, so it must be clamped to the vocabulary before + // allocation; unclamped, Integer.MAX_VALUE fails with OutOfMemoryError. The fixture has + // 8 rows, 3 of them special, so any request larger than the vocabulary returns the same + // 5 neighbors a topK of 8 would. + final List result = assertTimeoutPreemptively(Duration.ofSeconds(10), + () -> model.mostSimilar("king", Integer.MAX_VALUE)); + + assertEquals(model.vocabularySize() - 3, result.size()); + assertEquals("king", result.get(0).token()); + assertEquals(result, model.mostSimilar("king", model.vocabularySize())); + } + + @Test + void testMostSimilarOfZeroVectorQueryReturnsEmptyList(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertEquals(List.of(), model.mostSimilar("xyzzy", 3)); + } + + @Test + void testAnalogyFindsTheExactTarget(@TempDir Path dir) throws IOException { + // man is to king as woman is to ? Expected: queen (king - man + woman == queen exactly). + final StaticEmbeddingModel model = load(dir); + + final List result = model.analogy("man", "king", "woman", 1); + + assertEquals(1, result.size()); + assertEquals("queen", result.get(0).token()); + assertEquals(1.0, result.get(0).similarity(), 1e-5); + } + + @Test + void testAnalogyExcludesItsOwnInputTerms(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + // Only "queen" and "apple" remain eligible once man/king/woman and the special tokens are + // excluded, regardless of how close the raw analogy target vector is to the inputs. + final List result = model.analogy("man", "king", "woman", 4); + + assertEquals(2, result.size()); + assertFalse(result.stream().map(Neighbor::token) + .anyMatch(token -> List.of("man", "king", "woman").contains(token))); + } + + @Test + void testAnalogyToleratesEqualTerms(@TempDir Path dir) throws IOException { + // Repeating a term is legal: b - a + c with a == b is just c's vector, so with man and woman + // excluded the exactly collinear queen must win. + final StaticEmbeddingModel model = load(dir); + + final List result = model.analogy("man", "man", "woman", 2); + + assertEquals("queen", result.get(0).token()); + assertEquals(1.0, result.get(0).similarity(), 1e-5); + } + + @Test + void testAnalogyExclusionFoldsLikeEmbed(@TempDir Path dir) throws IOException { + // The exclusion folds terms through the model's own tokenizer, so on an uncased model a + // capitalized input excludes its lower-cased vocabulary row rather than handing it back. + final StaticEmbeddingModel model = load(dir); + + final List result = model.analogy("Man", "King", "Woman", 4); + + assertEquals(2, result.size()); + assertEquals("queen", result.get(0).token()); + assertFalse(result.stream().map(Neighbor::token) + .anyMatch(token -> List.of("man", "king", "woman").contains(token))); + } + + @Test + void testZeroVectorRowScoresZeroNotNaN(@TempDir Path dir) throws IOException { + // A non-special all-zero row has no direction; it must score exactly 0.0, not the NaN a + // naive 0/0 cosine would produce. + final Path vocab = dir.resolve("zero-vocab.txt"); + Files.write(vocab, List.of("[CLS]", "[SEP]", "[UNK]", "a", "zero")); + final float[][] rows = {{0f, 0f}, {0f, 0f}, {0f, 0f}, {1f, 0f}, {0f, 0f}}; + final Path tensors = dir.resolve("zero-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", rows)); + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(vocab, tensors, Casing.UNCASED, Normalization.NONE); + + final List result = model.mostSimilar("a", 5); + + assertEquals(2, result.size()); + assertEquals("a", result.get(0).token()); + assertEquals("zero", result.get(1).token()); + assertEquals(0.0, result.get(1).similarity()); + assertTrue(result.stream().allMatch(neighbor -> Double.isFinite(neighbor.similarity()))); + } + + @Test + void testMostSimilarRejectsInvalidArguments(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertThrows(IllegalArgumentException.class, () -> model.mostSimilar(null, 1)); + assertThrows(IllegalArgumentException.class, () -> model.mostSimilar("king", 0)); + assertThrows(IllegalArgumentException.class, () -> model.mostSimilar("king", -1)); + } + + @Test + void testAnalogyRejectsInvalidArguments(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertThrows(IllegalArgumentException.class, () -> model.analogy(null, "king", "woman", 1)); + assertThrows(IllegalArgumentException.class, () -> model.analogy("man", null, "woman", 1)); + assertThrows(IllegalArgumentException.class, () -> model.analogy("man", "king", null, 1)); + assertThrows(IllegalArgumentException.class, () -> model.analogy("man", "king", "woman", 0)); + } + + @Test + void testSimilarityRejectsNullArguments(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = load(dir); + + assertThrows(IllegalArgumentException.class, () -> model.similarity(null, "king")); + assertThrows(IllegalArgumentException.class, () -> model.similarity("king", null)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTermTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTermTest.java new file mode 100644 index 0000000000..3f42831e22 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTermTest.java @@ -0,0 +1,187 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * A model directory with a term table: term rows pool as single units where they match, the + * subword path is untouched everywhere else, terms appear as similarity-search neighbors, and a + * malformed or mismatched terms file fails loud at load time. + */ +class StaticEmbeddingModelTermTest { + + private static final List VOCABULARY = + List.of("[CLS]", "[SEP]", "[UNK]", "habeas", "corpus", "writ", "law"); + + /** + * The matrix rows: the three special tokens are zero, the content tokens have distinct + * directions, and the two term rows (habeas corpus, replevin) are distinct again. + */ + private static final float[][] ROWS = { + {0f, 0f}, // [CLS] + {0f, 0f}, // [SEP] + {0f, 0f}, // [UNK] + {1f, 0f}, // habeas + {0f, 1f}, // corpus + {2f, 0f}, // writ + {4f, 0f}, // law + {10f, 10f}, // term: habeas corpus + {5f, -5f}, // term: replevin + }; + + /** + * Writes a loadable WordPiece directory, optionally with the two term rows and their + * {@code terms.txt}, and loads it. + * + * @param dir The directory to write into. + * @param withTerms Whether to include the term rows and the terms file. + * @return The loaded model. + * @throws IOException Thrown if writing or loading fails. + */ + private static StaticEmbeddingModel model(Path dir, boolean withTerms) throws IOException { + Files.write(dir.resolve("vocab.txt"), VOCABULARY); + final int rows = withTerms ? ROWS.length : VOCABULARY.size(); + final float[][] matrix = new float[rows][]; + System.arraycopy(ROWS, 0, matrix, 0, rows); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + Files.writeString(dir.resolve("tokenizer_config.json"), "{\"do_lower_case\":true}"); + if (withTerms) { + Files.write(dir.resolve("terms.txt"), List.of("habeas corpus", "replevin")); + } + return StaticEmbeddingModel.load(dir); + } + + @Test + void testLoadsTheTermTable(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = model(dir, true); + assertEquals(VOCABULARY.size(), model.vocabularySize()); + assertEquals(2, model.termCount()); + } + + @Test + void testAMatchedTermPoolsItsSingleRow(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = model(dir, true); + assertArrayEquals(new float[] {10f, 10f}, model.embed("habeas corpus")); + // Case folding and punctuation between the words do not break the match. + assertArrayEquals(new float[] {10f, 10f}, model.embed("Habeas-Corpus!")); + // A single-word term matches ahead of its (absent) subword pieces. + assertArrayEquals(new float[] {5f, -5f}, model.embed("replevin")); + } + + @Test + void testTermAndPieceRowsPoolTogetherInTextOrder(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = model(dir, true); + // writ -> its piece row; "of" -> [UNK], skipped; "habeas corpus" -> the term row; + // law -> its piece row. Mean of (2,0), (10,10), (4,0). + assertArrayEquals(new float[] {16f / 3, 10f / 3}, model.embed("writ of habeas corpus law")); + } + + @Test + void testTextWithoutAMatchEmbedsExactlyLikeATermlessModel(@TempDir Path dir, + @TempDir Path termless) + throws IOException { + final StaticEmbeddingModel withTerms = model(dir, true); + final StaticEmbeddingModel without = model(termless, false); + // "habeas law" has both words in the vocabulary but matches no term: the two words are not + // adjacent words of any stored phrase. + assertArrayEquals(without.embed("habeas law"), withTerms.embed("habeas law")); + assertArrayEquals(without.embed("the writ, of law."), withTerms.embed("the writ, of law.")); + } + + @Test + void testATermlessModelZeroesWhatOnlyATermRowCouldEmbed(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel without = model(dir, false); + // Without the term table, "replevin" is out of vocabulary entirely. + assertArrayEquals(new float[] {0f, 0f}, without.embed("replevin")); + } + + @Test + void testTermsAreSimilarityNeighbors(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = model(dir, true); + final List neighbors = model.mostSimilar("replevin", 1); + assertEquals(1, neighbors.size()); + assertEquals("replevin", neighbors.get(0).token()); + assertEquals("habeas corpus", model.mostSimilar("habeas corpus", 1).get(0).token()); + } + + @Test + void testARowCountMismatchWithTermsFailsLoud(@TempDir Path dir) throws IOException { + Files.write(dir.resolve("vocab.txt"), VOCABULARY); + final float[][] matrix = new float[VOCABULARY.size()][]; + System.arraycopy(ROWS, 0, matrix, 0, VOCABULARY.size()); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", matrix)); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + Files.writeString(dir.resolve("tokenizer_config.json"), "{\"do_lower_case\":true}"); + Files.write(dir.resolve("terms.txt"), List.of("habeas corpus", "replevin")); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("plus 2 terms"), e.getMessage()); + } + + @Test + void testAMalformedTermsFileFailsLoud(@TempDir Path dir) throws IOException { + Files.write(dir.resolve("vocab.txt"), VOCABULARY); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + Files.writeString(dir.resolve("tokenizer_config.json"), "{\"do_lower_case\":true}"); + // Upper case is not the normalized form the matcher folds to. + Files.write(dir.resolve("terms.txt"), List.of("HABEAS CORPUS", "replevin")); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("normalized form"), e.getMessage()); + } + + @Test + void testASentencePieceDirectoryLoadsItsTermTable(@TempDir Path dir) throws IOException { + EmbeddingTestFixtures.writeSentencePieceDirectory(dir, List.of("lawbook")); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false}"); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + assertEquals(1, model.termCount()); + // The fixture's cell formula is row + d * 0.25, and the term owns the row after the + // vocabulary rows. + final float[] expected = new float[EmbeddingTestFixtures.SENTENCEPIECE_DIMENSION]; + for (int d = 0; d < expected.length; d++) { + expected[d] = model.vocabularySize() + d * 0.25f; + } + assertArrayEquals(expected, model.embed("Lawbook")); + } +} 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 new file mode 100644 index 0000000000..c721f3985a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java @@ -0,0 +1,477 @@ +/* + * 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 opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; +import opennlp.tools.embeddings.TextEmbedder; +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class StaticEmbeddingModelTest { + + // Fixture vocabulary: [CLS]=0, [SEP]=1, [UNK]=2, hello=3, world=4, cat=5. + private static final List VOCAB_TOKENS = + List.of("[CLS]", "[SEP]", "[UNK]", "hello", "world", "cat"); + private static final int DIMENSION = 3; + + // Row i is [i, i*10, i*100], so hand-computed expected pooled vectors are easy to verify. + private static final float[][] ROWS = { + {0f, 0f, 0f}, // [CLS] + {1f, 10f, 100f}, // [SEP] + {2f, 20f, 200f}, // [UNK] + {3f, 30f, 300f}, // hello + {4f, 40f, 400f}, // world + {5f, 50f, 500f}, // cat + }; + + private static Path writeVocab(Path dir) throws IOException { + final Path file = dir.resolve("vocab.txt"); + Files.write(file, VOCAB_TOKENS); + return file; + } + + private static Path writeSafetensors(Path dir, boolean withWeights) throws IOException { + 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. + 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)); + } + return file; + } + + private static Path writeSafetensorsF16(Path dir) throws IOException { + final Path file = dir.resolve("model.safetensors"); + SafetensorsTestFiles.write(file, "F16", SafetensorsTestFiles.matrix("embeddings", ROWS)); + return file; + } + + @Test + void testEmbedMeanPoolsWithoutWeights(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + final float[] result = model.embed("hello world"); + + // (hello + world) / 2 = ([3,30,300] + [4,40,400]) / 2 = [3.5, 35, 350] + assertArrayEquals(new float[] {3.5f, 35f, 350f}, result, 1e-5f); + } + + @Test + void testLoadsAnF16EmbeddingMatrix(@TempDir Path dir) throws IOException { + // model2vec writes float16 by default, so the loader must accept it and widen to float. + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensorsF16(dir), + Casing.UNCASED, Normalization.NONE); + + assertEquals(DIMENSION, model.dimension()); + // (hello + world) / 2 = [3.5, 35, 350]; the row values are all exact in IEEE half. + assertArrayEquals(new float[] {3.5f, 35f, 350f}, model.embed("hello world"), 1e-2f); + } + + @Test + void testLoadsAModelWhoseVocabularyDroppedTheFrameTokens(@TempDir Path dir) throws IOException { + // Model2Vec mean-pools content pieces and never frames, so it removes [CLS]/[SEP] from the + // distilled table, keeping only [PAD]/[UNK]. Such a table must still load; the loader caches + // the frame onto the unknown row and pooling skips it. The content rows below carry the same + // values as the framed fixture, so the embedding must match it piece for piece. + final List tokens = List.of("[PAD]", "[UNK]", "hello", "world", "cat"); + final float[][] rows = { + {9f, 9f, 9f}, // [PAD], never pooled + {8f, 8f, 8f}, // [UNK], never pooled + {3f, 30f, 300f}, // hello, same as the framed fixture's row + {4f, 40f, 400f}, // world, same as the framed fixture's row + {5f, 50f, 500f}, // cat, same as the framed fixture's row + }; + final Path vocab = dir.resolve("vocab.txt"); + Files.write(vocab, tokens); + final Path tensors = dir.resolve("model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", rows)); + + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(vocab, tensors, Casing.UNCASED, Normalization.NONE); + + // (hello + world) / 2, identical to testEmbedMeanPoolsWithoutWeights: the cached frame and + // any [UNK] are skipped, so only the two content pieces pool. + assertArrayEquals(new float[] {3.5f, 35f, 350f}, model.embed("hello world"), 1e-5f); + // "xyzzy" folds to [UNK] and is dropped, leaving just "cat". + assertArrayEquals(new float[] {5f, 50f, 500f}, model.embed("cat xyzzy"), 1e-5f); + // Text with no content pieces is a zero vector, not the frame or [UNK] vector. + assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed("xyzzy"), 1e-5f); + // The unknown row must never surface as a neighbor. + for (final Neighbor neighbor : model.mostSimilar("cat", 4)) { + assertTrue(!"[UNK]".equals(neighbor.token()) && !"[PAD]".equals(neighbor.token()), + "a special row leaked into neighbors: " + neighbor.token()); + } + } + + @Test + void testRejectsAWordPieceVocabularyWithoutUnknownToken(@TempDir Path dir) throws IOException { + final List tokens = List.of("[CLS]", "[SEP]", "hello", "world"); + final float[][] rows = {{0f, 0f, 0f}, {1f, 1f, 1f}, {2f, 2f, 2f}, {3f, 3f, 3f}}; + final Path vocab = dir.resolve("vocab.txt"); + Files.write(vocab, tokens); + final Path tensors = dir.resolve("model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", rows)); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocab, tensors, Casing.UNCASED, Normalization.NONE)); + assertTrue(e.getMessage().contains("[UNK]"), e.getMessage()); + } + + @Test + void testEmbedAppliesPerTokenWeightsButDividesByTokenCount(@TempDir Path dir) + throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, true), + Casing.UNCASED, Normalization.NONE); + + final float[] result = model.embed("hello world"); + + // hello has weight 2: (2*[3,30,300] + 1*[4,40,400]) / 2 (denominator is token COUNT, not + // the sum of weights) = ([6,60,600] + [4,40,400]) / 2 = [5, 50, 500] + assertArrayEquals(new float[] {5f, 50f, 500f}, result, 1e-5f); + } + + @Test + void testEmbedNormalizesToUnitLength(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.L2); + + final float[] result = model.embed("cat"); + + double normSquared = 0; + for (final float v : result) { + normSquared += (double) v * v; + } + assertEquals(1.0, Math.sqrt(normSquared), 1e-5); + // Direction preserved: cat's raw vector is [5, 50, 500], i.e. a positive multiple of + // [1, 10, 100]; the normalized result must be that same direction. + assertTrue(result[1] / result[0] > 9.9f && result[1] / result[0] < 10.1f); + } + + @Test + void testEmbedSkipsUnknownTokens(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + 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". + final float[] result = model.embed("cat xyzzy"); + + assertArrayEquals(new float[] {5f, 50f, 500f}, result, 1e-5f); + } + + @Test + void testEmbedOfTextWithNoInVocabularyTokensIsZeroVector(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed("xyzzy"), 1e-5f); + } + + @Test + void testEmbedOfEmptyTextIsZeroVectorNotAnError(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.L2); + + assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed(""), 1e-5f); + } + + @Test + void testEmbedOfWhitespaceOnlyTextIsZeroVector(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + // Whitespace-only text produces no content pieces at all, unlike unknown text, which still + // produces a (skipped) [UNK]; both must pool to the zero vector without dividing by zero. + assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed(" \t\n "), 1e-5f); + } + + @Test + void testEmbedSkipsSupplementaryPlaneTextAsUnknown(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + // An emoji is a supplementary-plane character (a surrogate pair in Java) no vocabulary + // piece covers; it must fold to [UNK] and be skipped, leaving just "cat" in the pool. + assertArrayEquals(new float[] {5f, 50f, 500f}, model.embed("cat \uD83D\uDE00"), 1e-5f); + } + + @Test + void testDimensionAndVocabularySizeAccessors(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + assertEquals(DIMENSION, model.dimension()); + assertEquals(VOCAB_TOKENS.size(), model.vocabularySize()); + } + + @Test + void testEmbedRejectsNullText(@TempDir Path dir) throws IOException { + final StaticEmbeddingModel model = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + assertThrows(IllegalArgumentException.class, () -> model.embed(null)); + } + + @Test + void testLoadRejectsNullArguments(@TempDir Path dir) throws IOException { + final Path vocab = writeVocab(dir); + final Path tensors = writeSafetensors(dir, false); + + assertThrows(IllegalArgumentException.class, + () -> StaticEmbeddingModel.load(null, tensors, Casing.UNCASED, Normalization.NONE)); + assertThrows(IllegalArgumentException.class, + () -> StaticEmbeddingModel.load(vocab, null, Casing.UNCASED, Normalization.NONE)); + } + + @Test + void testLoadRejectsVocabularySizeMismatch(@TempDir Path dir) throws IOException { + final Path shortVocab = dir.resolve("short-vocab.txt"); + Files.write(shortVocab, List.of("[CLS]", "[SEP]", "[UNK]")); + + // Malformed model content (files that disagree) is a checked InvalidFormatException, not + // an IllegalArgumentException; the latter is reserved for caller argument errors. + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(shortVocab, writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE)); + assertTrue(e.getMessage().contains("rows")); + } + + @Test + void testLoadRejectsANonFiniteMatrixValue(@TempDir Path dir) throws IOException { + // The distiller replaces non-finite teacher values with zero before writing, so a NaN in a + // loaded matrix marks a corrupt or foreign file. It must fail loud at load time: a NaN row + // defeats both the zero-norm guard and every similarity comparison downstream. + final float[][] rows = new float[ROWS.length][]; + for (int r = 0; r < ROWS.length; r++) { + rows[r] = ROWS[r].clone(); + } + rows[4][1] = Float.NaN; + final Path vocab = writeVocab(dir); + final Path nanTensors = dir.resolve("nan.safetensors"); + SafetensorsTestFiles.write(nanTensors, SafetensorsTestFiles.matrix("embeddings", rows)); + + final InvalidFormatException nan = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocab, nanTensors, Casing.UNCASED, Normalization.NONE)); + assertTrue(nan.getMessage().contains("row 4"), nan.getMessage()); + + // An infinity is just as corrupting and must be rejected the same way. + rows[4][1] = Float.POSITIVE_INFINITY; + final Path infiniteTensors = dir.resolve("infinite.safetensors"); + SafetensorsTestFiles.write(infiniteTensors, + SafetensorsTestFiles.matrix("embeddings", rows)); + + final InvalidFormatException infinite = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocab, infiniteTensors, + Casing.UNCASED, Normalization.NONE)); + assertTrue(infinite.getMessage().contains("row 4"), infinite.getMessage()); + } + + @Test + void testLoadRejectsWeightsSizeMismatch(@TempDir Path dir) throws IOException { + // A weights tensor sized for a different (smaller) vocabulary than the embedding matrix. + final Path file = dir.resolve("mismatched.safetensors"); + SafetensorsTestFiles.write(file, + SafetensorsTestFiles.matrix("embeddings", ROWS), + SafetensorsTestFiles.vector("weights", new float[] {1f})); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(writeVocab(dir), file, Casing.UNCASED, Normalization.NONE)); + assertTrue(e.getMessage().contains("weights")); + } + + // Writes the two JSON configuration files of a published model directory alongside the + // vocab/safetensors fixtures, with the shapes real releases use (extra fields, floats, + // nested objects, an explicit strip_accents null). + private static void writeConfigs(Path dir, String normalize, String doLowerCase) + throws IOException { + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"architectures\":[\"StaticModel\"]," + + "\"apply_pca\":256,\"normalize\":" + normalize + ",\"hidden_dim\":3}"); + Files.writeString(dir.resolve("tokenizer_config.json"), + "{\"added_tokens_decoder\":{\"0\":{\"content\":\"[PAD]\",\"special\":true}}," + + "\"do_lower_case\":" + doLowerCase + ",\"strip_accents\":null," + + "\"tokenizer_class\":\"BertTokenizer\"}"); + } + + @Test + void testLoadsFromAModelDirectory(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "true"); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + // Same fixture and switches as testEmbedMeanPoolsWithoutWeights, resolved from the configs + // this time; the upper-cased input additionally proves do_lower_case was picked up. + assertArrayEquals(new float[] {3.5f, 35f, 350f}, model.embed("HELLO WORLD"), 1e-5f); + } + + @Test + void testDirectoryLoadReadsCasedFromTheTokenizerConfig(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "false"); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + // do_lower_case=false maps to Casing.CASED: lower-case text still matches the vocabulary... + assertArrayEquals(new float[] {3.5f, 35f, 350f}, model.embed("hello world"), 1e-5f); + // ...but upper-case text is preserved as-is, matches no cased vocabulary entry, folds to + // the (skipped) [UNK], and pools to the zero vector instead of being lower-cased first. + assertArrayEquals(new float[] {0f, 0f, 0f}, model.embed("HELLO WORLD"), 1e-5f); + } + + @Test + void testDirectoryLoadReadsNormalizeFromTheConfig(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "true", "true"); + + final float[] result = StaticEmbeddingModel.load(dir).embed("cat"); + + double normSquared = 0; + for (final float v : result) { + normSquared += (double) v * v; + } + assertEquals(1.0, Math.sqrt(normSquared), 1e-5); + } + + @Test + void testDirectoryLoadRejectsNullAndNonDirectory(@TempDir Path dir) { + assertThrows(IllegalArgumentException.class, () -> StaticEmbeddingModel.load(null)); + assertThrows(IllegalArgumentException.class, + () -> StaticEmbeddingModel.load(dir.resolve("absent"))); + } + + @Test + void testDirectoryLoadNamesTheMissingFile(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + // no config.json, no tokenizer_config.json + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("config.json")); + assertTrue(e.getMessage().contains("explicit load overloads")); + } + + @Test + void testDirectoryLoadRejectsAConfigWithoutNormalize(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "true"); + Files.writeString(dir.resolve("config.json"), "{\"model_type\":\"model2vec\"}"); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("normalize")); + } + + @Test + void testDirectoryLoadRejectsAConfigDeclaringNonMeanPooling(@TempDir Path dir) + throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "true"); + // Only mean pooling is implemented; a third-party config declaring another pooling must + // fail loud instead of silently mean-pooling with the wrong semantics. + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false,\"pooling\":\"max\"}"); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("max"), e.getMessage()); + assertTrue(e.getMessage().contains("mean"), e.getMessage()); + } + + @Test + void testDirectoryLoadAcceptsTheDeclaredMeanPooling(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "true"); + // The pooling the distiller writes; declaring it explicitly must load like omitting it. + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":false,\"pooling\":\"mean\"}"); + + assertArrayEquals(new float[] {3.5f, 35f, 350f}, + StaticEmbeddingModel.load(dir).embed("hello world"), 1e-5f); + } + + @Test + void testDirectoryLoadRejectsContradictoryStripAccents(@TempDir Path dir) throws IOException { + writeVocab(dir); + writeSafetensors(dir, false); + writeConfigs(dir, "false", "true"); + Files.writeString(dir.resolve("tokenizer_config.json"), + "{\"do_lower_case\":true,\"strip_accents\":false}"); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> StaticEmbeddingModel.load(dir)); + assertTrue(e.getMessage().contains("strip_accents")); + } + + @Test + void testTextEmbedderSeamMatchesDirectUseAndBatches(@TempDir Path dir) throws IOException { + final TextEmbedder embedder = + StaticEmbeddingModel.load(writeVocab(dir), writeSafetensors(dir, false), + Casing.UNCASED, Normalization.NONE); + + // The CharSequence entry point produces the same vector as the String one, including for a + // CharSequence that is not a String. + assertArrayEquals(new float[] {3.5f, 35f, 350f}, + embedder.embed(new StringBuilder("hello world")), 1e-5f); + assertEquals(DIMENSION, embedder.dimension()); + + // The interface's default batch method returns one vector per input, in input order. + final float[][] vectors = embedder.embedAll(List.of("hello world", "cat")); + assertEquals(2, vectors.length); + assertArrayEquals(new float[] {3.5f, 35f, 350f}, vectors[0], 1e-5f); + assertArrayEquals(new float[] {5f, 50f, 500f}, vectors[1], 1e-5f); + + assertThrows(IllegalArgumentException.class, () -> embedder.embed(null)); + assertThrows(IllegalArgumentException.class, () -> embedder.embedAll(null)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingSearchExampleTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingSearchExampleTest.java new file mode 100644 index 0000000000..3538ed6608 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingSearchExampleTest.java @@ -0,0 +1,63 @@ +/* + * 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.Path; +import java.util.ArrayList; +import java.util.Comparator; +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.assertTrue; + +/** + * Pins the semantic-search listing documented in {@code embeddings.xml}: score every document + * against a query with {@link StaticEmbeddingModel#similarity(String, String)} and sort by + * descending score, asserting the exact ranking on the analogy fixture's known geometry. + */ +public class StaticEmbeddingSearchExampleTest { + + /** A scored document, as the manual's listing declares it. */ + record Scored(String document, double score) { + } + + @Test + void testRanksDocumentsByCosineSimilarityToTheQuery(@TempDir Path dir) throws IOException { + EmbeddingTestFixtures.writeAnalogyDirectory(dir); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + final String query = "king"; + final List documents = List.of("queen woman", "apple", "king man"); + + final List results = new ArrayList<>(); + for (final String document : documents) { + results.add(new Scored(document, model.similarity(query, document))); + } + results.sort(Comparator.comparingDouble(Scored::score).reversed()); + + // The fixture's mean-pooled vectors give distinct cosines to "king" ([3,3]): "king man" + // pools to [2.5,2] (0.994), "queen woman" to [1.5,3] (0.949), "apple" is [-3,-1] (-0.894). + assertEquals(List.of("king man", "queen woman", "apple"), + results.stream().map(Scored::document).toList()); + assertTrue(results.get(0).score() > results.get(1).score()); + assertTrue(results.get(2).score() < 0); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingUsageExampleTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingUsageExampleTest.java new file mode 100644 index 0000000000..ecd8513358 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingUsageExampleTest.java @@ -0,0 +1,96 @@ +/* + * 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.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.assertTrue; + +/** + * Pins the cookbook paths documented in {@code embeddings.xml}, mirroring its listings: load a + * model directory with {@link StaticEmbeddingModel#load(Path)}, embed a text, and call + * {@code similarity}, {@code mostSimilar}, and {@code analogy}; and load through the explicit + * WordPiece and SentencePiece overloads. + */ +public class StaticEmbeddingUsageExampleTest { + + @Test + void testEmbedSimilarityNeighborsAndAnalogy(@TempDir Path dir) throws IOException { + EmbeddingTestFixtures.writeAnalogyDirectory(dir); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + final float[] vector = model.embed("king"); + assertEquals(2, vector.length); + + assertEquals(1.0, model.similarity("king", "king"), 1e-5); + + final List neighbors = model.mostSimilar("king", 5); + assertTrue(!neighbors.isEmpty()); + assertEquals("king", neighbors.get(0).token()); + + final List analogy = model.analogy("man", "king", "woman", 1); + assertEquals(1, analogy.size()); + assertEquals("queen", analogy.get(0).token()); + } + + @Test + void testExplicitOverloads(@TempDir Path wordPieceDir, @TempDir Path sentencePieceDir) + throws IOException { + EmbeddingTestFixtures.writeAnalogyDirectory(wordPieceDir); + EmbeddingTestFixtures.writeSentencePieceDirectory(sentencePieceDir); + + // The manual's explicit WordPiece overload: the data files plus the two switches the + // model's configuration publishes. + final StaticEmbeddingModel model = StaticEmbeddingModel.load( + wordPieceDir.resolve("vocab.txt"), wordPieceDir.resolve("model.safetensors"), + StaticEmbeddingModel.Casing.UNCASED, + StaticEmbeddingModel.Normalization.L2); + assertEquals(2, model.dimension()); + assertUnitLength(model.embed("king")); + + // The manual's explicit SentencePiece overload: no casing switch, because the trained + // .model file carries the model's own text normalizer. + final StaticEmbeddingModel multilingual = StaticEmbeddingModel.loadSentencePiece( + sentencePieceDir.resolve("sentencepiece.bpe.model"), + sentencePieceDir.resolve("tokenizer.json"), + sentencePieceDir.resolve("model.safetensors"), + StaticEmbeddingModel.Normalization.L2); + assertEquals(EmbeddingTestFixtures.SENTENCEPIECE_DIMENSION, multilingual.dimension()); + assertUnitLength(multilingual.embed("a")); + } + + /** + * Asserts that a vector has unit L2 length, the visible effect of choosing + * {@code Normalization.L2} in the explicit overloads. + * + * @param vector The vector to measure. + */ + private static void assertUnitLength(float[] vector) { + double normSquared = 0; + for (final float v : vector) { + normSquared += (double) v * v; + } + assertEquals(1.0, Math.sqrt(normSquared), 1e-5); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TeacherTokenizerTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TeacherTokenizerTest.java new file mode 100644 index 0000000000..50028a3bb8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TeacherTokenizerTest.java @@ -0,0 +1,517 @@ +/* + * 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 org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The teacher tokenizer cleaning mirrors Model2Vec: unused tokens and special added tokens other + * than the unknown and pad tokens are dropped, the survivors are renumbered in their original id + * order, and the rewritten {@code tokenizer.json} carries the pruned vocabulary, the remapped + * unknown id, a null post-processor, and only the unknown/pad added tokens. + */ +class TeacherTokenizerTest { + + // A WordPiece teacher: the special tokens are added tokens, plus one [unused] row and one + // content row pair. The post-processor wraps sequences in [CLS]/[SEP] (ids 2 and 3). + private static final String WORDPIECE_TEACHER = + "{\"version\":\"1.0\"," + + "\"normalizer\":{\"type\":\"BertNormalizer\",\"lowercase\":true}," + + "\"added_tokens\":[" + + "{\"id\":0,\"content\":\"[PAD]\",\"special\":true}," + + "{\"id\":1,\"content\":\"[UNK]\",\"special\":true}," + + "{\"id\":2,\"content\":\"[CLS]\",\"special\":true}," + + "{\"id\":3,\"content\":\"[SEP]\",\"special\":true}," + + "{\"id\":4,\"content\":\"[MASK]\",\"special\":true}]," + + "\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":[{\"SpecialToken\":{\"id\":\"[CLS]\",\"type_id\":0}}," + + "{\"Sequence\":{\"id\":\"A\",\"type_id\":0}}," + + "{\"SpecialToken\":{\"id\":\"[SEP]\",\"type_id\":0}}]," + + "\"special_tokens\":{\"[CLS]\":{\"id\":\"[CLS]\",\"ids\":[2],\"tokens\":[\"[CLS]\"]}," + + "\"[SEP]\":{\"id\":\"[SEP]\",\"ids\":[3],\"tokens\":[\"[SEP]\"]}}}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[PAD]\":0,\"[UNK]\":1,\"[CLS]\":2,\"[SEP]\":3,\"[MASK]\":4," + + "\"hello\":5,\"[unused1]\":6,\"world\":7}}}"; + + // A Unigram teacher in the bge-m3 shape: , , , lead the vocabulary, + // trails it; all five are special added tokens. + private static final String UNIGRAM_TEACHER = + "{\"version\":\"1.0\"," + + "\"added_tokens\":[" + + "{\"id\":0,\"content\":\"\",\"special\":true}," + + "{\"id\":1,\"content\":\"\",\"special\":true}," + + "{\"id\":2,\"content\":\"\",\"special\":true}," + + "{\"id\":3,\"content\":\"\",\"special\":true}," + + "{\"id\":6,\"content\":\"\",\"special\":true}]," + + "\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":3,\"byte_fallback\":false," + + "\"vocab\":[[\"\",0.0],[\"\",0.0],[\"\",0.0],[\"\",0.0]," + + "[\"a\",-1.5],[\"b\",-2.5],[\"\",0.0]]}}"; + + private static Path write(Path dir, String name, String content) throws IOException { + final Path file = dir.resolve(name); + Files.writeString(file, content); + return file; + } + + @Test + void testWordpieceCleaningDropsSpecialsAndUnusedTokens(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", WORDPIECE_TEACHER); + write(dir, "tokenizer_config.json", "{\"do_lower_case\":true,\"pad_token\":\"[PAD]\"}"); + + final TeacherTokenizer tokenizer = + TeacherTokenizer.read(tokenizerJson, dir.resolve("tokenizer_config.json")); + + assertEquals(TeacherTokenizer.WORDPIECE, tokenizer.modelType()); + assertEquals(4, tokenizer.vocabularySize()); + assertArrayEquals(new int[] {0, 1, 5, 7}, tokenizer.keptOriginalIds()); + assertEquals(0, tokenizer.padTokenId()); + assertEquals("[UNK]", tokenizer.unkToken()); + assertEquals("[PAD]", tokenizer.padToken()); + // Each row is fed to the teacher as [CLS, token, SEP]. + assertArrayEquals(new long[] {2, 5, 3}, tokenizer.inputSequence(2)); + } + + @Test + void testWordpieceRewriteRenumbersTheSurvivors(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", WORDPIECE_TEACHER); + write(dir, "tokenizer_config.json", "{\"pad_token\":\"[PAD]\"}"); + final TeacherTokenizer tokenizer = + TeacherTokenizer.read(tokenizerJson, dir.resolve("tokenizer_config.json")); + + final Path cleaned = dir.resolve("cleaned.json"); + tokenizer.writeCleaned(cleaned); + + // The cleaned file parses again and names exactly the surviving rows in order (the pad + // token needs its tokenizer_config to be recognized, as in the teacher). + final TeacherTokenizer reread = + TeacherTokenizer.read(cleaned, dir.resolve("tokenizer_config.json")); + assertEquals(4, reread.vocabularySize()); + assertArrayEquals(new int[] {0, 1, 2, 3}, reread.keptOriginalIds()); + final String json = Files.readString(cleaned); + assertTrue(json.contains("\"post_processor\":null"), json); + assertTrue(json.contains("\"hello\":2"), json); + assertTrue(json.contains("\"world\":3"), json); + assertFalse(json.contains("[unused1]"), json); + assertFalse(json.contains("[MASK]"), json); + // The unk and pad added tokens remain, with Model2Vec's flag convention. + assertTrue(json.contains("{\"id\":0,\"content\":\"[PAD]\",\"single_word\":true," + + "\"lstrip\":true,\"rstrip\":true,\"normalized\":true,\"special\":true}"), json); + assertTrue(json.contains("{\"id\":1,\"content\":\"[UNK]\",\"single_word\":false," + + "\"lstrip\":false,\"rstrip\":false,\"normalized\":false,\"special\":true}"), json); + // Untouched sections survive byte for byte. + assertTrue(json.contains("\"normalizer\":{\"type\":\"BertNormalizer\",\"lowercase\":true}"), + json); + } + + @Test + void testUnigramCleaningKeepsPadAndUnkOnly(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", UNIGRAM_TEACHER); + write(dir, "tokenizer_config.json", "{\"pad_token\":\"\"}"); + + final TeacherTokenizer tokenizer = + TeacherTokenizer.read(tokenizerJson, dir.resolve("tokenizer_config.json")); + + assertEquals(TeacherTokenizer.UNIGRAM, tokenizer.modelType()); + assertEquals(4, tokenizer.vocabularySize()); + assertArrayEquals(new int[] {1, 3, 4, 5}, tokenizer.keptOriginalIds()); + assertEquals(1, tokenizer.padTokenId()); + assertEquals("", tokenizer.unkToken()); + // No post-processor, so the input sequence is the bare token. + assertArrayEquals(new long[] {4}, tokenizer.inputSequence(2)); + } + + @Test + void testUnigramRewriteRemapsUnkIdAndKeepsScores(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", UNIGRAM_TEACHER); + write(dir, "tokenizer_config.json", "{\"pad_token\":\"\"}"); + final TeacherTokenizer tokenizer = + TeacherTokenizer.read(tokenizerJson, dir.resolve("tokenizer_config.json")); + + final Path cleaned = dir.resolve("cleaned.json"); + tokenizer.writeCleaned(cleaned); + + // The loader's own Unigram reader must see the surviving rows in order. + assertEquals(List.of("", "", "a", "b"), TokenizerJsonVocab.rows(cleaned)); + final String json = Files.readString(cleaned); + assertTrue(json.contains("\"unk_id\":1"), json); + assertTrue(json.contains("[\"a\",-1.5]"), json); + assertTrue(json.contains("\"byte_fallback\":false"), json); + } + + @Test + void testUnigramWithoutPadTokenKeepsOnlyTheUnknownToken(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", UNIGRAM_TEACHER); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertNull(tokenizer.padToken()); + assertEquals(3, tokenizer.vocabularySize()); + assertArrayEquals(new int[] {3, 4, 5}, tokenizer.keptOriginalIds()); + } + + @Test + void testRejectsAnUnsupportedModelType(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"model\":{\"type\":\"BPE\",\"vocab\":{\"a\":0}}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(e.getMessage().contains("BPE"), e.getMessage()); + } + + @Test + void testRejectsANullTokenizerJsonFile() { + assertEquals("TokenizerJsonFile must not be null", assertThrows( + IllegalArgumentException.class, () -> TeacherTokenizer.read(null, null)).getMessage()); + } + + @Test + void testRejectsAMissingTokenizerJsonFile(@TempDir Path dir) { + final Path missing = dir.resolve("tokenizer.json"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> TeacherTokenizer.read(missing, null)); + assertEquals("File does not exist or is not a regular file: " + missing, e.getMessage()); + } + + /** + * The teacher must be rejected, not half-read, when it cannot describe a distilled table. Each + * case names the part of the contract it breaks and a fragment the message has to carry, so a + * silently weakened check shows up as a failing row rather than as a corrupt model directory. + */ + @ParameterizedTest + @CsvSource(delimiter = ';', value = { + "no model at all;{\"version\":\"1.0\"};has no model with a vocabulary", + "a model without a vocabulary;{\"model\":{\"type\":\"WordPiece\"}};" + + "has no model with a vocabulary", + "no unknown token;{\"model\":{\"type\":\"WordPiece\",\"vocab\":{\"a\":0}}};" + + "does not name an unknown token", + "an unknown token outside the vocabulary;{\"model\":{\"type\":\"WordPiece\"," + + "\"unk_token\":\"[UNK]\",\"vocab\":{\"a\":0}}};it is not in the vocabulary", + "a Unigram unk_id out of range;{\"model\":{\"type\":\"Unigram\",\"unk_id\":9," + + "\"vocab\":[[\"a\",0.0]]}};does not name an unknown token", + "vocabulary ids with a gap;{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\"," + + "\"vocab\":{\"a\":0,\"b\":2}}};not a gapless range", + "an unsupported post-processor;{\"post_processor\":{\"type\":\"ByteLevel\"}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "is not supported", + "trailing content;{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\"," + + "\"vocab\":{\"a\":0}}} junk;Trailing content"}) + void testRejectsATeacherItCannotDistill(String reason, String teacherJson, String messagePart, + @TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", teacherJson); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null), reason); + assertTrue(e.getMessage().contains(messagePart), + "a teacher with " + reason + " reported: " + e.getMessage()); + } + + /** + * The other shape a {@code TemplateProcessing} template takes: a single string, whose items are + * separated by whitespace rather than being a list of objects. + */ + @Test + void testReadsAStringTemplatePostProcessor(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\" $A \"," + + "\"special_tokens\":{\"\":{\"id\":\"\",\"ids\":[0]}," + + "\"\":{\"id\":\"\",\"ids\":[2]}}}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"\"," + + "\"vocab\":{\"\":0,\"\":1,\"\":2,\"\":3,\"a\":4}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + // Row 4 is 'a'; the string template wraps it the same way the structured form would. + assertArrayEquals(new long[] {0, 4, 2}, tokenizer.inputSequence(4)); + } + + /** + * A {@code BertProcessing} post-processor carries its wrapper as {@code cls}/{@code sep} token + * pairs instead of as a template, and the ids come straight from those pairs. + */ + @Test + void testReadsABertProcessingPostProcessor(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"post_processor\":{\"type\":\"BertProcessing\"," + + "\"cls\":[\"[CLS]\",2],\"sep\":[\"[SEP]\",3]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[PAD]\":0,\"[UNK]\":1,\"[CLS]\":2,\"[SEP]\":3,\"hello\":4}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertEquals(5, tokenizer.vocabularySize()); + assertArrayEquals(new long[] {2, 4, 3}, tokenizer.inputSequence(4)); + } + + /** A teacher without a post-processor feeds the bare token, with no wrapper ids. */ + @Test + void testANullPostProcessorAddsNoWrapper(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"post_processor\":null,\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0,\"hello\":1}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertArrayEquals(new long[] {1}, tokenizer.inputSequence(1)); + } + + /** + * A pad token the teacher's {@code tokenizer_config.json} names but the vocabulary does not have + * is not a row; it must not be kept, and the pad id falls back to 0. + */ + @Test + void testAPadTokenOutsideTheVocabularyIsIgnored(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", UNIGRAM_TEACHER); + write(dir, "tokenizer_config.json", "{\"pad_token\":\"\"}"); + + final TeacherTokenizer tokenizer = + TeacherTokenizer.read(tokenizerJson, dir.resolve("tokenizer_config.json")); + + assertEquals(0, tokenizer.padTokenId()); + assertArrayEquals(new int[] {3, 4, 5}, tokenizer.keptOriginalIds()); + } + + /** + * A template that names its special tokens without carrying a {@code special_tokens} table has to + * resolve those names through the vocabulary instead. + */ + @Test + void testATemplateWithoutASpecialTokenTableResolvesThroughTheVocabulary(@TempDir Path dir) + throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":[{\"SpecialToken\":{\"id\":\"[CLS]\",\"type_id\":0}}," + + "{\"Sequence\":{\"id\":\"A\",\"type_id\":0}}," + + "{\"SpecialToken\":{\"id\":\"[SEP]\",\"type_id\":0}}]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[PAD]\":0,\"[UNK]\":1,\"[CLS]\":2,\"[SEP]\":3,\"hello\":4}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertArrayEquals(new long[] {2, 4, 3}, tokenizer.inputSequence(4)); + } + + @Test + void testRejectsATemplateSpecialTokenThatResolvesNowhere(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":[{\"SpecialToken\":{\"id\":\"[BOS]\",\"type_id\":0}}," + + "{\"Sequence\":{\"id\":\"A\",\"type_id\":0}}]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0,\"hello\":1}}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(e.getMessage().contains("[BOS]"), e.getMessage()); + } + + @Test + void testRejectsAVocabularyIdUsedTwice(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0,\"b\":0}}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(e.getMessage().contains("assigned more than once"), e.getMessage()); + } + + @Test + void testAnExplicitlyNullUnkIdCountsAsNoUnknownToken(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"model\":{\"type\":\"Unigram\",\"unk_id\":null,\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(e.getMessage().contains("does not name an unknown token"), e.getMessage()); + } + + @ParameterizedTest + @ValueSource(strings = {"", " \n\t "}) + void testRejectsAnEmptyTokenizerJson(String content, @TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", content); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(e.getMessage().contains("Unexpected end of input"), e.getMessage()); + } + + /** + * The removal pattern is matched from the start of the token, so it drops a token that begins + * with an {@code [unusedN]} marker and keeps everything else, including a marker without digits + * and one that is not at the start. + */ + @ParameterizedTest + @CsvSource(delimiter = ';', value = { + "[unused0];1", + "[unused12];1", + "[unused7]tail;1", + "[unused];2", + "[unusedx];2", + "x[unused1];2", + "[UNUSED1];2"}) + void testUnusedTokenRemovalMatchesFromTheStartOnly(String token, int expectedSize, + @TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0,\"" + token + "\":1}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertEquals(expectedSize, tokenizer.vocabularySize()); + } + + /** + * Vocabulary entries are copied as raw spans, so a teacher's escapes reach the distilled file + * untouched and still decode to the pieces the loader resolves matrix rows by. + */ + @Test + void testUnicodeVocabularyEntriesSurviveTheRewrite(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"version\":\"1.0\",\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0,\"vocab\":[[\"\",0.0]," + + "[\"caf\\u00e9\",-1.0],[\"e\\u0301\",-2.0],[\"\\ud83d\\ude00\",-3.0]," + + "[\"a\",-4.0]]}}"); + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + final Path cleaned = dir.resolve("cleaned.json"); + tokenizer.writeCleaned(cleaned); + + assertEquals(5, tokenizer.vocabularySize()); + final String json = Files.readString(cleaned); + assertTrue(json.contains("[\"caf\\u00e9\",-1.0]"), json); + assertTrue(json.contains("[\"e\\u0301\",-2.0]"), json); + assertTrue(json.contains("[\"\\ud83d\\ude00\",-3.0]"), json); + // A precomposed letter, a base letter plus a combining acute, and a supplementary-plane + // character all decode to what the teacher declared. + assertEquals(List.of("", "caf\u00e9", "e\u0301", "\uD83D\uDE00", "a"), + TokenizerJsonVocab.rows(cleaned)); + } + + /** + * The added-token overlay is the only part of the rewrite that re-encodes a token string rather + * than copying its raw span, so it has to escape what JSON requires. + */ + @Test + void testTheAddedTokenOverlayEscapesTheUnknownTokenContent(@TempDir Path dir) throws IOException { + // The unknown token carries a backslash and a tab. + final String rawToken = "\"\""; + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"version\":\"1.0\"," + + "\"added_tokens\":[{\"id\":0,\"content\":" + rawToken + ",\"special\":true}]," + + "\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0," + + "\"vocab\":[[" + rawToken + ",0.0],[\"a\",-1.0]]}}"); + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + final Path cleaned = dir.resolve("cleaned.json"); + tokenizer.writeCleaned(cleaned); + + assertEquals(2, tokenizer.vocabularySize()); + final String json = Files.readString(cleaned); + assertTrue(json.contains("[" + rawToken + ",0.0]"), json); + assertTrue(json.contains("\"content\":\"\""), json); + } + + /** The rewrite emits only fields the teacher had, so an absent overlay stays absent. */ + @Test + void testATeacherWithoutAnAddedTokensSectionWritesNoOverlay(@TempDir Path dir) + throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"version\":\"1.0\",\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0," + + "\"vocab\":[[\"\",0.0],[\"a\",-1.0]]}}"); + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + final Path cleaned = dir.resolve("cleaned.json"); + tokenizer.writeCleaned(cleaned); + + assertFalse(Files.readString(cleaned).contains("added_tokens")); + assertEquals(List.of("", "a"), TokenizerJsonVocab.rows(cleaned)); + } + + /** + * The overlay is pruned by token content alone: the {@code special} flag is never read, so a + * plain vocabulary extension is dropped from the distilled table just like {@code [MASK]} is. + */ + @Test + void testEveryAddedTokenIsDroppedRegardlessOfItsSpecialFlag(@TempDir Path dir) + throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"version\":\"1.0\"," + + "\"added_tokens\":[{\"id\":1,\"content\":\"[UNK]\",\"special\":true}," + + "{\"id\":2,\"content\":\"covid\",\"special\":false}]," + + "\"post_processor\":null," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"hello\":0,\"[UNK]\":1,\"covid\":2}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertArrayEquals(new int[] {0, 1}, tokenizer.keptOriginalIds()); + } + + @Test + void testWriteCleanedRejectsANullFile(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", UNIGRAM_TEACHER); + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertEquals("File must not be null", assertThrows( + IllegalArgumentException.class, () -> tokenizer.writeCleaned(null)).getMessage()); + } + + @Test + void testTermInputSequenceMapsPieceStringsToOriginalIds(@TempDir Path dir) throws IOException { + final TeacherTokenizer tokenizer = TeacherTokenizer.read( + write(dir, "tokenizer.json", WORDPIECE_TEACHER), null); + + // hello and world map to their original ids, an unmapped piece falls to the unknown id, + // and the sequence is wrapped in the post-processor's [CLS]/[SEP] ids. + assertArrayEquals(new long[] {2, 5, 1, 7, 3}, + tokenizer.inputSequence(List.of("hello", "nope", "world"))); + assertEquals("Pieces must not be null", assertThrows(IllegalArgumentException.class, + () -> tokenizer.inputSequence((List) null)).getMessage()); + } + + @Test + void testReadsTheNormalizerLowercaseFlag(@TempDir Path dir) throws IOException { + assertEquals(Boolean.TRUE, TeacherTokenizer.read( + write(dir, "wordpiece.json", WORDPIECE_TEACHER), null).lowerCase()); + // The Unigram teacher states no normalizer, so the flag is unknown. + assertNull(TeacherTokenizer.read( + write(dir, "unigram.json", UNIGRAM_TEACHER), null).lowerCase()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermSegmenterTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermSegmenterTest.java new file mode 100644 index 0000000000..a52ff345e8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermSegmenterTest.java @@ -0,0 +1,131 @@ +/* + * 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.InputStream; +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.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The term segmenter's fidelity to the teacher's own tokenization: WordPiece casing and subword + * continuation, unknown-word fallback, delimiter removal, and the SentencePiece path through the + * teacher's trained model file. + */ +class TermSegmenterTest { + + // A WordPiece teacher whose vocabulary can subword-split "corpuses" into corpus + ##es. + private static final String WORDPIECE_TEACHER = + "{\"normalizer\":{\"type\":\"BertNormalizer\",\"lowercase\":true}," + + "\"post_processor\":null," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0,\"[CLS]\":1,\"[SEP]\":2," + + "\"habeas\":3,\"corpus\":4,\"##es\":5}}}"; + + private static TeacherTokenizer wordpieceTeacher(Path dir) throws IOException { + Files.writeString(dir.resolve(ModelFileNames.TOKENIZER_JSON), WORDPIECE_TEACHER); + return TeacherTokenizer.read(dir.resolve(ModelFileNames.TOKENIZER_JSON), null); + } + + @Test + void testSegmentsWithTheTeachersCasingAndSubwords(@TempDir Path dir) throws IOException { + final TermSegmenter segmenter = + TermSegmenter.forTeacher(wordpieceTeacher(dir), dir); + + assertEquals(List.of("habeas", "corpus"), segmenter.pieces("Habeas CORPUS")); + assertEquals(List.of("corpus", "##es"), segmenter.pieces("corpuses")); + } + + @Test + void testDropsTheEncodersDelimitersButKeepsTheUnknownPiece(@TempDir Path dir) + throws IOException { + final TermSegmenter segmenter = + TermSegmenter.forTeacher(wordpieceTeacher(dir), dir); + + // The wrapping [CLS]/[SEP] are the segmenter's own; an out-of-vocabulary word stays as the + // unknown piece, so the teacher still sees a position for it. + assertEquals(List.of("habeas", "[UNK]"), segmenter.pieces("habeas zzz")); + } + + @Test + void testWordpiecePiecesMapBackToTeacherInputIds(@TempDir Path dir) throws IOException { + final TeacherTokenizer teacher = wordpieceTeacher(dir); + final TermSegmenter segmenter = TermSegmenter.forTeacher(teacher, dir); + + // No post-processor, so the sequence is exactly the piece ids in the teacher's id space. + final long[] sequence = teacher.inputSequence(segmenter.pieces("habeas corpuses")); + assertArrayEquals(new long[] {3, 4, 5}, sequence); + } + + @Test + void testAUnigramTeacherSegmentsThroughItsTrainedModelFile(@TempDir Path dir) + throws IOException { + // The trained tiny SentencePiece fixture next to a matching Unigram tokenizer.json. + final byte[] modelBytes; + try (InputStream in = TermSegmenterTest.class + .getResourceAsStream(EmbeddingTestFixtures.TINY_UNIGRAM_RESOURCE)) { + modelBytes = in.readAllBytes(); + } + Files.write(dir.resolve("spiece.model"), modelBytes); + Files.writeString(dir.resolve(ModelFileNames.TOKENIZER_JSON), + "{\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0," + + "\"vocab\":[[\"\",0.0],[\"▁a\",-1.5],[\"a\",-2.0]]}}"); + final TeacherTokenizer teacher = TeacherTokenizer.read( + dir.resolve(ModelFileNames.TOKENIZER_JSON), null); + + final TermSegmenter segmenter = TermSegmenter.forTeacher(teacher, dir); + final List pieces = segmenter.pieces("a"); + + assertFalse(pieces.isEmpty()); + // Control pieces never appear; every piece is a string the trained model produced. + assertTrue(pieces.stream().noneMatch(p -> p.equals("") || p.equals("")), + pieces.toString()); + } + + @Test + void testAUnigramTeacherWithoutItsModelFileIsRejected(@TempDir Path dir) throws IOException { + Files.writeString(dir.resolve(ModelFileNames.TOKENIZER_JSON), + "{\"post_processor\":null," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0,\"vocab\":[[\"\",0.0]]}}"); + final TeacherTokenizer teacher = TeacherTokenizer.read( + dir.resolve(ModelFileNames.TOKENIZER_JSON), null); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> TermSegmenter.forTeacher(teacher, dir)); + assertTrue(e.getMessage().contains("SentencePiece"), e.getMessage()); + } + + @Test + void testNullArgumentsAreRejected(@TempDir Path dir) throws IOException { + final TeacherTokenizer teacher = wordpieceTeacher(dir); + assertThrows(IllegalArgumentException.class, () -> TermSegmenter.forTeacher(null, dir)); + assertThrows(IllegalArgumentException.class, () -> TermSegmenter.forTeacher(teacher, null)); + final TermSegmenter segmenter = TermSegmenter.forTeacher(teacher, dir); + assertThrows(IllegalArgumentException.class, () -> segmenter.pieces(null)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermTableTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermTableTest.java new file mode 100644 index 0000000000..aee76a3c90 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TermTableTest.java @@ -0,0 +1,155 @@ +/* + * 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; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.InvalidFormatException; + +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 term table's normalization contract, its validation of stored terms, and the greedy + * longest-first matching over word runs. + */ +class TermTableTest { + + private static final String SOURCE = "terms.txt"; + + private static TermTable table(String... terms) throws InvalidFormatException { + return TermTable.of(List.of(terms), 10, SOURCE); + } + + @ParameterizedTest + @CsvSource(delimiter = ';', value = { + "habeas corpus;habeas corpus", + "Habeas Corpus;habeas corpus", + "habeas-corpus!;habeas corpus", + "' writ OF Habeas ';writ of habeas", + "res judicata.;res judicata", + "42 USC 1983;42 usc 1983" + }) + void testNormalizeTermFoldsAndJoinsWordRuns(String raw, String expected) { + assertEquals(expected, TermTable.normalizeTerm(raw)); + } + + @ParameterizedTest + @ValueSource(strings = {"", " ", "&!.", "--"}) + void testNormalizeTermOfTextWithoutWordsIsEmpty(String raw) { + assertEquals("", TermTable.normalizeTerm(raw)); + } + + @Test + void testNormalizeTermFoldsSupplementaryPlaneLetters() { + // DESERET CAPITAL LETTER LONG I (U+10400) is a cased letter outside the BMP; its lower-case + // form is U+10428, one code point, so the word run survives the fold intact. + final String capital = new String(Character.toChars(0x10400)); + final String small = new String(Character.toChars(0x10428)); + assertEquals(small + "x", TermTable.normalizeTerm(capital + "x")); + } + + @Test + void testRejectsATermThatIsNotNormalized() { + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> table("HABEAS CORPUS")); + assertTrue(e.getMessage().contains("HABEAS CORPUS"), e.getMessage()); + assertTrue(e.getMessage().contains(SOURCE), e.getMessage()); + } + + @ParameterizedTest + @ValueSource(strings = {"", "habeas corpus", " habeas", "habeas "}) + void testRejectsMalformedTermForms(String term) { + assertThrows(InvalidFormatException.class, () -> table(term)); + } + + @Test + void testRejectsADuplicateTerm() { + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> table("habeas corpus", "habeas corpus")); + assertTrue(e.getMessage().contains("more than once"), e.getMessage()); + } + + @Test + void testRejectsNullArguments() { + assertThrows(IllegalArgumentException.class, () -> TermTable.of(null, 0, SOURCE)); + assertThrows(IllegalArgumentException.class, () -> TermTable.normalizeTerm(null)); + } + + @Test + void testTermsOwnRowsFromTheFirstRowOnward() throws InvalidFormatException { + final TermTable table = table("habeas corpus", "replevin"); + assertEquals(2, table.size()); + assertEquals("habeas corpus", table.term(10)); + assertEquals("replevin", table.term(11)); + assertThrows(IllegalArgumentException.class, () -> table.term(9)); + assertThrows(IllegalArgumentException.class, () -> table.term(12)); + } + + @Test + void testMatchesFoldCaseAndSpanPunctuation() throws InvalidFormatException { + final TermTable table = table("habeas corpus"); + final List matches = table.matches("The writ of Habeas-Corpus, granted."); + assertEquals(1, matches.size()); + assertEquals(10, matches.get(0).row()); + assertEquals("Habeas-Corpus", "The writ of Habeas-Corpus, granted." + .substring(matches.get(0).start(), matches.get(0).end())); + } + + @Test + void testTheLongestTermWinsAndConsumesItsWords() throws InvalidFormatException { + final TermTable table = table("habeas corpus", "writ of habeas corpus", "corpus"); + final List matches = table.matches("a writ of habeas corpus indeed"); + // The four-word term wins over both shorter terms, and its words are consumed: the inner + // "habeas corpus" and "corpus" do not match again. + assertEquals(1, matches.size()); + assertEquals(11, matches.get(0).row()); + } + + @Test + void testMatchingContinuesAfterAConsumedTerm() throws InvalidFormatException { + final TermTable table = table("habeas corpus", "replevin"); + final List matches = table.matches("habeas corpus then replevin"); + assertEquals(2, matches.size()); + assertEquals(10, matches.get(0).row()); + assertEquals(11, matches.get(1).row()); + assertTrue(matches.get(0).end() <= matches.get(1).start()); + } + + @Test + void testWordsSeparatedByOtherWordsDoNotMatchAPhrase() throws InvalidFormatException { + final TermTable table = table("habeas corpus"); + assertTrue(table.matches("habeas late corpus").isEmpty()); + } + + @Test + void testAnEmptyTableMatchesNothing() throws InvalidFormatException { + assertTrue(table().matches("habeas corpus").isEmpty()); + } + + @Test + void testMatchesRejectsNullText() throws InvalidFormatException { + final TermTable table = table("habeas corpus"); + assertThrows(IllegalArgumentException.class, () -> table.matches(null)); + } +} 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 new file mode 100644 index 0000000000..937a625630 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.java @@ -0,0 +1,183 @@ +/* + * 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 opennlp.tools.util.InvalidFormatException; + +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 {@code tokenizer.json} vocabulary contract: the Unigram {@code model.vocab} list order is + * the row order, {@code added_tokens} append or must agree, everything else is skipped, and any + * departure from the known shape fails loud. + */ +class TokenizerJsonVocabTest { + + @TempDir + private Path dir; + + private Path write(String json) throws IOException { + final Path file = dir.resolve("tokenizer.json"); + Files.writeString(file, json); + 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," + + "\"vocab\":[[\"\",0.0],[\"\",0.0],[\"\\u2581a\",-2.5],[\"b\",-3.0]]}}"); + + assertEquals(List.of("", "", "\u2581a", "b"), TokenizerJsonVocab.rows(file)); + } + + @Test + void testAddedTokenAtTheNextRowAppends() throws IOException { + final Path file = write("{\"added_tokens\":[{\"id\":2,\"content\":\"\"," + + "\"special\":true}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0],[\"b\",-1.0]]}}"); + + assertEquals(List.of("a", "b", ""), TokenizerJsonVocab.rows(file)); + } + + @Test + void testAddedTokenAtAnExistingRowMustAgree() throws IOException { + final Path agreeing = write("{\"added_tokens\":[{\"id\":0,\"content\":\"\"}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"\",0.0],[\"a\",-1.0]]}}"); + assertEquals(List.of("", "a"), TokenizerJsonVocab.rows(agreeing)); + + final Path contradicting = write("{\"added_tokens\":[{\"id\":0,\"content\":\"\"}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"\",0.0],[\"a\",-1.0]]}}"); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(contradicting)); + assertTrue(e.getMessage().contains("contradicts"), e.getMessage()); + } + + @Test + void testAddedTokenBeyondTheNextRowIsAGap() throws IOException { + final Path file = write("{\"added_tokens\":[{\"id\":5,\"content\":\"\"}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("gap"), e.getMessage()); + } + + @Test + void testAddedTokensAreOverlaidInIdOrderNotListOrder() throws IOException { + final Path file = write("{\"added_tokens\":[{\"id\":3,\"content\":\"y\"}," + + "{\"id\":2,\"content\":\"x\"}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0],[\"b\",-1.0]]}}"); + + assertEquals(List.of("a", "b", "x", "y"), TokenizerJsonVocab.rows(file)); + } + + @Test + void testSkipsUnrelatedSectionsAndDecodesEscapes() throws IOException { + final Path file = write("{\"version\":\"1.0\",\"truncation\":null," + + "\"normalizer\":{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"AAAA\"}," + + "\"pre_tokenizer\":[1,2,{\"a\":[true,false]}]," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0," + + "\"vocab\":[[\"\\\"quoted\\\"\",0.0],[\"tab\\there\",-1.0]]}}"); + + assertEquals(List.of("\"quoted\"", "tab\there"), TokenizerJsonVocab.rows(file)); + } + + @Test + void testRejectsANonUnigramModel() throws IOException { + final Path file = write("{\"model\":{\"type\":\"BPE\",\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("BPE"), e.getMessage()); + } + + @Test + void testRejectsAnObjectShapedVocab() throws IOException { + // The WordPiece/BPE tokenizer.json layout stores vocab as {piece: id}; ids in that shape + // are not list positions, so it must be refused rather than misread. + final Path file = write("{\"model\":{\"type\":\"Unigram\",\"vocab\":{\"a\":0,\"b\":1}}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("object"), e.getMessage()); + } + + @Test + void testRejectsAMissingVocab() throws IOException { + final Path noModel = write("{\"version\":\"1.0\"}"); + assertTrue(assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(noModel)).getMessage().contains("model.vocab")); + + final Path noVocab = write("{\"model\":{\"type\":\"Unigram\"}}"); + assertTrue(assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(noVocab)).getMessage().contains("model.vocab")); + } + + @Test + void testRejectsAnAddedTokenWithoutIdOrContent() throws IOException { + final Path file = write("{\"added_tokens\":[{\"content\":\"\"}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("id"), e.getMessage()); + } + + @Test + void testRejectsDuplicateTopLevelSections() throws IOException { + final Path file = write("{\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0]]}," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"b\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("more than once"), e.getMessage()); + } + + @Test + void testRejectsMalformedJson() throws IOException { + final Path file = write("{\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0]"); + + assertThrows(InvalidFormatException.class, () -> TokenizerJsonVocab.rows(file)); + } + + @Test + void testVocabularyEntryPointRejectsDuplicatePieces() throws IOException { + final Path file = write("{\"model\":{\"type\":\"Unigram\"," + + "\"vocab\":[[\"a\",0.0],[\"a\",-1.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> EmbeddingVocabulary.fromTokenizerJson(file)); + assertTrue(e.getMessage().contains("more than once"), e.getMessage()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/VectorIndexUsageExampleTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/VectorIndexUsageExampleTest.java new file mode 100644 index 0000000000..8977e2a6c9 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/VectorIndexUsageExampleTest.java @@ -0,0 +1,49 @@ +/* + * 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.Path; +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.embeddings.index.TurboQuantIndex; +import opennlp.embeddings.index.VectorIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** Pins the bounded in-memory vector search listing in {@code embeddings.xml}. */ +class VectorIndexUsageExampleTest { + + @Test + void testBuildFreezeAndQuery(@TempDir Path modelDirectory) throws IOException { + EmbeddingTestFixtures.writeAnalogyDirectory(modelDirectory); + + final StaticEmbeddingModel model = StaticEmbeddingModel.load(modelDirectory); + final VectorIndex index = new TurboQuantIndex(model.dimension(), 4, 42L); + + index.add("royal-article", model.embed("king queen")); + index.add("fruit-article", model.embed("apple")); + index.freeze(); + + final List hits = index.topK(model.embed("king"), 5); + assertEquals(2, hits.size()); + assertEquals("royal-article", hits.get(0).id()); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.java new file mode 100644 index 0000000000..c60e90cec0 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.java @@ -0,0 +1,84 @@ +/* + * 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.cmdline; + +import java.util.Set; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import opennlp.tools.cmdline.BasicCmdLineTool; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The command names the dispatcher offers and the help every tool produces. The names are the + * module's public surface (TRAINING.md and the manual quote them), so a class rename that changes + * a command has to fail here rather than in a user's shell. + */ +class CLITest { + + /** {@return the tools the dispatcher registers, as parameterized-test arguments} */ + static Stream tools() { + return Stream.of(new AssembleModelTool(), new DistillModelTool(), new QuantizeModelTool()); + } + + @Test + void testOffersExactlyTheModelCommands() { + assertEquals(Set.of("AssembleModel", "DistillModel", "QuantizeModel"), CLI.getToolNames()); + } + + @Test + void testTheToolNamesCannotBeModifiedByACaller() { + final Set names = CLI.getToolNames(); + + assertThrows(UnsupportedOperationException.class, () -> names.add("Other")); + } + + @ParameterizedTest + @MethodSource("tools") + void testEveryRegisteredToolDescribesItself(BasicCmdLineTool tool) { + assertTrue(CLI.getToolNames().contains(tool.getName()), + tool.getName() + " must be registered with the dispatcher"); + assertFalse(tool.getShortDescription().isBlank(), + tool.getName() + " must have a short description for the usage listing"); + assertTrue(tool.getHelp().contains(tool.getName()), tool.getHelp()); + } + + @Test + void testDistillHelpNamesEveryParameter() { + final String help = new DistillModelTool().getHelp(); + + assertTrue(help.contains("-teacher hf-id-or-path"), help); + assertTrue(help.contains("-out dir"), help); + // The optional parameters are bracketed, so a user can see they may be omitted. + assertTrue(help.contains("[-pcaDims "), help); + assertTrue(help.contains("[-terms "), help); + } + + @Test + void testAssembleHelpNamesItsParameter() { + final String help = new AssembleModelTool().getHelp(); + + assertTrue(help.contains("-modelDir dir"), help); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/FlatFloatIndexTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/FlatFloatIndexTest.java new file mode 100644 index 0000000000..7693c5da35 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/FlatFloatIndexTest.java @@ -0,0 +1,70 @@ +/* + * 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.index; + +import java.util.List; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * The exact index's scores: hand-computed cosine similarities, exact ordering, and the + * odd-dimension tail of the unrolled dot product. + */ +class FlatFloatIndexTest { + + @Test + void testScoresAreExactCosineSimilarities() { + final FlatFloatIndex index = new FlatFloatIndex(2); + index.add("x", new float[] {1f, 0f}); + index.add("y", new float[] {0f, 1f}); + index.add("diagonal", new float[] {1f, 1f}); + index.freeze(); + + final List hits = index.topK(new float[] {1f, 0f}, 3); + assertEquals(3, hits.size()); + assertEquals("x", hits.get(0).id()); + assertEquals(1.0, hits.get(0).score(), 1e-12); + assertEquals("diagonal", hits.get(1).id()); + assertEquals(Math.sqrt(0.5), hits.get(1).score(), 1e-12); + assertEquals("y", hits.get(2).id()); + assertEquals(0.0, hits.get(2).score(), 1e-12); + } + + @Test + void testAnOppositeVectorScoresMinusOne() { + final FlatFloatIndex index = new FlatFloatIndex(3); + index.add("opposite", new float[] {-2f, 0f, 0f}); + index.freeze(); + + assertEquals(-1.0, index.topK(new float[] {5f, 0f, 0f}, 1).get(0).score(), 1e-12); + } + + @Test + void testAnOddDimensionExercisesTheUnrolledTail() { + // Dimension 5 leaves one coordinate for the scalar tail after the 4-wide unroll. + final FlatFloatIndex index = new FlatFloatIndex(5); + index.add("v", new float[] {1f, 2f, 3f, 4f, 5f}); + index.freeze(); + + final double norm = Math.sqrt(1 + 4 + 9 + 16 + 25); + // Query along the last coordinate only: the dot is exactly the tail's contribution. + final float[] query = new float[] {0f, 0f, 0f, 0f, 2f}; + assertEquals(5.0 * 2 / (2 * norm), index.topK(query, 1).get(0).score(), 1e-12); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/TurboQuantIndexTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/TurboQuantIndexTest.java new file mode 100644 index 0000000000..dfc2c6bcf9 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/TurboQuantIndexTest.java @@ -0,0 +1,201 @@ +/* + * 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.index; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +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 opennlp.tools.util.InvalidFormatException; + +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 quantized index against the exact one: self-retrieval and recall on a deterministic + * random collection, the persisted round trip, and the malformed-directory contract. + */ +class TurboQuantIndexTest { + + private static final int DIMENSION = 64; + private static final int COUNT = 100; + private static final long SEED = 7; + + /** {@return the deterministic Gaussian test collection, one row per id} */ + private static float[][] collection() { + final Random random = new Random(SEED); + final float[][] vectors = new float[COUNT][DIMENSION]; + for (final float[] vector : vectors) { + for (int d = 0; d < DIMENSION; d++) { + vector[d] = (float) random.nextGaussian(); + } + } + return vectors; + } + + private static TurboQuantIndex quantized(float[][] vectors, int bits) { + final TurboQuantIndex index = new TurboQuantIndex(DIMENSION, bits, 42); + for (int i = 0; i < vectors.length; i++) { + index.add("v" + i, vectors[i]); + } + index.freeze(); + return index; + } + + @Test + void testSelfRetrievalSurvivesQuantization() { + final float[][] vectors = collection(); + final TurboQuantIndex index = quantized(vectors, 4); + for (int i = 0; i < vectors.length; i++) { + assertEquals("v" + i, index.topK(vectors[i], 1).get(0).id(), + "vector " + i + " must be its own nearest neighbor"); + } + } + + @ParameterizedTest + @ValueSource(ints = {2, 4}) + void testRecallAgainstTheExactIndex(int bits) { + final float[][] vectors = collection(); + final TurboQuantIndex quantized = quantized(vectors, bits); + final FlatFloatIndex exact = new FlatFloatIndex(DIMENSION); + for (int i = 0; i < vectors.length; i++) { + exact.add("v" + i, vectors[i]); + } + exact.freeze(); + + final Random random = new Random(SEED + 1); + double overlap = 0; + final int queries = 20; + final int k = 10; + for (int q = 0; q < queries; q++) { + final float[] query = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + query[d] = (float) random.nextGaussian(); + } + final Set truth = new HashSet<>(); + for (final VectorIndex.Hit hit : exact.topK(query, k)) { + truth.add(hit.id()); + } + for (final VectorIndex.Hit hit : quantized.topK(query, k)) { + if (truth.contains(hit.id())) { + overlap++; + } + } + } + final double recall = overlap / (queries * k); + // 4 bits tracks the exact ranking closely; 2 bits trades more recall for half the bytes. + assertTrue(recall >= (bits == 4 ? 0.85 : 0.6), + "recall@" + k + " at " + bits + " bits: " + recall); + } + + @Test + void testWriteReadRoundTripAnswersIdentically(@TempDir Path dir) throws IOException { + final float[][] vectors = collection(); + final TurboQuantIndex index = quantized(vectors, 4); + index.write(dir); + + final TurboQuantIndex reloaded = TurboQuantIndex.read(dir); + assertEquals(index.size(), reloaded.size()); + assertEquals(index.dimension(), reloaded.dimension()); + assertEquals(index.bits(), reloaded.bits()); + // The file stores the codes, scales, and seed, so a reloaded index scores identically. + assertEquals(index.topK(vectors[3], 5), reloaded.topK(vectors[3], 5)); + } + + @Test + void testWritingRequiresAFrozenNonEmptyIndex(@TempDir Path dir) { + final TurboQuantIndex building = new TurboQuantIndex(DIMENSION, 4, 42); + building.add("a", collection()[0]); + assertThrows(IllegalStateException.class, () -> building.write(dir)); + + final TurboQuantIndex empty = new TurboQuantIndex(DIMENSION, 4, 42); + empty.freeze(); + assertThrows(IllegalStateException.class, () -> empty.write(dir)); + } + + @ParameterizedTest + @ValueSource(ints = {2, 3, 4}) + void testBytesPerVectorIncludesPaddedCodesScaleAndNorm(int bits) { + final int unpaddedDimension = 65; + final TurboQuantIndex index = new TurboQuantIndex(unpaddedDimension, bits, 42); + final float[] vector = new float[unpaddedDimension]; + vector[0] = 1; + index.add("one", vector); + index.freeze(); + + final int paddedDimension = 128; + final double expected = (paddedDimension * bits + Byte.SIZE - 1) / Byte.SIZE + + 2 * Float.BYTES; + assertEquals(expected, index.bytesPerVector()); + } + + @Test + void testBytesPerVectorRequiresAFrozenNonEmptyIndex() { + final TurboQuantIndex building = new TurboQuantIndex(DIMENSION, 4, 42); + building.add("one", collection()[0]); + assertThrows(IllegalStateException.class, building::bytesPerVector); + + final TurboQuantIndex empty = new TurboQuantIndex(DIMENSION, 4, 42); + empty.freeze(); + assertThrows(IllegalStateException.class, empty::bytesPerVector); + } + + @Test + void testReadRejectsAMissingFile(@TempDir Path dir) { + assertThrows(IllegalArgumentException.class, () -> TurboQuantIndex.read(dir)); + } + + @Test + void testReadRejectsADuplicateId(@TempDir Path dir) throws IOException { + final TurboQuantIndex index = quantized(collection(), 4); + index.write(dir); + final List ids = Files.readAllLines(dir.resolve(TurboQuantIndex.IDS_FILE)); + ids.set(1, ids.get(0)); + Files.write(dir.resolve(TurboQuantIndex.IDS_FILE), ids); + + assertThrows(InvalidFormatException.class, () -> TurboQuantIndex.read(dir)); + } + + @Test + void testReadRejectsAnIdCountMismatch(@TempDir Path dir) throws IOException { + final TurboQuantIndex index = quantized(collection(), 4); + index.write(dir); + Files.write(dir.resolve(TurboQuantIndex.IDS_FILE), List.of("one-extra-id"), + StandardOpenOption.APPEND); + + final InvalidFormatException e = + assertThrows(InvalidFormatException.class, () -> TurboQuantIndex.read(dir)); + assertTrue(e.getMessage().contains("do not belong"), e.getMessage()); + } + + @ParameterizedTest + @ValueSource(ints = {0, 1, 5, 32}) + void testAnUnsupportedBitWidthIsRejected(int bits) { + assertThrows(IllegalArgumentException.class, () -> new TurboQuantIndex(DIMENSION, bits, 42)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/VectorIndexContractTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/VectorIndexContractTest.java new file mode 100644 index 0000000000..207143b06c --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/index/VectorIndexContractTest.java @@ -0,0 +1,160 @@ +/* + * 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.index; + +import java.util.List; +import java.util.function.IntFunction; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Named; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +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 {@link VectorIndex} lifecycle and validation contract, run against both implementations: + * build-phase validation, the freeze boundary, and the query result guarantees. + */ +class VectorIndexContractTest { + + private static final int DIMENSION = 8; + + static Stream indexes() { + return Stream.of( + Arguments.of(Named.>of("flat", FlatFloatIndex::new)), + Arguments.of(Named.>of("turboquant", + dimension -> new TurboQuantIndex(dimension, 4, 42)))); + } + + /** {@return a unit vector along the given axis} */ + private static float[] axis(int d) { + final float[] vector = new float[DIMENSION]; + vector[d] = 1f; + return vector; + } + + @ParameterizedTest + @MethodSource("indexes") + void testSelfRetrievalRanksTheIndexedVectorFirst(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + for (int d = 0; d < 4; d++) { + index.add("axis-" + d, axis(d)); + } + index.freeze(); + + for (int d = 0; d < 4; d++) { + final List hits = index.topK(axis(d), 2); + assertEquals(2, hits.size()); + assertEquals("axis-" + d, hits.get(0).id()); + assertTrue(hits.get(0).score() > 0.9, "self similarity: " + hits.get(0).score()); + assertTrue(hits.get(0).score() >= hits.get(1).score(), "ordering"); + } + } + + @ParameterizedTest + @MethodSource("indexes") + void testKBeyondTheSizeReturnsEverything(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("a", axis(0)); + index.add("b", axis(1)); + index.freeze(); + + assertEquals(2, index.topK(axis(0), 100).size()); + } + + @ParameterizedTest + @MethodSource("indexes") + void testAnEmptyIndexAnswersNoHits(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.freeze(); + assertEquals(0, index.size()); + assertTrue(index.topK(axis(0), 3).isEmpty()); + } + + @ParameterizedTest + @MethodSource("indexes") + void testAZeroQueryHasNoDirectionAndAnswersNoHits(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("a", axis(0)); + index.freeze(); + assertTrue(index.topK(new float[DIMENSION], 3).isEmpty()); + } + + @ParameterizedTest + @MethodSource("indexes") + void testAZeroVectorScoresZeroInsteadOfNaN(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("zero", new float[DIMENSION]); + index.freeze(); + + final List hits = index.topK(axis(0), 1); + assertEquals(1, hits.size()); + assertEquals(0.0, hits.get(0).score()); + } + + @ParameterizedTest + @MethodSource("indexes") + void testTheFreezeBoundaryIsEnforcedBothWays(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("a", axis(0)); + assertThrows(IllegalStateException.class, () -> index.topK(axis(0), 1)); + index.freeze(); + index.freeze(); + assertThrows(IllegalStateException.class, () -> index.add("b", axis(1))); + assertEquals(1, index.size()); + assertEquals(DIMENSION, index.dimension()); + } + + @ParameterizedTest + @MethodSource("indexes") + void testBuildPhaseValidation(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("a", axis(0)); + + assertThrows(IllegalArgumentException.class, () -> index.add(null, axis(1))); + assertThrows(IllegalArgumentException.class, () -> index.add(" ", axis(1))); + assertThrows(IllegalArgumentException.class, () -> index.add("b\nc", axis(1))); + assertThrows(IllegalArgumentException.class, () -> index.add("a", axis(1))); + assertThrows(IllegalArgumentException.class, () -> index.add("b", null)); + assertThrows(IllegalArgumentException.class, () -> index.add("b", new float[3])); + final float[] infinite = axis(1); + infinite[2] = Float.POSITIVE_INFINITY; + assertThrows(IllegalArgumentException.class, () -> index.add("b", infinite)); + } + + @ParameterizedTest + @MethodSource("indexes") + void testQueryValidation(IntFunction factory) { + final VectorIndex index = factory.apply(DIMENSION); + index.add("a", axis(0)); + index.freeze(); + + assertThrows(IllegalArgumentException.class, () -> index.topK(null, 1)); + assertThrows(IllegalArgumentException.class, () -> index.topK(new float[3], 1)); + assertThrows(IllegalArgumentException.class, () -> index.topK(axis(0), 0)); + } + + @ParameterizedTest + @MethodSource("indexes") + void testADimensionBelowOneIsRejected(IntFunction factory) { + assertThrows(IllegalArgumentException.class, () -> factory.apply(0)); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/resources/opennlp/embeddings/tiny-unigram.model b/opennlp-extensions/opennlp-embeddings/src/test/resources/opennlp/embeddings/tiny-unigram.model new file mode 100644 index 0000000000..b6e30611e4 Binary files /dev/null and b/opennlp-extensions/opennlp-embeddings/src/test/resources/opennlp/embeddings/tiny-unigram.model differ diff --git a/opennlp-extensions/opennlp-subword/pom.xml b/opennlp-extensions/opennlp-subword/pom.xml new file mode 100644 index 0000000000..b4d13c755e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/pom.xml @@ -0,0 +1,58 @@ + + + + + + 4.0.0 + + org.apache.opennlp + opennlp-extensions + 3.0.0-SNAPSHOT + + + opennlp-subword + jar + Apache OpenNLP :: Ext :: Subword + + + + org.apache.opennlp + opennlp-api + + + + org.junit.jupiter + junit-jupiter-api + test + + + + org.junit.jupiter + junit-jupiter-engine + test + + + + org.junit.jupiter + junit-jupiter-params + test + + + diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/BpeEncoder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/BpeEncoder.java new file mode 100644 index 0000000000..1c813d5314 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/BpeEncoder.java @@ -0,0 +1,232 @@ +/* + * 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.subword.sentencepiece; + +import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; + +/** + * Byte-pair-encoding segmentation: the normalized text starts as single characters (or + * user-defined symbols, which are frozen whole) and adjacent pairs merge greedily, highest piece + * score first, until no adjacent pair forms a vocabulary piece. + * + *

Only pieces of the normal, user-defined, and unused types participate in merges; a merge that + * lands on an unused piece is re-segmented back into its constituents.

+ */ +final class BpeEncoder implements Serializable { + + private static final long serialVersionUID = -57799941356582785L; + + private static final int MAX_RESEGMENT_DEPTH = 100; + + private final Map pieces; + private final float[] scores; + private final boolean[] unused; + private final boolean[] reserved; + private final int unkId; + private final PieceTrie userDefinedMatcher; + + /** + * Instantiates the encoder. + * + * @param pieces All pieces by content, mapping to their ids. + * @param scores The score of every piece, indexed by id. + * @param unused Whether each id has the unused piece type. + * @param reserved Whether each id is excluded from merging (any type other than + * normal, user-defined, or unused). + * @param unkId The id of the unknown piece. + * @param userDefinedMatcher Longest-match trie over user-defined symbols, or null when the + * model defines none. + */ + BpeEncoder(Map pieces, float[] scores, boolean[] unused, boolean[] reserved, + int unkId, PieceTrie userDefinedMatcher) { + this.pieces = pieces; + this.scores = scores; + this.unused = unused; + this.reserved = reserved; + this.unkId = unkId; + this.userDefinedMatcher = userDefinedMatcher; + } + + /** + * A candidate merge of the symbols at indices {@code left} and {@code right}; {@code size} is the + * merged byte length, used to detect staleness after either side has changed. + */ + private record Pair(int left, int right, float score, int size) { + } + + /** + * Segments normalized text. + * + * @param normalized The buffer holding the normalized UTF-8 bytes; must not be null. + * @param size The number of valid bytes in {@code normalized}. + * @return The segments covering all bytes, in text order. + * @throws IllegalArgumentException Thrown if {@code normalized} is null. + */ + List encode(byte[] normalized, int size) { + if (normalized == null) { + throw new IllegalArgumentException("normalized must not be null"); + } + if (size == 0) { + return List.of(); + } + + // The symbol list as index-linked ranges of the normalized bytes; merged-away symbols + // become empty ranges. Freeze flags travel as 0/1 bytes parallel to the ranges. + final IntBuilder fromB = new IntBuilder(size); + final IntBuilder toB = new IntBuilder(size); + final ByteBuilder freezeB = new ByteBuilder(size); + int position = 0; + while (position < size) { + int matched = 0; + if (userDefinedMatcher != null) { + matched = userDefinedMatcher.longestMatch(normalized, size, position); + } + final boolean frozen = matched > 0; + final int length = frozen ? matched + : Math.min(SentencePieceNormalizer.utf8Length(normalized[position]), + size - position); + fromB.append(position); + toB.append(position + length); + freezeB.append(frozen ? (byte) 1 : (byte) 0); + position += length; + } + final int symbolCount = freezeB.length(); + final int[] from = fromB.toArray(); + final int[] to = toB.toArray(); + final byte[] frozenFlags = freezeB.array(); + final int[] prev = new int[symbolCount]; + final int[] next = new int[symbolCount]; + final boolean[] freeze = new boolean[symbolCount]; + for (int i = 0; i < symbolCount; i++) { + prev[i] = i - 1; + next[i] = i + 1 < symbolCount ? i + 1 : -1; + freeze[i] = frozenFlags[i] != 0; + } + + // Higher score first; equal scores break towards the leftmost pair. + final PriorityQueue agenda = new PriorityQueue<>((a, b) -> { + final int byScore = Float.compare(b.score(), a.score()); + return byScore != 0 ? byScore : Integer.compare(a.left(), b.left()); + }); + // Merged piece content mapped back to its two constituents, for re-segmenting unused pieces. + final Map revMerge = new HashMap<>(); + + for (int left = 0; left + 1 < symbolCount; left++) { + maybeAddPair(normalized, from, to, freeze, left, left + 1, agenda, revMerge); + } + + while (!agenda.isEmpty()) { + final Pair top = agenda.poll(); + // Skips entries made stale by an earlier merge of either side. + if (from[top.left()] == to[top.left()] || from[top.right()] == to[top.right()] + || to[top.left()] - from[top.left()] + to[top.right()] - from[top.right()] + != top.size()) { + continue; + } + + // Replaces the pair with the merged symbol. + to[top.left()] = to[top.right()]; + next[top.left()] = next[top.right()]; + if (next[top.right()] >= 0) { + prev[next[top.right()]] = top.left(); + } + from[top.right()] = to[top.right()]; + + maybeAddPair(normalized, from, to, freeze, prev[top.left()], top.left(), agenda, revMerge); + maybeAddPair(normalized, from, to, freeze, top.left(), next[top.left()], agenda, revMerge); + } + + final List output = new ArrayList<>(symbolCount); + int consumed = 0; + for (int index = 0; index != -1; index = next[index]) { + final String piece = + new String(normalized, from[index], to[index] - from[index], StandardCharsets.UTF_8); + consumed = resegment(piece, consumed, 0, revMerge, output); + } + return output; + } + + /** + * Offers the adjacent symbol pair {@code (left, right)} as a merge candidate: the pair joins + * the agenda only when the concatenation is a mergeable vocabulary piece, and a merge landing + * on an unused piece is remembered in {@code revMerge} for later re-segmentation. + * + * @param normalized The buffer holding the normalized UTF-8 bytes. + * @param from Per symbol, the inclusive start offset in {@code normalized}. + * @param to Per symbol, the exclusive end offset in {@code normalized}. + * @param freeze Per symbol, whether it is a user-defined symbol excluded from merging. + * @param left The index of the left symbol, or {@code -1} for none. + * @param right The index of the right symbol, or {@code -1} for none. + * @param agenda The merge agenda to add to. + * @param revMerge The map from a merged piece to its two constituents. + */ + private void maybeAddPair(byte[] normalized, int[] from, int[] to, boolean[] freeze, + int left, int right, PriorityQueue agenda, + Map revMerge) { + if (left == -1 || right == -1 || freeze[left] || freeze[right]) { + return; + } + final String piece = + new String(normalized, from[left], to[right] - from[left], StandardCharsets.UTF_8); + final Integer id = pieces.get(piece); + if (id == null || id == unkId || reserved[id]) { + return; + } + agenda.add(new Pair(left, right, scores[id], to[right] - from[left])); + if (unused[id]) { + revMerge.put(piece, new String[] { + new String(normalized, from[left], to[left] - from[left], StandardCharsets.UTF_8), + new String(normalized, from[right], to[right] - from[right], StandardCharsets.UTF_8)}); + } + } + + /** + * Emits a symbol, splitting a piece of the unused type back into the pieces it was merged from. + * Positions are assigned by a running cursor; constituent byte lengths always sum to the merged + * length, so the cursor stays aligned with the normalized bytes. + * + * @param piece The piece content to emit. + * @param consumed The running byte cursor into the normalized text. + * @param depth The current recursion depth. + * @param revMerge The map from a merged piece to its two constituents. + * @param output The segment list to append to. + * @return The updated byte cursor. + */ + private int resegment(String piece, int consumed, int depth, Map revMerge, + List output) { + final Integer mapped = pieces.get(piece); + final int id = mapped == null ? unkId : mapped; + final int byteLength = piece.getBytes(StandardCharsets.UTF_8).length; + if (depth > MAX_RESEGMENT_DEPTH || !unused[id]) { + output.add(new Segment(consumed, consumed + byteLength, id)); + return consumed + byteLength; + } + final String[] parts = revMerge.get(piece); + if (parts == null) { + output.add(new Segment(consumed, consumed + byteLength, id)); + return consumed + byteLength; + } + consumed = resegment(parts[0], consumed, depth + 1, revMerge, output); + return resegment(parts[1], consumed, depth + 1, revMerge, output); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java new file mode 100644 index 0000000000..432f2104dc --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java @@ -0,0 +1,113 @@ +/* + * 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.subword.sentencepiece; + +import java.util.Arrays; + +/** A growable byte buffer supporting append, truncate, and suffix comparison. */ +final class ByteBuilder { + + /** The smallest backing array, so tiny requested capacities still grow geometrically. */ + private static final int MIN_CAPACITY = 16; + + private byte[] data; + private int length; + + /** + * Instantiates the buffer. + * + * @param capacity The initial capacity hint. + */ + ByteBuilder(int capacity) { + data = new byte[Math.max(capacity, MIN_CAPACITY)]; + } + + /** + * Appends one byte. + * + * @param b The byte to append. + */ + void append(byte b) { + if (length == data.length) { + data = Arrays.copyOf(data, grownLength()); + } + data[length++] = b; + } + + /** + * Appends a run of bytes. + * + * @param source The source array. + * @param from The inclusive start offset in {@code source}. + * @param count The number of bytes to append. + */ + void append(byte[] source, int from, int count) { + while (length + count > data.length) { + data = Arrays.copyOf(data, grownLength()); + } + System.arraycopy(source, from, data, length, count); + length += count; + } + + /** {@return the next backing-array length under the 1.5x growth policy} */ + private int grownLength() { + return data.length + (data.length >> 1); + } + + /** {@return the number of valid bytes} */ + int length() { + return length; + } + + /** + * Shrinks the valid length. + * + * @param newLength The new length, not negative and not greater than the current length. + * @throws IllegalArgumentException Thrown if {@code newLength} is negative or greater than the + * current length. + */ + void truncate(int newLength) { + if (newLength < 0 || newLength > length) { + throw new IllegalArgumentException( + "The new length " + newLength + " is outside [0, " + length + "]."); + } + length = newLength; + } + + /** + * Tests whether the valid bytes end with the given suffix. + * + * @param suffix The suffix to test. + * @return {@code true} when the buffer ends with {@code suffix}. + */ + boolean endsWith(byte[] suffix) { + if (length < suffix.length) { + return false; + } + return Arrays.equals(data, length - suffix.length, length, suffix, 0, suffix.length); + } + + /** {@return a trimmed copy of the valid bytes} */ + byte[] toArray() { + return Arrays.copyOf(data, length); + } + + /** {@return the backing array, valid up to {@link #length()}} */ + byte[] array() { + return data; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/DoubleArrayTrie.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/DoubleArrayTrie.java new file mode 100644 index 0000000000..5bed8a6852 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/DoubleArrayTrie.java @@ -0,0 +1,150 @@ +/* + * 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.subword.sentencepiece; + +import java.io.Serializable; + +/** + * Read-only lookup over a serialized Darts-clone double-array trie, the dictionary format + * embedded in a SentencePiece model's precompiled character map. + * + *

Each unit is one little-endian 32-bit word encoding a label, an offset to the unit's + * children, and a leaf flag; traversal XORs the offset with the next key byte. Only the longest + * prefix match is needed here, so this walks the byte key once and remembers the last accepting + * state. Out-of-range unit references, which a well-formed trie never produces, fail loudly + * rather than reading arbitrary memory.

+ * + *

This class is an independent re-implementation of the reader side of that format, written + * against the published double-array literature: the trie itself is Aoe's double-array, the unit + * encoding is the compact static variant of Yata et al., and the bit-9 offset extension is the + * two-kinds-of-offset scheme described by Kanda et al. Darts-clone, by the same Yata, is the + * reference implementation of the format.

+ * + * @see Darts-clone + * @see Aoe (1989): An Efficient Digital Search + * Algorithm by Using a Double-Array Structure + * @see Yata et al. (2007): A compact static + * double-array keeping character codes + * @see Kanda et al. (2023): Engineering faster + * double-array Aho-Corasick automata + */ +final class DoubleArrayTrie implements Serializable { + + private static final long serialVersionUID = -1572336116472261588L; + + // A non-leaf unit stores its transition label in the low 8 bits and the leaf flag in the sign + // bit. Key bytes are in [0, 255] with the sign bit clear, so comparing (unit & this mask) + // against a key byte both matches the label and rejects leaf units in one test. + private static final int LEAF_FLAG_AND_LABEL_MASK = 0x800000FF; + + // A leaf unit stores the key's value in its low 31 bits; the sign bit is the leaf flag. + private static final int LEAF_VALUE_MASK = 0x7FFFFFFF; + + // Bit 8 of a non-leaf unit marks that one of its children is a leaf holding this key's value. + private static final int HAS_LEAF_BIT = 8; + + private final int[] units; + + /** + * Wraps serialized trie units. + * + * @param data The bytes holding the units; must not be null. + * @param offset The offset of the first unit byte. + * @param length The number of bytes; must be a positive multiple of four. + * @throws IllegalArgumentException Thrown if {@code length} is not a positive multiple of four. + */ + DoubleArrayTrie(byte[] data, int offset, int length) { + if (length <= 0 || (length & 3) != 0) { + throw new IllegalArgumentException( + "The trie length " + length + " is not a positive multiple of four bytes."); + } + units = new int[length >> 2]; + for (int i = 0; i < units.length; i++) { + final int base = offset + (i << 2); + units[i] = (data[base] & 0xFF) | (data[base + 1] & 0xFF) << 8 + | (data[base + 2] & 0xFF) << 16 | (data[base + 3] & 0xFF) << 24; + } + } + + /** + * Finds the longest key that is a prefix of {@code key[from, to)}. + * + * @param key The byte key to match against; must not be null. + * @param from The inclusive start of the query window. + * @param to The exclusive end of the query window. + * @return {@code (value << 32) | matchedLength} for the longest match, or {@code -1} when no + * key matches. Values are non-negative, so the result is negative only on no-match. + * @throws IllegalArgumentException Thrown if the trie data references a unit outside its + * bounds, indicating corrupt data. + */ + long longestPrefixMatch(byte[] key, int from, int to) { + // The JVM's own bounds checks guard the walk; the catch below translates an out-of-range unit + // reference from corrupt data into a loud failure. + final int[] u = units; + try { + long result = -1; + int nodePos = 0; + int unit = u[0]; + nodePos ^= offset(unit); + for (int i = from; i < to; i++) { + final int b = key[i] & 0xFF; + nodePos ^= b; + unit = u[nodePos]; + if ((unit & LEAF_FLAG_AND_LABEL_MASK) != b) { + return result; + } + nodePos ^= offset(unit); + if (((unit >>> HAS_LEAF_BIT) & 1) == 1) { + final int value = u[nodePos] & LEAF_VALUE_MASK; + result = ((long) value << 32) | (i - from + 1); + } + } + return result; + } catch (ArrayIndexOutOfBoundsException e) { + throw new IllegalArgumentException( + "The trie references a unit outside its " + u.length + " units.", e); + } + } + + /** + * Tests whether any key starts with the given byte, which is exactly whether the root has a + * transition on it; used to precompute the first-byte gate of the normalizer scan. + * + * @param b The first key byte as an unsigned value. + * @return {@code true} if some key starts with {@code b}. + */ + boolean hasTransitionFromRoot(int b) { + final int root = units[0]; + final int nodePos = offset(root) ^ b; + if (nodePos < 0 || nodePos >= units.length) { + return false; + } + return (units[nodePos] & LEAF_FLAG_AND_LABEL_MASK) == b; + } + + /** + * Returns the offset from a unit to its children, as encoded by Darts-clone: bits 10 to 30 hold + * the raw offset, and bit 9 is an extension flag that scales it by 256 for far-away children. + * + * @param unit The unit word. + * @return The child offset. + */ + private static int offset(int unit) { + final int raw = unit >>> 10; + return (unit & 1 << 9) == 0 ? raw : raw << 8; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java new file mode 100644 index 0000000000..b06fb73209 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java @@ -0,0 +1,99 @@ +/* + * 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.subword.sentencepiece; + +import java.util.Arrays; + +/** A growable int buffer supporting append, indexed read, and truncate. */ +final class IntBuilder { + + /** The smallest backing array, so tiny requested capacities still grow geometrically. */ + private static final int MIN_CAPACITY = 16; + + private int[] data; + private int length; + + /** + * Instantiates the buffer. + * + * @param capacity The initial capacity hint. + */ + IntBuilder(int capacity) { + data = new int[Math.max(capacity, MIN_CAPACITY)]; + } + + /** + * Appends one value. + * + * @param value The value to append. + */ + void append(int value) { + if (length == data.length) { + data = Arrays.copyOf(data, grownLength()); + } + data[length++] = value; + } + + /** {@return the next backing-array length under the 1.5x growth policy} */ + private int grownLength() { + return data.length + (data.length >> 1); + } + + /** + * Reads a value by index. + * + * @param index An index in {@code [0, length())}. + * @return The value at {@code index}. + * @throws IndexOutOfBoundsException Thrown if {@code index} is out of range. + */ + int get(int index) { + if (index >= length) { + throw new IndexOutOfBoundsException("index " + index + " is outside [0, " + length + ")"); + } + return data[index]; + } + + /** {@return the number of valid values} */ + int length() { + return length; + } + + /** + * Shrinks the valid length. + * + * @param newLength The new length, not negative and not greater than the current length. + * @throws IllegalArgumentException Thrown if {@code newLength} is negative or greater than the + * current length. + */ + void truncate(int newLength) { + if (newLength < 0 || newLength > length) { + throw new IllegalArgumentException( + "The new length " + newLength + " is outside [0, " + length + "]."); + } + length = newLength; + } + + /** {@return a trimmed copy of the valid values} */ + int[] toArray() { + return Arrays.copyOf(data, length); + } + + /** {@return the backing array, valid up to {@link #length()}} */ + int[] array() { + return data; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ModelProtoReader.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ModelProtoReader.java new file mode 100644 index 0000000000..e9f4cb0bbd --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ModelProtoReader.java @@ -0,0 +1,416 @@ +/* + * 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.subword.sentencepiece; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; + +import opennlp.tools.util.InvalidFormatException; + +/** + * Reads the binary {@code ModelProto} serialization of a SentencePiece {@code .model} file. + * + *

The format is the standard protocol-buffer wire encoding of one flat message + * ({@code sentencepiece_model.proto}, Apache License 2.0). This reader walks the tag stream + * directly and keeps only the fields inference needs: the pieces with scores and types, the + * normalizer spec, the trainer-spec fields that change runtime behavior, and the embedded + * self-test samples. Unknown fields are skipped, and malformed input fails loudly.

+ * + * @see + * sentencepiece_model.proto + */ +final class ModelProtoReader { + + // Wire types of the protocol-buffer encoding. + private static final int WIRE_VARINT = 0; + private static final int WIRE_FIXED64 = 1; + private static final int WIRE_LEN = 2; + private static final int WIRE_FIXED32 = 5; + + // Field numbers of the ModelProto message in sentencepiece_model.proto. + private static final int FIELD_MODEL_PIECES = 1; + private static final int FIELD_MODEL_TRAINER_SPEC = 2; + private static final int FIELD_MODEL_NORMALIZER_SPEC = 3; + private static final int FIELD_MODEL_SELF_TEST_DATA = 4; + + // Field numbers of the ModelProto.SentencePiece sub-message. + private static final int FIELD_PIECE_PIECE = 1; + private static final int FIELD_PIECE_SCORE = 2; + private static final int FIELD_PIECE_TYPE = 3; + + // Field numbers of the TrainerSpec sub-message. + private static final int FIELD_TRAINER_MODEL_TYPE = 3; + private static final int FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX = 24; + private static final int FIELD_TRAINER_BYTE_FALLBACK = 35; + + // Field numbers of the NormalizerSpec sub-message. + private static final int FIELD_NORMALIZER_PRECOMPILED_CHARSMAP = 2; + private static final int FIELD_NORMALIZER_ADD_DUMMY_PREFIX = 3; + private static final int FIELD_NORMALIZER_REMOVE_EXTRA_WHITESPACES = 4; + private static final int FIELD_NORMALIZER_ESCAPE_WHITESPACES = 5; + + // Field numbers of the SelfTestData sub-message and its Sample entries. + private static final int FIELD_SELF_TEST_SAMPLES = 1; + private static final int FIELD_SAMPLE_INPUT = 1; + private static final int FIELD_SAMPLE_EXPECTED = 2; + + private final byte[] data; + private int pos; + + /** + * Prepares a reader positioned at the start of the given bytes; {@link #read(byte[])} drives + * the actual parse. + * + * @param data The raw bytes of a {@code .model} file. + */ + private ModelProtoReader(byte[] data) { + this.data = data; + } + + /** + * Parses a serialized {@code ModelProto}. + * + * @param data The raw bytes of a {@code .model} file; must not be null. + * @return The parsed model description. + * @throws IllegalArgumentException Thrown if {@code data} is null. + * @throws InvalidFormatException Thrown if the bytes are not a well-formed model. + */ + static RawModel read(byte[] data) throws InvalidFormatException { + if (data == null) { + throw new IllegalArgumentException("data must not be null"); + } + final ModelProtoReader reader = new ModelProtoReader(data); + final RawModel model = new RawModel(); + while (reader.pos < data.length) { + final long tag = reader.varint(); + switch (fieldOf(tag)) { + case FIELD_MODEL_PIECES -> reader.piece(model, reader.lenPayload(tag)); + case FIELD_MODEL_TRAINER_SPEC -> reader.trainerSpec(model, reader.lenPayload(tag)); + case FIELD_MODEL_NORMALIZER_SPEC -> reader.normalizerSpec(model, reader.lenPayload(tag)); + case FIELD_MODEL_SELF_TEST_DATA -> reader.selfTestData(model, reader.lenPayload(tag)); + default -> reader.skip(tag); + } + } + if (model.pieces.isEmpty()) { + throw new InvalidFormatException("The model defines no pieces."); + } + return model; + } + + /** + * Extracts the field number from a wire-format tag. + * + * @param tag The field tag. + * @return The field number. + */ + private static int fieldOf(long tag) { + return (int) (tag >>> 3); + } + + /** + * Extracts the wire type from a wire-format tag. + * + * @param tag The field tag. + * @return The wire type. + */ + private static int wireTypeOf(long tag) { + return (int) (tag & 7); + } + + /** + * Parses one {@code SentencePiece} sub-message and appends its piece, score, and type. + * + * @param model The model to append to. + * @param end The exclusive end offset of the sub-message payload. + * @throws InvalidFormatException Thrown if the sub-message is malformed or defines an empty + * piece or a non-finite score. + */ + private void piece(RawModel model, int end) throws InvalidFormatException { + String piece = null; + float score = 0; + int type = RawModel.TYPE_NORMAL; + while (pos < end) { + final long tag = varint(); + switch (fieldOf(tag)) { + case FIELD_PIECE_PIECE -> piece = utf8(lenPayload(tag)); + case FIELD_PIECE_SCORE -> score = fixed32Float(tag); + case FIELD_PIECE_TYPE -> type = (int) varintOf(tag); + default -> skip(tag); + } + } + if (piece == null || piece.isEmpty()) { + throw new InvalidFormatException( + "The model contains an empty piece at index " + model.pieces.size() + "."); + } + if (Float.isNaN(score) || Float.isInfinite(score)) { + throw new InvalidFormatException("The score of piece '" + piece + "' is not finite."); + } + model.pieces.add(piece); + model.scores.add(score); + model.types.add(type); + } + + /** + * Parses the {@code TrainerSpec} sub-message, keeping the fields that change runtime behavior. + * + * @param model The model to populate. + * @param end The exclusive end offset of the sub-message payload. + * @throws InvalidFormatException Thrown if the sub-message is malformed. + */ + private void trainerSpec(RawModel model, int end) throws InvalidFormatException { + while (pos < end) { + final long tag = varint(); + switch (fieldOf(tag)) { + case FIELD_TRAINER_MODEL_TYPE -> model.modelType = (int) varintOf(tag); + case FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX -> + model.treatWhitespaceAsSuffix = varintOf(tag) != 0; + case FIELD_TRAINER_BYTE_FALLBACK -> model.byteFallback = varintOf(tag) != 0; + default -> skip(tag); + } + } + } + + /** + * Parses the {@code NormalizerSpec} sub-message: the precompiled character map and the + * whitespace-handling flags. + * + * @param model The model to populate. + * @param end The exclusive end offset of the sub-message payload. + * @throws InvalidFormatException Thrown if the sub-message is malformed. + */ + private void normalizerSpec(RawModel model, int end) throws InvalidFormatException { + while (pos < end) { + final long tag = varint(); + switch (fieldOf(tag)) { + case FIELD_NORMALIZER_PRECOMPILED_CHARSMAP -> + model.precompiledCharsMap = bytes(lenPayload(tag)); + case FIELD_NORMALIZER_ADD_DUMMY_PREFIX -> model.addDummyPrefix = varintOf(tag) != 0; + case FIELD_NORMALIZER_REMOVE_EXTRA_WHITESPACES -> + model.removeExtraWhitespaces = varintOf(tag) != 0; + case FIELD_NORMALIZER_ESCAPE_WHITESPACES -> model.escapeWhitespaces = varintOf(tag) != 0; + default -> skip(tag); + } + } + } + + /** + * Parses the {@code SelfTestData} sub-message, collecting the input and expected-segmentation + * sample pairs. + * + * @param model The model to populate. + * @param end The exclusive end offset of the sub-message payload. + * @throws InvalidFormatException Thrown if the sub-message is malformed. + */ + private void selfTestData(RawModel model, int end) throws InvalidFormatException { + while (pos < end) { + final long tag = varint(); + if (fieldOf(tag) == FIELD_SELF_TEST_SAMPLES) { + final int sampleEnd = lenPayload(tag); + String input = null; + String expected = null; + while (pos < sampleEnd) { + final long sampleTag = varint(); + switch (fieldOf(sampleTag)) { + case FIELD_SAMPLE_INPUT -> input = utf8(lenPayload(sampleTag)); + case FIELD_SAMPLE_EXPECTED -> expected = utf8(lenPayload(sampleTag)); + default -> skip(sampleTag); + } + } + if (input != null && expected != null) { + model.selfTestInputs.add(input); + model.selfTestExpected.add(expected); + } + } else { + skip(tag); + } + } + } + + /** + * Reads the length prefix of a length-delimited field and returns the exclusive end offset of its + * payload. + * + * @param tag The field tag, whose wire type must be length-delimited. + * @return The exclusive end offset of the payload. + * @throws InvalidFormatException Thrown if the wire type is wrong or the length runs past the + * input. + */ + private int lenPayload(long tag) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_LEN) { + throw malformed("field " + fieldOf(tag) + " is not length-delimited"); + } + final long length = varint(); + if (length < 0 || pos + length > data.length) { + throw malformed("length " + length + " exceeds the remaining input"); + } + return pos + (int) length; + } + + /** + * Reads the varint value of a field after checking its wire type. + * + * @param tag The field tag, whose wire type must be varint. + * @return The decoded value. + * @throws InvalidFormatException Thrown if the wire type is wrong or the varint is malformed. + */ + private long varintOf(long tag) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_VARINT) { + throw malformed("field " + fieldOf(tag) + " is not a varint"); + } + return varint(); + } + + /** + * Reads the little-endian 32-bit float value of a field after checking its wire type. + * + * @param tag The field tag, whose wire type must be 32-bit. + * @return The decoded float. + * @throws InvalidFormatException Thrown if the wire type is wrong or the input is truncated. + */ + private float fixed32Float(long tag) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_FIXED32) { + throw malformed("field " + fieldOf(tag) + " is not a 32-bit value"); + } + if (pos + 4 > data.length) { + throw malformed("truncated 32-bit value"); + } + final int bits = (data[pos] & 0xFF) | (data[pos + 1] & 0xFF) << 8 + | (data[pos + 2] & 0xFF) << 16 | (data[pos + 3] & 0xFF) << 24; + pos += 4; + return Float.intBitsToFloat(bits); + } + + /** + * Decodes the bytes from the current position up to {@code end} as UTF-8, advancing past them. + * + * @param end The exclusive end offset of the payload. + * @return The decoded string. + */ + private String utf8(int end) { + final String s = new String(data, pos, end - pos, StandardCharsets.UTF_8); + pos = end; + return s; + } + + /** + * Copies the bytes from the current position up to {@code end}, advancing past them. + * + * @param end The exclusive end offset of the payload. + * @return The copied bytes. + */ + private byte[] bytes(int end) { + final byte[] b = new byte[end - pos]; + System.arraycopy(data, pos, b, 0, b.length); + pos = end; + return b; + } + + /** + * Reads a base-128 varint from the current position, advancing past it. + * + * @return The decoded value. + * @throws InvalidFormatException Thrown if the input ends mid-varint or the varint exceeds 64 + * bits. + */ + private long varint() throws InvalidFormatException { + long value = 0; + for (int shift = 0; shift < 64; shift += 7) { + if (pos >= data.length) { + throw malformed("truncated varint"); + } + final byte b = data[pos++]; + value |= (long) (b & 0x7F) << shift; + if ((b & 0x80) == 0) { + return value; + } + } + throw malformed("varint exceeds 64 bits"); + } + + /** + * Skips the value of an unrecognized field according to its wire type. + * + * @param tag The field tag. + * @throws InvalidFormatException Thrown if the wire type is unsupported or the value runs past + * the input. + */ + private void skip(long tag) throws InvalidFormatException { + switch (wireTypeOf(tag)) { + case WIRE_VARINT -> varint(); + case WIRE_FIXED64 -> advance(8); + case WIRE_LEN -> pos = lenPayload(tag); + case WIRE_FIXED32 -> advance(4); + default -> throw malformed("unsupported wire type " + wireTypeOf(tag)); + } + } + + /** + * Advances the position by a fixed number of bytes. + * + * @param count The number of bytes to skip. + * @throws InvalidFormatException Thrown if fewer than {@code count} bytes remain. + */ + private void advance(int count) throws InvalidFormatException { + if (pos + count > data.length) { + throw malformed("truncated field"); + } + pos += count; + } + + /** + * Creates the exception for malformed input, carrying the current byte position. + * + * @param detail A short description of what is malformed. + * @return The exception to throw. + */ + private InvalidFormatException malformed(String detail) { + return new InvalidFormatException( + "The model data is malformed at byte " + pos + ": " + detail + "."); + } + + /** The fields of a {@code ModelProto} that inference needs, with the proto's defaults. */ + static final class RawModel { + + static final int TYPE_NORMAL = 1; + static final int TYPE_UNKNOWN = 2; + static final int TYPE_CONTROL = 3; + static final int TYPE_USER_DEFINED = 4; + static final int TYPE_UNUSED = 5; + static final int TYPE_BYTE = 6; + + static final int MODEL_TYPE_UNIGRAM = 1; + static final int MODEL_TYPE_BPE = 2; + + final List pieces = new ArrayList<>(); + final List scores = new ArrayList<>(); + final List types = new ArrayList<>(); + + int modelType = MODEL_TYPE_UNIGRAM; + boolean byteFallback = false; + boolean treatWhitespaceAsSuffix = false; + + byte[] precompiledCharsMap = new byte[0]; + boolean addDummyPrefix = true; + boolean removeExtraWhitespaces = true; + boolean escapeWhitespaces = true; + + final List selfTestInputs = new ArrayList<>(); + final List selfTestExpected = new ArrayList<>(); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/PieceTrie.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/PieceTrie.java new file mode 100644 index 0000000000..42477251ab --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/PieceTrie.java @@ -0,0 +1,311 @@ +/* + * 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.subword.sentencepiece; + +import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Comparator; + +/** + * An immutable byte-level trie over vocabulary pieces, packed into flat arrays. + * + *

Encoding walks it one byte at a time ({@link #step(int, byte)}), so every piece that starts + * at a given input position is enumerated in one forward pass. Wide nodes dispatch through a + * 256-entry direct table and narrow nodes scan a short sorted label slice; both layouts enumerate + * identical transitions.

+ */ +final class PieceTrie implements Serializable { + + private static final long serialVersionUID = 30340094783102906L; + + /** The node id returned when no transition exists. */ + static final int DEAD = -1; + + // A node dispatches through a 256-entry slice of directPool when it has more children than + // this; otherwise a linear scan of the sorted label slice is used. + private static final int DIRECT_THRESHOLD = 8; + + /** The width of a wide node's direct dispatch slice. */ + private static final int DIRECT_TABLE_SIZE = 256; + + // Per node: the slice [childStart[n], childStart[n + 1]) of labels/childNodes, and the piece id + // accepted at the node, or -1. Wide nodes additionally index directPool at directStart[n]. + private final int[] childStart; + private final byte[] labels; + private final int[] childNodes; + private final int[] values; + private final int[] directStart; + private final int[] directPool; + + /** + * Wraps the packed arrays produced by {@link Builder} and derives the direct-dispatch tables + * for wide nodes. + * + * @param childStart Per node, the start of its edge slice; one trailing entry marks the end. + * @param labels The transition label of every edge. + * @param childNodes The target node of every edge, parallel to {@code labels}. + * @param values Per node, the accepted piece id, or {@code -1}. + */ + private PieceTrie(int[] childStart, byte[] labels, int[] childNodes, int[] values) { + this.childStart = childStart; + this.labels = labels; + this.childNodes = childNodes; + this.values = values; + this.directStart = new int[values.length]; + int wide = 0; + for (int node = 0; node < values.length; node++) { + if (childStart[node + 1] - childStart[node] > DIRECT_THRESHOLD) { + directStart[node] = wide * DIRECT_TABLE_SIZE; + wide++; + } else { + directStart[node] = -1; + } + } + this.directPool = new int[wide * DIRECT_TABLE_SIZE]; + Arrays.fill(directPool, DEAD); + for (int node = 0; node < values.length; node++) { + final int direct = directStart[node]; + if (direct >= 0) { + for (int edge = childStart[node]; edge < childStart[node + 1]; edge++) { + directPool[direct + (labels[edge] & 0xFF)] = childNodes[edge]; + } + } + } + } + + /** + * Builds a trie from pieces and their ids. + * + * @param pieces The UTF-8 bytes of each piece; must not be null or contain empty keys. + * @param ids The id stored for each piece, parallel to {@code pieces}. + * @return The packed trie. + * @throws IllegalArgumentException Thrown if a piece is defined more than once. + */ + static PieceTrie build(byte[][] pieces, int[] ids) { + final Integer[] order = new Integer[pieces.length]; + for (int i = 0; i < order.length; i++) { + order[i] = i; + } + Arrays.sort(order, Comparator.comparing(i -> pieces[i], Arrays::compareUnsigned)); + + // First pass counts nodes and edges, second pass fills the packed arrays; both walk the + // sorted keys with the same recursion, so the shapes agree by construction. + final Builder builder = new Builder(pieces, ids, order); + builder.count(0, pieces.length, 0); + builder.allocate(); + builder.fill(0, pieces.length, 0); + return new PieceTrie(builder.childStart, builder.labels, builder.childNodes, builder.values); + } + + /** {@return the root node id} */ + int root() { + return 0; + } + + /** + * Follows the transition labeled {@code b}. + * + * @param node The current node id. + * @param b The next key byte. + * @return The child node id, or {@link #DEAD} when no such transition exists. + */ + int step(int node, byte b) { + final int direct = directStart[node]; + if (direct >= 0) { + return directPool[direct + (b & 0xFF)]; + } + final int to = childStart[node + 1]; + for (int edge = childStart[node]; edge < to; edge++) { + if (labels[edge] == b) { + return childNodes[edge]; + } + } + return DEAD; + } + + /** + * Returns the piece id accepted at a node. + * + * @param node The node id. + * @return The id, or {@code -1} when the node accepts no piece. + */ + int value(int node) { + return values[node]; + } + + /** + * Returns the byte length of the longest piece in this trie that is a prefix of + * {@code input[from, inputLength)}. + * + * @param input The UTF-8 input buffer; must not be null. + * @param inputLength The number of valid bytes in {@code input}. + * @param from The offset to match from. + * @return The matched length in bytes, or zero when no piece matches. + */ + int longestMatch(byte[] input, int inputLength, int from) { + int node = root(); + int longest = 0; + for (int i = from; i < inputLength; i++) { + node = step(node, input[i]); + if (node == DEAD) { + break; + } + if (value(node) >= 0) { + longest = i - from + 1; + } + } + return longest; + } + + /** + * Creates the exception reported wherever a vocabulary piece turns out to be defined twice, + * keeping the message identical across all detection sites. + * + * @param piece The duplicated piece content. + * @return The exception to throw. + */ + static IllegalArgumentException duplicatePiece(String piece) { + return new IllegalArgumentException("The piece '" + piece + "' is defined more than once."); + } + + // Builds the packed form from keys sorted by unsigned byte order. Key ranges sharing a prefix + // are contiguous after the sort, so each recursion partitions its range by the byte at the + // current depth. + private static final class Builder { + + private final byte[][] pieces; + private final int[] ids; + private final Integer[] order; + + private int nodeCount; + private int edgeCount; + + private int[] childStart; + private byte[] labels; + private int[] childNodes; + private int[] values; + private int nextNode; + private int nextEdge; + + /** + * Prepares a builder over the keys and their sort order; {@link #count} and {@link #fill} + * perform the actual construction. + * + * @param pieces The UTF-8 bytes of each piece. + * @param ids The id stored for each piece, parallel to {@code pieces}. + * @param order The indices of {@code pieces} sorted by unsigned byte order. + */ + Builder(byte[][] pieces, int[] ids, Integer[] order) { + this.pieces = pieces; + this.ids = ids; + this.order = order; + } + + /** + * Counts the nodes and edges of the subtrie for the sorted key range {@code [from, to)} at the + * given depth. + * + * @param from The inclusive start index into {@code order}. + * @param to The exclusive end index into {@code order}. + * @param depth The byte depth this node partitions on. + * @throws IllegalArgumentException Thrown if a key is defined more than once. + */ + void count(int from, int to, int depth) { + nodeCount++; + int i = from; + if (i < to && pieces[order[i]].length == depth) { + i++; + // A second key ending at the same depth is a duplicate; the sort made them adjacent. + if (i < to && pieces[order[i]].length == depth) { + throw duplicatePiece(new String(pieces[order[i]], StandardCharsets.UTF_8)); + } + } + while (i < to) { + final byte label = pieces[order[i]][depth]; + int j = i; + while (j < to && pieces[order[j]][depth] == label) { + j++; + } + edgeCount++; + count(i, j, depth + 1); + i = j; + } + } + + /** Allocates the packed arrays to the node and edge counts gathered by {@link #count}. */ + void allocate() { + childStart = new int[nodeCount + 1]; + labels = new byte[edgeCount]; + childNodes = new int[edgeCount]; + values = new int[nodeCount]; + } + + /** + * Fills the packed arrays for the sorted key range {@code [from, to)} at the given depth and + * returns the node id assigned to it. + * + * @param from The inclusive start index into {@code order}. + * @param to The exclusive end index into {@code order}. + * @param depth The byte depth this node partitions on. + * @return The id of the node created for this range. + * @throws IllegalArgumentException Thrown if a key is defined more than once. + */ + int fill(int from, int to, int depth) { + final int node = nextNode++; + values[node] = -1; + int i = from; + if (i < to && pieces[order[i]].length == depth) { + if (values[node] != -1 || (i + 1 < to && pieces[order[i + 1]].length == depth)) { + throw duplicatePiece(new String(pieces[order[i]], StandardCharsets.UTF_8)); + } + values[node] = ids[order[i]]; + i++; + } + // Reserve this node's edge slice before recursing so siblings stay contiguous. + final int sliceStart = nextEdge; + int sliceCount = 0; + int scan = i; + while (scan < to) { + final byte label = pieces[order[scan]][depth]; + int j = scan; + while (j < to && pieces[order[j]][depth] == label) { + j++; + } + sliceCount++; + scan = j; + } + nextEdge += sliceCount; + childStart[node] = sliceStart; + childStart[node + 1] = nextEdge; + + int edge = sliceStart; + while (i < to) { + final byte label = pieces[order[i]][depth]; + int j = i; + while (j < to && pieces[order[j]][depth] == label) { + j++; + } + labels[edge] = label; + childNodes[edge] = fill(i, j, depth + 1); + edge++; + i = j; + } + return node; + } + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Segment.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Segment.java new file mode 100644 index 0000000000..7f603c6ac5 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Segment.java @@ -0,0 +1,27 @@ +/* + * 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.subword.sentencepiece; + +/** + * One encoded piece as a half-open byte range of the normalized text plus its vocabulary id. + * + * @param from The inclusive start offset in the normalized bytes. + * @param to The exclusive end offset in the normalized bytes. + * @param id The vocabulary id; the unknown id when no piece covers the range. + */ +record Segment(int from, int to, int id) { +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java new file mode 100644 index 0000000000..11ded3443c --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java @@ -0,0 +1,405 @@ +/* + * 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.subword.sentencepiece; + +import java.io.Serializable; + +import opennlp.tools.util.InvalidFormatException; + +/** + * The model-embedded text normalizer of a SentencePiece model, operating in UTF-8 byte space. + * + *

Normalization applies the model's precompiled character map (leftmost-longest replacement + * rules over UTF-8 prefixes), collapses and trims whitespace, optionally prepends the + * word-boundary marker, and escapes spaces to U+2581. Alongside the normalized bytes it produces + * {@code normToOrig}, mapping every normalized byte to the offset of the original byte chunk it + * was derived from, with one trailing entry for the end position; that map is what lets every + * downstream piece report an exact span of the caller's text.

+ */ +final class SentencePieceNormalizer implements Serializable { + + private static final long serialVersionUID = -3059745470932191300L; + + // U+2581 LOWER ONE EIGHTH BLOCK in UTF-8, the escaped form of a space. + static final byte[] SPACE_SYMBOL = {(byte) 0xE2, (byte) 0x96, (byte) 0x81}; + + // U+FFFD REPLACEMENT CHARACTER in UTF-8, emitted for a malformed byte. + private static final byte[] REPLACEMENT_CHAR = {(byte) 0xEF, (byte) 0xBF, (byte) 0xBD}; + + private final DoubleArrayTrie trie; + private final byte[] blob; + private final int replacementsFrom; + private final boolean addDummyPrefix; + private final boolean removeExtraWhitespaces; + private final boolean escapeWhitespaces; + private final boolean treatWhitespaceAsSuffix; + private final PieceTrie userDefinedMatcher; + // For each possible first byte, whether any character-map rule or user-defined symbol starts + // with it; a clear bit means normalizePrefix passes the byte through raw. + private final boolean[] ruleLead = new boolean[256]; + + /** + * Instantiates the normalizer. + * + * @param precompiledCharsMap The serialized character map; empty when the model has none. + * @param addDummyPrefix Whether a word-boundary marker is prepended. + * @param removeExtraWhitespaces Whether leading, trailing, and repeated whitespace collapses. + * @param escapeWhitespaces Whether spaces become U+2581. + * @param treatWhitespaceAsSuffix Whether the dummy marker is appended instead of prepended. + * @param userDefinedMatcher Longest-match trie over user-defined symbols that must pass + * through normalization untouched, or null when the model + * defines none. + * @throws InvalidFormatException Thrown if the character map is structurally invalid. + */ + SentencePieceNormalizer(byte[] precompiledCharsMap, boolean addDummyPrefix, + boolean removeExtraWhitespaces, boolean escapeWhitespaces, + boolean treatWhitespaceAsSuffix, PieceTrie userDefinedMatcher) + throws InvalidFormatException { + if (precompiledCharsMap.length == 0) { + trie = null; + blob = null; + replacementsFrom = 0; + } else { + // Layout: . + if (precompiledCharsMap.length <= 4) { + throw new InvalidFormatException("The precompiled character map is truncated."); + } + final long trieSize = (precompiledCharsMap[0] & 0xFFL) + | (precompiledCharsMap[1] & 0xFFL) << 8 + | (precompiledCharsMap[2] & 0xFFL) << 16 + | (precompiledCharsMap[3] & 0xFFL) << 24; + if (trieSize >= precompiledCharsMap.length - 4) { + throw new InvalidFormatException( + "The precompiled character map declares a trie of " + trieSize + + " bytes but only " + (precompiledCharsMap.length - 4) + " bytes follow."); + } + if (trieSize < 1024 || (trieSize & 0x3FF) != 0) { + throw new InvalidFormatException( + "The precompiled character map trie size " + trieSize + + " is not a positive multiple of 1024."); + } + if (precompiledCharsMap[precompiledCharsMap.length - 1] != 0) { + throw new InvalidFormatException( + "The precompiled character map replacement block is not null-terminated."); + } + trie = new DoubleArrayTrie(precompiledCharsMap, 4, (int) trieSize); + blob = precompiledCharsMap; + replacementsFrom = 4 + (int) trieSize; + } + this.addDummyPrefix = addDummyPrefix; + this.removeExtraWhitespaces = removeExtraWhitespaces; + this.escapeWhitespaces = escapeWhitespaces; + this.treatWhitespaceAsSuffix = treatWhitespaceAsSuffix; + this.userDefinedMatcher = userDefinedMatcher; + for (int b = 0; b < 256; b++) { + final boolean charsMapLead = trie != null && trie.hasTransitionFromRoot(b); + final boolean userDefinedLead = userDefinedMatcher != null + && userDefinedMatcher.step(userDefinedMatcher.root(), (byte) b) != PieceTrie.DEAD; + ruleLead[b] = charsMapLead || userDefinedLead; + } + } + + /** + * The normalized bytes plus the normalized-byte to original-byte offset map. The arrays are + * builder-backed and may be oversized; {@code length} bytes are valid, and the offset map + * holds {@code length + 1} entries. + */ + record Normalized(byte[] bytes, int length, int[] normToOrig) { + } + + /** + * One normalization step: {@code consumed} input bytes produced {@code data[from, to)}. The data + * array is the input itself (pass-through), the replacement blob, or the replacement character. + */ + private static final class Chunk { + + private byte[] data; + private int from; + private int to; + private int consumed; + + /** {@return whether this chunk is exactly one ASCII space byte} */ + boolean isSingleSpace() { + return to - from == 1 && data[from] == ' '; + } + } + + /** + * Normalizes UTF-8 input. + * + * @param input The buffer holding well-formed UTF-8 bytes; must not be null. + * @param inputLength The number of valid bytes in {@code input}. + * @return The normalized bytes with the offset map; the arrays are builder-backed, valid for + * {@code length} bytes and {@code length + 1} map entries. + */ + Normalized normalize(byte[] input, int inputLength) { + final ByteBuilder normalized = new ByteBuilder(inputLength + (inputLength >> 1) + 4); + final IntBuilder normToOrig = new IntBuilder(inputLength + (inputLength >> 1) + 5); + final Chunk chunk = new Chunk(); + + int from = 0; + int consumed = 0; + + // Ignores leading whitespace. + if (removeExtraWhitespaces) { + while (from < inputLength) { + normalizePrefix(input, inputLength, from, chunk); + if (!chunk.isSingleSpace()) { + break; + } + from += chunk.consumed; + consumed += chunk.consumed; + } + } + + // All input was whitespace. + if (from >= inputLength) { + normToOrig.append(consumed); + return new Normalized(normalized.array(), 0, normToOrig.array()); + } + + final byte[] spaceSymbol = escapeWhitespaces ? SPACE_SYMBOL : SINGLE_SPACE; + + if (!treatWhitespaceAsSuffix && addDummyPrefix) { + appendSpace(normalized, normToOrig, spaceSymbol, consumed); + } + + boolean isPrevSpace = removeExtraWhitespaces; + while (from < inputLength) { + final int lead = input[from] & 0xFF; + // An ASCII byte no rule starts with passes through raw: the chunk is the byte itself, no + // leading-space stripping applies, and it does not end in a space. + if (lead < 0x80 && lead != ' ' && !ruleLead[lead]) { + normalized.append(input[from]); + normToOrig.append(consumed); + consumed++; + from++; + isPrevSpace = false; + continue; + } + + normalizePrefix(input, inputLength, from, chunk); + int spFrom = chunk.from; + final int spTo = chunk.to; + final byte[] spData = chunk.data; + + // Removes leading spaces in the chunk if the previous chunk ended with whitespace. + while (isPrevSpace && spFrom < spTo && spData[spFrom] == ' ') { + spFrom++; + } + + if (spFrom < spTo) { + for (int n = spFrom; n < spTo; n++) { + if (spData[n] == ' ') { + appendSpace(normalized, normToOrig, spaceSymbol, consumed); + } else { + normalized.append(spData[n]); + normToOrig.append(consumed); + } + } + isPrevSpace = spData[spTo - 1] == ' '; + } + + consumed += chunk.consumed; + from += chunk.consumed; + if (!removeExtraWhitespaces) { + isPrevSpace = false; + } + } + + // Ignores trailing whitespace. + if (removeExtraWhitespaces) { + while (normalized.endsWith(spaceSymbol)) { + final int length = normalized.length() - spaceSymbol.length; + consumed = normToOrig.get(length); + normalized.truncate(length); + normToOrig.truncate(length); + } + } + + if (treatWhitespaceAsSuffix && addDummyPrefix) { + appendSpace(normalized, normToOrig, spaceSymbol, consumed); + } + + normToOrig.append(consumed); + if (normToOrig.length() != normalized.length() + 1) { + throw new IllegalStateException("The offset map has " + normToOrig.length() + + " entries for " + normalized.length() + " normalized bytes."); + } + return new Normalized(normalized.array(), normalized.length(), normToOrig.array()); + } + + private static final byte[] SINGLE_SPACE = {' '}; + + /** + * Appends the space symbol to the normalized output, mapping each of its bytes to the same + * original-byte offset. + * + * @param normalized The normalized-byte builder to append to. + * @param normToOrig The offset-map builder to append to. + * @param spaceSymbol The bytes of the (possibly escaped) space symbol. + * @param consumed The original-byte offset the symbol maps back to. + */ + private static void appendSpace(ByteBuilder normalized, IntBuilder normToOrig, + byte[] spaceSymbol, int consumed) { + normalized.append(spaceSymbol, 0, spaceSymbol.length); + for (int i = 0; i < spaceSymbol.length; i++) { + normToOrig.append(consumed); + } + } + + /** + * Fills the scratch with the normalized form of the longest applicable prefix of + * {@code input[from, inputLength)}: a user-defined symbol passes through raw, otherwise the + * longest character-map rule applies, otherwise one code point passes through raw (or becomes + * U+FFFD when the lead byte is malformed). + * + * @param input The UTF-8 input buffer. + * @param inputLength The number of valid bytes in {@code input}. + * @param from The offset to normalize from. + * @param chunk The scratch to fill. + */ + private void normalizePrefix(byte[] input, int inputLength, int from, Chunk chunk) { + if (userDefinedMatcher != null) { + final int matched = userDefinedMatcher.longestMatch(input, inputLength, from); + if (matched > 0) { + chunk.data = input; + chunk.from = from; + chunk.to = from + matched; + chunk.consumed = matched; + return; + } + } + + if (trie != null) { + final long match = trie.longestPrefixMatch(input, from, inputLength); + if (match >= 0) { + final int value = (int) (match >>> 32); + final int length = (int) (match & 0xFFFFFFFFL); + final int replacementFrom = replacementsFrom + value; + if (replacementFrom < blob.length) { + int replacementTo = replacementFrom; + while (blob[replacementTo] != 0) { + replacementTo++; + } + chunk.data = blob; + chunk.from = replacementFrom; + chunk.to = replacementTo; + chunk.consumed = length; + return; + } + } + } + + final int charLength = Math.min(utf8Length(input[from]), inputLength - from); + if (isMalformed(input, from, charLength)) { + chunk.data = REPLACEMENT_CHAR; + chunk.from = 0; + chunk.to = REPLACEMENT_CHAR.length; + chunk.consumed = 1; + return; + } + chunk.data = input; + chunk.from = from; + chunk.to = from + charLength; + chunk.consumed = charLength; + } + + /** + * Returns the byte length of a UTF-8 sequence from its lead byte; trail and malformed lead bytes + * report one byte. + * + * @param lead The lead byte. + * @return The sequence length in bytes, from one to four. + */ + static int utf8Length(byte lead) { + final int high = (lead & 0xFF) >>> 4; + if (high < 0xC) { + return 1; + } + return switch (high) { + case 0xC, 0xD -> 2; + case 0xE -> 3; + default -> 4; + }; + } + + /** + * Checks a single code point for well-formedness: correct trail-byte count and no unpaired + * surrogate or out-of-range value. + * + * @param input The UTF-8 input buffer. + * @param from The offset of the lead byte. + * @param length The candidate sequence length. + * @return {@code true} when the sequence is malformed. + */ + private static boolean isMalformed(byte[] input, int from, int length) { + if ((input[from] & 0x80) == 0) { + return false; + } + if ((input[from] & 0xC0) == 0x80 || length < utf8Length(input[from])) { + return true; + } + for (int i = from + 1; i < from + length; i++) { + if ((input[i] & 0xC0) != 0x80) { + return true; + } + } + final int codePoint = codePointAt(input, from, length); + return codePoint < 0 || (codePoint >= 0xD800 && codePoint <= 0xDFFF) || codePoint > 0x10FFFF + || length != minimalUtf8Length(codePoint); + } + + /** + * Decodes the code point of a UTF-8 sequence of the given length. + * + * @param input The UTF-8 input buffer. + * @param from The offset of the lead byte. + * @param length The sequence length in bytes, from one to four. + * @return The decoded code point. + */ + private static int codePointAt(byte[] input, int from, int length) { + return switch (length) { + case 1 -> input[from] & 0x7F; + case 2 -> (input[from] & 0x1F) << 6 | (input[from + 1] & 0x3F); + case 3 -> (input[from] & 0x0F) << 12 | (input[from + 1] & 0x3F) << 6 + | (input[from + 2] & 0x3F); + default -> (input[from] & 0x07) << 18 | (input[from + 1] & 0x3F) << 12 + | (input[from + 2] & 0x3F) << 6 | (input[from + 3] & 0x3F); + }; + } + + /** + * Returns the number of bytes the shortest UTF-8 encoding of a code point uses, which detects + * overlong encodings. + * + * @param codePoint The code point. + * @return The minimal encoding length in bytes, from one to four. + */ + private static int minimalUtf8Length(int codePoint) { + if (codePoint < 0x80) { + return 1; + } + if (codePoint < 0x800) { + return 2; + } + if (codePoint < 0x10000) { + return 3; + } + return 4; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceTokenizer.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceTokenizer.java new file mode 100644 index 0000000000..74caac0683 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceTokenizer.java @@ -0,0 +1,770 @@ +/* + * 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.subword.sentencepiece; + +import java.io.IOException; +import java.io.InputStream; +import java.io.ObjectInputFilter; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.IntUnaryOperator; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.tokenize.SubwordTokenizer; +import opennlp.tools.util.InvalidFormatException; +import opennlp.tools.util.normalizer.AlignedText; +import opennlp.tools.util.normalizer.Alignment; +import opennlp.tools.util.normalizer.OffsetAwareNormalizer; + +/** + * A {@link SubwordTokenizer} over a trained SentencePiece {@code .model} file, implemented purely + * in Java. The file carries the vocabulary with piece scores and types, the segmentation algorithm + * (unigram language model or byte-pair encoding), and the text normalizer, all of which this class + * runs. + * + *

Every piece carries the exact span of the caller's original text it came from, mapped back + * through the model's own normalizer, which is also exposed through {@link OffsetAwareNormalizer} + * for reuse outside tokenization.

+ * + *

Instances are immutable after loading and safe for concurrent use by multiple threads.

+ * + *

Beyond {@link #load(Path) loading} the native {@code .model} format, a tokenizer can be + * persisted with {@link #serialize(OutputStream)} and read back with + * {@link #deserialize(InputStream)}. Reads are guarded by an {@link java.io.ObjectInputFilter} + * that allow-lists only the classes reachable from a legitimate tokenizer graph and bounds graph + * depth, references, and array length.

+ * + * @see SentencePiece + * @see Kudo & Richardson (EMNLP 2018), + * "SentencePiece: A simple and language independent subword tokenizer and detokenizer for + * Neural Text Processing" + */ +@ThreadSafe +public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwareNormalizer { + + // Serializable through the OffsetAwareNormalizer contract. + private static final long serialVersionUID = -4472058014098085134L; + + /** The segmentation algorithm a model was trained with. */ + public enum Algorithm { + /** Unigram language model, decoded by best-path search. */ + UNIGRAM, + /** Byte-pair encoding, decoded by greedy highest-score merging. */ + BPE + } + + // Piece types of the model format. + private static final int TYPE_NORMAL = ModelProtoReader.RawModel.TYPE_NORMAL; + private static final int TYPE_UNKNOWN = ModelProtoReader.RawModel.TYPE_UNKNOWN; + private static final int TYPE_CONTROL = ModelProtoReader.RawModel.TYPE_CONTROL; + private static final int TYPE_USER_DEFINED = ModelProtoReader.RawModel.TYPE_USER_DEFINED; + private static final int TYPE_UNUSED = ModelProtoReader.RawModel.TYPE_UNUSED; + private static final int TYPE_BYTE = ModelProtoReader.RawModel.TYPE_BYTE; + + private static final int MAX_PIECE_LENGTH = 8000; + + private final Algorithm algorithm; + private final String[] pieces; + private final float[] scores; + private final int[] types; + private final int unkId; + private final boolean byteFallback; + private final Map mainPieces; + private final Map reservedPieces; + private final int[] byteToId; + private final SentencePieceNormalizer normalizer; + private final UnigramEncoder unigramEncoder; + private final BpeEncoder bpeEncoder; + private final List selfTestInputs; + private final List selfTestExpected; + + /** + * Validates a parsed model and derives the runtime structures: the piece maps, the byte-piece + * table, the normalizer, and the encoder matching the model's algorithm. + * + * @param model The parsed model description. + * @throws InvalidFormatException Thrown if the model is structurally invalid. + */ + private SentencePieceTokenizer(ModelProtoReader.RawModel model) throws InvalidFormatException { + final int count = model.pieces.size(); + pieces = model.pieces.toArray(new String[0]); + scores = new float[count]; + types = new int[count]; + for (int i = 0; i < count; i++) { + scores[i] = model.scores.get(i); + types[i] = model.types.get(i); + } + byteFallback = model.byteFallback; + algorithm = switch (model.modelType) { + case ModelProtoReader.RawModel.MODEL_TYPE_UNIGRAM -> Algorithm.UNIGRAM; + case ModelProtoReader.RawModel.MODEL_TYPE_BPE -> Algorithm.BPE; + default -> throw new InvalidFormatException( + "The model type " + model.modelType + " is not supported; only the unigram and BPE" + + " algorithms are."); + }; + + // Splits the vocabulary the way the reference does: pieces of the normal, user-defined, and + // unused types participate in segmentation, all others are reserved ids. + mainPieces = new HashMap<>(count * 2); + reservedPieces = new HashMap<>(); + final List userDefined = new ArrayList<>(); + byteToId = new int[256]; + Arrays.fill(byteToId, -1); + int foundUnkId = -1; + float minScore = Float.MAX_VALUE; + for (int i = 0; i < count; i++) { + final String piece = pieces[i]; + if (piece.length() >= MAX_PIECE_LENGTH) { + throw new InvalidFormatException("The piece with id " + i + " must be shorter than " + + MAX_PIECE_LENGTH + " characters."); + } + if (piece.indexOf(0) >= 0) { + throw new InvalidFormatException( + "The piece with id " + i + " contains a null character."); + } + final boolean isMain = + types[i] == TYPE_NORMAL || types[i] == TYPE_USER_DEFINED || types[i] == TYPE_UNUSED; + final Map target = + isMain || algorithm == Algorithm.BPE ? mainPieces : reservedPieces; + if (mainPieces.containsKey(piece) || reservedPieces.containsKey(piece)) { + throw new InvalidFormatException(PieceTrie.duplicatePiece(piece).getMessage()); + } + target.put(piece, i); + switch (types[i]) { + case TYPE_NORMAL -> minScore = Math.min(minScore, scores[i]); + case TYPE_USER_DEFINED -> userDefined.add(piece); + case TYPE_UNKNOWN -> { + if (foundUnkId >= 0) { + throw new InvalidFormatException("The model defines more than one unknown piece."); + } + foundUnkId = i; + } + case TYPE_BYTE -> { + if (!byteFallback) { + throw new InvalidFormatException("The model defines the byte piece '" + piece + + "' although byte fallback is disabled."); + } + final int b = parseBytePiece(piece); + if (b < 0) { + throw new InvalidFormatException("The byte piece '" + piece + "' is invalid."); + } + byteToId[b] = i; + } + default -> { + // CONTROL and UNUSED need no bookkeeping here. + } + } + } + if (foundUnkId < 0) { + throw new InvalidFormatException("The model defines no unknown piece."); + } + unkId = foundUnkId; + if (byteFallback) { + for (int b = 0; b < 256; b++) { + if (byteToId[b] < 0) { + throw new InvalidFormatException("The model enables byte fallback but defines no" + + " piece for byte " + b + "."); + } + } + } + + final PieceTrie userDefinedMatcher = + userDefined.isEmpty() ? null : trieOf(userDefined, id -> 0); + + normalizer = new SentencePieceNormalizer(model.precompiledCharsMap, model.addDummyPrefix, + model.removeExtraWhitespaces, model.escapeWhitespaces, model.treatWhitespaceAsSuffix, + userDefinedMatcher); + + final boolean[] unusedFlags = new boolean[count]; + final boolean[] userDefinedFlags = new boolean[count]; + final boolean[] reservedFlags = new boolean[count]; + for (int i = 0; i < count; i++) { + unusedFlags[i] = types[i] == TYPE_UNUSED; + userDefinedFlags[i] = types[i] == TYPE_USER_DEFINED; + reservedFlags[i] = types[i] != TYPE_NORMAL && types[i] != TYPE_USER_DEFINED + && types[i] != TYPE_UNUSED; + } + + if (algorithm == Algorithm.UNIGRAM) { + final List mainList = new ArrayList<>(mainPieces.size()); + final List mainIds = new ArrayList<>(mainPieces.size()); + for (int i = 0; i < count; i++) { + if (!reservedFlags[i]) { + mainList.add(pieces[i]); + mainIds.add(i); + } + } + final PieceTrie vocabulary = trieOf(mainList, mainIds::get); + unigramEncoder = new UnigramEncoder(vocabulary, scores, unusedFlags, userDefinedFlags, + minScore, unkId); + bpeEncoder = null; + } else { + unigramEncoder = null; + bpeEncoder = new BpeEncoder(mainPieces, scores, unusedFlags, reservedFlags, unkId, + userDefinedMatcher); + } + + selfTestInputs = List.copyOf(model.selfTestInputs); + selfTestExpected = List.copyOf(model.selfTestExpected); + } + + /** + * Builds a {@link PieceTrie} over the given pieces. + * + * @param pieceList The pieces to index. + * @param idOf Maps a piece's index in {@code pieceList} to the id the trie stores for it. + * @return The packed trie. + */ + private static PieceTrie trieOf(List pieceList, IntUnaryOperator idOf) { + final byte[][] keys = new byte[pieceList.size()][]; + final int[] ids = new int[pieceList.size()]; + for (int i = 0; i < keys.length; i++) { + keys[i] = pieceList.get(i).getBytes(StandardCharsets.UTF_8); + ids[i] = idOf.applyAsInt(i); + } + return PieceTrie.build(keys, ids); + } + + /** + * Loads a model from a file. + * + * @param modelFile The {@code .model} file to load; must not be null. + * @return The ready-to-use tokenizer. + * @throws IOException Thrown if the file cannot be read. + * @throws InvalidFormatException Thrown if the content is not a valid model. + * @throws IllegalArgumentException Thrown if {@code modelFile} is null. + */ + public static SentencePieceTokenizer load(Path modelFile) throws IOException { + if (modelFile == null) { + throw new IllegalArgumentException("modelFile must not be null"); + } + return new SentencePieceTokenizer(ModelProtoReader.read(Files.readAllBytes(modelFile))); + } + + /** + * Loads a model from a stream. The stream is read fully but not closed. + * + * @param in The stream positioned at the start of a {@code .model} serialization; must not be + * null. + * @return The ready-to-use tokenizer. + * @throws IOException Thrown if the stream cannot be read. + * @throws InvalidFormatException Thrown if the content is not a valid model. + * @throws IllegalArgumentException Thrown if {@code in} is null. + */ + public static SentencePieceTokenizer load(InputStream in) throws IOException { + if (in == null) { + throw new IllegalArgumentException("in must not be null"); + } + return new SentencePieceTokenizer(ModelProtoReader.read(in.readAllBytes())); + } + + /** + * Serializes this tokenizer to the given {@link OutputStream} using Java object serialization. + * The resulting stream can be read back with {@link #deserialize(InputStream)}. + * + * @param out The {@link OutputStream} to write to; must not be null. + * @throws IOException Thrown if IO errors occurred during serialization. + * @throws IllegalArgumentException Thrown if {@code out} is null. + */ + public void serialize(OutputStream out) throws IOException { + if (out == null) { + throw new IllegalArgumentException("out must not be null"); + } + try (ObjectOutputStream oos = new ObjectOutputStream(out)) { + oos.writeObject(this); + } + } + + /** + * Deserializes a {@link SentencePieceTokenizer} from the given {@link InputStream} using + * {@link DeserializationLimits#DEFAULT default} resource limits. + * + *

The stream is filtered via an {@link ObjectInputFilter} that allow-lists only the classes + * required to reconstruct a {@link SentencePieceTokenizer}, plus resource limits on graph depth, + * references, and array length. Foreign payloads are rejected with + * {@link java.io.InvalidClassException} before {@link ObjectInputStream#readObject()} + * returns.

+ * + *

Callers should still treat this method as defense-in-depth: only invoke it on streams from + * trusted sources. If the default limits are too tight for an unusually large model, use + * {@link #deserialize(InputStream, DeserializationLimits)} to supply higher limits. The class + * allow-list is intentionally not configurable; loosening it would defeat the purpose of the + * filter.

+ * + * @param in The {@link InputStream} to read from; must not be null. + * @return The reconstructed tokenizer. + * @throws IOException Thrown if IO errors occurred during deserialization, including + * {@link java.io.InvalidClassException} when the stream contains a class outside the + * allow-list or exceeds a resource limit. + * @throws ClassNotFoundException Thrown if required classes are not found. + * @throws IllegalArgumentException Thrown if {@code in} is null. + */ + public static SentencePieceTokenizer deserialize(InputStream in) + throws IOException, ClassNotFoundException { + return deserialize(in, DeserializationLimits.DEFAULT); + } + + /** + * Deserializes a {@link SentencePieceTokenizer} from the given {@link InputStream} using the + * supplied {@link DeserializationLimits resource limits}. + * + *

Use this overload when the {@link DeserializationLimits#DEFAULT default limits} reject a + * legitimate model, for example one with a very large vocabulary. The class allow-list applied + * to the stream is the same as for {@link #deserialize(InputStream)}; only the numeric limits + * change.

+ * + * @param in The {@link InputStream} to read from; must not be null. + * @param limits The {@link DeserializationLimits} to apply; must not be null. + * @return The reconstructed tokenizer. + * @throws IOException Thrown if IO errors occurred during deserialization, including + * {@link java.io.InvalidClassException} when the stream contains a class outside the + * allow-list or exceeds one of the supplied limits. + * @throws ClassNotFoundException Thrown if required classes are not found. + * @throws IllegalArgumentException Thrown if {@code in} or {@code limits} is null. + */ + public static SentencePieceTokenizer deserialize(InputStream in, DeserializationLimits limits) + throws IOException, ClassNotFoundException { + if (in == null) { + throw new IllegalArgumentException("in must not be null"); + } + if (limits == null) { + throw new IllegalArgumentException("limits must not be null"); + } + try (ObjectInputStream ois = new ObjectInputStream(in)) { + ois.setObjectInputFilter(buildFilter(limits)); + return (SentencePieceTokenizer) ois.readObject(); + } + } + + /** + * Resource limits applied by the {@link ObjectInputFilter} used by + * {@link SentencePieceTokenizer#deserialize(InputStream, DeserializationLimits)}. + * + *

The limits bound graph traversal regardless of the class allow-list and provide + * defense-in-depth against pathological streams. The {@linkplain #DEFAULT default values} are + * generous enough for typical production models; raise them only if a legitimate model is + * rejected.

+ * + * @param maxDepth Maximum object-graph nesting depth. Must be {@code > 0}. + * @param maxRefs Maximum number of internal references the stream may create. + * Must be {@code > 0}. + * @param maxArrayLength Maximum length of any single array allocation requested by the stream. + * Must be {@code > 0}. + */ + public record DeserializationLimits(long maxDepth, long maxRefs, long maxArrayLength) { + + /** + * Default limits. Sized so that models with vocabularies of several hundred thousand pieces + * round-trip while pathological streams stay bounded. + */ + public static final DeserializationLimits DEFAULT = + new DeserializationLimits(MAX_DEPTH_DEFAULT, MAX_REFS_DEFAULT, MAX_ARRAY_DEFAULT); + + /** + * Validates the limits. + * + * @throws IllegalArgumentException Thrown if any of {@code maxDepth}, {@code maxRefs}, or + * {@code maxArrayLength} is {@code <= 0}. + */ + public DeserializationLimits { + if (maxDepth <= 0) { + throw new IllegalArgumentException("maxDepth must be > 0"); + } + if (maxRefs <= 0) { + throw new IllegalArgumentException("maxRefs must be > 0"); + } + if (maxArrayLength <= 0) { + throw new IllegalArgumentException("maxArrayLength must be > 0"); + } + } + } + + private static final long MAX_DEPTH_DEFAULT = 64; + private static final long MAX_REFS_DEFAULT = 5_000_000; + private static final long MAX_ARRAY_DEFAULT = 10_000_000; + + // Allow-list of fully qualified class names that may appear in the serialized graph of a + // SentencePieceTokenizer. Anything else is rejected. + private static final Set ALLOWED_CLASSES = Set.of( + "opennlp.subword.sentencepiece.SentencePieceTokenizer", + "opennlp.subword.sentencepiece.SentencePieceTokenizer$Algorithm", + "opennlp.subword.sentencepiece.SentencePieceNormalizer", + "opennlp.subword.sentencepiece.UnigramEncoder", + "opennlp.subword.sentencepiece.BpeEncoder", + "opennlp.subword.sentencepiece.PieceTrie", + "opennlp.subword.sentencepiece.DoubleArrayTrie", + // JDK types used in field declarations. ObjectInputStream invokes the filter for every + // class descriptor in the inheritance chain, not only for the runtime class - so the + // abstract superclasses java.lang.Number (super of Integer) and java.lang.Enum (super of + // Algorithm) must be allow-listed even though no instance of either appears in the stream. + "java.lang.String", + "java.lang.Number", + "java.lang.Integer", + "java.lang.Enum", + "java.util.HashMap", + // HashMap.readObject() requests permission to allocate a Map.Entry[] before reading + // entries; the array type itself never appears as a value in the stream. + "java.util.Map$Entry", + // The unmodifiable lists created by List.copyOf serialize through the CollSer proxy, + // which requests an Object[] allocation for the elements, and the filter is also invoked + // for the concrete list class the proxy resolves to. + "java.util.CollSer", + "java.util.ImmutableCollections$List12", + "java.util.ImmutableCollections$ListN", + "java.lang.Object" + ); + + /** + * Builds the {@link ObjectInputFilter} enforcing the class allow-list and the given limits. + * + * @param limits The resource limits to enforce; never null here. + * @return The filter to install on the reading {@link ObjectInputStream}. + */ + private static ObjectInputFilter buildFilter(DeserializationLimits limits) { + return info -> { + if (info.depth() > limits.maxDepth() + || info.references() > limits.maxRefs() + || info.arrayLength() > limits.maxArrayLength()) { + return ObjectInputFilter.Status.REJECTED; + } + + final Class serialClass = info.serialClass(); + if (serialClass == null) { + return ObjectInputFilter.Status.UNDECIDED; + } + + Class componentType = serialClass; + while (componentType.isArray()) { + componentType = componentType.getComponentType(); + } + if (componentType.isPrimitive()) { + return ObjectInputFilter.Status.ALLOWED; + } + return ALLOWED_CLASSES.contains(componentType.getName()) + ? ObjectInputFilter.Status.ALLOWED + : ObjectInputFilter.Status.REJECTED; + }; + } + + /** {@inheritDoc} */ + @Override + public List encode(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final Utf8Text input = Utf8Text.of(text); + final SentencePieceNormalizer.Normalized normalized = + normalizer.normalize(input.bytes(), input.byteLength()); + final List segments = algorithm == Algorithm.UNIGRAM + ? unigramEncoder.encode(normalized.bytes(), normalized.length()) + : bpeEncoder.encode(normalized.bytes(), normalized.length()); + + final List out = new ArrayList<>(segments.size()); + final byte[] norm = normalized.bytes(); + final int[] normToOrig = normalized.normToOrig(); + + // Accumulates a run of adjacent unknown pieces into one, as the reference does, so a decoder + // sees a single unknown token per unknown region. + StringBuilder pendingUnk = null; + int pendingUnkStart = 0; + int pendingUnkEnd = 0; + + for (final Segment segment : segments) { + final boolean isUnk = segment.id() == unkId; + final boolean isControl = types[segment.id()] == TYPE_CONTROL; + // A non-unknown segment reuses its vocabulary string; only unknown segments need decoding. + final String piece = isUnk + ? new String(norm, segment.from(), segment.to() - segment.from(), StandardCharsets.UTF_8) + : pieces[segment.id()]; + + if (isControl) { + if (pendingUnk != null) { + out.add(new SubwordPiece(pendingUnk.toString(), unkId, pendingUnkStart, pendingUnkEnd)); + pendingUnk = null; + } + final int at = input.charOffset(normToOrig[segment.from()]); + out.add(new SubwordPiece(piece, segment.id(), at, at)); + continue; + } + + final int origBegin = input.charOffset(normToOrig[segment.from()]); + final int origEnd = input.charOffset(normToOrig[segment.to()]); + + if (isUnk && byteFallback) { + if (pendingUnk != null) { + out.add(new SubwordPiece(pendingUnk.toString(), unkId, pendingUnkStart, pendingUnkEnd)); + pendingUnk = null; + } + // Decomposes the unknown region into byte pieces; the last one carries the surface span. + for (int i = segment.from(); i < segment.to(); i++) { + final int b = norm[i] & 0xFF; + final boolean last = i == segment.to() - 1; + out.add(new SubwordPiece(BYTE_PIECES[b], byteToId[b], origBegin, + last ? origEnd : origBegin)); + } + } else if (isUnk) { + if (pendingUnk == null) { + pendingUnk = new StringBuilder(piece); + pendingUnkStart = origBegin; + } else { + pendingUnk.append(piece); + } + pendingUnkEnd = origEnd; + } else { + if (pendingUnk != null) { + out.add(new SubwordPiece(pendingUnk.toString(), unkId, pendingUnkStart, pendingUnkEnd)); + pendingUnk = null; + } + out.add(new SubwordPiece(piece, segment.id(), origBegin, origEnd)); + } + } + if (pendingUnk != null) { + out.add(new SubwordPiece(pendingUnk.toString(), unkId, pendingUnkStart, pendingUnkEnd)); + } + return out; + } + + /** {@inheritDoc} */ + @Override + public CharSequence normalize(CharSequence text) { + return normalizeAligned(text).normalized(); + } + + /** {@inheritDoc} */ + @Override + public AlignedText normalizeAligned(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final Utf8Text input = Utf8Text.of(text); + final SentencePieceNormalizer.Normalized result = + normalizer.normalize(input.bytes(), input.byteLength()); + final String normalized = new String(result.bytes(), 0, result.length(), + StandardCharsets.UTF_8); + final int[] normToOrig = result.normToOrig(); + final byte[] norm = result.bytes(); + + // Walks the normalized code points, grouping neighbors that came from the same original + // block into one replace run; gaps between blocks are deletions. + final Alignment.Builder builder = new Alignment.Builder(); + int cursor = 0; + int groupOrigStart = -1; + int groupOrigEnd = -1; + int groupChars = 0; + int b = 0; + final int normLength = result.length(); + while (b < normLength) { + final int byteLength = Math.min(SentencePieceNormalizer.utf8Length(norm[b]), + normLength - b); + final int origStart = input.charOffset(normToOrig[b]); + final int origEnd = input.charOffset(normToOrig[b + byteLength]); + final int chars = byteLength == 4 ? 2 : 1; + if (groupChars > 0 && origStart == groupOrigStart && origEnd == groupOrigEnd) { + groupChars += chars; + } else { + cursor = flushGroup(builder, cursor, groupOrigStart, groupOrigEnd, groupChars); + groupOrigStart = origStart; + groupOrigEnd = origEnd; + groupChars = chars; + } + b += byteLength; + } + cursor = flushGroup(builder, cursor, groupOrigStart, groupOrigEnd, groupChars); + if (cursor < input.charLength()) { + builder.replace(input.charLength() - cursor, 0); + } + return new AlignedText(text, normalized, builder.build(input.charLength())); + } + + /** + * Emits the pending alignment group as a replace run, preceded by a deletion for any original + * text skipped before it, and returns the advanced cursor. + * + * @param builder The alignment builder to append to. + * @param cursor The original-text offset reached so far. + * @param origStart The inclusive original-text start of the group. + * @param origEnd The exclusive original-text end of the group. + * @param chars The number of normalized chars in the group; zero flushes nothing. + * @return The original-text offset after the group. + */ + private static int flushGroup(Alignment.Builder builder, int cursor, int origStart, int origEnd, + int chars) { + if (chars == 0) { + return cursor; + } + if (origStart > cursor) { + builder.replace(origStart - cursor, 0); + } + builder.replace(origEnd - Math.max(origStart, cursor), chars); + return Math.max(origEnd, cursor); + } + + /** {@return the segmentation algorithm of the loaded model} */ + public Algorithm algorithm() { + return algorithm; + } + + /** {@return the number of pieces in the vocabulary} */ + public int vocabularySize() { + return pieces.length; + } + + /** + * Returns the piece string of an id. + * + * @param id A vocabulary id in {@code [0, vocabularySize())}. + * @return The piece string. + * @throws IllegalArgumentException Thrown if {@code id} is out of range. + */ + public String idToPiece(int id) { + checkId(id); + return pieces[id]; + } + + /** + * Returns the id of a piece string. + * + * @param piece The piece to look up; must not be null. + * @return The id, or the unknown id when the vocabulary does not contain the piece. + * @throws IllegalArgumentException Thrown if {@code piece} is null. + */ + public int pieceToId(String piece) { + if (piece == null) { + throw new IllegalArgumentException("piece must not be null"); + } + final Integer reserved = reservedPieces.get(piece); + if (reserved != null) { + return reserved; + } + return mainPieces.getOrDefault(piece, unkId); + } + + /** + * Returns the score of a piece. + * + * @param id A vocabulary id in {@code [0, vocabularySize())}. + * @return The score; a log-probability for unigram models, a merge rank for BPE models. + * @throws IllegalArgumentException Thrown if {@code id} is out of range. + */ + public float score(int id) { + checkId(id); + return scores[id]; + } + + /** {@return the id of the unknown piece} */ + public int unknownId() { + return unkId; + } + + /** + * Checks whether an id is the unknown piece. + * + * @param id A vocabulary id in {@code [0, vocabularySize())}. + * @return {@code true} for the unknown piece. + * @throws IllegalArgumentException Thrown if {@code id} is out of range. + */ + public boolean isUnknown(int id) { + checkId(id); + return types[id] == TYPE_UNKNOWN; + } + + /** + * Checks whether an id is a control piece. + * + * @param id A vocabulary id in {@code [0, vocabularySize())}. + * @return {@code true} for control pieces. + * @throws IllegalArgumentException Thrown if {@code id} is out of range. + */ + public boolean isControl(int id) { + checkId(id); + return types[id] == TYPE_CONTROL; + } + + /** + * Checks whether an id is a byte-fallback piece. + * + * @param id A vocabulary id in {@code [0, vocabularySize())}. + * @return {@code true} for byte pieces. + * @throws IllegalArgumentException Thrown if {@code id} is out of range. + */ + public boolean isByte(int id) { + checkId(id); + return types[id] == TYPE_BYTE; + } + + /** + * Verifies that an id is a valid vocabulary id. + * + * @param id The id to check. + * @throws IllegalArgumentException Thrown if {@code id} is outside {@code [0, vocabularySize())}. + */ + private void checkId(int id) { + if (id < 0 || id >= pieces.length) { + throw new IllegalArgumentException( + "The id " + id + " is outside [0, " + pieces.length + ")."); + } + } + + /** {@return the embedded self-test input samples} */ + List selfTestInputs() { + return selfTestInputs; + } + + /** {@return the embedded self-test expected segmentations} */ + List selfTestExpected() { + return selfTestExpected; + } + + // The prefix of a byte-fallback piece string; a full piece has the form "<0xAB>". + private static final String BYTE_PIECE_PREFIX = "<0x"; + + // "<0xAB>" piece strings for all byte values, as byte fallback emits them. + private static final String[] BYTE_PIECES = new String[256]; + + static { + final char[] hex = "0123456789ABCDEF".toCharArray(); + for (int b = 0; b < 256; b++) { + BYTE_PIECES[b] = BYTE_PIECE_PREFIX + hex[b >>> 4] + hex[b & 0xF] + ">"; + } + } + + /** + * Parses a byte-fallback piece string of the form {@code <0xAB>} into its byte value. + * + * @param piece The piece string. + * @return The byte value in {@code [0, 255]}, or {@code -1} when the string is not a byte piece. + */ + private static int parseBytePiece(String piece) { + if (piece.length() != 6 || !piece.startsWith(BYTE_PIECE_PREFIX) || piece.charAt(5) != '>') { + return -1; + } + final int high = Character.digit(piece.charAt(3), 16); + final int low = Character.digit(piece.charAt(4), 16); + return high < 0 || low < 0 ? -1 : (high << 4) | low; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java new file mode 100644 index 0000000000..5be10c2c84 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java @@ -0,0 +1,174 @@ +/* + * 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.subword.sentencepiece; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Viterbi segmentation under a unigram language model: of all ways to cover the normalized text + * with vocabulary pieces, it finds the one with the highest total log-probability. + * + *

Characters no piece covers fall back to the unknown id with a fixed penalty below the lowest + * piece score, and user-defined symbols receive a length-based bonus score so they always win.

+ */ +final class UnigramEncoder implements Serializable { + + private static final long serialVersionUID = 5648005733414803707L; + + private static final float UNK_PENALTY = 10.0f; + // The score of a user-defined symbol is this bonus per matched byte beyond the first instead + // of a trained log-probability, so longer user-defined matches always win the best path. + private static final float USER_DEFINED_LENGTH_BONUS = 0.1f; + private static final float SCORE_RESET_THRESHOLD = 100000.0f; + + private final PieceTrie trie; + private final float[] scores; + private final boolean[] unused; + private final boolean[] userDefined; + private final float unkScore; + private final int unkId; + + /** + * Instantiates the encoder. + * + * @param trie The vocabulary trie over all matchable pieces. + * @param scores The log-probability score of every piece, indexed by id. + * @param unused Whether each id has the unused piece type. + * @param userDefined Whether each id is a user-defined symbol. + * @param minScore The lowest score among normal pieces. + * @param unkId The id of the unknown piece. + */ + UnigramEncoder(PieceTrie trie, float[] scores, boolean[] unused, boolean[] userDefined, + float minScore, int unkId) { + this.trie = trie; + this.scores = scores; + this.unused = unused; + this.userDefined = userDefined; + this.unkScore = minScore - UNK_PENALTY; + this.unkId = unkId; + } + + /** + * Segments normalized text. + * + * @param normalized The buffer holding the normalized UTF-8 bytes; must not be null. + * @param size The number of valid bytes in {@code normalized}. + * @return The best-path segments covering all bytes, in text order. + * @throws IllegalArgumentException Thrown if {@code normalized} is null. + */ + List encode(byte[] normalized, int size) { + if (normalized == null) { + throw new IllegalArgumentException("normalized must not be null"); + } + if (size == 0) { + return List.of(); + } + + // The best path ending at each byte position (exclusive end), interleaved as + // [startsAt, scoreBits, id] triples; scores travel as raw float bits, a lossless round trip. + final int[] best = new int[3 * (size + 1)]; + for (int i = 0; i <= size; i++) { + best[3 * i] = -1; + } + best[1] = Float.floatToRawIntBits(0.0f); + + int startsAt = 0; + int maxFrontier = 0; + while (startsAt < size) { + float bestScoreTillHere = Float.intBitsToFloat(best[3 * startsAt + 1]); + if (bestScoreTillHere < -SCORE_RESET_THRESHOLD + || bestScoreTillHere > SCORE_RESET_THRESHOLD) { + // Re-bases accumulated scores to keep float precision on very long inputs; every + // reachable frontier position shifts by the same offset, so the argmax is unchanged. + final float offset = bestScoreTillHere; + for (int i = startsAt; i <= maxFrontier; i++) { + if (i == startsAt || best[3 * i] != -1) { + best[3 * i + 1] = Float.floatToRawIntBits( + Float.intBitsToFloat(best[3 * i + 1]) - offset); + } + } + bestScoreTillHere = 0.0f; + } + + boolean hasSingleNode = false; + final int mblen = Math.min(SentencePieceNormalizer.utf8Length(normalized[startsAt]), + size - startsAt); + + int node = trie.root(); + for (int keyPos = startsAt; keyPos < size; ) { + node = trie.step(node, normalized[keyPos]); + if (node == PieceTrie.DEAD) { + break; + } + keyPos++; + final int id = trie.value(node); + if (id < 0) { + continue; + } + if (unused[id]) { + continue; + } + maxFrontier = Math.max(maxFrontier, keyPos); + final int length = keyPos - startsAt; + // User-defined symbols receive a length bonus instead of a trained score. + final float score = userDefined[id] + ? USER_DEFINED_LENGTH_BONUS * (length - 1) : scores[id]; + final float candidate = score + bestScoreTillHere; + final int slot = 3 * keyPos; + if (best[slot] == -1 || candidate > Float.intBitsToFloat(best[slot + 1])) { + best[slot + 1] = Float.floatToRawIntBits(candidate); + best[slot] = startsAt; + best[slot + 2] = id; + } + if (!hasSingleNode && length == mblen) { + hasSingleNode = true; + } + } + + if (!hasSingleNode) { + final int end = startsAt + mblen; + maxFrontier = Math.max(maxFrontier, end); + final float candidate = unkScore + bestScoreTillHere; + final int slot = 3 * end; + if (best[slot] == -1 || candidate > Float.intBitsToFloat(best[slot + 1])) { + best[slot + 1] = Float.floatToRawIntBits(candidate); + best[slot] = startsAt; + best[slot + 2] = unkId; + } + } + + startsAt += mblen; + } + + final List results = new ArrayList<>(size / 4 + 1); + int endsAt = size; + while (endsAt > 0) { + final int from = best[3 * endsAt]; + if (from < 0) { + throw new IllegalStateException( + "The Viterbi path is broken at normalized byte " + endsAt + "."); + } + results.add(new Segment(from, endsAt, best[3 * endsAt + 2])); + endsAt = from; + } + Collections.reverse(results); + return results; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java new file mode 100644 index 0000000000..149b45be1e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java @@ -0,0 +1,136 @@ +/* + * 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.subword.sentencepiece; + +/** + * A caller's text encoded as UTF-8, keeping the map from every byte offset back to the UTF-16 + * offset it came from. + * + *

The pipeline runs in UTF-8 byte space, but the spans reported to the caller must be UTF-16 + * offsets into the original {@code CharSequence}; this map converts them. An unpaired surrogate, + * which UTF-8 cannot represent, is encoded as U+FFFD. Callers see those offsets only through + * {@link opennlp.tools.tokenize.SubwordPiece} spans.

+ */ +final class Utf8Text { + + private final byte[] bytes; + private final int byteLength; + // Null for pure-ASCII text, where byte offsets equal UTF-16 offsets. + private final int[] byteToChar; + private final int charLength; + + /** + * Wraps an encoded buffer and its offset map. + * + * @param bytes The UTF-8 buffer. + * @param byteLength The number of valid bytes in {@code bytes}. + * @param byteToChar The byte-to-UTF-16 offset map, or null for pure-ASCII text. + * @param charLength The length of the original text in UTF-16 units. + */ + private Utf8Text(byte[] bytes, int byteLength, int[] byteToChar, int charLength) { + this.bytes = bytes; + this.byteLength = byteLength; + this.byteToChar = byteToChar; + this.charLength = charLength; + } + + /** + * Encodes text. + * + * @param text The text to encode; must not be null. + * @return The encoded view. + */ + static Utf8Text of(CharSequence text) { + final int charLength = text.length(); + // The common case: pure ASCII, where the bytes are the chars and the map is the identity. + int ascii = 0; + while (ascii < charLength && text.charAt(ascii) < 0x80) { + ascii++; + } + if (ascii == charLength) { + final byte[] exact = new byte[charLength]; + for (int i = 0; i < charLength; i++) { + exact[i] = (byte) text.charAt(i); + } + return new Utf8Text(exact, charLength, null, charLength); + } + + final byte[] bytes = new byte[charLength * 3 + 1]; + final int[] byteToChar = new int[charLength * 3 + 2]; + int b = 0; + int c = 0; + while (c < charLength) { + int codePoint = text.charAt(c); + int charCount = 1; + if (Character.isHighSurrogate((char) codePoint) && c + 1 < charLength + && Character.isLowSurrogate(text.charAt(c + 1))) { + codePoint = Character.toCodePoint((char) codePoint, text.charAt(c + 1)); + charCount = 2; + } else if (Character.isSurrogate((char) codePoint)) { + // An unpaired surrogate has no UTF-8 form; U+FFFD keeps the encoding total. + codePoint = 0xFFFD; + } + final int start = b; + if (codePoint < 0x80) { + bytes[b++] = (byte) codePoint; + } else if (codePoint < 0x800) { + bytes[b++] = (byte) (0xC0 | codePoint >>> 6); + bytes[b++] = (byte) (0x80 | codePoint & 0x3F); + } else if (codePoint < 0x10000) { + bytes[b++] = (byte) (0xE0 | codePoint >>> 12); + bytes[b++] = (byte) (0x80 | codePoint >>> 6 & 0x3F); + bytes[b++] = (byte) (0x80 | codePoint & 0x3F); + } else { + bytes[b++] = (byte) (0xF0 | codePoint >>> 18); + bytes[b++] = (byte) (0x80 | codePoint >>> 12 & 0x3F); + bytes[b++] = (byte) (0x80 | codePoint >>> 6 & 0x3F); + bytes[b++] = (byte) (0x80 | codePoint & 0x3F); + } + for (int i = start; i < b; i++) { + byteToChar[i] = c; + } + c += charCount; + } + byteToChar[b] = charLength; + return new Utf8Text(bytes, b, byteToChar, charLength); + } + + /** {@return the UTF-8 buffer; valid up to {@link #byteLength()}} */ + byte[] bytes() { + return bytes; + } + + /** {@return the number of valid bytes in {@link #bytes()}} */ + int byteLength() { + return byteLength; + } + + /** {@return the length of the original text in UTF-16 units} */ + int charLength() { + return charLength; + } + + /** + * Maps a byte offset to the UTF-16 offset of the character containing it. + * + * @param byteOffset An offset in {@code [0, bytes().length]}. + * @return The UTF-16 offset; the text length for the end offset. + */ + int charOffset(int byteOffset) { + return byteToChar == null ? byteOffset : byteToChar[byteOffset]; + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/PieceTrieTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/PieceTrieTest.java new file mode 100644 index 0000000000..adaf2a292e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/PieceTrieTest.java @@ -0,0 +1,117 @@ +/* + * 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.subword.sentencepiece; + +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Random; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * The trie's transition function held against a naive map-backed reference over randomized + * vocabularies, so the hybrid direct-table and linear-scan node layouts are proven to enumerate + * identical transitions; the encoder's correctness rests on that equivalence. + */ +class PieceTrieTest { + + @Test + void testStepsMatchAMapBackedReferenceOverRandomVocabularies() { + final Random random = new Random(7); + for (int round = 0; round < 20; round++) { + // Piece count crosses the direct-table threshold in both directions, so wide and narrow + // roots are both exercised. + final int pieceCount = 2 + random.nextInt(60); + final Set keys = new HashSet<>(); + while (keys.size() < pieceCount) { + final int length = 1 + random.nextInt(5); + final StringBuilder key = new StringBuilder(); + for (int i = 0; i < length; i++) { + key.append((char) ('a' + random.nextInt(random.nextBoolean() ? 26 : 4))); + } + keys.add(key.toString()); + } + final byte[][] pieces = new byte[keys.size()][]; + final int[] ids = new int[keys.size()]; + final Map reference = new HashMap<>(); + int index = 0; + for (final String key : keys) { + pieces[index] = key.getBytes(StandardCharsets.UTF_8); + ids[index] = index; + reference.put(key, index); + index++; + } + final PieceTrie trie = PieceTrie.build(pieces, ids); + + // Every walk over random query strings must accept exactly the reference's prefixes. + for (int query = 0; query < 200; query++) { + final int length = 1 + random.nextInt(8); + final StringBuilder text = new StringBuilder(); + for (int i = 0; i < length; i++) { + text.append((char) ('a' + random.nextInt(6))); + } + int node = trie.root(); + for (int i = 0; i < length; i++) { + node = trie.step(node, (byte) text.charAt(i)); + final String prefix = text.substring(0, i + 1); + final boolean anyKeyHasPrefix = + keys.stream().anyMatch(k -> k.startsWith(prefix)); + if (node == PieceTrie.DEAD) { + assertFalse(anyKeyHasPrefix, "dead end despite live prefix: " + prefix); + break; + } + final Integer expected = reference.get(prefix); + assertEquals(expected == null ? -1 : expected, trie.value(node), + "value mismatch at prefix: " + prefix); + } + } + } + } + + @Test + void testWideRootDispatchesAllByteValues() { + // 200 distinct single-byte pieces force the direct-table layout at the root. + final byte[][] pieces = new byte[200][]; + final int[] ids = new int[200]; + for (int i = 0; i < 200; i++) { + pieces[i] = new byte[] {(byte) (i + 20)}; + ids[i] = i; + } + final PieceTrie trie = PieceTrie.build(pieces, ids); + for (int b = 0; b < 256; b++) { + final int node = trie.step(trie.root(), (byte) b); + if (b >= 20 && b < 220) { + assertEquals(b - 20, trie.value(node)); + } else { + assertEquals(PieceTrie.DEAD, node); + } + } + } + + @Test + void testDuplicatePiecesFailLoudly() { + final byte[][] pieces = {{'a'}, {'a'}}; + assertThrows(IllegalArgumentException.class, () -> PieceTrie.build(pieces, new int[] {0, 1})); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceAlignmentTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceAlignmentTest.java new file mode 100644 index 0000000000..7ac28e67ba --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceAlignmentTest.java @@ -0,0 +1,145 @@ +/* + * 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.subword.sentencepiece; + +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.util.Span; +import opennlp.tools.util.normalizer.AlignedText; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Exercises the {@code OffsetAwareNormalizer} view: the model normalizer's output must align + * back to the original text exactly, including through whitespace collapsing, character-map + * replacements, and supplementary characters. + */ +class SentencePieceAlignmentTest { + + private static SentencePieceTokenizer unigram() { + return SentencePieceFixtures.tokenizer("tiny-unigram"); + } + + @ParameterizedTest + @ValueSource(strings = {"Hello world", " Hello world ", "Hello world.\nSecond line", + "3.14159 x 42", "family emoji", " leading and trailing ", "a", "", " "}) + void testAlignedMatchesUnaligned(String input) { + final AlignedText aligned = unigram().normalizeAligned(input); + assertEquals(unigram().normalize(input).toString(), aligned.normalizedString()); + assertEquals(input, aligned.original().toString()); + assertEquals(aligned.normalizedString().length(), aligned.alignment().normalizedLength()); + assertEquals(input.length(), aligned.alignment().originalLength()); + } + + @Test + void testWordMapsBackThroughCollapsedWhitespace() { + final String input = " Hello world "; + final AlignedText aligned = unigram().normalizeAligned(input); + final String normalized = aligned.normalizedString(); + + final int at = normalized.indexOf("world"); + final Span original = aligned.toOriginalSpan(at, at + "world".length()); + assertEquals("world", input.substring(original.getStart(), original.getEnd())); + } + + @Test + void testLigatureReplacementMapsToItsSourceCharacter() { + // The character map expands the single ligature to two letters; both normalized letters + // must map back to the one original character. + final String input = cp(0xFB01) + "nancial"; + final AlignedText aligned = unigram().normalizeAligned(input); + final String normalized = aligned.normalizedString(); + + final int at = normalized.indexOf("fi"); + assertTrue(at >= 0, "the character map must expand the ligature, got " + normalized); + final Span original = aligned.toOriginalSpan(at, at + 2); + assertEquals(0, original.getStart()); + assertEquals(1, original.getEnd()); + } + + @Test + void testSupplementaryCharacterSpansUseUtf16Units() { + final String input = "I love " + new String(Character.toChars(0x1F355)) + " pizza"; + final List pieces = unigram().encode(input); + + SubwordPiece pizzaSlice = null; + for (final SubwordPiece piece : pieces) { + if (piece.piece().contains(new String(Character.toChars(0x1F355)))) { + pizzaSlice = piece; + } + } + assertTrue(pizzaSlice != null, "the emoji must surface as a piece, got " + pieces); + assertEquals(7, pizzaSlice.start()); + assertEquals(9, pizzaSlice.end()); + assertEquals(new String(Character.toChars(0x1F355)), + input.substring(pizzaSlice.start(), pizzaSlice.end())); + } + + @Test + void testEverySpanIsWithinTheOriginalText() { + final String input = "quotes " + cp(0x201C) + "fancy" + cp(0x201D) + " and " + + cp(0x2018) + "single" + cp(0x2019) + " " + cp(0x2014) + " dash"; + for (final SubwordPiece piece : unigram().encode(input)) { + assertTrue(piece.start() >= 0 && piece.end() <= input.length(), + "span " + piece + " must lie inside the input"); + assertTrue(piece.start() <= piece.end(), "span " + piece + " must not be inverted"); + } + } + + @Test + void testSpansAreMonotonicAndAdjacent() { + final String input = "The quick brown fox jumps over the lazy dog."; + int previousEnd = 0; + for (final SubwordPiece piece : unigram().encode(input)) { + assertTrue(piece.start() >= previousEnd || piece.start() == piece.end(), + "piece " + piece + " must not step back before " + previousEnd); + previousEnd = Math.max(previousEnd, piece.end()); + } + assertEquals(input.length(), previousEnd, "the last span must reach the end of the input"); + } + + @Test + void testUnpairedSurrogateIsDeterministic() { + final String input = "a" + (char) 0xD83C + "b"; + final List first = unigram().encode(input); + final List second = unigram().encode(input); + assertEquals(first, second); + int covered = 0; + for (final SubwordPiece piece : first) { + covered = Math.max(covered, piece.end()); + } + assertEquals(input.length(), covered); + } + + private static String cp(int codePoint) { + return new String(Character.toChars(codePoint)); + } + + @Test + void testNullInputsFailLoudly() { + assertThrows(IllegalArgumentException.class, () -> unigram().encode(null)); + assertThrows(IllegalArgumentException.class, () -> unigram().normalizeAligned(null)); + assertThrows(IllegalArgumentException.class, () -> unigram().normalize(null)); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.java new file mode 100644 index 0000000000..a2c06d6177 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.java @@ -0,0 +1,184 @@ +/* + * 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.subword.sentencepiece; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Stream; + +import opennlp.tools.tokenize.SubwordPiece; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +/** + * Shared support for the bundled test models and for the tab-separated parity fixture files + * produced by the {@code gen_fixtures.py} and {@code gen_real_fixtures.py} scripts in the test + * resources: one line per input, holding the input, the expected piece count, four columns per + * expected piece (content, id, start, end), and the expected normalized form. + */ +final class SentencePieceFixtures { + + /** The file name suffix of a bundled or downloaded SentencePiece model. */ + static final String MODEL_SUFFIX = ".model"; + + /** The file name suffix of the fixture file belonging to a model. */ + static final String FIXTURES_SUFFIX = ".fixtures.tsv"; + + private static final Map LOADED = new ConcurrentHashMap<>(); + + private SentencePieceFixtures() { + } + + /** + * The bundled models, one per algorithm and normalizer variant the reader must handle. + * + * @return The model names, usable as {@code @MethodSource} arguments. + */ + static Stream models() { + return Stream.of("tiny-unigram", "tiny-unigram-bytefb", "tiny-bpe", "tiny-unigram-identity", + "tiny-unigram-suffix"); + } + + /** + * Loads a bundled model from the test resources, caching the result so the parsing cost is paid + * once per model across all test classes. + * + * @param model The bundled model name, without the {@link #MODEL_SUFFIX} suffix. + * @return The loaded tokenizer, shared by all callers. + */ + static SentencePieceTokenizer tokenizer(String model) { + return LOADED.computeIfAbsent(model, name -> { + try (InputStream in = + SentencePieceFixtures.class.getResourceAsStream(name + MODEL_SUFFIX)) { + assertNotNull(in, "missing test resource " + name + MODEL_SUFFIX); + return SentencePieceTokenizer.load(in); + } catch (IOException e) { + throw new IllegalStateException(e); + } + }); + } + + /** + * Reads the fixture file belonging to a bundled model. + * + * @param model The bundled model name, without the {@link #MODEL_SUFFIX} suffix. + * @return The parsed fixtures in file order. + * @throws IOException Thrown if the fixture resource cannot be read. + */ + static List fixtures(String model) throws IOException { + try (InputStream in = + SentencePieceFixtures.class.getResourceAsStream(model + FIXTURES_SUFFIX)) { + assertNotNull(in, "missing test resource " + model + FIXTURES_SUFFIX); + return read(new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))); + } + } + + /** + * One parsed fixture line: an input with the piece sequence and normalized form the reference + * implementation produced for it. + * + * @param input The text to encode. + * @param pieces The expected pieces with ids and original-text spans, in text order. + * @param normalized The expected normalized form of {@code input}. + */ + record Fixture(String input, List pieces, String normalized) { + } + + /** + * Reads all fixture lines from a reader. + * + * @param reader The reader positioned at the start of a fixture file; must not be null. + * @return The parsed fixtures in file order. + * @throws IOException Thrown if the reader fails. + */ + static List read(BufferedReader reader) throws IOException { + final List fixtures = new ArrayList<>(); + String line; + while ((line = reader.readLine()) != null) { + final String[] cols = line.split("\t", -1); + final String input = unescape(cols[0]); + final int count = Integer.parseInt(cols[1]); + final List pieces = new ArrayList<>(count); + for (int i = 0; i < count; i++) { + pieces.add(new SubwordPiece(unescape(cols[2 + i * 4]), + Integer.parseInt(cols[3 + i * 4]), Integer.parseInt(cols[4 + i * 4]), + Integer.parseInt(cols[5 + i * 4]))); + } + fixtures.add(new Fixture(input, pieces, unescape(cols[2 + count * 4]))); + } + return fixtures; + } + + /** + * Asserts that a tokenizer reproduces one fixture exactly: the piece sequence with ids and + * spans, and the normalized form. + * + * @param tokenizer The tokenizer under test. + * @param fixture The expected encoding. + * @param context A prefix for failure messages that identifies the model and input. + */ + static void assertFixture(SentencePieceTokenizer tokenizer, Fixture fixture, String context) { + final List actual = tokenizer.encode(fixture.input()); + assertEquals(fixture.pieces().size(), actual.size(), + context + " piece count; got " + actual); + for (int i = 0; i < actual.size(); i++) { + final SubwordPiece expected = fixture.pieces().get(i); + final SubwordPiece got = actual.get(i); + assertEquals(expected.piece(), got.piece(), context + " piece " + i); + assertEquals(expected.id(), got.id(), context + " id of piece " + i); + assertEquals(expected.start(), got.start(), context + " start of piece " + i); + assertEquals(expected.end(), got.end(), context + " end of piece " + i); + } + assertEquals(fixture.normalized(), tokenizer.normalize(fixture.input()).toString(), + context + " normalized form"); + } + + /** + * Reverses the fixture files' escaping of tab, newline, carriage return, and backslash. + * + * @param s The escaped column content. + * @return The unescaped text. + * @throws IllegalArgumentException Thrown if an unknown escape sequence occurs. + */ + static String unescape(String s) { + final StringBuilder out = new StringBuilder(s.length()); + for (int i = 0; i < s.length(); i++) { + final char c = s.charAt(i); + if (c == '\\' && i + 1 < s.length()) { + i++; + switch (s.charAt(i)) { + case 't' -> out.append('\t'); + case 'n' -> out.append('\n'); + case 'r' -> out.append('\r'); + case '\\' -> out.append('\\'); + default -> throw new IllegalArgumentException("bad escape in fixture: " + s); + } + } else { + out.append(c); + } + } + return out.toString(); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java new file mode 100644 index 0000000000..5f340eda3a --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java @@ -0,0 +1,238 @@ +/* + * 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.subword.sentencepiece; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Fail-loud behavior on malformed models, plus the concurrency guarantee: one loaded tokenizer + * must produce identical results from many threads. + */ +class SentencePieceModelValidationTest { + + @Test + void testNullAndEmptyInputFailLoudly() { + assertThrows(IllegalArgumentException.class, + () -> SentencePieceTokenizer.load((Path) null)); + assertThrows(IllegalArgumentException.class, + () -> SentencePieceTokenizer.load((InputStream) null)); + assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(new byte[0]))); + } + + @Test + void testGarbageBytesFailLoudly() { + final byte[] garbage = "this is not a model file at all".getBytes(StandardCharsets.UTF_8); + assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(garbage))); + } + + @Test + void testTruncatedModelFailsLoudly() throws IOException { + final byte[] whole = readModel(); + final byte[] truncated = Arrays.copyOf(whole, whole.length / 3); + assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(truncated))); + } + + @Test + void testUnsupportedModelTypeFailsLoudly() { + // A minimal well-formed model claiming the WORD algorithm (model_type = 3). + final byte[] model = minimalModel(3); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(model))); + assertTrue(e.getMessage().contains("not supported"), e.getMessage()); + } + + @Test + void testMissingUnknownPieceFailsLoudly() { + final byte[] model = minimalModelWithoutUnk(); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(model))); + assertTrue(e.getMessage().contains("unknown piece"), e.getMessage()); + } + + @Test + void testMalformedPrecompiledCharsMapFailsLoudly() { + // A well-formed proto whose normalizer spec carries a truncated precompiled character map; + // load() must report it as an invalid model, like every other malformed model content. + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + writePiece(out, "", 2); + writePiece(out, "a", 1); + // normalizer_spec { precompiled_charsmap = <3 bytes> } + out.write(0x1A); + out.write(5); + out.write(0x12); + out.write(3); + out.writeBytes(new byte[] {1, 2, 3}); + final byte[] model = out.toByteArray(); + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream(model))); + assertTrue(e.getMessage().contains("character map"), e.getMessage()); + } + + @Test + void testConcurrentEncodingIsConsistent() throws Exception { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer("tiny-unigram"); + final String[] inputs = { + "The quick brown fox jumps over the lazy dog.", + "tokenization and segmentation", + " Hello world ", + "water running walked faster apple book work play"}; + final List> expected = new ArrayList<>(); + for (final String input : inputs) { + expected.add(tokenizer.encode(input)); + } + + final ExecutorService pool = Executors.newFixedThreadPool(8); + try { + final List> futures = new ArrayList<>(); + for (int t = 0; t < 8; t++) { + futures.add(pool.submit((Callable) () -> { + for (int round = 0; round < 500; round++) { + for (int i = 0; i < inputs.length; i++) { + if (!expected.get(i).equals(tokenizer.encode(inputs[i]))) { + return false; + } + } + } + return true; + })); + } + for (final Future future : futures) { + assertTrue(future.get(), "concurrent encoding must match single-threaded results"); + } + } finally { + pool.shutdownNow(); + } + } + + @Test + void testVocabularyAccessors() { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer("tiny-unigram"); + assertEquals(300, tokenizer.vocabularySize()); + assertEquals(SentencePieceTokenizer.Algorithm.UNIGRAM, tokenizer.algorithm()); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + final String piece = tokenizer.idToPiece(id); + if (!tokenizer.isUnknown(id) && !tokenizer.isControl(id)) { + assertEquals(id, tokenizer.pieceToId(piece), "round trip of piece '" + piece + "'"); + } + } + assertEquals(tokenizer.unknownId(), tokenizer.pieceToId("definitely-not-in-the-vocabulary")); + assertThrows(IllegalArgumentException.class, () -> tokenizer.idToPiece(-1)); + assertThrows(IllegalArgumentException.class, + () -> tokenizer.idToPiece(tokenizer.vocabularySize())); + assertThrows(IllegalArgumentException.class, () -> tokenizer.pieceToId(null)); + } + + @Test + void testScoresAreFiniteAndRangeChecked() { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer("tiny-unigram"); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + assertTrue(Float.isFinite(tokenizer.score(id)), "score of piece " + id); + } + assertThrows(IllegalArgumentException.class, () -> tokenizer.score(-1)); + assertThrows(IllegalArgumentException.class, + () -> tokenizer.score(tokenizer.vocabularySize())); + } + + @Test + void testBytePiecesExistOnlyInByteFallbackModels() { + final SentencePieceTokenizer byteFallback = + SentencePieceFixtures.tokenizer("tiny-unigram-bytefb"); + int bytePieces = 0; + for (int id = 0; id < byteFallback.vocabularySize(); id++) { + if (byteFallback.isByte(id)) { + bytePieces++; + assertTrue(byteFallback.idToPiece(id).startsWith("<0x"), + "byte piece " + id + " is " + byteFallback.idToPiece(id)); + } + } + assertEquals(256, bytePieces, "byte fallback defines one piece per byte value"); + + final SentencePieceTokenizer plain = SentencePieceFixtures.tokenizer("tiny-unigram"); + for (int id = 0; id < plain.vocabularySize(); id++) { + assertFalse(plain.isByte(id), "piece " + id + " must not be a byte piece"); + } + assertThrows(IllegalArgumentException.class, () -> plain.isByte(-1)); + } + + private static byte[] readModel() throws IOException { + try (InputStream in = + SentencePieceModelValidationTest.class.getResourceAsStream("tiny-unigram.model")) { + return in.readAllBytes(); + } + } + + // Hand-encodes a minimal ModelProto: three pieces (, , ) and a trainer spec with + // the requested model type. + private static byte[] minimalModel(int modelType) { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + writePiece(out, "", 2); + writePiece(out, "", 3); + writePiece(out, "", 3); + writePiece(out, "a", 1); + // trainer_spec { model_type = } + out.write(0x12); + out.write(2); + out.write(0x18); + out.write(modelType); + return out.toByteArray(); + } + + private static byte[] minimalModelWithoutUnk() { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + writePiece(out, "a", 1); + writePiece(out, "b", 1); + return out.toByteArray(); + } + + private static void writePiece(ByteArrayOutputStream out, String piece, int type) { + final byte[] utf8 = piece.getBytes(StandardCharsets.UTF_8); + // pieces { piece = ; score = 0.0; type = } as nested length-delimited field 1. + final int inner = 2 + utf8.length + 2; + out.write(0x0A); + out.write(inner); + out.write(0x0A); + out.write(utf8.length); + out.writeBytes(utf8); + out.write(0x18); + out.write(type); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java new file mode 100644 index 0000000000..f6361a93b2 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java @@ -0,0 +1,66 @@ +/* + * 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.subword.sentencepiece; + +import java.io.IOException; +import java.util.List; +import java.util.StringJoiner; + +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Asserts exact parity with the reference implementation: for every fixture input, the pieces, + * ids, original-text spans, and the normalized form must equal what the reference produced for + * the same bundled model. The fixtures were generated by the {@code gen_fixtures.py} script in + * the test resources against the sentencepiece Python package. + */ +class SentencePieceParityTest { + + @ParameterizedTest + @MethodSource("opennlp.subword.sentencepiece.SentencePieceFixtures#models") + void testFixtureParity(String model) throws IOException { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer(model); + int lines = 0; + for (final SentencePieceFixtures.Fixture fixture : SentencePieceFixtures.fixtures(model)) { + lines++; + SentencePieceFixtures.assertFixture(tokenizer, fixture, + model + " input <" + fixture.input() + ">"); + } + assertTrue(lines >= 30, "the fixture file must not be empty or truncated"); + } + + @ParameterizedTest + @MethodSource("opennlp.subword.sentencepiece.SentencePieceFixtures#models") + void testEmbeddedSelfTestSamples(String model) { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer(model); + final List inputs = tokenizer.selfTestInputs(); + final List expected = tokenizer.selfTestExpected(); + assertTrue(!inputs.isEmpty(), "the tiny models embed self-test samples"); + for (int i = 0; i < inputs.size(); i++) { + final StringJoiner joined = new StringJoiner(" "); + for (final String piece : tokenizer.encodeToPieces(inputs.get(i))) { + joined.add(piece); + } + assertEquals(expected.get(i), joined.toString(), + model + " self-test sample <" + inputs.get(i) + ">"); + } + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java new file mode 100644 index 0000000000..1174f8652c --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java @@ -0,0 +1,77 @@ +/* + * 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.subword.sentencepiece; + +import java.io.BufferedReader; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +/** + * Opt-in parity check against real pre-trained models, which are downloaded rather than bundled. + * + *

Point {@code -Dopennlp.subword.eval.dir} at a directory holding {@code .model} files + * with sibling {@code .fixtures.tsv} files generated by the {@code gen_real_fixtures.py} + * script from the test resources; every model found is asserted piece for piece. Without the + * property the test is skipped.

+ */ +class SentencePieceRealModelEvalTest { + + @Test + void testRealModelParity() throws IOException { + final String dir = System.getProperty("opennlp.subword.eval.dir"); + assumeTrue(dir != null && !dir.isBlank(), + "set -Dopennlp.subword.eval.dir to run the real-model parity check"); + + int models = 0; + try (Stream files = Files.list(Path.of(dir))) { + for (final Path model : files + .filter(f -> f.toString().endsWith(SentencePieceFixtures.MODEL_SUFFIX)) + .sorted().toList()) { + final String path = model.toString(); + final Path fixtures = Path.of( + path.substring(0, path.length() - SentencePieceFixtures.MODEL_SUFFIX.length()) + + SentencePieceFixtures.FIXTURES_SUFFIX); + assumeTrue(Files.exists(fixtures), "no fixtures for " + model.getFileName()); + models++; + assertModel(model, fixtures); + } + } + assertTrue(models > 0, "the eval directory contains no models"); + } + + private static void assertModel(Path modelPath, Path fixturesPath) throws IOException { + final SentencePieceTokenizer tokenizer = SentencePieceTokenizer.load(modelPath); + final List fixtures; + try (BufferedReader reader = Files.newBufferedReader(fixturesPath, StandardCharsets.UTF_8)) { + fixtures = SentencePieceFixtures.read(reader); + } + for (final SentencePieceFixtures.Fixture fixture : fixtures) { + SentencePieceFixtures.assertFixture(tokenizer, fixture, + modelPath.getFileName() + " input <" + fixture.input() + ">"); + } + assertTrue(fixtures.size() >= 30, modelPath.getFileName() + " fixtures must not be truncated"); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java new file mode 100644 index 0000000000..e5bb42d148 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java @@ -0,0 +1,156 @@ +/* + * 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.subword.sentencepiece; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.InvalidClassException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.ArrayList; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertIterableEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Asserts the {@code Serializable} contract inherited through + * {@code opennlp.tools.util.normalizer.CharSequenceNormalizer}: a tokenizer round-tripped + * through Java object serialization must encode and normalize exactly like the original. + * Also asserts the guarded read path of + * {@link SentencePieceTokenizer#deserialize(InputStream)}: foreign payloads and streams + * exceeding the resource limits are rejected before materialisation. + */ +class SentencePieceTokenizerSerializationTest { + + private static final String[] INPUTS = { + "", + "The quick brown fox jumps over the lazy dog.", + " Hello world ", + "tokenization and segmentation", + "caf\u00e9 na\u00efve \u4e2d\u6587" + }; + + @ParameterizedTest + @MethodSource("opennlp.subword.sentencepiece.SentencePieceFixtures#models") + void testRoundTripPreservesEncoding(String model) throws IOException, ClassNotFoundException { + final SentencePieceTokenizer original = SentencePieceFixtures.tokenizer(model); + + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + out.writeObject(original); + } + final SentencePieceTokenizer copy; + try (ObjectInputStream in = + new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { + copy = (SentencePieceTokenizer) in.readObject(); + } + + assertEquals(original.algorithm(), copy.algorithm(), model + " algorithm"); + assertEquals(original.vocabularySize(), copy.vocabularySize(), model + " vocabulary size"); + for (final String input : INPUTS) { + final String context = model + " input <" + input + ">"; + assertIterableEquals(original.encode(input), copy.encode(input), context + " pieces"); + assertEquals(original.normalize(input).toString(), copy.normalize(input).toString(), + context + " normalized form"); + } + } + + /** + * Serializes the tokenizer of the given fixture model through + * {@link SentencePieceTokenizer#serialize(java.io.OutputStream)}. + * + * @param model The fixture model name. + * @return The serialized bytes. + * @throws IOException Thrown if the model cannot be read or serialized. + */ + private static byte[] serialized(String model) throws IOException { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + SentencePieceFixtures.tokenizer(model).serialize(bytes); + return bytes.toByteArray(); + } + + @ParameterizedTest + @MethodSource("opennlp.subword.sentencepiece.SentencePieceFixtures#models") + void testGuardedDeserializePreservesEncoding(String model) + throws IOException, ClassNotFoundException { + final SentencePieceTokenizer original = SentencePieceFixtures.tokenizer(model); + final SentencePieceTokenizer copy = + SentencePieceTokenizer.deserialize(new ByteArrayInputStream(serialized(model))); + + assertEquals(original.algorithm(), copy.algorithm(), model + " algorithm"); + for (final String input : INPUTS) { + final String context = model + " input <" + input + ">"; + assertIterableEquals(original.encode(input), copy.encode(input), context + " pieces"); + } + } + + /** + * Verifies that a stream whose top-level object is not on the allow-list is rejected + * before it is materialised, even though its classes are harmless JDK types. + */ + @Test + void testForeignPayloadIsRejected() throws IOException { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + final ArrayList foreign = new ArrayList<>(); + foreign.add("not a tokenizer"); + out.writeObject(foreign); + } + assertThrows(InvalidClassException.class, () -> + SentencePieceTokenizer.deserialize(new ByteArrayInputStream(bytes.toByteArray()))); + } + + /** + * Verifies that a legitimate stream is rejected when it exceeds the supplied resource + * limits, so the limits bound the graph regardless of the class allow-list. + */ + @Test + void testStreamExceedingLimitsIsRejected() throws IOException { + final byte[] legitimate = serialized("tiny-unigram"); + final SentencePieceTokenizer.DeserializationLimits tight = + new SentencePieceTokenizer.DeserializationLimits(1, 1, 1); + assertThrows(InvalidClassException.class, () -> + SentencePieceTokenizer.deserialize(new ByteArrayInputStream(legitimate), tight)); + } + + /** + * Verifies that null arguments are rejected with {@link IllegalArgumentException} at the + * API boundary. + */ + @Test + void testNullArgumentsAreRejected() throws IOException { + final SentencePieceTokenizer tokenizer = SentencePieceFixtures.tokenizer("tiny-unigram"); + assertThrows(IllegalArgumentException.class, () -> tokenizer.serialize(null)); + assertThrows(IllegalArgumentException.class, () -> + SentencePieceTokenizer.deserialize(null)); + assertThrows(IllegalArgumentException.class, () -> + SentencePieceTokenizer.deserialize(new ByteArrayInputStream(new byte[0]), null)); + assertThrows(IllegalArgumentException.class, () -> + new SentencePieceTokenizer.DeserializationLimits(0, 1, 1)); + assertThrows(IllegalArgumentException.class, () -> + new SentencePieceTokenizer.DeserializationLimits(1, 0, 1)); + assertThrows(IllegalArgumentException.class, () -> + new SentencePieceTokenizer.DeserializationLimits(1, 1, 0)); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceUsageExampleTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceUsageExampleTest.java new file mode 100644 index 0000000000..b1bd4cc9ee --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceUsageExampleTest.java @@ -0,0 +1,69 @@ +/* + * 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.subword.sentencepiece; + +import java.io.IOException; +import java.io.InputStream; +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 opennlp.tools.tokenize.SubwordPiece; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Pins the cookbook path documented in {@code tokenizer.xml}: load a + * {@link SentencePieceTokenizer} from a {@code .model} file, encode text to pieces with + * original offsets, and obtain id arrays. + */ +class SentencePieceUsageExampleTest { + + @Test + void testLoadEncodeAndEncodeToIds(@TempDir Path dir) throws IOException { + final Path modelFile = dir.resolve("spiece.model"); + try (InputStream in = SentencePieceUsageExampleTest.class + .getResourceAsStream("tiny-unigram.model")) { + assertNotNull(in, "missing test resource tiny-unigram.model"); + Files.copy(in, modelFile); + } + + final SentencePieceTokenizer tokenizer = SentencePieceTokenizer.load(modelFile); + final String text = "hello world"; + final List pieces = tokenizer.encode(text); + assertFalse(pieces.isEmpty()); + for (final SubwordPiece piece : pieces) { + assertTrue(piece.id() >= 0); + assertTrue(piece.start() >= 0); + assertTrue(piece.end() <= text.length()); + // Control or whitespace pieces may report an empty span (start == end). + assertTrue(piece.start() <= piece.end()); + } + + final int[] ids = tokenizer.encodeToIds(text); + assertEquals(pieces.size(), ids.length); + for (int i = 0; i < ids.length; i++) { + assertEquals(pieces.get(i).id(), ids[i]); + } + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/README.md b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/README.md new file mode 100644 index 0000000000..6b8cbc39db --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/README.md @@ -0,0 +1,100 @@ + + +# SentencePiece parity fixtures + +Tiny trained `.model` files and matching `.fixtures.tsv` files used by +`SentencePieceParityTest` and related tests. They are **not** third-party +pretrained models: they are generated in-tree from `corpus.txt` plus a short +multilingual add-on list in `gen_fixtures.py`, using the reference +[sentencepiece](https://github.com/google/sentencepiece) Python package. + +The expected outputs in the TSVs come from the reference implementation, not +from the Java code under test, so the parity tests stay independent of the +implementation they check. + +## Regenerating the tiny models + +From this directory (or any directory; pass absolute paths as needed): + +```bash +python3 -m venv .venv +source .venv/bin/activate +pip install sentencepiece +python gen_fixtures.py corpus.txt . +``` + +That trains each model listed in `MODELS` inside `gen_fixtures.py` +(`tiny-unigram`, `tiny-unigram-bytefb`, `tiny-bpe`, `tiny-unigram-identity`, +`tiny-unigram-suffix`) and writes: + +- `.model` — SentencePiece binary model +- `.fixtures.tsv` — expected pieces, ids, UTF-16 spans, and normalized form +- `corpus-full.txt` — training corpus (`corpus.txt` plus multilingual lines) + +Pin the `sentencepiece` package version you used if regenerating for a PR, so +reviewers can reproduce the same bytes. + +## Validating the Java implementation + +To verify parity end to end, regenerate the fixtures as above, then run the +test suite from the repository root: + +```bash +./mvnw -pl opennlp-extensions/opennlp-subword -am test +``` + +`SentencePieceParityTest` asserts every fixture line piece for piece, span +for span, against `SentencePieceTokenizer`, for each bundled tiny model. + +To additionally validate against real published models, generate fixtures for +a directory of pre-trained `*.model` files and point the eval test at it: + +```bash +source .venv/bin/activate +python gen_real_fixtures.py /path/to/models +./mvnw -pl opennlp-extensions/opennlp-subword -am test \ + -Dopennlp.subword.eval.dir=/path/to/models +``` + +`SentencePieceRealModelEvalTest` is skipped unless +`opennlp.subword.eval.dir` is set. + +## Real-model fixtures (optional, not bundled) + +`gen_real_fixtures.py` writes the same TSV format for any directory of +pre-trained `*.model` files (no training). It reuses the escaping helpers and +input list from `gen_fixtures.py`: + +```bash +source .venv/bin/activate # same venv as above +python gen_real_fixtures.py /path/to/models +``` + +Real models and their TSVs are not checked into this tree; the script is for +local eval against published SentencePiece models. + +## Fixture TSV format + +Tab-separated, with backslash escapes (`\\`, `\t`, `\n`, `\r`): + +```text +esc(input) TAB pieceCount TAB [esc(piece) TAB id TAB begin TAB end]... TAB esc(normalized) +``` + +`begin` / `end` are UTF-16 code-unit offsets into the original input (Java +`String` indexing). diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/corpus.txt b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/corpus.txt new file mode 100644 index 0000000000..d578654315 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/corpus.txt @@ -0,0 +1,66 @@ +The quick brown fox jumps over the lazy dog. +Apache OpenNLP is a machine learning based toolkit for the processing of natural language text. +It supports the most common NLP tasks, such as tokenization, sentence segmentation, and named entity extraction. +Subword tokenization decomposes words into smaller units drawn from a fixed vocabulary. +The unigram language model selects the segmentation with the highest total log probability. +Byte pair encoding merges the most frequent adjacent symbol pairs until no merge applies. +A sentence piece model carries its own text normalizer inside the model file. +Character offsets should always point back into the original text. +Whitespace is escaped with a special marker so word boundaries survive segmentation. +Numbers like 3.14159 and 42 and 1024 appear in ordinary text. +Punctuation, quotes, and dashes are folded by the normalizer! +Questions? Answers! Ellipses... and (parentheses) too. +The cafe served naive patrons a souffle with creme fraiche. +Internationalization and localization are long words. +Antidisestablishmentarianism remains one of the longest English words. +She sells seashells by the seashore. +Peter Piper picked a peck of pickled peppers. +How much wood would a woodchuck chuck if a woodchuck could chuck wood? +The rain in Spain stays mainly in the plain. +To be or not to be, that is the question. +All happy families are alike; each unhappy family is unhappy in its own way. +It was the best of times, it was the worst of times. +Call me Ishmael. +In the beginning was the word. +The world is everything that is the case. +Language models assign probabilities to sequences of tokens. +Retrieval systems rank documents by similarity to a query. +Embeddings map text into dense vector spaces. +Search engines combine lexical and semantic signals. +The tokenizer must be fast, deterministic, and thread safe. +Model files are loaded once and shared across threads. +Tests must prove parity with the reference implementation. +Offsets are measured in code units of the original encoding. +The normalizer collapses repeated whitespace into one marker. +A leading marker separates words that start a sentence. +Unknown characters fall back to a penalty score. +Byte fallback decomposes unknown characters into byte pieces. +User defined symbols are never split by the tokenizer. +Control symbols never appear in encoded output. +The vocabulary maps each piece to an integer identifier. +Scores are log probabilities in the unigram model. +Merge ranks order the byte pair encoding agenda. +The trie enumerates every piece that starts at a position. +Dynamic programming finds the best path in one pass. +Backtracking recovers the winning segmentation. +The agenda is a priority queue ordered by score. +Stale entries are skipped when their symbols have merged. +Un texto corto en espanol para variar el corpus. +Un petit texte en francais pour la diversite. +Ein kurzer deutscher Satz steht auch hier. +Ancora una frase italiana per completezza. +Tokenization quality depends on the training corpus. +The model was trained on a tiny corpus for testing only. +Nothing in this file is quoted from any external work. +Short lines help. +One. +Two words. +Three little words. +water water water water water +running runner ran runs +walked walking walker walks +faster fastest fast +apple apples applesauce +book books bookshelf bookstore +work works worked working worker +play plays played playing player diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_fixtures.py b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_fixtures.py new file mode 100644 index 0000000000..658902e48f --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_fixtures.py @@ -0,0 +1,139 @@ +# 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. + +"""Trains the tiny SentencePiece test models and generates the parity fixtures. + +Run inside a venv with the sentencepiece package installed: + python gen_fixtures.py + +Every fixture line is tab-separated with backslash escaping (\\\\, \\t, \\n, \\r): + esc(input) TAB pieceCount TAB [esc(piece) TAB id TAB begin TAB end]... TAB esc(normalized) +Offsets are UTF-16 code-unit offsets into the original input, matching Java string indexing. +""" +import sys +import sentencepiece as spm + +MULTILINGUAL = [ + "Le café coûte trois euros à Paris.", + "Der Straßenname ändert sich häufig.", + "Ça va très bien, merci beaucoup.", + "El niño pequeño come una manzana.", + "Привет мир и всем добро.", + "東京タワーに登りました。", + "日本語の文章も少しあります。", + "안녕하세요 세계입니다.", + "你好世界这是中文。", + "I love \U0001f355 and \U0001f1e9\U0001f1ea a lot!", + "Emoji test \U0001f600 \U0001f680 ❤️ done.", +] + +INPUTS = [ + "", + " ", + " ", + "a", + "Hello world", + " Hello world ", + "Hello world.\nSecond line\ttabbed", + "The quick brown fox jumps over the lazy dog.", + "tokenization and segmentation", + "Antidisestablishmentarianism", + "water running walked faster apple book work play", + "3.14159 x 42 = 1024?", + "!!!???...", + "(parentheses) and [brackets] and {braces}", + "café naïve fiancé résumé", + "financial fluid", + "① ⑪ ㋿ ＫＡＴＡＫＡＮＡ", + "ｶﾀｶﾅ half width", + "東京タワーへ行きました", + "日本語とEnglish混在", + "Привет мир", + "안녕하세요 세계", + "你好，世界！", + "I love \U0001f355 pizza", + "flags \U0001f1e9\U0001f1ea \U0001f1fa\U0001f1f8 end", + "family \U0001f469‍\U0001f469‍\U0001f467‍\U0001f466 emoji", + "zero​width and non breaking", + "quotes “fancy” and ‘single’ — dash", + " the [URL] token", + "a b[URL]c", + "control tokens inline", + "https://example.com/path?q=1&x=2", + "UPPER lower MiXeD case", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "Ω≈ç√∫˜µ≤", + "مرحبا بالعالم", + " leading and trailing ", + "\ttab\tstart", + "newline\n\n\nruns", + "mid spaces collapse", +] + +MODELS = { + "tiny-unigram": dict(model_type="unigram", vocab_size=300), + "tiny-unigram-bytefb": dict(model_type="unigram", vocab_size=600, byte_fallback=True, + character_coverage=0.995), + "tiny-bpe": dict(model_type="bpe", vocab_size=300), + "tiny-unigram-identity": dict(model_type="unigram", vocab_size=300, + normalization_rule_name="identity"), + "tiny-unigram-suffix": dict(model_type="unigram", vocab_size=300, + treat_whitespace_as_suffix=True), +} + + +def esc(s): + return (s.replace("\\", "\\\\").replace("\t", "\\t") + .replace("\n", "\\n").replace("\r", "\\r")) + + +def utf16_offset(text, codepoint_offset): + return len(text[:codepoint_offset].encode("utf-16-le")) // 2 + + +def main(corpus, outdir): + full_corpus = outdir + "/corpus-full.txt" + with open(corpus, encoding="utf-8") as f: + lines = f.read().splitlines() + lines += MULTILINGUAL + with open(full_corpus, "w", encoding="utf-8") as f: + f.write("\n".join(lines) + "\n") + + for name, opts in MODELS.items(): + spm.SentencePieceTrainer.Train( + input=full_corpus, + model_prefix=outdir + "/" + name, + hard_vocab_limit=False, + character_coverage=opts.pop("character_coverage", 1.0), + user_defined_symbols=["", "[URL]"], + self_test_sample_size=10, + **opts, + ) + sp = spm.SentencePieceProcessor(model_file=outdir + "/" + name + ".model") + with open(outdir + "/" + name + ".fixtures.tsv", "w", encoding="utf-8") as out: + for text in INPUTS: + proto = sp.EncodeAsImmutableProto(text) + cols = [esc(text), str(len(proto.pieces))] + for piece in proto.pieces: + cols += [esc(piece.piece), str(piece.id), + str(utf16_offset(text, piece.begin)), + str(utf16_offset(text, piece.end))] + cols.append(esc(sp.Normalize(text))) + out.write("\t".join(cols) + "\n") + print(name, "vocab", sp.GetPieceSize()) + + +if __name__ == "__main__": + main(sys.argv[1], sys.argv[2]) diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_real_fixtures.py b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_real_fixtures.py new file mode 100644 index 0000000000..9b27c7a11e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/gen_real_fixtures.py @@ -0,0 +1,75 @@ +# 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. + +"""Generates parity fixtures for pre-trained real-world models (no training). + +Usage: python gen_real_fixtures.py +Reads every *.model in the directory and writes a sibling *.fixtures.tsv in the same +escaped-TSV format as gen_fixtures.py, over a larger and messier input list. +""" +import glob +import os +import sys +import sentencepiece as spm +from gen_fixtures import INPUTS, esc, utf16_offset + +EXTRA = [ + "The Transformer architecture revolutionized natural language processing in 2017.", + "supercalifragilisticexpialidocious and pneumonoultramicroscopicsilicovolcanoconiosis", + "e=mc^2, F=ma, and a^2+b^2=c^2 are famous equations.", + "Mixed scripts: English, 日本語, 한국어, русский, and العربية together.", + "Prices: $19.99, €25,50, £12, ¥1500, and ₹999.", + "C++ and C# and F# are programming languages; so is Java.", + "def encode(text): return sp.encode(text, out_type=str)", + "SELECT * FROM documents WHERE score > 0.5 ORDER BY rank;", + "The 2024 Summer Olympics were held in Paris, France.", + "COVID-19 vaccines use mRNA technology (Pfizer-BioNTech, Moderna).", + "Email me at test.user+tag@example.co.uk or call +1 (555) 010-9999.", + "10,000 steps a day keeps the doctor away... allegedly!", + "The naive resume of the fiancee included a cafe visit.", + "¿Dónde está la biblioteca? ¡Allí está!", + "Smørrebrød og æbleskiver er danske specialiteter.", + "Zażółć gęślą jaźń is a Polish pangram.", + "Đây là tiếng Việt với nhiều dấu.", + "今日はいい天気ですね。明日も晴れるといいな。", + "北京和上海都是大城市。", + "한국의 수도는 서울입니다.", + "\U0001f9d1‍\U0001f4bb codes while \U0001f9d1‍\U0001f373 cooks \U0001f35c!", + "
line separator and 
paragraph separator lurk here", + "BOM at the start of this sentence", + "tabs\tand\ttabs\tand\ttabs", + "CRLF\r\nline endings\r\nhappen", +] + + +def main(model_dir): + for model_path in sorted(glob.glob(os.path.join(model_dir, "*.model"))): + name = os.path.splitext(model_path)[0] + sp = spm.SentencePieceProcessor(model_file=model_path) + with open(name + ".fixtures.tsv", "w", encoding="utf-8") as out: + for text in INPUTS + EXTRA: + proto = sp.EncodeAsImmutableProto(text) + cols = [esc(text), str(len(proto.pieces))] + for piece in proto.pieces: + cols += [esc(piece.piece), str(piece.id), + str(utf16_offset(text, piece.begin)), + str(utf16_offset(text, piece.end))] + cols.append(esc(sp.Normalize(text))) + out.write("\t".join(cols) + "\n") + print(os.path.basename(name), "vocab", sp.GetPieceSize()) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.fixtures.tsv b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.fixtures.tsv new file mode 100644 index 0000000000..e3560148f9 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.fixtures.tsv @@ -0,0 +1,40 @@ + 0 + 0 + 0 +a 1 ▁a 8 0 1 ▁a +Hello world 7 ▁ 177 0 0 H 246 0 1 el 48 1 3 l 186 3 4 o 183 4 5 ▁wor 38 5 9 ld 129 9 11 ▁Hello▁world + Hello world 7 ▁ 177 1 1 H 246 1 2 el 48 2 4 l 186 4 5 o 183 5 6 ▁wor 38 6 12 ld 129 12 14 ▁Hello▁world +Hello world.\nSecond line\ttabbed 19 ▁ 177 0 0 H 246 0 1 el 48 1 3 l 186 3 4 o 183 4 5 ▁wor 38 5 9 ld 129 9 11 . 193 11 12 ▁S 64 12 14 ec 51 14 16 on 16 16 18 d 189 18 19 ▁l 28 19 21 in 7 21 23 e 178 23 24 ▁t 5 24 26 ab 85 26 28 b 199 28 29 ed 24 29 31 ▁Hello▁world.▁Second▁line▁tabbed +The quick brown fox jumps over the lazy dog. 27 ▁The 50 0 3 ▁qu 81 3 6 ic 61 6 8 k 196 8 9 ▁b 23 9 11 ro 41 11 13 wn 90 13 15 ▁f 25 15 17 o 183 17 18 x 203 18 19 ▁ 177 19 20 j 220 20 21 u 192 21 22 mp 108 22 24 s 182 24 25 ▁o 45 25 27 v 200 27 28 er 6 28 30 ▁the 19 30 34 ▁l 28 34 36 a 179 36 37 z 202 37 38 y 198 38 39 ▁d 42 39 41 o 183 41 42 g 194 42 43 . 193 43 44 ▁The▁quick▁brown▁fox▁jumps▁over▁the▁lazy▁dog. +tokenization and segmentation 4 ▁tokeniz 170 0 7 ation 47 7 12 ▁and 59 12 16 ▁segmentation 171 16 29 ▁tokenization▁and▁segmentation +Antidisestablishmentarianism 16 ▁A 76 0 1 n 181 1 2 t 180 2 3 id 176 3 5 is 31 5 7 est 57 7 10 ab 85 10 12 l 186 12 13 is 31 13 15 h 187 15 16 ment 83 16 20 ar 18 20 22 i 184 22 23 an 40 23 25 is 31 25 27 m 191 27 28 ▁Antidisestablishmentarianism +water running walked faster apple book work play 15 ▁water 124 0 5 ▁r 93 5 7 un 39 7 9 ning 154 9 13 ▁wal 160 13 17 k 196 17 18 ed 24 18 20 ▁fast 162 20 25 er 6 25 27 ▁app 99 27 31 le 107 31 33 ▁boo 155 33 37 k 196 37 38 ▁work 167 38 43 ▁play 123 43 48 ▁water▁running▁walked▁faster▁apple▁book▁work▁play +3.14159 x 42 = 1024? 21 ▁ 177 0 0 3 242 0 1 . 193 1 2 1 215 2 3 4 216 3 4 1 215 4 5 5 243 5 6 9 244 6 7 ▁ 177 7 8 x 203 8 9 ▁ 177 9 10 4 216 10 11 2 224 11 12 ▁ 177 12 13 = 0 13 14 ▁ 177 14 15 1 215 15 16 0 241 16 17 2 224 17 18 4 216 18 19 ? 225 19 20 ▁3.14159▁x▁42▁=▁1024? +!!!???... 10 ▁ 177 0 0 ! 214 0 1 ! 214 1 2 ! 214 2 3 ? 225 3 4 ? 225 4 5 ? 225 5 6 . 193 6 7 . 193 7 8 . 193 8 9 ▁!!!???... +(parentheses) and [brackets] and {braces} 25 ▁ 177 0 0 ( 239 0 1 p 190 1 2 ar 18 2 4 ent 35 4 7 he 9 7 9 ses 114 9 12 ) 240 12 13 ▁and 59 13 17 ▁ 177 17 18 [ 0 18 19 b 199 19 20 r 185 20 21 ack 112 21 24 et 86 24 26 s 182 26 27 ] 0 27 28 ▁and 59 28 32 ▁ 177 32 33 { 0 33 34 b 199 34 35 r 185 35 36 ac 29 36 38 es 11 38 40 } 0 40 41 ▁(parentheses)▁and▁[brackets]▁and▁{braces} +café naïve fiancé résumé 18 ▁c 26 0 1 af 173 1 3 é 254 3 4 ▁n 49 4 6 a 179 6 7 ï 0 7 8 ve 109 8 10 ▁f 25 10 12 i 184 12 13 an 40 13 15 c 188 15 16 é 254 16 17 ▁r 93 17 19 é 254 19 20 s 182 20 21 u 192 21 22 m 191 22 23 é 254 23 24 ▁café▁naïve▁fiancé▁résumé +financial fluid 10 ▁f 25 0 0 in 7 0 2 an 40 2 4 c 188 4 5 i 184 5 6 al 20 6 8 ▁f 25 8 9 l 186 9 10 u 192 10 11 id 176 11 13 ▁financial▁fluid +① ⑪ ㋿ ＫＡＴＡＫＡＮＡ 16 ▁ 177 0 0 1 215 0 1 ▁ 177 1 2 1 215 2 2 1 215 2 3 ▁ 177 3 4 令和 0 4 5 ▁ 177 5 6 K 0 6 7 A 207 7 8 T 201 8 9 A 207 9 10 K 0 10 11 A 207 11 12 N 212 12 13 A 207 13 14 ▁1▁11▁令和▁KATAKANA +ｶﾀｶﾅ half width 11 ▁ 177 0 0 カ 0 0 1 タ 268 1 2 カナ 0 2 4 ▁h 110 4 6 al 20 6 8 f 197 8 9 ▁w 14 9 11 id 176 11 13 t 180 13 14 h 187 14 15 ▁カタカナ▁half▁width +東京タワーへ行きました 10 ▁ 177 0 0 東 280 0 1 京 273 1 2 タ 268 2 3 ワ 269 3 4 ー 270 4 5 へ行き 0 5 8 ま 235 8 9 し 234 9 10 た 264 10 11 ▁東京タワーへ行きました +日本語とEnglish混在 12 ▁ 177 0 0 日 277 0 1 本 279 1 2 語 284 2 3 と 0 3 4 E 208 4 5 n 181 5 6 g 194 6 7 l 186 7 8 is 31 8 10 h 187 10 11 混在 0 11 13 ▁日本語とEnglish混在 +Привет мир 11 ▁ 177 0 0 П 256 0 1 р 222 1 2 и 221 2 3 в 230 3 4 е 231 4 5 т 260 5 6 ▁ 177 6 7 м 232 7 8 и 221 8 9 р 222 9 10 ▁Привет▁мир +안녕하세요 세계 9 ▁ 177 0 0 안 290 0 1 녕 287 1 2 하 293 2 3 세 238 3 4 요 291 4 5 ▁ 177 5 6 세 238 6 7 계 286 7 8 ▁안녕하세요▁세계 +你好，世界！ 7 ▁ 177 0 0 你 274 0 1 好 275 1 2 , 204 2 3 世 271 3 4 界 281 4 5 ! 214 5 6 ▁你好,世界! +I love 🍕 pizza 9 ▁I 92 0 1 ▁lo 96 1 4 ve 109 4 6 ▁ 177 6 7 🍕 297 7 9 ▁p 15 9 11 iz 54 11 13 z 202 13 14 a 179 14 15 ▁I▁love▁🍕▁pizza +flags 🇩🇪 🇺🇸 end 12 ▁f 25 0 1 l 186 1 2 a 179 2 3 g 194 3 4 s 182 4 5 ▁ 177 5 6 🇩 295 6 8 🇪 296 8 10 ▁ 177 10 11 🇺🇸 0 11 15 ▁en 149 15 18 d 189 18 19 ▁flags▁🇩🇪▁🇺🇸▁end +family 👩‍👩‍👧‍👦 emoji 12 ▁f 25 0 1 am 104 1 3 il 53 3 5 y 198 5 6 ▁ 177 6 7 👩‍👩‍👧‍👦 0 7 18 ▁ 177 18 19 e 178 19 20 m 191 20 21 o 183 21 22 j 220 22 23 i 184 23 24 ▁family▁👩‍👩‍👧‍👦▁emoji +zero​width and non breaking 16 ▁ 177 0 0 z 202 0 1 er 6 1 3 o 183 3 4 ▁w 14 4 6 id 176 6 8 t 180 8 9 h 187 9 10 ▁and 59 10 14 ▁n 49 14 16 on 16 16 18 ▁b 23 18 20 re 32 20 22 a 179 22 23 k 196 23 24 ing 27 24 27 ▁zero▁width▁and▁non▁breaking +quotes “fancy” and ‘single’ — dash 22 ▁qu 81 0 2 ot 130 2 4 es 11 4 6 ▁ 177 6 7 “ 0 7 8 f 197 8 9 an 40 9 11 c 188 11 12 y 198 12 13 ” 0 13 14 ▁and 59 14 18 ▁ 177 18 19 ‘ 0 19 20 s 182 20 21 ing 27 21 24 le 107 24 26 ’ 0 26 27 ▁ 177 27 28 — 0 28 29 ▁d 42 29 31 as 30 31 33 h 187 33 34 ▁quotes▁“fancy”▁and▁‘single’▁—▁dash + the [URL] token 7 ▁ 177 0 0 3 0 6 ▁the 19 6 10 ▁ 177 10 11 [URL] 4 11 16 ▁to 43 16 19 ken 95 19 22 ▁▁the▁[URL]▁token +a b[URL]c 6 ▁a 8 0 1 ▁ 177 1 2 3 2 8 b 199 8 9 [URL] 4 9 14 c 188 14 15 ▁a▁b[URL]c +control tokens inline 20 ▁c 26 0 1 on 16 1 3 t 180 3 4 ro 41 4 6 l 186 6 7 ▁ 177 7 8 < 0 8 9 s 182 9 10 > 0 10 11 ▁to 43 11 14 ken 95 14 17 s 182 17 18 ▁ 177 18 19 0 22 23 ▁in 37 23 26 l 186 26 27 in 7 27 29 e 178 29 30 ▁control▁▁tokens▁▁inline +https://example.com/path?q=1&x=2 26 ▁h 110 0 1 t 180 1 2 t 180 2 3 p 190 3 4 s 182 4 5 :// 0 5 8 ex 52 8 10 am 104 10 12 p 190 12 13 le 107 13 15 . 193 15 16 c 188 16 17 o 183 17 18 m 191 18 19 / 0 19 20 p 190 20 21 at 17 21 23 h 187 23 24 ? 225 24 25 q 205 25 26 = 0 26 27 1 215 27 28 & 0 28 29 x 203 29 30 = 0 30 31 2 224 31 32 ▁https://example.com/path?q=1&x=2 +UPPER lower MiXeD case 17 ▁U 135 0 1 P 210 1 2 P 210 2 3 E 208 3 4 R 248 4 5 ▁lo 96 5 8 w 195 8 9 er 6 9 11 ▁ 177 11 12 M 227 12 13 i 184 13 14 X 0 14 15 e 178 15 16 D 226 16 17 ▁c 26 17 19 as 30 19 21 e 178 21 22 ▁UPPER▁lower▁MiXeD▁case +aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa 40 ▁a 8 0 1 a 179 1 2 a 179 2 3 a 179 3 4 a 179 4 5 a 179 5 6 a 179 6 7 a 179 7 8 a 179 8 9 a 179 9 10 a 179 10 11 a 179 11 12 a 179 12 13 a 179 13 14 a 179 14 15 a 179 15 16 a 179 16 17 a 179 17 18 a 179 18 19 a 179 19 20 a 179 20 21 a 179 21 22 a 179 22 23 a 179 23 24 a 179 24 25 a 179 25 26 a 179 26 27 a 179 27 28 a 179 28 29 a 179 29 30 a 179 30 31 a 179 31 32 a 179 32 33 a 179 33 34 a 179 34 35 a 179 35 36 a 179 36 37 a 179 37 38 a 179 38 39 a 179 39 40 ▁aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +Ω≈ç√∫˜µ≤ 4 ▁ 177 0 0 Ω≈ç√∫ 0 0 5 ▁ 177 5 5 ̃μ≤ 0 5 8 ▁Ω≈ç√∫▁̃μ≤ +مرحبا بالعالم 4 ▁ 177 0 0 مرحبا 0 0 5 ▁ 177 5 6 بالعالم 0 6 13 ▁مرحبا▁بالعالم + leading and trailing 9 ▁l 28 2 3 e 178 3 4 ad 172 4 6 ing 27 6 9 ▁and 59 9 13 ▁t 5 13 15 ra 62 15 17 il 53 17 19 ing 27 19 22 ▁leading▁and▁trailing +\ttab\tstart 6 ▁t 5 1 2 ab 85 2 4 ▁s 13 4 6 t 180 6 7 ar 18 7 9 t 180 9 10 ▁tab▁start +newline\n\n\nruns 9 ▁n 49 0 1 e 178 1 2 w 195 2 3 l 186 3 4 in 7 4 6 e 178 6 7 ▁r 93 7 11 un 39 11 13 s 182 13 14 ▁newline▁runs +mid spaces collapse 11 ▁m 22 0 1 id 176 1 3 ▁s 13 3 7 pac 147 7 10 es 11 10 12 ▁c 26 12 16 ol 69 16 18 l 186 18 19 ap 127 19 21 s 182 21 22 e 178 22 23 ▁mid▁spaces▁collapse diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.model b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.model new file mode 100644 index 0000000000..8b6f22eb82 Binary files /dev/null and b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-bpe.model differ diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.fixtures.tsv b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.fixtures.tsv new file mode 100644 index 0000000000..b8f9733dba --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.fixtures.tsv @@ -0,0 +1,40 @@ + 0 + 0 + 0 +a 1 ▁a 266 0 1 ▁a +Hello world 7 ▁ 261 0 0 H 552 0 1 e 264 1 2 ll 412 2 4 o 274 4 5 ▁wor 427 5 9 ld 317 9 11 ▁Hello▁world + Hello world 7 ▁ 261 1 1 H 552 1 2 e 264 2 3 ll 412 3 5 o 274 5 6 ▁wor 427 6 12 ld 317 12 14 ▁Hello▁world +Hello world.\nSecond line\ttabbed 21 ▁ 261 0 0 H 552 0 1 e 264 1 2 ll 412 2 4 o 274 4 5 ▁wor 427 5 9 ld 317 9 11 . 262 11 12 ▁S 296 12 14 e 264 14 15 c 411 15 16 o 274 16 17 nd 284 17 19 ▁l 492 19 21 ine 415 21 24 ▁ 261 24 25 t 269 25 26 a 279 26 27 b 344 27 28 b 344 28 29 ed 267 29 31 ▁Hello▁world.▁Second▁line▁tabbed +The quick brown fox jumps over the lazy dog. 26 ▁Th 277 0 2 e 264 2 3 ▁qu 346 3 6 i 282 6 7 ck 307 7 9 ▁b 309 9 11 r 288 11 12 ow 306 12 14 n 268 14 15 ▁fo 359 15 18 x 421 18 19 ▁ 261 19 20 j 381 20 21 um 529 21 23 p 352 23 24 s 263 24 25 ▁o 513 25 27 ve 366 27 29 r 288 29 30 ▁the 265 30 34 ▁la 339 34 37 z 351 37 38 y 275 38 39 ▁do 464 39 42 g 356 42 43 . 262 43 44 ▁The▁quick▁brown▁fox▁jumps▁over▁the▁lazy▁dog. +tokenization and segmentation 4 ▁tokenization 461 0 12 ▁a 266 12 14 nd 284 14 16 ▁segmentation 335 16 29 ▁tokenization▁and▁segmentation +Antidisestablishmentarianism 15 ▁An 408 0 2 ti 523 2 4 d 273 4 5 is 278 5 7 est 323 7 10 a 279 10 11 b 344 11 12 lish 454 12 16 ment 508 16 20 ar 353 20 22 i 282 22 23 a 279 23 24 n 268 24 25 is 278 25 27 m 320 27 28 ▁Antidisestablishmentarianism +water running walked faster apple book work play 15 ▁water 312 0 5 ▁runn 467 5 10 ing 272 10 13 ▁walk 348 13 18 ed 267 18 20 ▁fast 358 20 25 er 270 25 27 ▁app 498 27 31 le 405 31 33 ▁b 309 33 35 o 274 35 36 o 274 36 37 k 354 37 38 ▁work 303 38 43 ▁play 313 43 48 ▁water▁running▁walked▁faster▁apple▁book▁work▁play +3.14159 x 42 = 1024? 21 ▁ 261 0 0 3 543 0 1 . 262 1 2 1 371 2 3 4 372 3 4 1 371 4 5 5 549 5 6 9 550 6 7 ▁ 261 7 8 x 421 8 9 ▁ 261 9 10 4 372 10 11 2 434 11 12 ▁ 261 12 13 <0x3D> 66 13 14 ▁ 261 14 15 1 371 15 16 0 548 16 17 2 434 17 18 4 372 18 19 ? 435 19 20 ▁3.14159▁x▁42▁=▁1024? +!!!???... 10 ▁ 261 0 0 ! 370 0 1 ! 370 1 2 ! 370 2 3 ? 435 3 4 ? 435 4 5 ? 435 5 6 . 262 6 7 . 262 7 8 . 262 8 9 ▁!!!???... +(parentheses) and [brackets] and {braces} 27 ▁ 261 0 0 ( 546 0 1 p 352 1 2 are 286 2 5 n 268 5 6 th 289 6 8 e 264 8 9 ses 414 9 12 ) 547 12 13 ▁a 266 13 15 nd 284 15 17 ▁ 261 17 18 <0x5B> 96 18 19 b 344 19 20 r 288 20 21 ack 327 21 24 e 264 24 25 ts 311 25 27 <0x5D> 98 27 28 ▁a 266 28 30 nd 284 30 32 ▁ 261 32 33 <0x7B> 128 33 34 b 344 34 35 ra 409 35 37 ces 478 37 40 <0x7D> 130 40 41 ▁(parentheses)▁and▁[brackets]▁and▁{braces} +café naïve fiancé résumé 19 ▁caf 484 0 3 é 557 3 4 ▁ 261 4 5 n 268 5 6 a 279 6 7 <0xC3> 200 7 7 <0xAF> 180 7 8 ve 366 8 10 ▁fi 518 10 13 a 279 13 14 n 268 14 15 c 411 15 16 é 557 16 17 ▁ 261 17 18 r 288 18 19 é 557 19 20 s 263 20 21 um 529 21 23 é 557 23 24 ▁café▁naïve▁fiancé▁résumé +financial fluid 13 ▁fi 518 0 1 n 268 1 2 a 279 2 3 n 268 3 4 c 411 4 5 i 282 5 6 al 281 6 8 ▁ 261 8 9 f 337 9 9 l 319 9 10 u 271 10 11 i 282 11 12 d 273 12 13 ▁financial▁fluid +① ⑪ ㋿ ＫＡＴＡＫＡＮＡ 21 ▁ 261 0 0 1 371 0 1 ▁ 261 1 2 1 371 2 2 1 371 2 3 ▁ 261 3 4 <0xE4> 233 4 4 <0xBB> 192 4 4 <0xA4> 169 4 4 <0xE5> 234 4 4 <0x92> 151 4 4 <0x8C> 145 4 5 ▁ 261 5 6 <0x4B> 80 6 7 A 596 7 8 T 599 8 9 A 596 9 10 <0x4B> 80 10 11 A 596 11 12 N 542 12 13 A 596 13 14 ▁1▁11▁令和▁KATAKANA +ｶﾀｶﾅ half width 17 ▁ 261 0 0 <0xE3> 232 0 0 <0x82> 135 0 0 <0xAB> 176 0 1 タ 565 1 2 <0xE3> 232 2 2 <0x82> 135 2 2 <0xAB> 176 2 3 <0xE3> 232 3 3 <0x83> 136 3 3 <0x8A> 143 3 4 ▁h 512 4 6 al 281 6 8 f 337 8 9 ▁wi 314 9 12 d 273 12 13 th 289 13 15 ▁カタカナ▁half▁width +東京タワーへ行きました 18 ▁ 261 0 0 東 571 0 1 京 568 1 2 タ 565 2 3 ワ 580 3 4 ー 581 4 5 <0xE3> 232 5 5 <0x81> 134 5 5 <0xB8> 189 5 6 <0xE8> 237 6 6 <0xA1> 166 6 6 <0x8C> 145 6 7 <0xE3> 232 7 7 <0x81> 134 7 7 <0x8D> 146 7 8 ま 587 8 9 し 450 9 10 た 564 10 11 ▁東京タワーへ行きました +日本語とEnglish混在 18 ▁ 261 0 0 日 583 0 1 本 585 1 2 <0xE8> 237 2 2 <0xAA> 175 2 2 <0x9E> 163 2 3 <0xE3> 232 3 3 <0x81> 134 3 3 <0xA8> 173 3 4 E 595 4 5 ng 328 5 7 lish 454 7 11 <0xE6> 235 11 11 <0xB7> 188 11 11 <0xB7> 188 11 12 <0xE5> 234 12 12 <0x9C> 161 12 12 <0xA8> 173 12 13 ▁日本語とEnglish混在 +Привет мир 11 ▁ 261 0 0 П 559 0 1 р 382 1 2 и 373 2 3 в 437 3 4 е 438 4 5 т 561 5 6 ▁ 261 6 7 м 439 7 8 и 373 8 9 р 382 9 10 ▁Привет▁мир +안녕하세요 세계 19 ▁ 261 0 0 <0xEC> 241 0 0 <0x95> 154 0 0 <0x88> 141 0 1 <0xEB> 240 1 1 <0x85> 138 1 1 <0x95> 154 1 2 <0xED> 242 2 2 <0x95> 154 2 2 <0x98> 157 2 3 세 432 3 4 <0xEC> 241 4 4 <0x9A> 159 4 4 <0x94> 153 4 5 ▁ 261 5 6 세 432 6 7 <0xEA> 239 7 7 <0xB3> 184 7 7 <0x84> 137 7 8 ▁안녕하세요▁세계 +你好，世界！ 7 ▁ 261 0 0 你 569 0 1 好 570 1 2 , 280 2 3 世 566 3 4 界 572 4 5 ! 370 5 6 ▁你好,世界! +I love 🍕 pizza 12 ▁I 329 0 1 ▁lo 305 1 4 ve 366 4 6 ▁ 261 6 7 <0xF0> 245 7 7 <0x9F> 164 7 7 <0x8D> 146 7 7 <0x95> 154 7 9 ▁p 486 9 11 i 282 11 12 z 351 12 13 za 540 13 15 ▁I▁love▁🍕▁pizza +flags 🇩🇪 🇺🇸 end 26 ▁ 261 0 0 f 337 0 1 l 319 1 2 a 279 2 3 g 356 3 4 s 263 4 5 ▁ 261 5 6 <0xF0> 245 6 6 <0x9F> 164 6 6 <0x87> 140 6 6 <0xA9> 174 6 8 <0xF0> 245 8 8 <0x9F> 164 8 8 <0x87> 140 8 8 <0xAA> 175 8 10 ▁ 261 10 11 <0xF0> 245 11 11 <0x9F> 164 11 11 <0x87> 140 11 11 <0xBA> 191 11 13 <0xF0> 245 13 13 <0x9F> 164 13 13 <0x87> 140 13 13 <0xB8> 189 13 15 ▁en 318 15 18 d 273 18 19 ▁flags▁🇩🇪▁🇺🇸▁end +family 👩‍👩‍👧‍👦 emoji 34 ▁famil 446 0 5 y 275 5 6 ▁ 261 6 7 <0xF0> 245 7 7 <0x9F> 164 7 7 <0x91> 150 7 7 <0xA9> 174 7 9 <0xE2> 231 9 9 <0x80> 133 9 9 <0x8D> 146 9 10 <0xF0> 245 10 10 <0x9F> 164 10 10 <0x91> 150 10 10 <0xA9> 174 10 12 <0xE2> 231 12 12 <0x80> 133 12 12 <0x8D> 146 12 13 <0xF0> 245 13 13 <0x9F> 164 13 13 <0x91> 150 13 13 <0xA7> 172 13 15 <0xE2> 231 15 15 <0x80> 133 15 15 <0x8D> 146 15 16 <0xF0> 245 16 16 <0x9F> 164 16 16 <0x91> 150 16 16 <0xA6> 171 16 18 ▁ 261 18 19 e 264 19 20 m 320 20 21 o 274 21 22 j 381 22 23 i 282 23 24 ▁family▁👩‍👩‍👧‍👦▁emoji +zero​width and non breaking 16 ▁ 261 0 0 z 351 0 1 er 270 1 3 o 274 3 4 ▁wi 314 4 7 d 273 7 8 th 289 8 10 ▁a 266 10 12 nd 284 12 14 ▁no 489 14 17 n 268 17 18 ▁b 309 18 20 re 349 20 22 a 279 22 23 k 354 23 24 ing 272 24 27 ▁zero▁width▁and▁non▁breaking +quotes “fancy” and ‘single’ — dash 33 ▁quote 456 0 5 s 263 5 6 ▁ 261 6 7 <0xE2> 231 7 7 <0x80> 133 7 7 <0x9C> 161 7 8 f 337 8 9 a 279 9 10 n 268 10 11 c 411 11 12 y 275 12 13 <0xE2> 231 13 13 <0x80> 133 13 13 <0x9D> 162 13 14 ▁a 266 14 16 nd 284 16 18 ▁ 261 18 19 <0xE2> 231 19 19 <0x80> 133 19 19 <0x98> 157 19 20 s 263 20 21 ing 272 21 24 le 405 24 26 <0xE2> 231 26 26 <0x80> 133 26 26 <0x99> 158 26 27 ▁ 261 27 28 <0xE2> 231 28 28 <0x80> 133 28 28 <0x94> 153 28 29 ▁d 368 29 31 as 420 31 33 h 308 33 34 ▁quotes▁“fancy”▁and▁‘single’▁—▁dash + the [URL] token 8 ▁ 261 0 0 3 0 6 ▁the 265 6 10 ▁ 261 10 11 [URL] 4 11 16 ▁to 304 16 19 k 354 19 20 en 350 20 22 ▁▁the▁[URL]▁token +a b[URL]c 6 ▁a 266 0 1 ▁ 261 1 2 3 2 8 b 344 8 9 [URL] 4 9 14 c 411 14 15 ▁a▁b[URL]c +control tokens inline 20 ▁co 326 0 2 n 268 2 3 tro 429 3 6 l 319 6 7 ▁ 261 7 8 <0x3C> 65 8 9 s 263 9 10 <0x3E> 67 10 11 ▁to 304 11 14 k 354 14 15 en 350 15 17 s 263 17 18 ▁ 261 18 19 <0x3C> 65 19 20 <0x2F> 52 20 21 s 263 21 22 <0x3E> 67 22 23 ▁in 276 23 26 l 319 26 27 ine 415 27 30 ▁control▁▁tokens▁▁inline +https://example.com/path?q=1&x=2 27 ▁h 512 0 1 t 269 1 2 t 269 2 3 p 352 3 4 s 263 4 5 <0x3A> 63 5 6 <0x2F> 52 6 7 <0x2F> 52 7 8 e 264 8 9 x 421 9 10 a 279 10 11 mple 485 11 15 . 262 15 16 c 411 16 17 om 475 17 19 <0x2F> 52 19 20 p 352 20 21 a 279 21 22 th 289 22 24 ? 435 24 25 q 598 25 26 <0x3D> 66 26 27 1 371 27 28 <0x26> 43 28 29 x 421 29 30 <0x3D> 66 30 31 2 434 31 32 ▁https://example.com/path?q=1&x=2 +UPPER lower MiXeD case 17 ▁ 261 0 0 U 593 0 1 P 425 1 2 P 425 2 3 E 595 3 4 R 544 4 5 ▁lo 305 5 8 w 419 8 9 er 270 9 11 ▁M 442 11 13 i 282 13 14 <0x58> 93 14 15 e 264 15 16 D 589 16 17 ▁c 338 17 19 as 420 19 21 e 264 21 22 ▁UPPER▁lower▁MiXeD▁case +aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa 40 ▁a 266 0 1 a 279 1 2 a 279 2 3 a 279 3 4 a 279 4 5 a 279 5 6 a 279 6 7 a 279 7 8 a 279 8 9 a 279 9 10 a 279 10 11 a 279 11 12 a 279 12 13 a 279 13 14 a 279 14 15 a 279 15 16 a 279 16 17 a 279 17 18 a 279 18 19 a 279 19 20 a 279 20 21 a 279 21 22 a 279 22 23 a 279 23 24 a 279 24 25 a 279 25 26 a 279 26 27 a 279 27 28 a 279 28 29 a 279 29 30 a 279 30 31 a 279 31 32 a 279 32 33 a 279 33 34 a 279 34 35 a 279 35 36 a 279 36 37 a 279 37 38 a 279 38 39 a 279 39 40 ▁aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +Ω≈ç√∫˜µ≤ 22 ▁ 261 0 0 <0xCE> 211 0 0 <0xA9> 174 0 1 <0xE2> 231 1 1 <0x89> 142 1 1 <0x88> 141 1 2 <0xC3> 200 2 2 <0xA7> 172 2 3 <0xE2> 231 3 3 <0x88> 141 3 3 <0x9A> 159 3 4 <0xE2> 231 4 4 <0x88> 141 4 4 <0xAB> 176 4 5 ▁ 261 5 5 <0xCC> 209 5 5 <0x83> 136 5 6 <0xCE> 211 6 6 <0xBC> 193 6 7 <0xE2> 231 7 7 <0x89> 142 7 7 <0xA4> 169 7 8 ▁Ω≈ç√∫▁̃μ≤ +مرحبا بالعالم 26 ▁ 261 0 0 <0xD9> 222 0 0 <0x85> 138 0 1 <0xD8> 221 1 1 <0xB1> 182 1 2 <0xD8> 221 2 2 <0xAD> 178 2 3 <0xD8> 221 3 3 <0xA8> 173 3 4 <0xD8> 221 4 4 <0xA7> 172 4 5 ▁ 261 5 6 <0xD8> 221 6 6 <0xA8> 173 6 7 <0xD8> 221 7 7 <0xA7> 172 7 8 <0xD9> 222 8 8 <0x84> 137 8 9 <0xD8> 221 9 9 <0xB9> 190 9 10 <0xD8> 221 10 10 <0xA7> 172 10 11 <0xD9> 222 11 11 <0x84> 137 11 12 <0xD9> 222 12 12 <0x85> 138 12 13 ▁مرحبا▁بالعالم + leading and trailing 10 ▁lea 499 2 5 ding 501 5 9 ▁a 266 9 11 nd 284 11 13 ▁ 261 13 14 t 269 14 15 ra 409 15 17 i 282 17 18 l 319 18 19 ing 272 19 22 ▁leading▁and▁trailing +\ttab\tstart 5 ▁ 261 1 1 t 269 1 2 a 279 2 3 b 344 3 4 ▁start 453 4 10 ▁tab▁start +newline\n\n\nruns 11 ▁ 261 0 0 n 268 0 1 e 264 1 2 w 419 2 3 l 319 3 4 ine 415 4 7 ▁ 261 7 10 r 288 10 11 u 271 11 12 n 268 12 13 s 263 13 14 ▁newline▁runs +mid spaces collapse 12 ▁ 261 0 0 m 320 0 1 i 282 1 2 d 273 2 3 ▁ 261 3 6 space 389 6 11 s 263 11 12 ▁co 326 12 17 ll 412 17 19 a 279 19 20 p 352 20 21 se 325 21 23 ▁mid▁spaces▁collapse diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.model b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.model new file mode 100644 index 0000000000..6548571f49 Binary files /dev/null and b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-bytefb.model differ diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.fixtures.tsv b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.fixtures.tsv new file mode 100644 index 0000000000..3677b86a84 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.fixtures.tsv @@ -0,0 +1,40 @@ + 0 + 0 + 0 +a 1 ▁a 10 0 1 ▁a +Hello world 7 ▁ 5 0 0 H 241 0 1 e 8 1 2 ll 105 2 4 o 17 4 5 ▁wor 160 5 9 ld 74 9 11 ▁Hello▁world + Hello world 7 ▁ 5 1 1 H 241 1 2 e 8 2 3 ll 105 3 5 o 17 5 6 ▁wor 160 6 12 ld 74 12 14 ▁Hello▁world +Hello world.\nSecond line\ttabbed 23 ▁ 5 0 0 H 241 0 1 e 8 1 2 ll 105 2 4 o 17 4 5 ▁wor 160 5 9 ld 74 9 11 . 7 11 12 \n 0 12 13 S 297 13 14 e 8 14 15 c 38 15 16 o 17 16 17 nd 24 17 19 ▁ 5 19 20 l 30 20 21 ine 147 21 24 \t 0 24 25 t 11 25 26 a 13 26 27 b 45 27 28 b 45 28 29 ed 18 29 31 ▁Hello▁world.\nSecond▁line\ttabbed +The quick brown fox jumps over the lazy dog. 28 ▁Th 25 0 2 e 8 2 3 ▁qu 69 3 6 i 15 6 7 ck 43 7 9 ▁b 47 9 11 r 23 11 12 ow 60 12 14 n 9 14 15 ▁fo 97 15 18 x 108 18 19 ▁ 5 19 20 j 115 20 21 u 14 21 22 m 26 22 23 p 27 23 24 s 6 24 25 ▁ 5 25 26 o 17 26 27 ve 102 27 29 r 23 29 30 ▁the 12 30 34 ▁la 88 34 37 z 54 37 38 y 19 38 39 ▁do 159 39 42 g 48 42 43 . 7 43 44 ▁The▁quick▁brown▁fox▁jumps▁over▁the▁lazy▁dog. +tokenization and segmentation 4 ▁tokenization 200 0 12 ▁a 10 12 14 nd 24 14 16 ▁segmentation 80 16 29 ▁tokenization▁and▁segmentation +Antidisestablishmentarianism 14 ▁An 123 0 2 t 11 2 3 i 15 3 4 d 33 4 5 is 22 5 7 est 63 7 10 a 13 10 11 b 45 11 12 lish 193 12 16 ment 209 16 20 aria 221 20 24 n 9 24 25 is 22 25 27 m 26 27 28 ▁Antidisestablishmentarianism +water running walked faster apple book work play 15 ▁water 70 0 5 ▁runn 211 5 10 ing 20 10 13 ▁walk 87 13 18 ed 18 18 20 ▁fast 85 20 25 er 16 25 27 ▁appl 214 27 32 e 8 32 33 ▁b 47 33 35 o 17 35 36 o 17 36 37 k 66 37 38 ▁work 57 38 43 ▁play 71 43 48 ▁water▁running▁walked▁faster▁apple▁book▁work▁play +3.14159 x 42 = 1024? 21 ▁ 5 0 0 3 225 0 1 . 7 1 2 1 109 2 3 4 110 3 4 1 109 4 5 5 238 5 6 9 239 6 7 ▁ 5 7 8 x 108 8 9 ▁ 5 9 10 4 110 10 11 2 169 11 12 ▁ 5 12 13 = 0 13 14 ▁ 5 14 15 1 109 15 16 0 282 16 17 2 169 17 18 4 110 18 19 ? 170 19 20 ▁3.14159▁x▁42▁=▁1024? +!!!???... 10 ▁ 5 0 0 ! 113 0 1 ! 113 1 2 ! 113 2 3 ? 170 3 4 ? 170 4 5 ? 170 5 6 . 7 6 7 . 7 7 8 . 7 8 9 ▁!!!???... +(parentheses) and [brackets] and {braces} 27 ▁ 5 0 0 ( 236 0 1 p 27 1 2 are 34 2 5 n 9 5 6 th 36 6 8 e 8 8 9 ses 150 9 12 ) 237 12 13 ▁a 10 13 15 nd 24 15 17 ▁ 5 17 18 [ 0 18 19 b 45 19 20 r 23 20 21 ack 89 21 24 e 8 24 25 ts 101 25 27 ] 0 27 28 ▁a 10 28 30 nd 24 30 32 ▁ 5 32 33 { 0 33 34 b 45 34 35 ra 152 35 37 ces 216 37 40 } 0 40 41 ▁(parentheses)▁and▁[brackets]▁and▁{braces} +café naïve fiancé résumé 21 ▁ 5 0 0 ca 104 0 2 f 41 2 3 é 247 3 4 ▁ 5 4 5 n 9 5 6 a 13 6 7 ï 0 7 8 ve 102 8 10 ▁fi 210 10 13 a 13 13 14 n 9 14 15 c 38 15 16 é 247 16 17 ▁ 5 17 18 r 23 18 19 é 247 19 20 s 6 20 21 u 14 21 22 m 26 22 23 é 247 23 24 ▁café▁naïve▁fiancé▁résumé +financial fluid 13 ▁ 5 0 0 fi 0 0 1 n 9 1 2 a 13 2 3 n 9 3 4 c 38 4 5 i 15 5 6 al 21 6 8 ▁ 5 8 9 fl 0 9 10 u 14 10 11 i 15 11 12 d 33 12 13 ▁financial▁fluid +① ⑪ ㋿ ＫＡＴＡＫＡＮＡ 8 ▁ 5 0 0 ① 0 0 1 ▁ 5 1 2 ⑪ 0 2 3 ▁ 5 3 4 ㋿ 0 4 5 ▁ 5 5 6 ＫＡＴＡＫＡＮＡ 0 6 14 ▁①▁⑪▁㋿▁ＫＡＴＡＫＡＮＡ +ｶﾀｶﾅ half width 9 ▁ 5 0 0 ｶﾀｶﾅ 0 0 4 ▁ 5 4 5 h 32 5 6 al 21 6 8 f 41 8 9 ▁wi 73 9 12 d 33 12 13 th 36 13 15 ▁ｶﾀｶﾅ▁half▁width +東京タワーへ行きました 10 ▁ 5 0 0 東 231 0 1 京 230 1 2 タ 227 2 3 ワ 228 3 4 ー 229 4 5 へ行き 0 5 8 ま 287 8 9 し 179 9 10 た 256 10 11 ▁東京タワーへ行きました +日本語とEnglish混在 9 ▁ 5 0 0 日 264 0 1 本 266 1 2 語 269 2 3 と 0 3 4 E 295 4 5 ng 86 5 7 lish 193 7 11 混在 0 11 13 ▁日本語とEnglish混在 +Привет мир 11 ▁ 5 0 0 П 249 0 1 р 112 1 2 и 111 2 3 в 172 3 4 е 173 4 5 т 252 5 6 ▁ 5 6 7 м 174 7 8 и 111 8 9 р 112 9 10 ▁Привет▁мир +안녕하세요 세계 9 ▁ 5 0 0 안 234 0 1 녕 233 1 2 하 235 2 3 세 181 3 4 요 274 4 5 ▁ 5 5 6 세 181 6 7 계 271 7 8 ▁안녕하세요▁세계 +你好，世界！ 7 ▁ 5 0 0 你 261 0 1 好 262 1 2 ， 0 2 3 世 259 3 4 界 267 4 5 ！ 0 5 6 ▁你好，世界！ +I love 🍕 pizza 12 ▁ 5 0 0 I 294 0 1 ▁lo 53 1 4 ve 102 4 6 ▁ 5 6 7 🍕 279 7 9 ▁ 5 9 10 p 27 10 11 i 15 11 12 z 54 12 13 z 54 13 14 a 13 14 15 ▁I▁love▁🍕▁pizza +flags 🇩🇪 🇺🇸 end 14 ▁ 5 0 0 f 41 0 1 l 30 1 2 a 13 2 3 g 48 3 4 s 6 4 5 ▁ 5 5 6 🇩 277 6 8 🇪 278 8 10 ▁ 5 10 11 🇺🇸 0 11 15 ▁ 5 15 16 e 8 16 17 nd 24 17 19 ▁flags▁🇩🇪▁🇺🇸▁end +family 👩‍👩‍👧‍👦 emoji 10 ▁famil 185 0 5 y 19 5 6 ▁ 5 6 7 👩‍👩‍👧‍👦 0 7 18 ▁ 5 18 19 e 8 19 20 m 26 20 21 o 17 21 22 j 115 22 23 i 15 23 24 ▁family▁👩‍👩‍👧‍👦▁emoji +zero​width and non breaking 22 ▁ 5 0 0 z 54 0 1 er 16 1 3 o 17 3 4 ​ 0 4 5 w 78 5 6 i 15 6 7 d 33 7 8 th 36 8 10 ▁a 10 10 12 nd 24 12 14 ▁ 5 14 15 n 9 15 16 o 17 16 17 n 9 17 18   0 18 19 b 45 19 20 r 23 20 21 e 8 21 22 a 13 22 23 k 66 23 24 ing 20 24 27 ▁zero​width▁and▁non breaking +quotes “fancy” and ‘single’ — dash 24 ▁quote 198 0 5 s 6 5 6 ▁ 5 6 7 “ 0 7 8 f 41 8 9 a 13 9 10 n 9 10 11 c 38 11 12 y 19 12 13 ” 0 13 14 ▁a 10 14 16 nd 24 16 18 ▁ 5 18 19 ‘ 0 19 20 s 6 20 21 ing 20 21 24 le 107 24 26 ’ 0 26 27 ▁ 5 27 28 — 0 28 29 ▁d 100 29 31 a 13 31 32 s 6 32 33 h 32 33 34 ▁quotes▁“fancy”▁and▁‘single’▁—▁dash + the [URL] token 9 ▁ 5 0 0 3 0 6 ▁the 12 6 10 ▁ 5 10 11 [URL] 4 11 16 ▁to 51 16 19 k 66 19 20 e 8 20 21 n 9 21 22 ▁▁the▁[URL]▁token +a b[URL]c 6 ▁a 10 0 1 ▁ 5 1 2 3 2 8 b 45 8 9 [URL] 4 9 14 c 38 14 15 ▁a▁b[URL]c +control tokens inline 20 ▁co 99 0 2 n 9 2 3 tro 128 3 6 l 30 6 7 ▁ 5 7 8 < 0 8 9 s 6 9 10 > 0 10 11 ▁to 51 11 14 k 66 14 15 e 8 15 16 n 9 16 17 s 6 17 18 ▁ 5 18 19 0 22 23 ▁in 35 23 26 l 30 26 27 ine 147 27 30 ▁control▁▁tokens▁▁inline +https://example.com/path?q=1&x=2 29 ▁ 5 0 0 h 32 0 1 t 11 1 2 t 11 2 3 p 27 3 4 s 6 4 5 :// 0 5 8 e 8 8 9 x 108 9 10 a 13 10 11 m 26 11 12 p 27 12 13 le 107 13 15 . 7 15 16 c 38 16 17 o 17 17 18 m 26 18 19 / 0 19 20 p 27 20 21 a 13 21 22 th 36 22 24 ? 170 24 25 q 298 25 26 = 0 26 27 1 109 27 28 & 0 28 29 x 108 29 30 = 0 30 31 2 169 31 32 ▁https://example.com/path?q=1&x=2 +UPPER lower MiXeD case 19 ▁ 5 0 0 U 293 0 1 P 156 1 2 P 156 2 3 E 295 3 4 R 242 4 5 ▁lo 53 5 8 w 78 8 9 er 16 9 11 ▁ 5 11 12 M 288 12 13 i 15 13 14 X 0 14 15 e 8 15 16 D 289 16 17 ▁ 5 17 18 ca 104 18 20 s 6 20 21 e 8 21 22 ▁UPPER▁lower▁MiXeD▁case +aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa 40 ▁a 10 0 1 a 13 1 2 a 13 2 3 a 13 3 4 a 13 4 5 a 13 5 6 a 13 6 7 a 13 7 8 a 13 8 9 a 13 9 10 a 13 10 11 a 13 11 12 a 13 12 13 a 13 13 14 a 13 14 15 a 13 15 16 a 13 16 17 a 13 17 18 a 13 18 19 a 13 19 20 a 13 20 21 a 13 21 22 a 13 22 23 a 13 23 24 a 13 24 25 a 13 25 26 a 13 26 27 a 13 27 28 a 13 28 29 a 13 29 30 a 13 30 31 a 13 31 32 a 13 32 33 a 13 33 34 a 13 34 35 a 13 35 36 a 13 36 37 a 13 37 38 a 13 38 39 a 13 39 40 ▁aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +Ω≈ç√∫˜µ≤ 2 ▁ 5 0 0 Ω≈ç√∫˜µ≤ 0 0 8 ▁Ω≈ç√∫˜µ≤ +مرحبا بالعالم 4 ▁ 5 0 0 مرحبا 0 0 5 ▁ 5 5 6 بالعالم 0 6 13 ▁مرحبا▁بالعالم + leading and trailing 13 ▁ 5 2 2 le 107 2 4 a 13 4 5 d 33 5 6 ing 20 6 9 ▁a 10 9 11 nd 24 11 13 ▁ 5 13 14 t 11 14 15 ra 152 15 17 i 15 17 18 l 30 18 19 ing 20 19 22 ▁leading▁and▁trailing +\ttab\tstart 9 ▁ 5 0 0 \t 0 0 1 t 11 1 2 a 13 2 3 b 45 3 4 \t 0 4 5 st 46 5 7 ar 50 7 9 t 11 9 10 ▁\ttab\tstart +newline\n\n\nruns 11 ▁ 5 0 0 n 9 0 1 e 8 1 2 w 78 2 3 l 30 3 4 ine 147 4 7 \n\n\n 0 7 10 r 23 10 11 u 14 11 12 n 9 12 13 s 6 13 14 ▁newline\n\n\nruns +mid spaces collapse 13 ▁ 5 0 0 m 26 0 1 i 15 1 2 d 33 2 3 ▁ 5 3 6 space 127 6 11 s 6 11 12 ▁co 99 12 17 ll 105 17 19 a 13 19 20 p 27 20 21 s 6 21 22 e 8 22 23 ▁mid▁spaces▁collapse diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.model b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.model new file mode 100644 index 0000000000..54c3482777 Binary files /dev/null and b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-identity.model differ diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.fixtures.tsv b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.fixtures.tsv new file mode 100644 index 0000000000..6dbd8aad46 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.fixtures.tsv @@ -0,0 +1,40 @@ + 0 + 0 + 0 +a 1 a▁ 21 0 1 a▁ +Hello world 9 H 242 0 1 e 20 1 2 l 15 2 3 lo 59 3 5 ▁ 5 5 6 wor 71 6 9 l 15 9 10 d 16 10 11 ▁ 5 11 11 Hello▁world▁ + Hello world 9 H 242 1 2 e 20 2 3 l 15 3 4 lo 59 4 6 ▁ 5 6 9 wor 71 9 12 l 15 12 13 d 16 13 14 ▁ 5 14 14 Hello▁world▁ +Hello world.\nSecond line\ttabbed 23 H 242 0 1 e 20 1 2 l 15 2 3 lo 59 3 5 ▁ 5 5 6 wor 71 6 9 l 15 9 10 d 16 10 11 .▁ 6 11 13 S 44 13 14 e 20 14 15 co 72 15 17 n 13 17 18 d 16 18 19 ▁ 5 19 20 l 15 20 21 in 36 21 23 e▁ 18 23 25 t 9 25 26 a 14 26 27 b 31 27 28 b 31 28 29 ed▁ 17 29 31 Hello▁world.▁Second▁line▁tabbed▁ +The quick brown fox jumps over the lazy dog. 29 The▁ 28 0 4 q 298 4 5 u 19 5 6 i 8 6 7 ck▁ 214 7 10 b 31 10 11 r 33 11 12 own▁ 65 12 16 f 35 16 17 o 10 17 18 x 149 18 19 ▁ 5 19 20 j 109 20 21 u 19 21 22 m 22 22 23 p 27 23 24 s▁ 12 24 26 o 10 26 27 v 51 27 28 er▁ 26 28 31 the▁ 11 31 35 la 98 35 37 z 78 37 38 y 37 38 39 ▁ 5 39 40 d 16 40 41 o 10 41 42 g 38 42 43 .▁ 6 43 44 The▁quick▁brown▁fox▁jumps▁over▁the▁lazy▁dog.▁ +tokenization and segmentation 8 t 9 0 1 o 10 1 2 k 41 2 3 en 39 3 5 ization▁ 150 5 13 and▁ 34 13 17 segmentation 82 17 29 ▁ 5 29 29 tokenization▁and▁segmentation▁ +Antidisestablishmentarianism 16 A 68 0 1 nti 137 1 4 d 16 4 5 i 8 5 6 s 7 6 7 est 47 7 10 a 14 10 11 b 31 11 12 lish 190 12 16 ment 229 16 20 aria 224 20 24 n 13 24 25 i 8 25 26 s 7 26 27 m 22 27 28 ▁ 5 28 28 Antidisestablishmentarianism▁ +water running walked faster apple book work play 19 water▁ 64 0 6 runn 179 6 10 ing▁ 30 10 14 walk 87 14 18 ed▁ 17 18 21 fast 88 21 25 er▁ 26 25 28 app 126 28 31 l 15 31 32 e▁ 18 32 34 b 31 34 35 o 10 35 36 o 10 36 37 k 41 37 38 ▁ 5 38 39 work 53 39 43 ▁ 5 43 44 play 66 44 48 ▁ 5 48 48 water▁running▁walked▁faster▁apple▁book▁work▁play▁ +3.14159 x 42 = 1024? 21 3 238 0 1 . 97 1 2 1 101 2 3 4 102 3 4 1 101 4 5 5 239 5 6 9 240 6 7 ▁ 5 7 8 x 149 8 9 ▁ 5 9 10 4 102 10 11 2 155 11 12 ▁ 5 12 13 = 0 13 14 ▁ 5 14 15 1 101 15 16 0 234 16 17 2 155 17 18 4 102 18 19 ? 296 19 20 ▁ 5 20 20 3.14159▁x▁42▁=▁1024?▁ +!!!???... 9 ! 297 0 1 ! 297 1 2 ! 297 2 3 ? 296 3 4 ? 296 4 5 ? 296 5 6 . 97 6 7 . 97 7 8 .▁ 6 8 9 !!!???...▁ +(parentheses) and [brackets] and {braces} 26 ( 236 0 1 par 143 1 4 en 39 4 6 the 142 6 9 se 58 9 11 s 7 11 12 ) 237 12 13 ▁ 5 13 14 and▁ 34 14 18 [ 0 18 19 b 31 19 20 ra 76 20 22 c 25 22 23 k 41 23 24 e 20 24 25 t 9 25 26 s 7 26 27 ] 0 27 28 ▁ 5 28 29 and▁ 34 29 33 { 0 33 34 b 31 34 35 ra 76 35 37 ces 218 37 40 } 0 40 41 ▁ 5 41 41 (parentheses)▁and▁[brackets]▁and▁{braces}▁ +café naïve fiancé résumé 20 ca 73 0 2 f 35 2 3 é 247 3 4 ▁ 5 4 5 na 96 5 7 ï 0 7 8 ve▁ 70 8 11 fi 74 11 13 a 14 13 14 n 13 14 15 c 25 15 16 é 247 16 17 ▁ 5 17 18 r 33 18 19 é 247 19 20 s 7 20 21 u 19 21 22 m 22 22 23 é 247 23 24 ▁ 5 24 24 café▁naïve▁fiancé▁résumé▁ +financial fluid 13 fi 74 0 1 na 96 1 3 n 13 3 4 c 25 4 5 i 8 5 6 al 23 6 8 ▁ 5 8 9 f 35 9 9 l 15 9 10 u 19 10 11 i 8 11 12 d 16 12 13 ▁ 5 13 13 financial▁fluid▁ +① ⑪ ㋿ ＫＡＴＡＫＡＮＡ 16 1 101 0 1 ▁ 5 1 2 1 101 2 2 1 101 2 3 ▁ 5 3 4 令和 0 4 5 ▁ 5 5 6 K 0 6 7 A 68 7 8 T 62 8 9 A 68 9 10 K 0 10 11 A 68 11 12 N 151 12 13 A 68 13 14 ▁ 5 14 14 1▁11▁令和▁KATAKANA▁ +ｶﾀｶﾅ half width 14 カ 0 0 1 タ 275 1 2 カナ 0 2 4 ▁ 5 4 5 h 24 5 6 al 23 6 8 f 35 8 9 ▁ 5 9 10 w 48 10 11 i 8 11 12 d 16 12 13 t 9 13 14 h 24 14 15 ▁ 5 15 15 カタカナ▁half▁width▁ +東京タワーへ行きました 10 東 284 0 1 京 280 1 2 タ 275 2 3 ワ 276 3 4 ー 277 4 5 へ行き 0 5 8 ま 294 8 9 し 177 9 10 た 254 10 11 ▁ 5 11 11 東京タワーへ行きました▁ +日本語とEnglish混在 10 日 256 0 1 本 257 1 2 語 258 2 3 と 0 3 4 E 144 4 5 n 13 5 6 g 38 6 7 lish 190 7 11 混在 0 11 13 ▁ 5 13 13 日本語とEnglish混在▁ +Привет мир 11 П 248 0 1 р 107 1 2 и 106 2 3 в 163 3 4 е 164 4 5 т 251 5 6 ▁ 5 6 7 м 165 7 8 и 106 8 9 р 107 9 10 ▁ 5 10 10 Привет▁мир▁ +안녕하세요 세계 9 안 233 0 1 녕 259 1 2 하 261 2 3 세 174 3 4 요 260 4 5 ▁ 5 5 6 세 174 6 7 계 271 7 8 ▁ 5 8 8 안녕하세요▁세계▁ +你好，世界！ 7 你 281 0 1 好 282 1 2 , 299 2 3 世 278 3 4 界 285 4 5 ! 297 5 6 ▁ 5 6 6 你好,世界!▁ +I love 🍕 pizza 11 I 145 0 1 ▁ 5 1 2 lo 59 2 4 ve▁ 70 4 7 🍕 265 7 9 ▁ 5 9 10 p 27 10 11 i 8 11 12 z 78 12 13 z 78 13 14 a▁ 21 14 15 I▁love▁🍕▁pizza▁ +flags 🇩🇪 🇺🇸 end 12 f 35 0 1 la 98 1 3 g 38 3 4 s▁ 12 4 6 🇩 263 6 8 🇪 264 8 10 ▁ 5 10 11 🇺🇸 0 11 15 ▁ 5 15 16 en 39 16 18 d 16 18 19 ▁ 5 19 19 flags▁🇩🇪▁🇺🇸▁end▁ +family 👩‍👩‍👧‍👦 emoji 11 famil 168 0 5 y 37 5 6 ▁ 5 6 7 👩‍👩‍👧‍👦 0 7 18 ▁ 5 18 19 e 20 19 20 m 22 20 21 o 10 21 22 j 109 22 23 i 8 23 24 ▁ 5 24 24 family▁👩‍👩‍👧‍👦▁emoji▁ +zero​width and non breaking 18 z 78 0 1 er 50 1 3 o 10 3 4 ▁ 5 4 5 w 48 5 6 i 8 6 7 d 16 7 8 t 9 8 9 h 24 9 10 ▁ 5 10 11 and▁ 34 11 15 n 13 15 16 on▁ 132 16 19 b 31 19 20 re 40 20 22 a 14 22 23 k 41 23 24 ing▁ 30 24 27 zero▁width▁and▁non▁breaking▁ +quotes “fancy” and ‘single’ — dash 25 quote 159 0 5 s▁ 12 5 7 “ 0 7 8 f 35 8 9 a 14 9 10 n 13 10 11 c 25 11 12 y 37 12 13 ” 0 13 14 ▁ 5 14 15 and▁ 34 15 19 ‘ 0 19 20 s 7 20 21 in 36 21 23 g 38 23 24 l 15 24 25 e 20 25 26 ’ 0 26 27 ▁ 5 27 28 — 0 28 29 ▁ 5 29 30 d 16 30 31 as 77 31 33 h 24 33 34 ▁ 5 34 34 quotes▁“fancy”▁and▁‘single’▁—▁dash▁ + the [URL] token 10 3 0 6 ▁ 5 6 7 the▁ 11 7 11 [URL] 4 11 16 ▁ 5 16 17 t 9 17 18 o 10 18 19 k 41 19 20 en 39 20 22 ▁ 5 22 22 ▁the▁[URL]▁token▁ +a b[URL]c 6 a▁ 21 0 2 3 2 8 b 31 8 9 [URL] 4 9 14 c 25 14 15 ▁ 5 15 15 a▁b[URL]c▁ +control tokens inline 21 co 72 0 2 n 13 2 3 tro 119 3 6 l▁ 61 6 8 < 0 8 9 s 7 9 10 > 0 10 11 ▁ 5 11 12 t 9 12 13 o 10 13 14 k 41 14 15 en 39 15 17 s▁ 12 17 19 0 22 23 ▁ 5 23 24 in 36 24 26 l 15 26 27 in 36 27 29 e▁ 18 29 30 control▁▁tokens▁▁inline▁ +https://example.com/path?q=1&x=2 24 h 24 0 1 t 9 1 2 t 9 2 3 p 27 3 4 s 7 4 5 :// 0 5 8 ex 120 8 10 a 14 10 11 mple 194 11 15 . 97 15 16 com 134 16 19 / 0 19 20 pa 56 20 22 t 9 22 23 h 24 23 24 ? 296 24 25 q 298 25 26 = 0 26 27 1 101 27 28 & 0 28 29 x 149 29 30 = 0 30 31 2 155 31 32 ▁ 5 32 32 https://example.com/path?q=1&x=2▁ +UPPER lower MiXeD case 18 U 230 0 1 P 80 1 2 P 80 2 3 E 144 3 4 R 235 4 5 ▁ 5 5 6 lo 59 6 8 w 48 8 9 er▁ 26 9 12 M 154 12 13 i 8 13 14 X 0 14 15 e 20 15 16 D 152 16 17 ▁ 5 17 18 ca 73 18 20 s 7 20 21 e▁ 18 21 22 UPPER▁lower▁MiXeD▁case▁ +aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa 40 a 14 0 1 a 14 1 2 a 14 2 3 a 14 3 4 a 14 4 5 a 14 5 6 a 14 6 7 a 14 7 8 a 14 8 9 a 14 9 10 a 14 10 11 a 14 11 12 a 14 12 13 a 14 13 14 a 14 14 15 a 14 15 16 a 14 16 17 a 14 17 18 a 14 18 19 a 14 19 20 a 14 20 21 a 14 21 22 a 14 22 23 a 14 23 24 a 14 24 25 a 14 25 26 a 14 26 27 a 14 27 28 a 14 28 29 a 14 29 30 a 14 30 31 a 14 31 32 a 14 32 33 a 14 33 34 a 14 34 35 a 14 35 36 a 14 36 37 a 14 37 38 a 14 38 39 a▁ 21 39 40 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa▁ +Ω≈ç√∫˜µ≤ 4 Ω≈ç√∫ 0 0 5 ▁ 5 5 5 ̃μ≤ 0 5 8 ▁ 5 8 8 Ω≈ç√∫▁̃μ≤▁ +مرحبا بالعالم 4 مرحبا 0 0 5 ▁ 5 5 6 بالعالم 0 6 13 ▁ 5 13 13 مرحبا▁بالعالم▁ + leading and trailing 11 l 15 2 3 e 20 3 4 a 14 4 5 d 16 5 6 ing▁ 30 6 10 and▁ 34 10 14 t 9 14 15 ra 76 15 17 i 8 17 18 l 15 18 19 ing▁ 30 19 22 leading▁and▁trailing▁ +\ttab\tstart 6 t 9 1 2 a 14 2 3 b 31 3 4 ▁ 5 4 5 start 187 5 10 ▁ 5 10 10 tab▁start▁ +newline\n\n\nruns 10 n 13 0 1 e 20 1 2 w 48 2 3 l 15 3 4 in 36 4 6 e▁ 18 6 10 r 33 10 11 u 19 11 12 n 13 12 13 s▁ 12 13 14 newline▁runs▁ +mid spaces collapse 12 m 22 0 1 i 8 1 2 d 16 2 3 ▁ 5 3 6 space 117 6 11 s▁ 12 11 15 co 72 15 17 l 15 17 18 la 98 18 20 p 27 20 21 s 7 21 22 e▁ 18 22 23 mid▁spaces▁collapse▁ diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.model b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.model new file mode 100644 index 0000000000..984acd070e Binary files /dev/null and b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram-suffix.model differ diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.fixtures.tsv b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.fixtures.tsv new file mode 100644 index 0000000000..5eeffe47e5 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.fixtures.tsv @@ -0,0 +1,40 @@ + 0 + 0 + 0 +a 1 ▁a 10 0 1 ▁a +Hello world 7 ▁ 5 0 0 H 241 0 1 e 8 1 2 ll 105 2 4 o 17 4 5 ▁wor 160 5 9 ld 74 9 11 ▁Hello▁world + Hello world 7 ▁ 5 1 1 H 241 1 2 e 8 2 3 ll 105 3 5 o 17 5 6 ▁wor 160 6 12 ld 74 12 14 ▁Hello▁world +Hello world.\nSecond line\ttabbed 23 ▁ 5 0 0 H 241 0 1 e 8 1 2 ll 105 2 4 o 17 4 5 ▁wor 160 5 9 ld 74 9 11 . 7 11 12 ▁ 5 12 13 S 297 13 14 e 8 14 15 c 38 15 16 o 17 16 17 nd 24 17 19 ▁ 5 19 20 l 30 20 21 ine 147 21 24 ▁ 5 24 25 t 11 25 26 a 13 26 27 b 45 27 28 b 45 28 29 ed 18 29 31 ▁Hello▁world.▁Second▁line▁tabbed +The quick brown fox jumps over the lazy dog. 28 ▁Th 25 0 2 e 8 2 3 ▁qu 69 3 6 i 15 6 7 ck 43 7 9 ▁b 47 9 11 r 23 11 12 ow 60 12 14 n 9 14 15 ▁fo 97 15 18 x 108 18 19 ▁ 5 19 20 j 115 20 21 u 14 21 22 m 26 22 23 p 27 23 24 s 6 24 25 ▁ 5 25 26 o 17 26 27 ve 102 27 29 r 23 29 30 ▁the 12 30 34 ▁la 88 34 37 z 54 37 38 y 19 38 39 ▁do 159 39 42 g 48 42 43 . 7 43 44 ▁The▁quick▁brown▁fox▁jumps▁over▁the▁lazy▁dog. +tokenization and segmentation 4 ▁tokenization 200 0 12 ▁a 10 12 14 nd 24 14 16 ▁segmentation 80 16 29 ▁tokenization▁and▁segmentation +Antidisestablishmentarianism 14 ▁An 123 0 2 t 11 2 3 i 15 3 4 d 33 4 5 is 22 5 7 est 63 7 10 a 13 10 11 b 45 11 12 lish 193 12 16 ment 209 16 20 aria 221 20 24 n 9 24 25 is 22 25 27 m 26 27 28 ▁Antidisestablishmentarianism +water running walked faster apple book work play 15 ▁water 70 0 5 ▁runn 211 5 10 ing 20 10 13 ▁walk 87 13 18 ed 18 18 20 ▁fast 85 20 25 er 16 25 27 ▁appl 214 27 32 e 8 32 33 ▁b 47 33 35 o 17 35 36 o 17 36 37 k 66 37 38 ▁work 57 38 43 ▁play 71 43 48 ▁water▁running▁walked▁faster▁apple▁book▁work▁play +3.14159 x 42 = 1024? 21 ▁ 5 0 0 3 225 0 1 . 7 1 2 1 109 2 3 4 110 3 4 1 109 4 5 5 238 5 6 9 239 6 7 ▁ 5 7 8 x 108 8 9 ▁ 5 9 10 4 110 10 11 2 169 11 12 ▁ 5 12 13 = 0 13 14 ▁ 5 14 15 1 109 15 16 0 282 16 17 2 169 17 18 4 110 18 19 ? 170 19 20 ▁3.14159▁x▁42▁=▁1024? +!!!???... 10 ▁ 5 0 0 ! 113 0 1 ! 113 1 2 ! 113 2 3 ? 170 3 4 ? 170 4 5 ? 170 5 6 . 7 6 7 . 7 7 8 . 7 8 9 ▁!!!???... +(parentheses) and [brackets] and {braces} 27 ▁ 5 0 0 ( 236 0 1 p 27 1 2 are 34 2 5 n 9 5 6 th 36 6 8 e 8 8 9 ses 150 9 12 ) 237 12 13 ▁a 10 13 15 nd 24 15 17 ▁ 5 17 18 [ 0 18 19 b 45 19 20 r 23 20 21 ack 89 21 24 e 8 24 25 ts 101 25 27 ] 0 27 28 ▁a 10 28 30 nd 24 30 32 ▁ 5 32 33 { 0 33 34 b 45 34 35 ra 152 35 37 ces 216 37 40 } 0 40 41 ▁(parentheses)▁and▁[brackets]▁and▁{braces} +café naïve fiancé résumé 21 ▁ 5 0 0 ca 104 0 2 f 41 2 3 é 247 3 4 ▁ 5 4 5 n 9 5 6 a 13 6 7 ï 0 7 8 ve 102 8 10 ▁fi 210 10 13 a 13 13 14 n 9 14 15 c 38 15 16 é 247 16 17 ▁ 5 17 18 r 23 18 19 é 247 19 20 s 6 20 21 u 14 21 22 m 26 22 23 é 247 23 24 ▁café▁naïve▁fiancé▁résumé +financial fluid 13 ▁fi 210 0 1 n 9 1 2 a 13 2 3 n 9 3 4 c 38 4 5 i 15 5 6 al 21 6 8 ▁ 5 8 9 f 41 9 9 l 30 9 10 u 14 10 11 i 15 11 12 d 33 12 13 ▁financial▁fluid +① ⑪ ㋿ ＫＡＴＡＫＡＮＡ 16 ▁ 5 0 0 1 109 0 1 ▁ 5 1 2 1 109 2 2 1 109 2 3 ▁ 5 3 4 令和 0 4 5 ▁ 5 5 6 K 0 6 7 A 296 7 8 T 299 8 9 A 296 9 10 K 0 10 11 A 296 11 12 N 224 12 13 A 296 13 14 ▁1▁11▁令和▁KATAKANA +ｶﾀｶﾅ half width 11 ▁ 5 0 0 カ 0 0 1 タ 227 1 2 カナ 0 2 4 ▁ 5 4 5 h 32 5 6 al 21 6 8 f 41 8 9 ▁wi 73 9 12 d 33 12 13 th 36 13 15 ▁カタカナ▁half▁width +東京タワーへ行きました 10 ▁ 5 0 0 東 231 0 1 京 230 1 2 タ 227 2 3 ワ 228 3 4 ー 229 4 5 へ行き 0 5 8 ま 287 8 9 し 179 9 10 た 256 10 11 ▁東京タワーへ行きました +日本語とEnglish混在 9 ▁ 5 0 0 日 264 0 1 本 266 1 2 語 269 2 3 と 0 3 4 E 295 4 5 ng 86 5 7 lish 193 7 11 混在 0 11 13 ▁日本語とEnglish混在 +Привет мир 11 ▁ 5 0 0 П 249 0 1 р 112 1 2 и 111 2 3 в 172 3 4 е 173 4 5 т 252 5 6 ▁ 5 6 7 м 174 7 8 и 111 8 9 р 112 9 10 ▁Привет▁мир +안녕하세요 세계 9 ▁ 5 0 0 안 234 0 1 녕 233 1 2 하 235 2 3 세 181 3 4 요 274 4 5 ▁ 5 5 6 세 181 6 7 계 271 7 8 ▁안녕하세요▁세계 +你好，世界！ 7 ▁ 5 0 0 你 261 0 1 好 262 1 2 , 31 2 3 世 259 3 4 界 267 4 5 ! 113 5 6 ▁你好,世界! +I love 🍕 pizza 12 ▁ 5 0 0 I 294 0 1 ▁lo 53 1 4 ve 102 4 6 ▁ 5 6 7 🍕 279 7 9 ▁ 5 9 10 p 27 10 11 i 15 11 12 z 54 12 13 z 54 13 14 a 13 14 15 ▁I▁love▁🍕▁pizza +flags 🇩🇪 🇺🇸 end 14 ▁ 5 0 0 f 41 0 1 l 30 1 2 a 13 2 3 g 48 3 4 s 6 4 5 ▁ 5 5 6 🇩 277 6 8 🇪 278 8 10 ▁ 5 10 11 🇺🇸 0 11 15 ▁ 5 15 16 e 8 16 17 nd 24 17 19 ▁flags▁🇩🇪▁🇺🇸▁end +family 👩‍👩‍👧‍👦 emoji 10 ▁famil 185 0 5 y 19 5 6 ▁ 5 6 7 👩‍👩‍👧‍👦 0 7 18 ▁ 5 18 19 e 8 19 20 m 26 20 21 o 17 21 22 j 115 22 23 i 15 23 24 ▁family▁👩‍👩‍👧‍👦▁emoji +zero​width and non breaking 19 ▁ 5 0 0 z 54 0 1 er 16 1 3 o 17 3 4 ▁wi 73 4 7 d 33 7 8 th 36 8 10 ▁a 10 10 12 nd 24 12 14 ▁ 5 14 15 n 9 15 16 o 17 16 17 n 9 17 18 ▁b 47 18 20 r 23 20 21 e 8 21 22 a 13 22 23 k 66 23 24 ing 20 24 27 ▁zero▁width▁and▁non▁breaking +quotes “fancy” and ‘single’ — dash 24 ▁quote 198 0 5 s 6 5 6 ▁ 5 6 7 “ 0 7 8 f 41 8 9 a 13 9 10 n 9 10 11 c 38 11 12 y 19 12 13 ” 0 13 14 ▁a 10 14 16 nd 24 16 18 ▁ 5 18 19 ‘ 0 19 20 s 6 20 21 ing 20 21 24 le 107 24 26 ’ 0 26 27 ▁ 5 27 28 — 0 28 29 ▁d 100 29 31 a 13 31 32 s 6 32 33 h 32 33 34 ▁quotes▁“fancy”▁and▁‘single’▁—▁dash + the [URL] token 9 ▁ 5 0 0 3 0 6 ▁the 12 6 10 ▁ 5 10 11 [URL] 4 11 16 ▁to 51 16 19 k 66 19 20 e 8 20 21 n 9 21 22 ▁▁the▁[URL]▁token +a b[URL]c 6 ▁a 10 0 1 ▁ 5 1 2 3 2 8 b 45 8 9 [URL] 4 9 14 c 38 14 15 ▁a▁b[URL]c +control tokens inline 20 ▁co 99 0 2 n 9 2 3 tro 128 3 6 l 30 6 7 ▁ 5 7 8 < 0 8 9 s 6 9 10 > 0 10 11 ▁to 51 11 14 k 66 14 15 e 8 15 16 n 9 16 17 s 6 17 18 ▁ 5 18 19 0 22 23 ▁in 35 23 26 l 30 26 27 ine 147 27 30 ▁control▁▁tokens▁▁inline +https://example.com/path?q=1&x=2 29 ▁ 5 0 0 h 32 0 1 t 11 1 2 t 11 2 3 p 27 3 4 s 6 4 5 :// 0 5 8 e 8 8 9 x 108 9 10 a 13 10 11 m 26 11 12 p 27 12 13 le 107 13 15 . 7 15 16 c 38 16 17 o 17 17 18 m 26 18 19 / 0 19 20 p 27 20 21 a 13 21 22 th 36 22 24 ? 170 24 25 q 298 25 26 = 0 26 27 1 109 27 28 & 0 28 29 x 108 29 30 = 0 30 31 2 169 31 32 ▁https://example.com/path?q=1&x=2 +UPPER lower MiXeD case 19 ▁ 5 0 0 U 293 0 1 P 156 1 2 P 156 2 3 E 295 3 4 R 242 4 5 ▁lo 53 5 8 w 78 8 9 er 16 9 11 ▁ 5 11 12 M 288 12 13 i 15 13 14 X 0 14 15 e 8 15 16 D 289 16 17 ▁ 5 17 18 ca 104 18 20 s 6 20 21 e 8 21 22 ▁UPPER▁lower▁MiXeD▁case +aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa 40 ▁a 10 0 1 a 13 1 2 a 13 2 3 a 13 3 4 a 13 4 5 a 13 5 6 a 13 6 7 a 13 7 8 a 13 8 9 a 13 9 10 a 13 10 11 a 13 11 12 a 13 12 13 a 13 13 14 a 13 14 15 a 13 15 16 a 13 16 17 a 13 17 18 a 13 18 19 a 13 19 20 a 13 20 21 a 13 21 22 a 13 22 23 a 13 23 24 a 13 24 25 a 13 25 26 a 13 26 27 a 13 27 28 a 13 28 29 a 13 29 30 a 13 30 31 a 13 31 32 a 13 32 33 a 13 33 34 a 13 34 35 a 13 35 36 a 13 36 37 a 13 37 38 a 13 38 39 a 13 39 40 ▁aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +Ω≈ç√∫˜µ≤ 4 ▁ 5 0 0 Ω≈ç√∫ 0 0 5 ▁ 5 5 5 ̃μ≤ 0 5 8 ▁Ω≈ç√∫▁̃μ≤ +مرحبا بالعالم 4 ▁ 5 0 0 مرحبا 0 0 5 ▁ 5 5 6 بالعالم 0 6 13 ▁مرحبا▁بالعالم + leading and trailing 13 ▁ 5 2 2 le 107 2 4 a 13 4 5 d 33 5 6 ing 20 6 9 ▁a 10 9 11 nd 24 11 13 ▁ 5 13 14 t 11 14 15 ra 152 15 17 i 15 17 18 l 30 18 19 ing 20 19 22 ▁leading▁and▁trailing +\ttab\tstart 5 ▁ 5 1 1 t 11 1 2 a 13 2 3 b 45 3 4 ▁start 191 4 10 ▁tab▁start +newline\n\n\nruns 11 ▁ 5 0 0 n 9 0 1 e 8 1 2 w 78 2 3 l 30 3 4 ine 147 4 7 ▁ 5 7 10 r 23 10 11 u 14 11 12 n 9 12 13 s 6 13 14 ▁newline▁runs +mid spaces collapse 13 ▁ 5 0 0 m 26 0 1 i 15 1 2 d 33 2 3 ▁ 5 3 6 space 127 6 11 s 6 11 12 ▁co 99 12 17 ll 105 17 19 a 13 19 20 p 27 20 21 s 6 21 22 e 8 22 23 ▁mid▁spaces▁collapse diff --git a/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.model b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.model new file mode 100644 index 0000000000..b6e30611e4 Binary files /dev/null and b/opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/tiny-unigram.model differ diff --git a/opennlp-extensions/pom.xml b/opennlp-extensions/pom.xml index 9afcd3fe3c..4fed4c405e 100644 --- a/opennlp-extensions/pom.xml +++ b/opennlp-extensions/pom.xml @@ -38,8 +38,10 @@ + opennlp-embeddings opennlp-morfologik opennlp-spellcheck + opennlp-subword opennlp-uima diff --git a/pom.xml b/pom.xml index 53aa93d56d..9811276bd4 100644 --- a/pom.xml +++ b/pom.xml @@ -204,6 +204,12 @@ test-jar + + opennlp-embeddings + ${project.groupId} + ${project.version} + + opennlp-morfologik ${project.groupId} @@ -216,6 +222,12 @@ ${project.version} + + opennlp-subword + ${project.groupId} + ${project.version} + + opennlp-uima ${project.groupId} diff --git a/rat-excludes b/rat-excludes index 5a5d86b90c..90b6a69803 100644 --- a/rat-excludes +++ b/rat-excludes @@ -70,3 +70,14 @@ src/main/resources/opennlp/tools/tokenize/uax29/WordBreakProperty.txt src/main/resources/opennlp/tools/tokenize/uax29/ExtendedPictographic.txt src/main/resources/opennlp/tools/util/normalizer/confusables.txt src/test/resources/opennlp/tools/tokenize/uax29/WordBreakTest.txt + + +src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx + +src/test/resources/opennlp/subword/sentencepiece/*.model +src/test/resources/opennlp/subword/sentencepiece/*.fixtures.tsv +src/test/resources/opennlp/subword/sentencepiece/corpus.txt + +src/test/resources/opennlp/embeddings/tiny-unigram.model + +dev/embeddings/parity/sentences.txt