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/generate_test_teacher.py b/dev/embeddings/generate_test_teacher.py new file mode 100644 index 0000000000..8a43f518ee --- /dev/null +++ b/dev/embeddings/generate_test_teacher.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# 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. + +"""Print ONNX test constants for EmbeddingTestFixtures. + +Run with: uv run --with onnx==1.19.0 python dev/embeddings/generate_test_teacher.py +Python is required for regeneration, not for Maven tests. The graphs use original +numeric tables with no trained parameters or external model files. +""" + +import base64 +import textwrap + +from onnx import TensorProto, checker, helper + + +def print_model(name, nodes, tensors, dimension): + """Validate a graph and print a Java base64 constant.""" + graph = helper.make_graph( + nodes, + name, + [helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "tokens"])], + [helper.make_tensor_value_info( + "last_hidden_state", TensorProto.FLOAT, ["batch", "tokens", dimension])], + tensors, + ) + model = helper.make_model(graph, ir_version=8, opset_imports=[helper.make_opsetid("", 13)]) + checker.check_model(model, full_check=True) + chunks = textwrap.wrap(base64.b64encode(model.SerializeToString()).decode("ascii"), 76) + print(f" private static final String {name} =") + for index, chunk in enumerate(chunks): + prefix = " " if index == 0 else " + " + suffix = ";" if index == len(chunks) - 1 else "" + print(f'{prefix}"{chunk}"{suffix}') + + +def main(): + """Generate the lookup graph and a graph with a variable output dimension.""" + # PAD, UNK, CLS, SEP, coffee, espresso, tea, history. + table = [ + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 1, + 0, 0, 0, -1, + 3, 0, 0, 0, + 2, 1, 0, 0, + -1, 2, 0, 0, + -1, -2, 0, 0, + ] + print_model("LOOKUP_TEACHER_ONNX", [ + helper.make_node("Gather", ["table", "input_ids"], ["last_hidden_state"], axis=0), + ], [helper.make_tensor("table", TensorProto.FLOAT, [8, 4], table)], 4) + + print_model("VARIABLE_DIMENSION_ONNX", [ + helper.make_node("Cast", ["input_ids"], ["as_float"], to=TensorProto.FLOAT), + helper.make_node("Unsqueeze", ["as_float", "axes"], ["states"]), + helper.make_node("Shape", ["input_ids"], ["input_shape"]), + helper.make_node("Gather", ["input_shape", "batch_axis"], ["batch_size"], axis=0), + helper.make_node("Concat", ["ones", "batch_size"], ["repeats"], axis=0), + helper.make_node("Tile", ["states", "repeats"], ["last_hidden_state"]), + ], [ + helper.make_tensor("axes", TensorProto.INT64, [1], [2]), + helper.make_tensor("batch_axis", TensorProto.INT64, [1], [0]), + helper.make_tensor("ones", TensorProto.INT64, [2], [1, 1]), + ], "hidden") + + +if __name__ == "__main__": + main() 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..1c3253d315 --- /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 text into a fixed-length vector. + * + *

Unlike {@link opennlp.tools.util.wordvector.WordVectorTable}, which looks up a stored vector + * for one word, this interface accepts a sentence, paragraph, or document.

+ * + *

Thread safety is implementation specific.

+ * + *

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++) { + final CharSequence text = texts.get(i); + if (text == null) { + throw new IllegalArgumentException("texts[" + i + "] must not be null"); + } + vectors[i] = embed(text); + } + return vectors; + } + + /** {@return the dimension of every vector this embedder produces} */ + int dimension(); +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java index dedfb236c7..9e620c3f5c 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java @@ -51,7 +51,8 @@ * vocabulary, because each emitted piece must have an id. Vocabulary entries starting with * {@code ##} are continuation pieces and can match only after the first piece of a word.

* - *

Lower casing applies the Unicode full case mapping, including the {@code Final_Sigma} + *

Lower casing applies the + * Unicode full case mapping, including the {@code Final_Sigma} * context, so a word-final Greek capital sigma becomes U+03C2 as in the reference * implementation.

* diff --git a/opennlp-api/src/test/java/opennlp/tools/embeddings/TextEmbedderTest.java b/opennlp-api/src/test/java/opennlp/tools/embeddings/TextEmbedderTest.java new file mode 100644 index 0000000000..875f90f023 --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/embeddings/TextEmbedderTest.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.tools.embeddings; + +import java.util.Arrays; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +class TextEmbedderTest { + + private final TextEmbedder permissiveEmbedder = new TextEmbedder() { + @Override + public float[] embed(CharSequence text) { + return new float[] {1f}; + } + + @Override + public int dimension() { + return 1; + } + }; + + @Test + void testEmbedAllRejectsNullList() { + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> permissiveEmbedder.embedAll(null)); + + assertEquals("texts must not be null", exception.getMessage()); + } + + @Test + void testEmbedAllRejectsNullElement() { + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> permissiveEmbedder.embedAll(Arrays.asList("first", null))); + + assertEquals("texts[1] must not be null", exception.getMessage()); + } +} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/Tokens.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/Tokens.java index f5c105fe20..7cffd636b9 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/Tokens.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/Tokens.java @@ -23,7 +23,7 @@ * @param tokens The tokens themselves. * @param ids The token IDs as retrieved from the vocabulary. * @param mask The token mask. (Typically all 1.) - * @param types The token types. (Typically all 1.) + * @param types The segment IDs, all 0 for a single input segment. */ public record Tokens(String[] tokens, long[] ids, long[] mask, long[] types) { 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 0cf0b3ced1..48e1f70d1a 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,16 +20,22 @@ import java.io.File; import java.io.IOException; import java.nio.LongBuffer; +import java.util.ArrayList; 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; /** * Facilitates the generation of sentence vectors using @@ -53,9 +59,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 @@ -91,6 +106,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); } @@ -100,10 +116,15 @@ public SentenceVectorsDL(final File model, final File vocabulary, final boolean * @param sentence The input sentence. * @return The sentence vector. * + * @throws IllegalArgumentException Thrown if {@code sentence} is {@code null}. * @throws OrtException Thrown if an error occurs during inference. */ public float[] getVectors(final String sentence) throws OrtException { + if (sentence == null) { + throw new IllegalArgumentException("sentence must not be null"); + } + final Tokens tokens = encodeTokens(sentence); final Map inputs = new HashMap<>(); @@ -129,4 +150,155 @@ 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[" + i + "] must not be null"); + } + encoded[i] = encodeTokens(text); + 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; + } + } diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/AbstractDLChunkingTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/AbstractDLChunkingTest.java index 13957ea9be..922ee8ca74 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/AbstractDLChunkingTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/AbstractDLChunkingTest.java @@ -27,7 +27,7 @@ /** * Model-free tests for {@link AbstractDL#whitespaceChunks(String, int, int)} and - * {@link AbstractDL#whitespaceChunkSpans(String, int, int)}, the shared tokenize-and-chunk seam + * {@link AbstractDL#whitespaceChunkSpans(String, int, int)}, the shared tokenize-and-chunk path * used by both {@code NameFinderDL} and {@code DocumentCategorizerDL}. */ public class AbstractDLChunkingTest { 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 a4c9f05a29..c0677df348 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 @@ -35,17 +35,6 @@ public class CreateTokenizerTest { - private static final class TestDL extends AbstractDL { - - private TestDL(Map vocab) { - super(null, null, vocab, true); - } - - private Tokens encode(String text) { - return encodeTokens(text); - } - } - private static Map bertVocab() { final Map vocab = new HashMap<>(); vocab.put(WordpieceTokenizer.BERT_CLS_TOKEN, 0); @@ -151,7 +140,7 @@ void testDlEncodingPreservesVocabularyIds() { WordpieceTokenizer.BERT_UNK_TOKEN, 999, "hello", 42); - final Tokens tokens = new TestDL(vocab).encode("Hello"); + final Tokens tokens = new ModelFreeDL(vocab, true).encode("Hello"); assertArrayEquals(new String[] {"[CLS]", "hello", "[SEP]"}, tokens.tokens()); assertArrayEquals(new long[] {101, 42, 205}, tokens.ids()); diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/ModelFreeDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/ModelFreeDL.java new file mode 100644 index 0000000000..e21b3ad938 --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/ModelFreeDL.java @@ -0,0 +1,54 @@ +/* + * 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; + +import java.util.Map; + +/** + * An {@link AbstractDL} without an ONNX environment or session, so the shared encoder can be + * exercised on a plain vocabulary. Only {@link #encode(CharSequence)} may be called; anything + * reaching the absent session fails. + */ +final class ModelFreeDL extends AbstractDL { + + /** + * Creates an encoder over the given vocabulary. + * + * @param vocab The token-to-id map; must not be {@code null} and must contain the special + * tokens the vocabulary's model family requires. + * @param lowerCase {@code true} to lower case and strip accents, as for an uncased model. + * + * @throws IllegalArgumentException Thrown if {@code vocab} is {@code null} or lacks a + * required special token. + */ + ModelFreeDL(Map vocab, boolean lowerCase) { + super(null, null, vocab, lowerCase); + } + + /** + * Encodes text through {@link AbstractDL#encodeTokens(CharSequence)}. + * + * @param text The text to encode; must not be {@code null}. + * @return The model input arrays. + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + Tokens encode(CharSequence text) { + return encodeTokens(text); + } +} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/SharedDlEncodingTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/SharedDlEncodingTest.java new file mode 100644 index 0000000000..a4949010ad --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/SharedDlEncodingTest.java @@ -0,0 +1,353 @@ +/* + * 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; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.stream.Stream; + +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; + +import opennlp.tools.tokenize.WordpieceTokenizer; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Checks the complete token arrays created by the shared deep-learning encoder. */ +public class SharedDlEncodingTest { + + private static final int CONCURRENT_CALLS = 16; + private static final int CONCURRENT_WORKERS = 4; + private static final long WAIT_SECONDS = 10; + + private static final Map BERT_TOKEN_IDS = Map.ofEntries( + Map.entry(WordpieceTokenizer.BERT_CLS_TOKEN, 101), + Map.entry(WordpieceTokenizer.BERT_SEP_TOKEN, Integer.MAX_VALUE), + Map.entry(WordpieceTokenizer.BERT_UNK_TOKEN, 900_001), + Map.entry("hello", 42), + Map.entry("world", 1_500_000_000), + Map.entry("play", 0), + Map.entry("##ing", 2_147_483_646), + Map.entry("cafe", 88), + Map.entry("Caf\u00E9", 89), + Map.entry("\u03C3\u03BF\u03C6\u03BF\u03C2", 90), + Map.entry("\u03A3\u039F\u03A6\u039F\u03A3", 91)); + + private static final Map ROBERTA_TOKEN_IDS = Map.of( + WordpieceTokenizer.ROBERTA_CLS_TOKEN, Integer.MAX_VALUE, + WordpieceTokenizer.ROBERTA_SEP_TOKEN, 2, + WordpieceTokenizer.ROBERTA_UNK_TOKEN, 800_000_000, + "hello", 0, + "world", 500, + "play", 10, + "##ing", 1_500_000_000); + + private static final Map ROBERTA_BERT_UNKNOWN_TOKEN_IDS = Map.of( + WordpieceTokenizer.ROBERTA_CLS_TOKEN, 71, + WordpieceTokenizer.ROBERTA_SEP_TOKEN, 72, + WordpieceTokenizer.BERT_UNK_TOKEN, Integer.MAX_VALUE, + "hello", 73); + + private static final EncodingExpectation BERT_HELLO_EXPECTATION = new EncodingExpectation( + BERT_TOKEN_IDS, true, "Hello", + new String[] {"[CLS]", "hello", "[SEP]"}, + new long[] {101, 42, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0}); + + private static final EncodingExpectation BERT_MIXED_EXPECTATION = new EncodingExpectation( + BERT_TOKEN_IDS, true, "Hello playing rabbit", + new String[] {"[CLS]", "hello", "play", "##ing", "[UNK]", "[SEP]"}, + new long[] {101, 42, 0, 2_147_483_646L, 900_001, 2_147_483_647L}, + new long[] {1, 1, 1, 1, 1, 1}, + new long[] {0, 0, 0, 0, 0, 0}); + + private static final EncodingExpectation BERT_UNKNOWN_EXPECTATION = new EncodingExpectation( + BERT_TOKEN_IDS, true, "rabbit", + new String[] {"[CLS]", "[UNK]", "[SEP]"}, + new long[] {101, 900_001, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0}); + + private static final EncodingExpectation BERT_WORDPIECES_EXPECTATION = + new EncodingExpectation( + BERT_TOKEN_IDS, true, "Playing", + new String[] {"[CLS]", "play", "##ing", "[SEP]"}, + new long[] {101, 0, 2_147_483_646L, 2_147_483_647L}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0}); + + private static final List CONCURRENT_EXPECTATIONS = List.of( + BERT_HELLO_EXPECTATION, + BERT_UNKNOWN_EXPECTATION, + BERT_WORDPIECES_EXPECTATION); + + /** Expected output for one model-free encoding case. */ + private record EncodingExpectation( + Map tokenIds, + boolean lowerCase, + CharSequence input, + String[] tokens, + long[] ids, + long[] mask, + long[] types) { + } + + /** + * Confirms the token strings, ids, attention mask, and token types. + * + * @param name The case name. + * @param expected The fixture and expected arrays. + */ + @ParameterizedTest(name = "{0}") + @MethodSource("encodingCases") + void testEncodingArrays(String name, EncodingExpectation expected) { + final Tokens actual = new ModelFreeDL(expected.tokenIds(), expected.lowerCase()) + .encode(expected.input()); + + assertEncoding(expected, actual); + } + + /** Confirms null text is rejected at the instance encoder boundary. */ + @Test + void testRejectsNullText() { + final ModelFreeDL encoder = new ModelFreeDL(BERT_TOKEN_IDS, true); + + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> encoder.encode(null)); + assertEquals("text must not be null", exception.getMessage()); + } + + /** Confirms that successive calls return independent arrays. */ + @Test + void testRepeatedCallsReturnIndependentArrays() { + final ModelFreeDL encoder = new ModelFreeDL(BERT_TOKEN_IDS, true); + final Tokens first = encoder.encode("Hello"); + final Tokens second = encoder.encode("Hello"); + + assertNotSame(first.tokens(), second.tokens()); + assertNotSame(first.ids(), second.ids()); + assertNotSame(first.mask(), second.mask()); + assertNotSame(first.types(), second.types()); + + first.tokens()[1] = "changed"; + first.ids()[1] = -1; + first.mask()[1] = 0; + first.types()[1] = 1; + + assertEncoding(BERT_HELLO_EXPECTATION, second); + assertEncoding(BERT_HELLO_EXPECTATION, encoder.encode("Hello")); + } + + /** + * Confirms bounded concurrent calls produce complete independent results. + * + * @throws InterruptedException Thrown if the test thread is interrupted. + * @throws ExecutionException Thrown if an encoding task fails. + * @throws TimeoutException Thrown if an encoding task exceeds the wait limit. + */ + @Test + void testConcurrentCallsReturnIndependentArrays() + throws InterruptedException, ExecutionException, TimeoutException { + final ModelFreeDL encoder = new ModelFreeDL(BERT_TOKEN_IDS, true); + final ExecutorService executor = Executors.newFixedThreadPool(CONCURRENT_WORKERS); + final CountDownLatch start = new CountDownLatch(1); + final List> futures = new ArrayList<>(); + try { + for (int i = 0; i < CONCURRENT_CALLS; i++) { + final EncodingExpectation expected = + CONCURRENT_EXPECTATIONS.get(i % CONCURRENT_EXPECTATIONS.size()); + futures.add(executor.submit(() -> { + start.await(); + return encoder.encode(expected.input()); + })); + } + start.countDown(); + + final List completed = new ArrayList<>(); + for (int i = 0; i < futures.size(); i++) { + final Future future = futures.get(i); + final Tokens actual = future.get(WAIT_SECONDS, TimeUnit.SECONDS); + assertEncoding(CONCURRENT_EXPECTATIONS.get(i % CONCURRENT_EXPECTATIONS.size()), actual); + for (final Tokens prior : completed) { + assertNotSame(prior.tokens(), actual.tokens()); + assertNotSame(prior.ids(), actual.ids()); + assertNotSame(prior.mask(), actual.mask()); + assertNotSame(prior.types(), actual.types()); + } + completed.add(actual); + } + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(WAIT_SECONDS, TimeUnit.SECONDS)); + } + } + + /** + * Supplies model-free encoding fixtures. + * + * @return The named fixtures. + */ + private static Stream encodingCases() { + return Stream.of( + Arguments.of("bert-empty", new EncodingExpectation(BERT_TOKEN_IDS, true, "", + new String[] {"[CLS]", "[SEP]"}, + new long[] {101, 2_147_483_647L}, + new long[] {1, 1}, + new long[] {0, 0})), + Arguments.of("bert-ascii-whitespace", new EncodingExpectation( + BERT_TOKEN_IDS, true, " \t\n", + new String[] {"[CLS]", "[SEP]"}, + new long[] {101, 2_147_483_647L}, + new long[] {1, 1}, + new long[] {0, 0})), + Arguments.of("bert-unicode-whitespace", new EncodingExpectation( + BERT_TOKEN_IDS, true, "\u00A0\u2028\u3000", + new String[] {"[CLS]", "[SEP]"}, + new long[] {101, 2_147_483_647L}, + new long[] {1, 1}, + new long[] {0, 0})), + Arguments.of("bert-known", BERT_HELLO_EXPECTATION), + Arguments.of("bert-known-sequence", new EncodingExpectation( + BERT_TOKEN_IDS, true, "Hello WORLD", + new String[] {"[CLS]", "hello", "world", "[SEP]"}, + new long[] {101, 42, 1_500_000_000L, 2_147_483_647L}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0})), + Arguments.of("bert-unicode-separator", new EncodingExpectation( + BERT_TOKEN_IDS, true, "Hello\u00A0WORLD", + new String[] {"[CLS]", "hello", "world", "[SEP]"}, + new long[] {101, 42, 1_500_000_000L, 2_147_483_647L}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0})), + Arguments.of("bert-unknown", BERT_UNKNOWN_EXPECTATION), + Arguments.of("bert-known-unknown", new EncodingExpectation( + BERT_TOKEN_IDS, true, "Hello rabbit", + new String[] {"[CLS]", "hello", "[UNK]", "[SEP]"}, + new long[] {101, 42, 900_001, 2_147_483_647L}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0})), + Arguments.of("bert-wordpieces", BERT_WORDPIECES_EXPECTATION), + Arguments.of("bert-mixed", BERT_MIXED_EXPECTATION), + Arguments.of("bert-string-builder", new EncodingExpectation( + BERT_TOKEN_IDS, true, new StringBuilder("Hello"), + new String[] {"[CLS]", "hello", "[SEP]"}, + new long[] {101, 42, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-cased-accent", new EncodingExpectation( + BERT_TOKEN_IDS, false, "Caf\u00E9", + new String[] {"[CLS]", "Caf\u00E9", "[SEP]"}, + new long[] {101, 89, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-cased-miss", new EncodingExpectation( + BERT_TOKEN_IDS, false, "Hello", + new String[] {"[CLS]", "[UNK]", "[SEP]"}, + new long[] {101, 900_001, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-precomposed-accent", new EncodingExpectation( + BERT_TOKEN_IDS, true, "CAF\u00C9", + new String[] {"[CLS]", "cafe", "[SEP]"}, + new long[] {101, 88, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-decomposed-accent", new EncodingExpectation( + BERT_TOKEN_IDS, true, "Cafe\u0301", + new String[] {"[CLS]", "cafe", "[SEP]"}, + new long[] {101, 88, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-final-sigma", new EncodingExpectation( + BERT_TOKEN_IDS, true, "\u03A3\u039F\u03A6\u039F\u03A3", + new String[] {"[CLS]", "\u03C3\u03BF\u03C6\u03BF\u03C2", "[SEP]"}, + new long[] {101, 90, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("bert-cased-sigma", new EncodingExpectation( + BERT_TOKEN_IDS, false, "\u03A3\u039F\u03A6\u039F\u03A3", + new String[] {"[CLS]", "\u03A3\u039F\u03A6\u039F\u03A3", "[SEP]"}, + new long[] {101, 91, 2_147_483_647L}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("roberta-empty", new EncodingExpectation( + ROBERTA_TOKEN_IDS, true, "", + new String[] {"", ""}, + new long[] {2_147_483_647L, 2}, + new long[] {1, 1}, + new long[] {0, 0})), + Arguments.of("roberta-known", new EncodingExpectation( + ROBERTA_TOKEN_IDS, true, "Hello", + new String[] {"", "hello", ""}, + new long[] {2_147_483_647L, 0, 2}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("roberta-unknown", new EncodingExpectation( + ROBERTA_TOKEN_IDS, true, "rabbit", + new String[] {"", "", ""}, + new long[] {2_147_483_647L, 800_000_000, 2}, + new long[] {1, 1, 1}, + new long[] {0, 0, 0})), + Arguments.of("roberta-mixed", new EncodingExpectation( + ROBERTA_TOKEN_IDS, true, "Hello rabbit WORLD", + new String[] {"", "hello", "", "world", ""}, + new long[] {2_147_483_647L, 0, 800_000_000, 500, 2}, + new long[] {1, 1, 1, 1, 1}, + new long[] {0, 0, 0, 0, 0})), + Arguments.of("roberta-wordpieces", new EncodingExpectation( + ROBERTA_TOKEN_IDS, true, "Playing", + new String[] {"", "play", "##ing", ""}, + new long[] {2_147_483_647L, 10, 1_500_000_000L, 2}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0})), + Arguments.of("roberta-bert-unknown", new EncodingExpectation( + ROBERTA_BERT_UNKNOWN_TOKEN_IDS, true, "Hello rabbit", + new String[] {"", "hello", "[UNK]", ""}, + new long[] {71, 73, 2_147_483_647L, 72}, + new long[] {1, 1, 1, 1}, + new long[] {0, 0, 0, 0}))); + } + + /** + * Compares all model input arrays. + * + * @param expected The expected arrays. + * @param actual The encoded arrays. + */ + private void assertEncoding(EncodingExpectation expected, Tokens actual) { + assertArrayEquals(expected.tokens(), actual.tokens()); + assertArrayEquals(expected.ids(), actual.ids()); + assertArrayEquals(expected.mask(), actual.mask()); + assertArrayEquals(expected.types(), actual.types()); + } +} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/doccat/DocumentCategorizerDLTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/doccat/DocumentCategorizerDLTest.java index 90c9abcead..1c7bcc4da5 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/doccat/DocumentCategorizerDLTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/doccat/DocumentCategorizerDLTest.java @@ -57,7 +57,7 @@ private static DocumentCategorizerDL categorizerWithoutSession() { } @Test - void testCategorizeFailsLoudlyWhenInferenceFails() { + void testCategorizePropagatesInferenceFailure() { final IllegalStateException e = assertThrows(IllegalStateException.class, () -> categorizerWithoutSession().categorize(new String[] {"hello world"})); @@ -66,7 +66,7 @@ void testCategorizeFailsLoudlyWhenInferenceFails() { } @Test - void testScoreMapsFailLoudlyWhenInferenceFails() { + void testScoreMapsPropagateInferenceFailure() { final DocumentCategorizerDL categorizer = categorizerWithoutSession(); assertThrows(IllegalStateException.class, () -> @@ -108,7 +108,7 @@ void testConstructorRejectsNullInferenceOptions() { @Test void testSoftmaxRejectsNaNLogit() { - // A NaN logit would otherwise poison the whole distribution into NaN scores; fail loudly instead. + // A NaN logit would otherwise turn the whole distribution into NaN scores. final IllegalStateException e = assertThrows(IllegalStateException.class, () -> DocumentCategorizerDL.softmax(new float[] {0f, Float.NaN, 0f})); assertTrue(e.getMessage().contains("NaN"), e.getMessage()); @@ -118,7 +118,7 @@ void testSoftmaxRejectsNaNLogit() { void testSoftmaxRejectsInfiniteLogit() { // A +Infinity logit (not NaN, so it slips past an isNaN-only guard) poisons the distribution too: // max becomes +Inf, so value - max is Inf - Inf == NaN, every exp() is NaN, and categorize() would - // silently return all-NaN scores. It must fail loud like the NaN case. -Infinity is non-finite too. + // return all-NaN scores. It is invalid for the same reason as NaN. -Infinity is non-finite too. final IllegalStateException pos = assertThrows(IllegalStateException.class, () -> DocumentCategorizerDL.softmax(new float[] {0f, Float.POSITIVE_INFINITY, 0f})); assertTrue(pos.getMessage().contains("non-finite") || pos.getMessage().contains("Infinity"), @@ -181,7 +181,7 @@ void testLogitsFromOutputDispatchesOnModelShape() { } @Test - void testLogitsFromOutputFailsLoudlyOnNullAndUnexpectedType() { + void testLogitsFromOutputRejectsNullAndUnexpectedType() { // A null or otherwise-shaped model output is a contract violation, not an "inference failed". final IllegalStateException onNull = assertThrows(IllegalStateException.class, () -> DocumentCategorizerDL.logitsFromOutput(null)); @@ -191,7 +191,7 @@ void testLogitsFromOutputFailsLoudlyOnNullAndUnexpectedType() { } @Test - void testRequireMatchingCategoryCountFailsLoudlyOnMismatch() { + void testRequireMatchingCategoryCountRejectsMismatch() { // A distribution whose length differs from the configured category count means the model and // the categorizer configuration do not match; the matching case passes the array through. final double[] ok = {0.5, 0.5}; 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..f1f72a3621 --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLEmbedderTest.java @@ -0,0 +1,135 @@ +/* + * 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 interface 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 the same for every input. + 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); + + assertEquals("text must not be null", assertThrows(IllegalArgumentException.class, + () -> embedder.embed(null)).getMessage()); + assertEquals("sentence must not be null", assertThrows(IllegalArgumentException.class, + () -> vectors.getVectors(null)).getMessage()); + assertEquals("texts must not be null", assertThrows(IllegalArgumentException.class, + () -> embedder.embedAll(null)).getMessage()); + } + } + + /** + * 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); + assertEquals("texts[1] must not be null", assertThrows(IllegalArgumentException.class, + () -> vectors.embedAll(Arrays.asList("hello", null))).getMessage()); + } + } +} 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-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..32e374712f --- /dev/null +++ b/opennlp-docs/src/docbkx/embeddings.xml @@ -0,0 +1,295 @@ + + + + + + + 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 WordPiece or SentencePiece tokenization. The module loads Model2Vec + layouts and can distill a sentence-transformer into the same flat-table form. + Embedding uses table lookups, mean pooling, and optional normalization, without a + model forward pass or 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 for semantic similarity, deduplication, candidate retrieval, clustering, or + classifier features when a static table suits the task. Use a contextual model when + the task depends on distinguishing word senses in context. + + + OpenNLP also supports contextual, ONNX-backed sentence vectors in the + opennlp-dl module. Contextual models preserve word-sense context at a + higher inference cost. 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), 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 contains vocab.txt, + model.safetensors, config.json, and + tokenizer_config.json; a Model2Vec Unigram model contains + tokenizer.json, model.safetensors, and + config.json. Its JSON file contains the Unigram vocabulary, scores, + normalizer, and pre-tokenizer. Separate-file 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);]]> + + + 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 contains the model's text + normalizer. + + + + + + Matrix rows are resolved by piece string, never by tokenizer id, because the two + files of a separate-file SentencePiece model may order and offset their ids differently. + A poolable piece with no matrix row is rejected during loading. A self-contained + Model2Vec Unigram layout loads directly from its tokenizer.json. + + + 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, and similarity reports 0 for them. + +
+ +
+ Semantic Search + + A common 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());]]> + + + 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. Any index that accepts float vectors can store these embeddings. + +
+ +
+ 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, row count, dimension reduction, and variance retained by PCA. The summary is + printed only after the directory loads successfully. + + + The teacher must return the same vector length for all batches, including + phrase batches with different sequence lengths. An inconsistent length causes + an IllegalArgumentException before output files are written. + + + Java applications can distill a local teacher and load the saved table directly: + + related = model.mostSimilar("coffee espresso", 3);]]> + + This small example requests 2 PCA components. Select the dimension using + retrieval quality on application data. The executable + ModelDistillerExampleTest uses an original ONNX lookup table to + check inference, PCA, weighting, saved files and search. That test is not a + language-quality evaluation. + + + -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. A model without a term file uses subword pieces throughout. + 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. A missing file is + reported with its expected name. The tool 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. The size + depends on the bit width and on padding each row to the next power of two. For example, + a 500,000-row table with 300 dimensions uses 600 MB as 32-bit floats. At 4 bits, its + rows are padded to 512 dimensions and the quantized file uses about 136 MB, including + its per-row scales and norms. The implementation uses a randomized Hadamard transform, + a standard-normal Lloyd-Max scalar grid, and a least-squares scale per row. This is an + MSE-oriented variant of + TurboQuant. It does not include + the residual QJL stage used by the paper's unbiased inner-product estimator. + + + 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 stored inside the quantized file. A directory that still contains both files is + rejected at load time, so the deployment must choose which matrix is authoritative. + Embedding, similarity, and mostSimilar use the same API as the + float model. Fewer bits reduce the file size and increase quantization error. + +
+ +
+ Loading and Inference Benchmarks + + The optional jmh Maven profile compiles + StaticEmbeddingModelBenchmark. Run the class with the module's test + classpath and -t 1 -prof gc to collect timing and allocation statistics. + Use -p modelDir=/path/to/model for a local model. The default + synthetic fixture is a generated 29,528 by 256 F32 table and requires + no downloads. The module README provides build and launch commands. + + + Embedding throughput and allocated bytes are reported per input text. Top-10 search + uses one query per operation. Model loading reports milliseconds and allocated bytes + per load, including file access, decoding and construction. Fixture generation is + outside timing. Repeated loads can use the operating system's file cache, so the + loading result does not describe disk startup. Allocated bytes include temporary + objects and do not describe retained or peak memory. + +
+ +
+ The safetensors Reader + + Weights are read with a small reader for the safetensors format. It parses a JSON + header and raw tensor bytes without deserializing Java or Python objects. 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 0761fc95ff..93a9045434 100644 --- a/opennlp-docs/src/docbkx/opennlp.xml +++ b/opennlp-docs/src/docbkx/opennlp.xml @@ -121,6 +121,7 @@ under the License. + diff --git a/opennlp-extensions/opennlp-embeddings/README.md b/opennlp-extensions/opennlp-embeddings/README.md new file mode 100644 index 0000000000..65efbf7ded --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/README.md @@ -0,0 +1,248 @@ + + +# OpenNLP Static Embeddings + +Turn text into embedding vectors from a static (non-contextual) table: a per-token vector matrix plus WordPiece or SentencePiece tokenization. The module loads [Model2Vec](https://github.com/MinishLab/model2vec) layouts and can distill a sentence-transformer into the same flat-table form. SentencePiece support permits multilingual tables distilled from encoders in the [XLM-RoBERTa](https://arxiv.org/abs/1911.02116) family. Embedding uses JVM table lookups and arithmetic, without a model forward pass or native runtime. + +OpenNLP also supports contextual ONNX models, which preserve word-sense context at a higher inference cost. Both embedding methods 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 self-contained Model2Vec Unigram model carries `tokenizer.json`, `model.safetensors`, and `config.json`. A separate-file SentencePiece model uses the Unigram layout and adds its trained `sentencepiece.bpe.model`, `spiece.model`, or `tokenizer.model`. The tokenizer and pooling switches are read from the model's own configuration. One loaded model is immutable and thread-safe, so it can serve every thread of an application. + +A multilingual SentencePiece table can compare text in the languages covered by its teacher: + +```java +double crossLingual = model.similarity( + "The weather is beautiful today", "今天天气很好"); +``` + +## When to use it + +Use static embeddings when throughput and deployment simplicity matter: semantic similarity, deduplication, candidate retrieval before a reranker, clustering, or classifier features. Use a contextual model 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* instead of the tokenizer's numeric id. The two files of a SentencePiece model may order or offset their ids differently, so string lookup keeps them aligned. Loading rejects a poolable piece with no matrix row. Unknown pieces are omitted, and text with no known pieces embeds to a zero vector. +3. **Weight and pool.** Per-token weights (when present) multiply into the running sum, which is divided by the number of pooled tokens. This is the pooling rule used by Model2Vec tables. +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: Unigram" --> SPM{"trained .model present?"} + SPM -- "yes" --> SEP["load separate-file SentencePiece
(its own normalizer)"] + SPM -- "no" --> SELF["load self-contained tokenizer.json
(normalizer and scores)"] + SEP --> TJ["tokenizer.json vocab
names the matrix rows"] + SELF --> TJ + 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 small [safetensors](https://github.com/huggingface/safetensors) reader. It parses a JSON header and raw tensor bytes, without deserializing Java or Python objects. Tensor data streams directly into the decoded array. A decoded tensor is limited to the maximum Java array length, about 2.1 billion float elements. + +## 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 +``` + +`SubwordTokenizer` provides one piece stream for the WordPiece encoder in `opennlp-api` and the pure-JVM SentencePiece implementation in `opennlp-subword`. `TextEmbedder` provides one embedding contract for this static implementation and the contextual ONNX implementation in `opennlp-dl`. + +## Performance + +A static table avoids a model forward pass. `StaticEmbeddingModelBenchmark` uses the Java +Microbenchmark Harness (JMH) to report loading time, embedding throughput and top-10 search +throughput. Add `-prof gc` for allocation statistics. + +Compile the benchmarks and run the fixture tests from the repository root: + +```sh +./mvnw -pl opennlp-extensions/opennlp-embeddings -am -Pjmh \ + -Dopennlp.forkCount=1 -Dtest=StaticEmbeddingModelBenchmarkTest \ + -Dsurefire.failIfNoSpecifiedTests=false clean test +``` + +With the `jmh` Maven profile enabled in an IDE, run +`opennlp.embeddings.StaticEmbeddingModelBenchmark.main` using the module's test classpath. +For a command-line launch, use that test classpath with `java -cp`: + +```sh +java -cp "$JMH_CLASSPATH" opennlp.embeddings.StaticEmbeddingModelBenchmark \ + -t 1 -prof gc -rf json -rff embeddings-jmh.json +``` + +Set `JMH_CLASSPATH` to the compiled test and main classes plus the test dependencies, including +the reactor modules. Use current reactor classes, not older snapshot JARs. The default +`modelDir=synthetic` generates a 29,528 by 256 F32 table without downloads. To use a local model, +add `-p modelDir=/path/to/model`. The program accepts JMH options and uses forked JVMs by default. + +`embed` reports operations and `gc.alloc.rate.norm` bytes per input text, although each +invocation processes a batch of 5 texts. `mostSimilarTop10` reports them per query. `load` +reports milliseconds and allocated bytes per model load, including file access, decoding and +construction. Fixture generation is outside timing. Repeated loads can use the operating +system's file cache; this is not a disk-startup or peak-memory measurement. The allocation +statistic includes temporary objects, not just the retained model. + +`embed()` tokenizes and pools only the rows used by the input. `mostSimilar()` scans every matrix row, so its cost grows with the vocabulary. Use a vector index when a full scan is too expensive. The harness in `dev/embeddings/parity/` compares single-thread speed and vector output with the Model2Vec Python implementation. Run it with the model and hardware used for deployment. + +## Quantizing a model + +`QuantizeModel` converts `model.safetensors` to a 2, 3, or 4-bit matrix: + +```text +opennlp-embeddings QuantizeModel -modelDir /path/to/model-directory -bits 4 +``` + +The command writes `model.quantized` and reports its size and sampled reconstruction cosine. +Delete `model.safetensors` before loading the quantized model. A directory containing both matrix +files is rejected. The tokenizer, configuration, and optional `terms.txt` stay unchanged. + +The format uses a randomized Hadamard transform, a Gaussian Lloyd-Max grid, and a scale for each +row. It is an MSE-oriented variant of +[TurboQuant](https://arxiv.org/abs/2504.19874) and does not implement the paper's residual QJL +estimator. + +## 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 + +For a small corpus, rank documents directly with `similarity`: + +```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"); + +String query = "home espresso machine"; + +IntStream.range(0, docs.size()) + .boxed() + .sorted(Comparator.comparingDouble( + (Integer i) -> model.similarity(query, docs.get(i))).reversed()) + .forEach(i -> System.out.println(docs.get(i))); +``` + +For a larger corpus, embed each document once, store the vectors in a vector index, and embed each query with the same model. Any index that accepts float vectors, including a Hierarchical Navigable Small World (HNSW) index, can store these embeddings. + +## 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 `sentencepiece.bpe.model` from the teacher repository into the same directory. A missing file is reported during loading. + +## 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. +- Input with no known pieces embeds to a zero vector. Decide whether that represents "no signal" for the application. + +## Testing distillation + +The executable [ModelDistillerExampleTest](src/test/java/opennlp/embeddings/ModelDistillerExampleTest.java) +tests local ONNX inference, PCA, weighting, saved-model loading and search with original numeric +fixtures. It runs without a downloaded model or Python. To regenerate the ONNX constants, run +`uv run --with onnx==1.19.0 python dev/embeddings/generate_test_teacher.py` from the repository root. +The fixture tests data flow, not language quality. + +## 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..0d1de66fe1 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/TRAINING.md @@ -0,0 +1,93 @@ + + +# Distilling a Model for OpenNLP Static Embeddings + +The `DistillModel` command produces a static embedding table from a sentence-transformer teacher. It follows the [Model2Vec](https://github.com/MinishLab/model2vec) pipeline: run the teacher's ONNX graph over its vocabulary, apply principal component analysis (PCA) and Zipf weighting, then write a flat per-token matrix. This is a single pass over the vocabulary, without a training corpus or optimization loop. + +## 1. Distill the teacher + +``` +opennlp-embeddings DistillModel -teacher BAAI/bge-m3 -out bge-m3-static -pcaDims 256 +``` + +`-teacher` is a Hugging Face model id or a local directory. A remote teacher is cached under `~/.cache/opennlp-embeddings/-` with its `tokenizer.json`, optional `tokenizer_config.json`, `onnx/model.onnx`, optional `onnx/model.onnx_data`, and SentencePiece model when required. `-pcaDims` defaults to 256. The command assembles the output directory and verifies it with `StaticEmbeddingModel.load` before printing its summary. + +The command replaces model and tokenizer files produced by an earlier distillation in the output directory. Unrelated files remain. An interrupted run may leave an incomplete output directory, so rerun the command before loading it. + +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` controls the output vector width and defaults to 256. A larger value increases the model's memory, disk, and inference cost. Evaluate retrieval or classification quality on the target task before changing it. + +## 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")); + +double crossLingual = model.similarity( + "The weather is beautiful today", "今天天气很好"); +double unrelated = model.similarity( + "The weather is beautiful today", "quarterly earnings missed"); + +List neighbors = model.mostSimilar("coffee", 5); +``` + +The reference Python flow lives in `dev/embeddings/distill_bge_m3.py`. The `dev/embeddings/parity/` harness embeds the same text with both implementations and reports vector differences and single-thread speed. Compare independently distilled tables by similarities and rankings because PCA bases can differ. + +## 4. Quantize the matrix + +Quantization reduces the matrix size after distillation or assembly: + +```text +opennlp-embeddings QuantizeModel -modelDir bge-m3-static -bits 4 +``` + +The command writes `model.quantized` and verifies the written file against sampled source rows. +Remove `model.safetensors` to select the quantized matrix. Keep the tokenizer, configuration, and +term files in the directory. + +## 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 + +Check the teacher model's license before publishing a distilled table. Record the exact teacher revision and retain any attribution or redistribution terms that apply to derived weights. diff --git a/opennlp-extensions/opennlp-embeddings/pom.xml b/opennlp-extensions/opennlp-embeddings/pom.xml new file mode 100644 index 0000000000..8a62d5e146 --- /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 + + + default-testCompile + 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..ebe3a1fd08 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmark.java @@ -0,0 +1,241 @@ +/* + * 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 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.OperationsPerInvocation; +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.CommandLineOptions; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; + +/** + * Benchmarks model loading, embedding and nearest-neighbor queries. + * + *

The default {@code modelDir=synthetic} uses a generated 29,528 by 256 table without + * downloads. Use {@code -p modelDir=/models/table} for a local model. Fixture creation is + * outside the timed operations. Loading includes file access, decoding and model construction; + * repeated loads can use the operating system's file cache. Add {@code -prof gc} for allocation + * statistics.

+ */ +@BenchmarkMode(Mode.Throughput) +@OutputTimeUnit(TimeUnit.SECONDS) +@Warmup(iterations = 5, time = 2) +@Measurement(iterations = 10, time = 2) +@Fork(2) +@Threads(1) +public class StaticEmbeddingModelBenchmark { + + /** The synthetic-fixture selector; any other value is treated as a model directory path. */ + private static final String SYNTHETIC = "synthetic"; + + 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", + }; + + static final List SENTENCES = List.of( + "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."); + + /** Model files shared by the benchmark threads. */ + @State(Scope.Benchmark) + public static class ModelFiles { + + /** + * 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; + + Path directory; + private Path tempDir; + + /** + * Prepares a local model directory without loading the model. + * + * @throws IOException If fixture creation fails. + */ + @Setup(Level.Trial) + public void prepare() throws IOException { + if (SYNTHETIC.equals(modelDir)) { + tempDir = Files.createTempDirectory("opennlp-embeddings-jmh"); + directory = tempDir; + writeVocab(); + writeSafetensors(); + Files.writeString(directory.resolve(ModelFileNames.CONFIG), "{\"normalize\":true}"); + Files.writeString(directory.resolve(ModelFileNames.TOKENIZER_CONFIG), "{\"do_lower_case\":true}"); + } else { + directory = Path.of(modelDir); + } + } + + /** + * Deletes generated fixtures, not user model files. + * + * @throws IOException If fixture deletion fails. + */ + @TearDown(Level.Trial) + public void cleanup() throws IOException { + if (tempDir != null) { + Files.deleteIfExists(tempDir.resolve(ModelFileNames.VOCABULARY)); + Files.deleteIfExists(tempDir.resolve(ModelFileNames.SAFETENSORS)); + Files.deleteIfExists(tempDir.resolve(ModelFileNames.CONFIG)); + Files.deleteIfExists(tempDir.resolve(ModelFileNames.TOKENIZER_CONFIG)); + Files.deleteIfExists(tempDir); + } + } + + /** + * Writes WordPiece tokens for the generated table. + * + * @throws IOException If the file cannot be written. + */ + private void writeVocab() 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)); + } + Files.write(directory.resolve(ModelFileNames.VOCABULARY), tokens); + } + + /** + * Writes deterministic F32 embedding values. + * + * @throws IOException If the file cannot be written. + */ + private void writeSafetensors() throws IOException { + final Random random = new Random(42); + final float[] values = new float[VOCAB_SIZE * DIMENSION]; + for (int i = 0; i < values.length; i++) { + values[i] = (random.nextFloat() - 0.5f) * 2f; + } + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + new SafetensorsTestFiles.Tensor("embeddings", new int[] {VOCAB_SIZE, DIMENSION}, values)); + } + } + + /** A loaded model shared by inference threads. */ + @State(Scope.Benchmark) + public static class ModelState { + + StaticEmbeddingModel model; + + /** + * Loads the model before inference timing starts. + * + * @param files The prepared model directory. + * @throws IllegalArgumentException If the path is not a model directory. + * @throws IOException If the model cannot be loaded. + */ + @Setup(Level.Trial) + public void load(ModelFiles files) throws IOException { + model = StaticEmbeddingModel.load(files.directory); + } + } + + /** + * Loads a new model instance per operation. + * + * @param files The prepared model directory. + * @return The loaded model. + * @throws IllegalArgumentException If the path is not a model directory. + * @throws IOException If the model cannot be loaded. + */ + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public StaticEmbeddingModel load(ModelFiles files) throws IOException { + return StaticEmbeddingModel.load(files.directory); + } + + /** + * Embeds a batch with throughput and allocation expressed per input text. + * + * @param state The loaded model. + * @param blackhole Receives the output vectors. + */ + @Benchmark + @OperationsPerInvocation(5) + public void embed(ModelState state, Blackhole blackhole) { + for (final String sentence : SENTENCES) { + blackhole.consume(state.model.embed(sentence)); + } + } + + /** + * Searches the model for 10 nearest tokens per operation. + * + * @param state The loaded model. + * @param blackhole Receives the search results. + */ + @Benchmark + public void mostSimilarTop10(ModelState state, Blackhole blackhole) { + blackhole.consume(state.model.mostSimilar(SENTENCES.get(0), 10)); + } + + /** + * Runs the embedding benchmarks with JMH command-line options. + * + * @param args JMH options, such as {@code -t 1 -prof gc}. + * @throws Exception If option parsing or benchmark execution fails. + */ + public static void main(String[] args) throws Exception { + final Options opt = new OptionsBuilder() + .parent(new CommandLineOptions(args)) + .include(StaticEmbeddingModelBenchmark.class.getSimpleName()) + .shouldFailOnError(true) + .build(); + new Runner(opt).run(); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmarkTest.java b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmarkTest.java new file mode 100644 index 0000000000..8fc854b3d7 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/jmh/java/opennlp/embeddings/StaticEmbeddingModelBenchmarkTest.java @@ -0,0 +1,196 @@ +/* + * 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 java.util.concurrent.TimeUnit; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OperationsPerInvocation; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.infra.Blackhole; +import org.openjdk.jmh.runner.BenchmarkList; +import org.openjdk.jmh.runner.BenchmarkListEntry; + +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.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests benchmark units, model reloading and fixture cleanup. */ +class StaticEmbeddingModelBenchmarkTest { + + private static final String QUERY = "fox"; + + @TempDir + Path modelDirectory; + + /** + * Checks that throughput and allocation are reported per text, not per batch. + * + * @throws NoSuchMethodException If the benchmark method cannot be found. + */ + @Test + void testEmbedOperationCount() throws NoSuchMethodException { + final OperationsPerInvocation operations = StaticEmbeddingModelBenchmark.class + .getMethod("embed", StaticEmbeddingModelBenchmark.ModelState.class, Blackhole.class) + .getAnnotation(OperationsPerInvocation.class); + assertNotNull(operations, "The batch requires an explicit operation count"); + assertEquals(StaticEmbeddingModelBenchmark.SENTENCES.size(), operations.value()); + } + + /** + * Checks that generated JMH metadata reflects the current benchmark source. + * + * @throws IOException If the generated metadata cannot be read. + */ + @Test + void testGeneratedMetadata() throws IOException { + try (InputStream metadata = StaticEmbeddingModelBenchmark.class + .getResourceAsStream("/META-INF/BenchmarkList")) { + assertNotNull(metadata, "JMH annotation processing must generate benchmark metadata"); + final List entries = BenchmarkList.readBenchmarkList(metadata).stream() + .filter(entry -> entry.getUserClassQName() + .equals(StaticEmbeddingModelBenchmark.class.getName())) + .toList(); + final String benchmarkName = StaticEmbeddingModelBenchmark.class.getName(); + assertEquals(List.of(benchmarkName + ".embed", benchmarkName + ".load", + benchmarkName + ".mostSimilarTop10"), + entries.stream().map(BenchmarkListEntry::getUsername).sorted().toList()); + for (final BenchmarkListEntry entry : entries) { + assertEquals(StaticEmbeddingModelBenchmark.class.getAnnotation(Threads.class).value(), + entry.getThreads().get()); + assertArrayEquals(new String[] {"synthetic"}, entry.getParams().get().get("modelDir")); + if (entry.getUsername().endsWith(".embed")) { + assertEquals(StaticEmbeddingModelBenchmark.SENTENCES.size(), + entry.getOperationsPerInvocation().get()); + } + if (entry.getUsername().endsWith(".load")) { + assertEquals(TimeUnit.MILLISECONDS, entry.getTimeUnit().get()); + assertEquals(Mode.AverageTime, entry.getMode()); + } else { + assertEquals(TimeUnit.SECONDS, entry.getTimeUnit().get()); + assertEquals(Mode.Throughput, entry.getMode()); + } + } + } + } + + /** + * Loads the generated table through the public directory API and removes the fixture. + * + * @throws IOException If fixture creation or loading fails. + */ + @Test + void testSyntheticModel() throws IOException { + final var files = new StaticEmbeddingModelBenchmark.ModelFiles(); + files.modelDir = "synthetic"; + try { + files.prepare(); + final var state = new StaticEmbeddingModelBenchmark.ModelState(); + state.load(files); + final StaticEmbeddingModel model = state.model; + assertEquals(29_528, model.vocabularySize()); + assertEquals(256, model.dimension()); + assertEquals(10, model.mostSimilar(QUERY, 10).size()); + final StaticEmbeddingModel reloaded = new StaticEmbeddingModelBenchmark().load(files); + assertNotSame(model, reloaded); + for (final String text : StaticEmbeddingModelBenchmark.SENTENCES) { + final float[] vector = model.embed(text); + assertArrayEquals(vector, reloaded.embed(text)); + double norm = 0; + for (final float value : vector) { + norm += (double) value * value; + } + assertEquals(1, norm, 1e-6, text); + } + } finally { + files.cleanup(); + } + assertFalse(Files.exists(files.directory)); + files.cleanup(); + } + + /** + * Checks fresh file loading without changing or deleting a supplied directory. + * + * @throws IOException If fixture creation or loading fails. + */ + @Test + void testUserModelReloadAndCleanup() throws IOException { + Files.write(modelDirectory.resolve(ModelFileNames.VOCABULARY), List.of("[UNK]", QUERY)); + Files.writeString(modelDirectory.resolve(ModelFileNames.CONFIG), "{\"normalize\":false}"); + Files.writeString(modelDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG), "{\"do_lower_case\":true}"); + writeMatrix(new float[] {1, 2}); + final var files = new StaticEmbeddingModelBenchmark.ModelFiles(); + files.modelDir = modelDirectory.toString(); + files.prepare(); + final var benchmark = new StaticEmbeddingModelBenchmark(); + final StaticEmbeddingModel original = benchmark.load(files); + assertArrayEquals(new float[] {1, 2}, original.embed(QUERY)); + writeMatrix(new float[] {3, 4}); + final byte[] tensorBytes = Files.readAllBytes(modelDirectory.resolve(ModelFileNames.SAFETENSORS)); + final StaticEmbeddingModel updated = benchmark.load(files); + assertNotSame(original, updated); + assertArrayEquals(new float[] {3, 4}, updated.embed(QUERY)); + assertArrayEquals(new float[] {1, 2}, original.embed(QUERY)); + files.cleanup(); + assertTrue(Files.isDirectory(modelDirectory)); + assertEquals(List.of("[UNK]", QUERY), + Files.readAllLines(modelDirectory.resolve(ModelFileNames.VOCABULARY))); + assertEquals("{\"normalize\":false}", Files.readString(modelDirectory.resolve(ModelFileNames.CONFIG))); + assertEquals("{\"do_lower_case\":true}", + Files.readString(modelDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG))); + assertArrayEquals(tensorBytes, Files.readAllBytes(modelDirectory.resolve(ModelFileNames.SAFETENSORS))); + } + + /** + * Checks that a missing model directory fails without creating input files. + * + * @throws IOException If setup or cleanup fails. + */ + @Test + void testMissingUserModel() throws IOException { + final Path missing = modelDirectory.resolve("missing-model"); + final var files = new StaticEmbeddingModelBenchmark.ModelFiles(); + files.modelDir = missing.toString(); + files.prepare(); + assertThrows(IllegalArgumentException.class, () -> new StaticEmbeddingModelBenchmark().load(files)); + files.cleanup(); + assertFalse(Files.exists(missing)); + } + + /** + * Writes a small embedding table for the supplied-directory test. + * + * @param vector The vector for the query token. + * @throws IOException If the file cannot be written. + */ + private void writeMatrix(float[] vector) throws IOException { + SafetensorsTestFiles.write(modelDirectory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", new float[][] {new float[2], vector})); + } +} 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..b7164e8b25 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingTable.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; + +/** + * 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, double[])} + * accumulates into a vector of {@link #pooledLength()}, and {@link #finishPooling(double[])} + * maps the accumulated vector to original space once per pooled result. Scoring mirrors this: + * {@link #prepareQuery(double[])} maps a query into the working space once, and + * {@link #dot(int, double[])} scores every row against the prepared query there. The working + * space must preserve norms and dot products, so cosine math is space-independent.

+ * + *

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 double-precision accumulator, of length {@link #pooledLength()}. + */ + void addRow(int row, float weight, double[] sum); + + /** + * Maps an accumulated vector to original space. Called once per pooled result. The result stays + * in double precision so normalization can run before conversion to the public float vector. + * + * @param sum The double-precision accumulator, of length {@link #pooledLength()}. + * @return The double-precision pooled vector in original space, of length + * {@link #dimension()}. + */ + double[] finishPooling(double[] sum); + + /** + * Maps an original-space query into this table's working space, once per scan. + * + * @param query The query, of length {@link #dimension()}. Not modified. + * @return The double-precision prepared query, of length {@link #pooledLength()}. + */ + double[] prepareQuery(double[] 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(double[])}. + * @return The dot product. + */ + double dot(int row, double[] 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..2ee1ed9152 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/EmbeddingVocabulary.java @@ -0,0 +1,182 @@ +/* + * 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; + private final Set specialRows; + + /** Holds the parsed piece-to-row and row-to-piece views; built by the {@code from*} factories. */ + private EmbeddingVocabulary(Map idByToken, List tokenById, + Set specialRows) { + this.idByToken = idByToken; + this.tokenById = tokenById; + this.specialRows = specialRows; + } + + /** + * 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); + final TokenizerJsonVocab.Result result = TokenizerJsonVocab.read(file); + return fromLines(result.rows(), file.toString(), result.specialRows()); + } + + /** + * 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 { + return fromLines(lines, sourceName, Set.of()); + } + + /** Builds a vocabulary and records rows declared as special tokens. */ + private static EmbeddingVocabulary fromLines(List lines, String sourceName, + Set specialRows) + 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), + Set.copyOf(specialRows)); + } + + /** {@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(); + } + + /** {@return the rows declared as special tokens by the vocabulary source} */ + Set specialRows() { + return specialRows; + } + + /** + * 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..eb7e6e677b --- /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..8c0449d2b4 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/FloatEmbeddingTable.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; + +/** + * A float, row-major {@link EmbeddingTable}. Its working space is the original vector space, and + * row norms are computed when the table is constructed. + */ +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, double[] 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] += (double) values[base + d] * weight; + } + } + } + + /** {@inheritDoc} */ + @Override + public double[] finishPooling(double[] sum) { + return sum; + } + + /** {@inheritDoc} */ + @Override + public double[] prepareQuery(double[] query) { + return query.clone(); + } + + /** {@inheritDoc} */ + @Override + public double dot(int row, double[] preparedQuery) { + final int base = row * dimension; + 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..17dbe1bcdb --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/GaussianQuantizer.java @@ -0,0 +1,227 @@ +/* + * 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; + +/** + * A scalar quantizer for standard-normal values: {@code 2^bits} representation levels minimizing + * mean squared error over {@code N(0,1)}, with encoding by nearest level. A + * {@link HadamardRotation} and per-row scaling let {@link QuantizedEmbeddingMatrix} use this one + * fixed grid for every row. + * + *

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.

+ * + *

A quantized file stores its grid, and reading rebuilds the quantizer from the stored levels + * through {@link #fromLevels(float[])}. Decoding therefore uses the serialized levels instead of + * regenerating them.

+ * + *

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 when the value lies in + // (thresholds[i-1], thresholds[i]], with the outermost intervals unbounded. + private final double[] thresholds; + + /** + * Constructs a validated grid for {@link #forBits(int)} and + * {@link #fromLevels(float[])}. + * + * @param levels The representation levels, ascending. + */ + private GaussianQuantizer(float[] levels) { + this.levels = levels; + this.thresholds = new double[levels.length - 1]; + for (int i = 0; i < thresholds.length; i++) { + thresholds[i] = ((double) levels[i] + levels[i + 1]) / 2.0; + } + } + + /** + * {@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); + } + // Start with levels spaced across the central mass. + 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..65eec880ed --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HadamardRotation.java @@ -0,0 +1,186 @@ +/* + * 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. The transform reduces coordinate concentration + * before {@link GaussianQuantizer} encodes each coordinate independently. + * + *

The transform is orthonormal, so it preserves norms and inner products within floating-point + * rounding): two vectors rotated with the same instance have the same dot product as the + * originals, which lets similarity calculations remain in rotated space without an 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 double 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 = 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(double[] 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(double[] vector) { + requirePaddedLength(vector); + walshHadamard(vector); + for (int i = 0; i < paddedDimension; i++) { + if (flip[i]) { + vector[i] = -vector[i]; + } + } + } + + /** + * Checks that a vector has 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(double[] 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(double[] 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 double a = vector[i]; + final double 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 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..bf9e08ac5e --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/HuggingFaceModelCache.java @@ -0,0 +1,818 @@ +/* + * 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.channels.FileChannel; +import java.nio.channels.FileLock; +import java.nio.channels.OverlappingFileLockException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +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; + +/** + * Downloads the files needed to distill a Hugging Face model and keeps a verified local cache. + * Each ref is resolved to one commit, and each downloaded file must match the digest reported by + * the hub. Missing optional files are omitted. + * + *

{@value #REVISION_FILE} and {@value #FILES_FILE} describe a complete snapshot. Writers for + * the same teacher are serialized with a filesystem lock.

+ */ +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 percent-encoded form of a slash inside one URI path segment. */ + private static final String ENCODED_SLASH = "%2F"; + + /** 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 file listing every repository artifact present in a completed cache snapshot. */ + static final String FILES_FILE = ".opennlp-files"; + + /** The suffix of the temporary file a download streams into before it is moved into place. */ + private static final String DOWNLOAD_SUFFIX = ".download"; + + /** The suffix of the lock file that serializes writers to one cache directory. */ + private static final String LOCK_SUFFIX = ".opennlp-lock"; + + /** Delay before retrying a file lock held by another thread in this JVM. */ + private static final long LOCK_RETRY_MILLIS = 10L; + + /** 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}. A relative path containing a {@code ..} segment is rejected as + * ambiguous; pass an absolute or normalized path instead. + * @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 (!isAmbiguousRelativePath(local) && 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)); + Files.createDirectories(cacheRoot); + final Path lockFile = cache.resolveSibling(cache.getFileName() + LOCK_SUFFIX); + try (FileChannel channel = FileChannel.open(lockFile, + StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { + final FileLock cacheLock = acquireCacheLock(channel); + try (cacheLock) { + return resolveLocked(cache, hubBase, modelId, requestedRevision, listener); + } + } + } + + /** + * Resolves one teacher while holding its cache lock. + * + * @param cache The teacher's cache directory. + * @param hubBase The hub's base URL. + * @param modelId The hub model id. + * @param requestedRevision The requested revision, or {@code null} for the default branch. + * @param listener The progress listener; may be {@code null}. + * @return The completed cache directory. + * @throws IOException Thrown if the snapshot cannot be resolved, downloaded, or verified. + */ + private static Path resolveLocked(Path cache, String hubBase, String modelId, + String requestedRevision, + ModelDistiller.ProgressListener listener) throws IOException { + final String pinned = pinnedRevision(cache); + if (pinned != null && hasCompleteSnapshot(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 failed download therefore leaves no completion record, and the next attempt verifies + // each existing file before reusing it. + Files.deleteIfExists(cache.resolve(REVISION_FILE)); + Files.deleteIfExists(cache.resolve(FILES_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); + } + writeFileRecord(cache); + Files.writeString(cache.resolve(REVISION_FILE), commit + System.lineSeparator(), + StandardCharsets.UTF_8); + return cache; + } + + /** + * Acquires an exclusive cache lock, waiting when another thread or process holds it. + * + * @param channel The lock-file channel. + * @return The acquired lock. + * @throws IOException Thrown if the lock cannot be acquired or the thread is interrupted. + */ + private static FileLock acquireCacheLock(FileChannel channel) throws IOException { + while (true) { + try { + return channel.lock(); + } catch (OverlappingFileLockException e) { + try { + Thread.sleep(LOCK_RETRY_MILLIS); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while waiting for the model cache lock", interrupted); + } + } + } + } + + /** + * {@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 the cache contains exactly the known files recorded for its snapshot} + * + * @param cache The cache directory. + */ + private static boolean hasCompleteSnapshot(Path cache) { + final List recorded; + try { + recorded = Files.readAllLines(cache.resolve(FILES_FILE), StandardCharsets.UTF_8); + } catch (IOException e) { + return false; + } + if (recorded.size() != recorded.stream().distinct().count()) { + return false; + } + for (final String file : recorded) { + if (!REQUIRED_FILES.contains(file) && !OPTIONAL_FILES.contains(file)) { + return false; + } + } + for (final String file : REQUIRED_FILES) { + if (!recorded.contains(file) || !Files.isRegularFile(cache.resolve(file))) { + return false; + } + } + for (final String file : OPTIONAL_FILES) { + if (recorded.contains(file) != Files.isRegularFile(cache.resolve(file))) { + return false; + } + } + return true; + } + + /** + * Records the repository files present after a verified download. + * + * @param cache The completed cache directory. + * @throws IOException Thrown if the record cannot be written. + */ + private static void writeFileRecord(Path cache) throws IOException { + final List files = new ArrayList<>(REQUIRED_FILES); + for (final String file : OPTIONAL_FILES) { + if (Files.isRegularFile(cache.resolve(file))) { + files.add(file); + } + } + Files.write(cache.resolve(FILES_FILE), files, StandardCharsets.UTF_8); + } + + /** + * 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") + + "; every downloaded file must be attributable 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 + encodeRevision(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 and validates the digest the hub publishes for a file. + * + * @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, so the file cannot be verified"); + } + String hex = header.trim(); + if (hex.length() >= 2 && hex.charAt(0) == '"' && hex.charAt(hex.length() - 1) == '"') { + hex = hex.substring(1, hex.length() - 1); + } + 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, so the file cannot be verified"); + } + return new Digest(form, hex); + } + + /** + * {@return the value the original hub response 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 final response. This method + * reads only the original response, so a redirect target cannot supply the verification + * value.

+ * + * @param response The response, at the end of its redirect chain. + * @param name The header name. + */ + private static String originHeader(HttpResponse response, String name) { + HttpResponse origin = response; + while (origin.previousResponse().isPresent()) { + origin = origin.previousResponse().get(); + } + return origin.headers().firstValue(name).orElse(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; + } + + /** + * {@return whether {@code path} is relative and contains a {@code ..} segment} + * + *

Windows collapses {@code ..} lexically, without checking that the segment before it + * exists, so {@code BAAI/..} denotes the working directory there even when {@code BAAI} does + * not exist; POSIX resolves the same string to nothing. Refusing the shape keeps a misspelled + * hub id from silently naming a directory the caller did not mean, and keeps the rejection + * identical on every platform.

+ * + * @param path The teacher reference as a path. + */ + private static boolean isAmbiguousRelativePath(Path path) { + if (path.isAbsolute()) { + return false; + } + for (final Path segment : path) { + if ("..".equals(segment.toString())) { + return true; + } + } + return false; + } + + /** + * 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 made from one or more slash-delimited parts 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 && !isRevision(revision))) { + return null; + } + return new TeacherReference(organization + "/" + model, revision); + } + + /** + * {@return whether a revision contains one or more non-empty slash-delimited reference parts} + * + * @param revision The revision to check. + */ + private static boolean isRevision(String revision) { + int start = 0; + for (int i = 0; i <= revision.length(); i++) { + if (i == revision.length() || revision.charAt(i) == '/') { + if (!isReferencePart(revision.substring(start, i))) { + return false; + } + start = i + 1; + } + } + return true; + } + + /** + * {@return a revision encoded as one URI path segment} + * + * @param revision The validated revision. + */ + private static String encodeRevision(String revision) { + return revision.replace("/", ENCODED_SLASH); + } + + /** + * {@return whether a reference part is one or more ASCII word characters, dots, or dashes, + * excluding the path segments {@code .} and {@code ..}} + * + * @param part The part to check. + */ + private static boolean isReferencePart(String part) { + if (part.isEmpty() || ".".equals(part) || "..".equals(part)) { + 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, distinguished by + * hex length. + */ + 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..253c3abd2e --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/JsonCursor.java @@ -0,0 +1,409 @@ +/* + * 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 shared by the small JSON readers in this package. It parses scalar values and can skip + * one value of any type. Each reader handles its expected input structure. Malformed input raises + * an {@link InvalidFormatException} that includes the input name and offset. + */ +final class JsonCursor { + + private static final int MAX_NESTING_DEPTH = 128; + + 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()) { + final char c = text.charAt(position); + if (c != ' ' && c != '\t' && c != '\n' && c != '\r') { + return; + } + 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 if (c <= 0x1F) { + throw malformed("Unescaped control character in a string"); + } 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; + // JSON escape digits are limited to the ASCII hexadecimal characters. + int value = 0; + for (int i = 0; i < 4; i++) { + final int digit = hexadecimalValue(hex.charAt(i)); + if (digit < 0) { + throw malformed("Malformed \\u escape sequence: " + hex); + } + value = (value << 4) | digit; + } + return (char) value; + } + + /** {@return the value of an ASCII hexadecimal digit, or {@code -1} for another character} */ + private int hexadecimalValue(char c) { + if (c >= '0' && c <= '9') { + return c - '0'; + } + if (c >= 'a' && c <= 'f') { + return c - 'a' + 10; + } + if (c >= 'A' && c <= 'F') { + return c - 'A' + 10; + } + return -1; + } + + /** + * Skips one JSON number, holding it to the grammar (optional minus, digits, optional fraction, + * optional signed exponent). This validation also applies to skipped fields. + */ + private void skipNumber() throws InvalidFormatException { + skipIntegerPart(); + if (position < text.length() && text.charAt(position) == '.') { + position++; + if (position >= text.length() || !isAsciiDigit(text.charAt(position))) { + throw malformed("Malformed number: digit expected after the decimal point"); + } + while (position < text.length() && isAsciiDigit(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() || !isAsciiDigit(text.charAt(position))) { + throw malformed("Malformed number: digit expected in the exponent"); + } + while (position < text.length() && isAsciiDigit(text.charAt(position))) { + position++; + } + } + } + + /** + * Skips the optional sign and integer part of a JSON number. + * + * @throws InvalidFormatException Thrown if the integer part is absent or has a leading zero. + */ + private void skipIntegerPart() throws InvalidFormatException { + if (peek() == '-') { + position++; + } + if (position >= text.length() || !isAsciiDigit(text.charAt(position))) { + throw malformed("Malformed number"); + } + if (text.charAt(position) == '0') { + position++; + if (position < text.length() && isAsciiDigit(text.charAt(position))) { + throw malformed("Malformed number: leading zeros are not allowed"); + } + return; + } + while (position < text.length() && isAsciiDigit(text.charAt(position))) { + position++; + } + } + + /** {@return whether {@code c} is an ASCII decimal digit} */ + private boolean isAsciiDigit(char c) { + return c >= '0' && c <= '9'; + } + + /** + * {@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; + skipIntegerPart(); + 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), allowing a + * reader to ignore unknown fields. + */ + void skipValue() throws InvalidFormatException { + skipValue(0); + } + + /** + * Skips one JSON value at the given container depth. + * + * @param depth The number of enclosing arrays and objects. + * @throws InvalidFormatException Thrown if the value is malformed or nested too deeply. + */ + private void skipValue(int depth) throws InvalidFormatException { + skipWhitespace(); + final char c = peek(); + if (c == '"') { + parseString(); + } else if (c == '[') { + requireContainerDepth(depth); + position++; + skipWhitespace(); + if (peek() != ']') { + while (true) { + skipValue(depth + 1); + 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 == '{') { + requireContainerDepth(depth); + position++; + skipWhitespace(); + if (peek() != '}') { + while (true) { + skipWhitespace(); + parseString(); + skipWhitespace(); + expect(':'); + skipValue(depth + 1); + 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 + "'"); + } + } + + /** + * Rejects a container whose contents would exceed the nesting limit. + * + * @param depth The number of enclosing arrays and objects. + * @throws InvalidFormatException Thrown at the nesting limit. + */ + private void requireContainerDepth(int depth) throws InvalidFormatException { + if (depth >= MAX_NESTING_DEPTH) { + throw malformed("JSON nesting depth exceeds " + MAX_NESTING_DEPTH); + } + } + + /** + * {@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..43b9335f72 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Model2VecUnigramTokenizer.java @@ -0,0 +1,811 @@ +/* + * 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; +import opennlp.tools.util.Span; +import opennlp.tools.util.normalizer.AlignedText; +import opennlp.tools.util.normalizer.Alignment; + +/** + * Runs a supported Hugging Face Unigram tokenizer through OpenNLP's SentencePiece decoder. + * Normalization occurs in memory, and returned offsets refer to the original input text. + */ +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; + + /** + * Creates a tokenizer from validated configuration. + * + * @param parsed The tokenizer configuration. + * @throws IOException Thrown if an internal SentencePiece model cannot be loaded. + */ + 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. + * + * @param tokenizerJson The Hugging Face tokenizer configuration. + * @return A tokenizer for the configuration. + * @throws IllegalArgumentException Thrown if {@code tokenizerJson} is null or not a regular file. + * @throws IOException Thrown if the configuration cannot be read or loaded. + */ + 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"); + } + final String original = text.toString(); + AlignedText aligned = normalizer.normalizeAligned(original); + for (NormalizationOperation operation : operations) { + final AlignedText next = operation.applyAligned(aligned.normalizedString()); + aligned = new AlignedText(original, next.normalized(), + aligned.alignment().andThen(next.alignment())); + } + final List normalizedPieces = segmenter.encode(aligned.normalized()); + final List originalPieces = new ArrayList<>(normalizedPieces.size()); + for (SubwordPiece piece : normalizedPieces) { + final Span span = aligned.toOriginalSpan(piece.start(), piece.end()); + originalPieces.add( + new SubwordPiece(piece.piece(), piece.id(), span.getStart(), span.getEnd())); + } + return originalPieces; + } + + /** {@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); + } + + /** Reads and validates the tokenizer configuration. */ + 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(); + final Set fields = new HashSet<>(); + if (cursor.peek() != '}') { + while (true) { + final String key = uniqueKey(cursor, fields, "top-level"); + 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); + } + + /** Reads the Unigram model object. */ + private static ParsedModel parseModel(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + int unknownId = -1; + boolean byteFallback = false; + List pieces = null; + final Set fields = new HashSet<>(); + while (cursor.peek() != '}') { + final String key = uniqueKey(cursor, fields, "model"); + 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(); + if (cursor.peek() == '}') { + throw cursor.malformed("Trailing comma in model object"); + } + } 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); + } + + /** Reads the ordered Unigram vocabulary. */ + 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(); + if (cursor.peek() == ']') { + throw cursor.malformed("Trailing comma in model vocabulary"); + } + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after a vocabulary entry"); + } else { + return pieces; + } + } + cursor.consume(); + return pieces; + } + + /** Reads the tokenizer normalizer. */ + private static ParsedNormalizer parseNormalizer(JsonCursor cursor) + throws InvalidFormatException { + final NormalizerBuilder builder = new NormalizerBuilder(); + parseNormalizerObject(cursor, builder); + return new ParsedNormalizer(builder.precompiledCharsMap, builder.operations); + } + + /** Adds one normalizer object to {@code builder}. */ + 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; + final Set fields = new HashSet<>(); + while (cursor.peek() != '}') { + final String key = uniqueKey(cursor, fields, "normalizer"); + 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"); + } + if (!builder.operations.isEmpty()) { + throw cursor.malformed( + "Precompiled normalizer must precede other normalization steps"); + } + 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 + "'"); + } + } + + /** Reads the children of a Sequence normalizer. */ + 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(); + if (cursor.peek() == ']') { + throw cursor.malformed("Trailing comma in normalizer array"); + } + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after a normalizer"); + } else { + return children; + } + } + cursor.consume(); + return children; + } + + /** Reads a Replace normalizer pattern. */ + 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); + } + + /** Creates a supported Replace operation. */ + 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 (pattern.value().isEmpty()) { + throw cursor.malformed("Replace normalizer has an unsupported empty literal pattern"); + } + 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"); + } + + /** Reads the pre-tokenizer and reports whether it is the supported Metaspace form. */ + private static boolean parsePreTokenizer(JsonCursor cursor) throws InvalidFormatException { + cursor.expect('{'); + cursor.skipWhitespace(); + String type = null; + String replacement = null; + String prependScheme = null; + boolean split = true; + final Set fields = new HashSet<>(); + while (cursor.peek() != '}') { + final String key = uniqueKey(cursor, fields, "pre-tokenizer"); + 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; + } + + /** Reads tokenizer rows that carry added-token metadata. */ + private static List parseAddedTokens(JsonCursor cursor) + throws InvalidFormatException { + cursor.expect('['); + cursor.skipWhitespace(); + final List tokens = new ArrayList<>(); + final Set tokenIds = new HashSet<>(); + while (cursor.peek() != ']') { + cursor.expect('{'); + cursor.skipWhitespace(); + int id = -1; + String content = null; + boolean special = false; + final Set fields = new HashSet<>(); + while (cursor.peek() != '}') { + final String key = uniqueKey(cursor, fields, "added token"); + 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"); + } + if (!tokenIds.add(id)) { + throw cursor.malformed("added token id " + id + " occurs more than once"); + } + tokens.add(new AddedToken(id, content, special)); + cursor.skipWhitespace(); + final char next = cursor.consume(); + if (next == ',') { + cursor.skipWhitespace(); + if (cursor.peek() == ']') { + throw cursor.malformed("Trailing comma in added-token array"); + } + } else if (next != ']') { + throw cursor.malformed("Expected ',' or ']' after an added token"); + } else { + return tokens; + } + } + cursor.consume(); + return tokens; + } + + /** + * Reads an object field name and rejects a name already seen in that object. + * + * @param cursor The JSON cursor positioned at the field name. + * @param fields The names already read from the current object. + * @param object A short object name for error messages. + * @return The field name. + * @throws InvalidFormatException Thrown if the field name occurs more than once. + */ + private static String uniqueKey(JsonCursor cursor, Set fields, String object) + throws InvalidFormatException { + final String key = cursor.parseString(); + if (!fields.add(key)) { + throw cursor.malformed(object + " field '" + key + "' occurs more than once"); + } + return key; + } + + /** Converts a non-negative JSON integer to an {@code int}. */ + 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; + } + + /** {@return whether {@code piece} has the ASCII form {@code <0xNN>}} */ + 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 isAsciiHexDigit(piece.charAt(3)) && isAsciiHexDigit(piece.charAt(4)); + } + + /** {@return whether {@code c} is an ASCII hexadecimal digit} */ + private static boolean isAsciiHexDigit(char c) { + return c >= '0' && c <= '9' || c >= 'A' && c <= 'F' || c >= 'a' && c <= 'f'; + } + + /** Encodes the supplied tokenizer data as a SentencePiece model. */ + 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 { + + /** + * Applies this operation and records how its output maps to {@code input}. + * + * @param input The text produced by the previous normalization stage. + * @return The operation result and its alignment to {@code input}. + */ + AlignedText applyAligned(String input); + } + + private record SurroundOperation(String literal) implements NormalizationOperation { + /** {@inheritDoc} */ + @Override + public AlignedText applyAligned(String input) { + final StringBuilder out = new StringBuilder(input.length() + 8); + final Alignment.Builder alignment = new Alignment.Builder(input.length() + 8); + int cursor = 0; + while (cursor < input.length()) { + if (input.startsWith(literal, cursor)) { + appendMarker(out, alignment); + out.append(literal); + alignment.equal(literal.length()); + appendMarker(out, alignment); + cursor += literal.length(); + } else { + final int codePoint = input.codePointAt(cursor); + out.appendCodePoint(codePoint); + final int width = Character.charCount(codePoint); + alignment.equal(width); + cursor += width; + } + } + return new AlignedText(input, out.toString(), alignment.build(input.length())); + } + } + + private enum CollapseOperation implements NormalizationOperation { + INSTANCE; + + /** {@inheritDoc} */ + @Override + public AlignedText applyAligned(String input) { + final StringBuilder out = new StringBuilder(input.length()); + final Alignment.Builder alignment = new Alignment.Builder(input.length()); + int cursor = 0; + while (cursor < input.length()) { + final int codePoint = input.codePointAt(cursor); + if (codePoint == MARKER_CHAR) { + final int start = cursor; + do { + cursor++; + } while (cursor < input.length() && input.charAt(cursor) == MARKER_CHAR); + out.append(MARKER_CHAR); + alignment.replace(cursor - start, 1); + } else { + final int width = Character.charCount(codePoint); + out.appendCodePoint(codePoint); + alignment.equal(width); + cursor += width; + } + } + return new AlignedText(input, out.toString(), alignment.build(input.length())); + } + } + + private record StripOperation(boolean left, boolean right) implements NormalizationOperation { + /** {@inheritDoc} */ + @Override + public AlignedText applyAligned(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--; + } + } + if (start == 0 && end == input.length()) { + return identity(input); + } + final Alignment.Builder alignment = new Alignment.Builder(end - start); + alignment.replace(start, 0); + alignment.equal(end - start); + alignment.replace(input.length() - end, 0); + return new AlignedText(input, input.substring(start, end), alignment.build(input.length())); + } + } + + /** + * Creates an aligned identity result. + * + * @param input The text used as both sides of the result. + * @return An identity alignment over {@code input}. + */ + private static AlignedText identity(String input) { + return new AlignedText(input, input, + new Alignment.Builder(input.length()).equal(input.length()).build(input.length())); + } + + /** + * Appends a metaspace marker when the output does not already end with one. + * + * @param out The normalized output. + * @param alignment The alignment being built for {@code out}. + */ + private static void appendMarker(StringBuilder out, Alignment.Builder alignment) { + if (out.isEmpty() || out.charAt(out.length() - 1) != MARKER_CHAR) { + out.append(MARKER_CHAR); + alignment.replace(0, 1); + } + } + + 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(); + + /** Writes an embedded message field. */ + void message(int field, byte[] value) { + bytesField(field, value); + } + + /** Writes a UTF-8 string field. */ + void string(int field, String value) { + bytesField(field, value.getBytes(StandardCharsets.UTF_8)); + } + + /** Writes a length-delimited field. */ + void bytesField(int field, byte[] value) { + varint((long) field << 3 | 2); + varint(value.length); + out.writeBytes(value); + } + + /** Writes an integer field. */ + void varintField(int field, long value) { + varint((long) field << 3); + varint(value); + } + + /** Writes a 32-bit floating-point field. */ + 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); + } + + /** Writes an unsigned variable-length integer. */ + void varint(long value) { + long remaining = value; + while ((remaining & ~0x7fL) != 0) { + out.write((int) (remaining & 0x7f) | 0x80); + remaining >>>= 7; + } + out.write((int) remaining); + } + + /** {@return the encoded message} */ + 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..06d59b2bf9 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelAssembler.java @@ -0,0 +1,393 @@ +/* + * 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 SentencePiece tokenizer family used 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 separateSentencePiece = ModelFileNames.firstRegularFile(modelDirectory, + ModelFileNames.SENTENCEPIECE_MODELS) != null; + return new Result(separateSentencePiece ? 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; + boolean seenModel = false; + boolean seenNormalizer = 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 (seenModel) { + throw cursor.malformed("Field 'model' appears more than once"); + } + seenModel = true; + final ModelSection model = parseModel(cursor); + modelType = model.type(); + orderedVocabulary = model.orderedVocabulary(); + } + case "normalizer" -> { + if (seenNormalizer) { + throw cursor.malformed("Field 'normalizer' appears more than once"); + } + seenNormalizer = true; + 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; + boolean seenType = false; + boolean seenVocabulary = false; + 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)) { + if (seenType) { + throw cursor.malformed("Field 'model.type' appears more than once"); + } + seenType = true; + type = cursor.parseString(); + } else if ("vocab".equals(key)) { + if (seenVocabulary) { + throw cursor.malformed("Field 'model.vocab' appears more than once"); + } + seenVocabulary = true; + if (cursor.peek() == '{') { + orderedVocabulary = parseVocabularyDictionary(cursor); + } else { + cursor.skipValue(); + } + } 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<>(); + final Map idByToken = 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(); + final Long previousId = idByToken.putIfAbsent(token, id); + if (previousId != null) { + throw cursor.malformed("Vocabulary token '" + token + "' is assigned more than once, " + + "at ids " + previousId + " and " + id); + } + 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..c74662692f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelDistiller.java @@ -0,0 +1,602 @@ +/* + * 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.LinkOption; +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 with the randomized SVD in + * {@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.

+ * + *

The teacher must produce a consistent vector length across batches. A change in length + * causes an {@link IllegalArgumentException} before output files are written.

+ * + *

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. Files produced + * or derived by distillation are replaced; unrelated files remain. A + * failure part way through leaves an incomplete output directory. 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 and output are the same directory, 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 and output are the same directory, 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); + if (Files.exists(outputDirectory) && Files.isSameFile(teacherDirectory, outputDirectory)) { + throw new IllegalArgumentException("outputDirectory must differ from teacherDirectory"); + } + 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(); + 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); + removeDerivedArtifacts(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()) { + // The current matrix has no term rows, so an existing terms file cannot describe it. + 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); + } + + /** + * Removes tokenizer files derived or copied by an earlier distillation. They must describe the + * same tokenizer as the matrix and {@code tokenizer.json} written by the current run. + * + * @param outputDirectory The model output directory. + * @throws IOException Thrown if an old artifact cannot be removed. + */ + private static void removeDerivedArtifacts(Path outputDirectory) throws IOException { + Files.deleteIfExists(outputDirectory.resolve(ModelFileNames.VOCABULARY)); + Files.deleteIfExists(outputDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG)); + for (final String name : ModelFileNames.SENTENCEPIECE_MODELS) { + Files.deleteIfExists(outputDirectory.resolve(name)); + } + } + + /** + * 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 (int termIndex = 0; termIndex < terms.size(); termIndex++) { + final String term = terms.get(termIndex); + if (term == null) { + throw new IllegalArgumentException("terms[" + termIndex + "] must not be 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 (Files.exists(outputDirectory, LinkOption.NOFOLLOW_LINKS) + && !Files.isDirectory(outputDirectory)) { + throw new IllegalArgumentException( + "outputDirectory must be a directory or not exist: " + outputDirectory); + } + 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\": " + + jsonString(String.valueOf(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 branch or tag may later identify different model files, so the output records the exact + * teacher revision used for the distillation.

+ * + * @param teacherDirectory The teacher's directory. + */ + private static String teacherRevisionField(Path teacherDirectory) { + final String revision = HuggingFaceModelCache.pinnedRevision(teacherDirectory); + return revision == null ? "" : " \"teacher_revision\": " + jsonString(revision) + ",\n"; + } + + /** + * {@return {@code value} as a JSON string literal} + * + *

Package-private for tests: the characters that need escaping are illegal in file names + * on Windows, so a teacher directory cannot carry them there and the escaper is exercised + * directly instead.

+ * + * @param value The value to quote and escape. + */ + static String jsonString(String value) { + final StringBuilder json = new StringBuilder(value.length() + 2).append('"'); + for (int i = 0; i < value.length(); i++) { + final char c = value.charAt(i); + switch (c) { + case '"' -> json.append("\\\""); + case '\\' -> json.append("\\\\"); + case '\b' -> json.append("\\b"); + case '\f' -> json.append("\\f"); + case '\n' -> json.append("\\n"); + case '\r' -> json.append("\\r"); + case '\t' -> json.append("\\t"); + default -> { + if (c < 0x20) { + json.append("\\u00") + .append(Character.forDigit(c >>> 4, 16)) + .append(Character.forDigit(c & 0x0f, 16)); + } else { + json.append(c); + } + } + } + } + return json.append('"').toString(); + } + + /** + * 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..17a82a292e --- /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}. Every model directory has + * one matrix file, either {@link #SAFETENSORS} or {@link #QUANTIZED}. A WordPiece directory also + * has {@link #CONFIG}, {@link #VOCABULARY}, and {@link #TOKENIZER_CONFIG}; a Unigram directory has + * {@link #CONFIG} and {@link #TOKENIZER_JSON}. Separate-file SentencePiece directories may additionally + * have 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 contains 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 containing 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 containing 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..29cacb0ae1 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelQuantizer.java @@ -0,0 +1,169 @@ +/* + * 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; +import opennlp.tools.util.java.Experimental; + +/** + * 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. Delete the safetensors before loading the quantized + * deployment; a model directory containing both matrix files is ambiguous and is rejected. + * + *

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.

+ * + *

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

+ */ +@Experimental +public final class ModelQuantizer { + + // At most this many rows enter the verification sample, evenly strided so it is + // deterministic and spans the row range. + private static final int VERIFICATION_SAMPLE_CAP = 1024; + + /** Not instantiable. */ + private ModelQuantizer() { + } + + /** + * Statistics for a completed quantization. + * + * @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 included. + * @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 Statistics for the written quantized matrix. + * @throws IllegalArgumentException Thrown if an argument is invalid. + * @throws InvalidFormatException Thrown if the directory has no safetensors file or its + * tensors do not have the required shapes. + * @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); + } + GaussianQuantizer.requireSupportedBits(bits); + 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 boolean hasWeights = tensors.tensorNames() + .contains(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME); + if (hasWeights && tensors.tensorInfo(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME) + .shape().length != 1) { + throw new InvalidFormatException("Tensor '" + + StaticEmbeddingModel.WEIGHTS_TENSOR_NAME + "' in " + safetensorsFile + + " must be 1-D"); + } + 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 (hasWeights) { + weights = tensors.readFloats(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME); + if (weights.length != rowCount) { + throw new InvalidFormatException("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 sampleCount = Math.min(rowCount, VERIFICATION_SAMPLE_CAP); + int sampled = 0; + int nonZero = 0; + double cosineSum = 0; + for (int sample = 0; sample < sampleCount; sample++) { + final int row = sampleCount == 1 ? 0 + : (int) ((long) sample * (rowCount - 1) / (sampleCount - 1)); + 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); + if (denominator == 0) { + return Double.NaN; + } + return Math.max(-1.0, Math.min(1.0, 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..a695f66c0f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/Neighbor.java @@ -0,0 +1,50 @@ +/* + * 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 matrix row's text: a tokenizer piece or a term-table entry. + * @param similarity Cosine similarity to the query vector, in {@code [-1, 1]}. + */ +@Experimental +public record Neighbor(String token, double similarity) { + + /** + * Creates a search result. + * + * @param token The token or term text. + * @param similarity Cosine similarity in {@code [-1, 1]}. + * @throws IllegalArgumentException Thrown if {@code token} is {@code null}, or + * {@code similarity} is non-finite or outside {@code [-1, 1]}. + */ + public Neighbor { + if (token == null) { + throw new IllegalArgumentException("token must not be null"); + } + if (!Double.isFinite(similarity) || similarity < -1.0 || similarity > 1.0) { + throw new IllegalArgumentException( + "similarity must be finite and within [-1, 1], got " + 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..3b2dccd4d4 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/OnnxTeacherEncoder.java @@ -0,0 +1,410 @@ +/* + * 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.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; +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 accepts {@code input_ids}, optional {@code attention_mask}, and optional + * {@code token_type_ids}. The encoder supplies only the inputs declared by the graph. + * The pooled output is the mean of the rank-3 float {@code last_hidden_state} output over all + * sequence positions. When that name is absent, the graph must have one rank-3 float output. The + * attention mask is all ones because a batch is not 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 optional attention-mask input. */ + 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; + + /** The conventional name of a transformer encoder's last hidden state. */ + private static final String LAST_HIDDEN_STATE = "last_hidden_state"; + + private final OrtEnvironment environment; + private final OrtSession session; + private final OnnxJavaType inputIdsType; + private final OnnxJavaType attentionMaskType; + private final OnnxJavaType tokenTypeIdsType; + private final String hiddenStateOutput; + private final AtomicBoolean closed = new AtomicBoolean(); + + /** Expected vector length after a successful inference. */ + private int hiddenDimension = -1; + + /** Holds the open session; created by {@link #load(Path)}. */ + private OnnxTeacherEncoder(OrtEnvironment environment, OrtSession session, + OnnxJavaType inputIdsType, OnnxJavaType attentionMaskType, + OnnxJavaType tokenTypeIdsType, + String hiddenStateOutput) { + this.environment = environment; + this.session = session; + this.inputIdsType = inputIdsType; + this.attentionMaskType = attentionMaskType; + this.tokenTypeIdsType = tokenTypeIdsType; + 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 named {@code last_hidden_state} + * or one unambiguous rank-3 float tensor output. + * @return The encoder. + * @throws IllegalArgumentException Thrown if the file is missing, the graph has unsupported + * inputs, a sequence input has the wrong rank or element type, the hidden-state output is + * missing or ambiguous, 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); + } + // Close the session if graph validation fails. + 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 Set unsupportedInputs = new TreeSet<>(session.getInputNames()); + unsupportedInputs.removeAll(Set.of(INPUT_IDS, ATTENTION_MASK, TOKEN_TYPE_IDS)); + if (!unsupportedInputs.isEmpty()) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + + " declares unsupported inputs: " + unsupportedInputs); + } + final Map inputInfo = session.getInputInfo(); + final OnnxJavaType inputIdsType = integerInputType(inputInfo, INPUT_IDS, onnxFile); + final OnnxJavaType attentionMaskType = inputInfo.containsKey(ATTENTION_MASK) + ? integerInputType(inputInfo, ATTENTION_MASK, onnxFile) : null; + final OnnxJavaType tokenTypeIdsType = inputInfo.containsKey(TOKEN_TYPE_IDS) + ? integerInputType(inputInfo, TOKEN_TYPE_IDS, onnxFile) : null; + + final Map outputInfo = session.getOutputInfo(); + final String hiddenStateOutput; + if (outputInfo.containsKey(LAST_HIDDEN_STATE)) { + if (!isHiddenState(outputInfo.get(LAST_HIDDEN_STATE))) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " declares '" + + LAST_HIDDEN_STATE + "', but it is not a rank-3 FLOAT tensor"); + } + hiddenStateOutput = LAST_HIDDEN_STATE; + } else { + final List candidates = new ArrayList<>(); + for (final Map.Entry output : outputInfo.entrySet()) { + if (isHiddenState(output.getValue())) { + candidates.add(output.getKey()); + } + } + if (candidates.isEmpty()) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " has no rank-3 float " + + "tensor output (a last hidden state) to pool (outputs: " + + outputInfo.keySet() + ")"); + } + if (candidates.size() > 1) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + + " has multiple rank-3 FLOAT outputs and none is named '" + LAST_HIDDEN_STATE + + "': " + candidates); + } + hiddenStateOutput = candidates.get(0); + } + return new OnnxTeacherEncoder(environment, session, inputIdsType, attentionMaskType, + tokenTypeIdsType, 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; + } + } + + /** + * Reads and validates one integer sequence input. + * + * @param inputInfo The graph's input metadata. + * @param name The input name. + * @param onnxFile The graph file, for error messages. + * @return The input's INT32 or INT64 element type. + */ + private static OnnxJavaType integerInputType(Map inputInfo, String name, + Path onnxFile) { + final NodeInfo node = inputInfo.get(name); + if (node == null || !(node.getInfo() instanceof TensorInfo tensorInfo)) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " input '" + name + + "' must be a tensor"); + } + if (tensorInfo.getShape().length != 2) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " input '" + name + + "' must have rank 2, but has rank " + tensorInfo.getShape().length); + } + if (tensorInfo.type != OnnxJavaType.INT32 && tensorInfo.type != OnnxJavaType.INT64) { + throw new IllegalArgumentException("ONNX graph " + onnxFile + " input '" + name + + "' must be INT32 or INT64, but is " + tensorInfo.type); + } + return tensorInfo.type; + } + + /** + * Checks whether output metadata describes a last hidden state. + * + * @param node The output metadata. + * @return {@code true} for a rank-three FLOAT tensor. + */ + private static boolean isHiddenState(NodeInfo node) { + return node.getInfo() instanceof TensorInfo tensorInfo + && tensorInfo.type == OnnxJavaType.FLOAT + && tensorInfo.getShape().length == HIDDEN_STATE_RANK; + } + + /** + * 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 and must not contain a null or empty sequence. + * @return The pooled vectors, {@code [batchSize][hiddenDimension]}. + * @throws IllegalArgumentException Thrown if the batch is empty, contains a null or empty + * sequence, is ragged, the vector length changes between batches, 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"); + } + for (int i = 0; i < batch.length; i++) { + if (batch[i] == null) { + throw new IllegalArgumentException("batch[" + i + "] must not be null"); + } + if (batch[i].length == 0) { + throw new IllegalArgumentException("batch[" + i + "] must not be 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"); + } + } + try { + final Map inputs = new HashMap<>(); + OnnxTensor mask = null; + OnnxTensor tokenTypeIds = null; + try (OnnxTensor inputIds = createIntegerTensor(batch, inputIdsType, INPUT_IDS)) { + inputs.put(INPUT_IDS, inputIds); + if (attentionMaskType != null) { + final long[][] attentionMask = new long[batch.length][sequenceLength]; + for (final long[] row : attentionMask) { + Arrays.fill(row, 1L); + } + mask = createIntegerTensor(attentionMask, attentionMaskType, ATTENTION_MASK); + inputs.put(ATTENTION_MASK, mask); + } + if (tokenTypeIdsType != null) { + tokenTypeIds = createIntegerTensor(new long[batch.length][sequenceLength], + tokenTypeIdsType, TOKEN_TYPE_IDS); + 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(); + validateOutputShape(hidden, batch.length, sequenceLength); + final float[][] pooled = new float[batch.length][]; + for (int i = 0; i < batch.length; i++) { + final float[][] states = hidden[i]; + final double[] sum = new double[states[0].length]; + for (final float[] state : states) { + for (int d = 0; d < sum.length; d++) { + sum[d] += state[d]; + } + } + final float[] mean = new float[sum.length]; + for (int d = 0; d < sum.length; d++) { + final double meanValue = sum[d] / states.length; + mean[d] = Double.isNaN(meanValue) ? 0 : (float) meanValue; + } + pooled[i] = mean; + } + return pooled; + } finally { + if (tokenTypeIds != null) { + tokenTypeIds.close(); + } + if (mask != null) { + mask.close(); + } + } + } + } catch (OrtException e) { + throw new IllegalArgumentException("ONNX forward pass failed: " + e.getMessage(), e); + } + } + + /** + * Creates an INT32 or INT64 tensor for a graph input. + * + * @param values The input values. + * @param type The element type declared by the graph. + * @param name The input name, for range errors. + * @return The created tensor. + * @throws IllegalArgumentException Thrown if an INT32 input value is outside its range. + * @throws OrtException Thrown if ONNX Runtime cannot create the tensor. + */ + private OnnxTensor createIntegerTensor(long[][] values, OnnxJavaType type, String name) + throws OrtException { + if (type == OnnxJavaType.INT64) { + return OnnxTensor.createTensor(environment, values); + } + final int[][] converted = new int[values.length][]; + for (int row = 0; row < values.length; row++) { + converted[row] = new int[values[row].length]; + for (int column = 0; column < values[row].length; column++) { + final long value = values[row][column]; + if (value < Integer.MIN_VALUE || value > Integer.MAX_VALUE) { + throw new IllegalArgumentException(name + "[" + row + "][" + column + "] value " + + value + " does not fit an INT32 tensor"); + } + converted[row][column] = (int) value; + } + } + return OnnxTensor.createTensor(environment, converted); + } + + /** + * Validates output dimensions and requires a consistent vector length between batches. + * + * @param hidden The last hidden state. + * @param batchSize The input batch size. + * @param sequenceLength The input sequence length. + * @throws IllegalArgumentException Thrown if output dimensions are inconsistent. + */ + private void validateOutputShape(float[][][] hidden, int batchSize, int sequenceLength) { + if (hidden.length != batchSize) { + throw new IllegalArgumentException("ONNX output '" + hiddenStateOutput + + "' batch dimension " + hidden.length + " does not match input batch dimension " + + batchSize); + } + int dimension = -1; + for (int row = 0; row < hidden.length; row++) { + final float[][] states = hidden[row]; + if (states.length != sequenceLength) { + throw new IllegalArgumentException("ONNX output '" + hiddenStateOutput + "' row " + row + + " has sequence dimension " + states.length + ", expected " + sequenceLength); + } + for (int position = 0; position < states.length; position++) { + if (dimension < 0) { + dimension = states[position].length; + if (dimension == 0) { + throw new IllegalArgumentException("ONNX output '" + hiddenStateOutput + + "' has an empty hidden dimension"); + } + } else if (states[position].length != dimension) { + throw new IllegalArgumentException("ONNX output '" + hiddenStateOutput + "' row " + + row + " position " + position + " has hidden dimension " + + states[position].length + ", expected " + dimension); + } + } + } + if (hiddenDimension >= 0 && dimension != hiddenDimension) { + throw new IllegalArgumentException("ONNX output '" + hiddenStateOutput + + "' has hidden dimension " + dimension + ", expected " + hiddenDimension + + " from the initial batch"); + } + hiddenDimension = dimension; + } + + /** + * Closes the native session; calling this more than once is a no-op after the first call. + * + *

The shared {@link OrtEnvironment} remains open because {@link + * OrtEnvironment#getEnvironment()} returns a process-wide singleton. The atomic guard prevents + * a second {@link OrtSession#close()} call.

+ */ + @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..e63c68d081 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedEmbeddingMatrix.java @@ -0,0 +1,752 @@ +/* + * 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; +import opennlp.tools.util.java.Experimental; + +/** + * An embedding matrix quantized to {@code 2}-{@code 4} bits per dimension using an MSE-oriented + * variant of TurboQuant (Zandieh, Daliri, Hadian, Mirrokni, + * TurboQuant: Online Vector Quantization with + * Near-optimal Distortion Rate). Each row is transformed by a seeded + * {@link HadamardRotation}, and each rotated coordinate is encoded with the + * {@link GaussianQuantizer} grid for the selected bit width. A row decodes to a per-row scale + * times its grid levels; the scale is fitted by least squares. + * + *

The cited algorithm uses a dense random rotation and a dimension-specific coordinate + * distribution in its MSE stage. This implementation substitutes the fast Hadamard transform and + * a standard-normal grid. It also omits the paper's residual QJL stage, so it does not claim the + * paper's unbiased inner-product estimator.

+ * + *

The storage is {@code bits} per padded dimension plus two doubles per row for the + * fitted scale and decoded norm. The rotation pads each row to the next power of two.

+ * + *

Rows live in rotated space. 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, double[])} and applying {@link #toOriginal(double[])} + * 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, double[])}, without + * leaving rotated space. {@link #decodeRow(int)} exists for callers that need one original-space + * row and for measuring reconstruction quality.

+ * + *

The file format stores the grid levels and rotation seed. The reader therefore uses the + * same decoder parameters that wrote the file.

+ * + *

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 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; + + // "ONQ2": OpenNLP quantized matrix, format 2. Version 2 stores scales and norms as doubles. + private static final int MAGIC = 0x4F4E5132; + + 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 double[] 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 during quantization + // time (one inverse rotation per row) and stored in the file rather than recomputed from the + // codes on load. + private final double[] decodedNorms; + // Optional per-row pooling weights stored 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; + + /** + * Constructs state validated by {@link #quantize} or {@link #read}. + */ + private QuantizedEmbeddingMatrix(int rowCount, int dimension, int bits, long seed, + GaussianQuantizer quantizer, double[] scales, byte[] codes, + double[] 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 double[] scales = new double[rowCount]; + final byte[] codes = new byte[rowCount * rowBytes]; + final double[] decodedNorms = new double[rowCount]; + final double[] rotated = new double[paddedDimension]; + final double[] decoded = new double[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 double standardized = rotated[i] * standardize; + final int code = quantizer.encode((float) 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] = 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 largestDecodedMagnitude = 0; + for (int d = 0; d < dimension; d++) { + largestDecodedMagnitude = Math.max(largestDecodedMagnitude, Math.abs(decoded[d])); + } + double decodedAdjustment = 1.0; + if (largestDecodedMagnitude > Float.MAX_VALUE) { + // Quantization can overshoot the finite float range even when every source coordinate is + // finite. Scale all reconstructed coordinates instead of clipping them individually, so + // decoding, stored norms, and rotated-space dot products continue to describe one row. + decodedAdjustment = Float.MAX_VALUE / largestDecodedMagnitude; + scales[row] *= decodedAdjustment; + } + double decodedSumOfSquares = 0; + for (int d = 0; d < dimension; d++) { + final float value = finiteFloat(decoded[d] * decodedAdjustment); + decodedSumOfSquares += (double) value * value; + } + decodedNorms[row] = Math.sqrt(decodedSumOfSquares); + } + return new QuantizedEmbeddingMatrix(rowCount, dimension, bits, seed, quantizer, scales, + codes, decodedNorms, null); + } + + /** + * {@return a copy of this matrix with per-row pooling weights} The file stores the weights + * 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} when absent. 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 has + * 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, double[])}. + * + * @param vector The original-space vector. Must not be {@code null} and must have length + * {@link #dimension()}. + * @return A new double-precision array of length {@link #paddedDimension()} containing the + * rotated vector. + * @throws IllegalArgumentException Thrown if {@code vector} is {@code null} or has the wrong + * length. + */ + public double[] 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 double[] padded = new double[paddedDimension]; + for (int d = 0; d < dimension; d++) { + padded[d] = vector[d]; + } + rotation.rotate(padded); + return padded; + } + + /** + * Rotates an internal double-precision query without narrowing analogy results to floats. + * + * @param query The query, of length {@link #dimension()}. Not modified. + * @return A new rotated query of length {@link #paddedDimension()}. + */ + double[] rotateQuery(double[] query) { + final double[] padded = Arrays.copyOf(query, paddedDimension); + 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, double[])}; 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 double-precision array of length {@link #dimension()} containing the + * original-space vector. + * @throws IllegalArgumentException Thrown if {@code rotated} is {@code null} or has the wrong + * length. + */ + public double[] toOriginal(double[] 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 double[] copy = rotated.clone(); + rotation.inverse(copy); + return Arrays.copyOf(copy, dimension); + } + + /** + * 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, double[] 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 double 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, within floating-point 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, double[] 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 += 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 + * during quantization 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 performs 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 double[] rotated = new double[paddedDimension]; + addRowRotated(row, 1f, rotated); + return toFloatVector(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 double scale : scales) { + data.writeDouble(scale); + } + for (final double decodedNorm : decodedNorms) { + data.writeDouble(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 the decoder. + * + * @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 has an unsupported version, sizes that + * conflict with the file length, an invalid grid, invalid row metadata, or trailing bytes. + * @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 has 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(), e); + } + 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 paddedDimension; + final int rowBytes; + try { + paddedDimension = HadamardRotation.paddedDimension(dimension); + rowBytes = rowByteCount(paddedDimension, bits); + requireStorableSize(rowCount, rowBytes); + } catch (IllegalArgumentException e) { + throw new InvalidFormatException(file + " declares a matrix this reader cannot " + + "store: " + e.getMessage(), e); + } + // Check the minimum length before allocating row-sized arrays. + final long declaredBytes = 28L + 4L * levelCount + (16L + 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 has 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(), e); + } + double maximumLevelMagnitude = 0; + for (final float level : levels) { + maximumLevelMagnitude = Math.max(maximumLevelMagnitude, Math.abs(level)); + } + // Each inverse-transform output can sum paddedDimension decoded coordinates before the + // normalization factor is applied. + final double maximumDecodableScale = + Double.MAX_VALUE / paddedDimension / maximumLevelMagnitude; + final double[] scales = new double[rowCount]; + for (int row = 0; row < rowCount; row++) { + scales[row] = data.readDouble(); + if (!Double.isFinite(scales[row]) || scales[row] < 0 + || scales[row] >= maximumDecodableScale) { + throw new InvalidFormatException(file + " has an invalid scale for row " + row + + ": " + scales[row]); + } + } + // A finite float vector's L2 norm cannot exceed its L1 norm. + final double maximumDecodedNorm = (double) Float.MAX_VALUE * dimension; + final double[] decodedNorms = new double[rowCount]; + for (int row = 0; row < rowCount; row++) { + decodedNorms[row] = data.readDouble(); + if (!Double.isFinite(decodedNorms[row]) || decodedNorms[row] < 0 + || decodedNorms[row] > maximumDecodedNorm) { + throw new InvalidFormatException(file + " has an invalid decoded norm for row " + + row + ": " + decodedNorms[row]); + } + if (scales[row] == 0.0 && decodedNorms[row] > 0.0) { + throw new InvalidFormatException(file + " has zero scale but positive decoded norm for " + + "row " + row + ": " + decodedNorms[row]); + } + } + float[] poolingWeights = null; + final int poolingWeightsFlag = data.readUnsignedByte(); + if (poolingWeightsFlag > 1) { + throw new InvalidFormatException(file + " has invalid pooling-weight flag " + + poolingWeightsFlag + "; expected 0 or 1"); + } + if (poolingWeightsFlag == 1) { + if (declaredBytes + 4L * rowCount > fileSize) { + throw new InvalidFormatException(file + " declares per-row pooling weights, " + + "needing at least " + (declaredBytes + 4L * rowCount) + " bytes in total, but " + + "has 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} A 3-bit code may cross a byte boundary, so this reads a second byte + * when required. + * + * @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)); + } + } + + /** + * Converts a finite double to float, saturating values outside the finite float range. + * + * @param value The value to convert. + * @return The converted value. + */ + private static float finiteFloat(double value) { + if (value > Float.MAX_VALUE) { + return Float.MAX_VALUE; + } + if (value < -Float.MAX_VALUE) { + return -Float.MAX_VALUE; + } + return (float) value; + } + + /** + * Converts a double-precision vector to finite floats. + * + * @param values The vector to convert. + * @return The converted vector. + */ + private float[] toFloatVector(double[] values) { + final float[] result = new float[values.length]; + for (int i = 0; i < values.length; i++) { + result[i] = finiteFloat(values[i]); + } + return result; + } +} 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..b3ef71c1ad --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/QuantizedTableAdapter.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; + +/** + * Adapts a {@link QuantizedEmbeddingMatrix} to {@link EmbeddingTable}. Pooling and scoring use the + * matrix's rotated vector space. + */ +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, double[] sum) { + matrix.addRowRotated(row, weight, sum); + } + + /** {@inheritDoc} */ + @Override + public double[] finishPooling(double[] sum) { + return matrix.toOriginal(sum); + } + + /** {@inheritDoc} */ + @Override + public double[] prepareQuery(double[] query) { + return matrix.rotateQuery(query); + } + + /** {@inheritDoc} */ + @Override + public double dot(int row, double[] 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..2dcd295b2b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/RandomizedPca.java @@ -0,0 +1,617 @@ +/* + * 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 to approximate the dominant subspace of transformer token embeddings. + * + *

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. A fixed seed makes the Java calculation reproducible.

+ * + *

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..ca5a63daa6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsFile.java @@ -0,0 +1,410 @@ +/* + * 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.CharacterCodingException; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +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; +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 between those operations is rejected.

+ * + *

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; + + private static final long MAX_HEADER_SIZE = 100_000_000L; + + /** 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; + + // Array allocation limits are slightly below Integer.MAX_VALUE and vary by JVM; 8 is the + // commonly reserved headroom. + 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) { + throw new InvalidFormatException( + "File " + file + " declares a negative header length: " + headerLength); + } + if (headerLength > MAX_HEADER_SIZE) { + throw new InvalidFormatException("File " + file + " declares a header length of " + + headerLength + " bytes, which exceeds the safetensors limit of " + + MAX_HEADER_SIZE + " bytes"); + } + if (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"); + } + final ByteBuffer headerBytes = ByteBuffer.allocate((int) headerLength); + readFully(channel, headerBytes, HEADER_LENGTH_PREFIX_BYTES, file); + final String headerJson; + try { + headerJson = StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(headerBytes.flip()).toString(); + } catch (CharacterCodingException e) { + throw new InvalidFormatException( + "File " + file + " does not contain a valid UTF-8 header", e); + } + 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"); + } + } + final List tensorsByOffset = new ArrayList<>(parsed.tensors()); + tensorsByOffset.sort(Comparator.comparingLong(TensorInfo::dataOffsetBegin) + .thenComparingLong(TensorInfo::dataOffsetEnd)); + long expectedOffset = 0; + for (final TensorInfo tensor : tensorsByOffset) { + if (tensor.dataOffsetBegin() != expectedOffset) { + throw new InvalidFormatException("File " + file + " tensor '" + tensor.name() + + "' begins at data offset " + tensor.dataOffsetBegin() + " instead of " + + expectedOffset + "; tensor ranges must be contiguous and non-overlapping"); + } + expectedOffset = tensor.dataOffsetEnd(); + } + if (expectedOffset != dataLength) { + throw new InvalidFormatException("File " + file + " declares " + expectedOffset + + " bytes of tensor data but its data section has " + dataLength + " bytes"); + } + 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; + try { + elementCount = info.elementCount(); + } catch (IllegalArgumentException e) { + throw new InvalidFormatException(e.getMessage(), e); + } + 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). It does not guess from the tensor name. + * + * @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..3d5c62ed37 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsHeaderParser.java @@ -0,0 +1,272 @@ +/* + * 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. Input outside that structure is rejected. + */ +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(); + boolean metadataSeen = false; + 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)) { + if (metadataSeen) { + throw cursor.malformed("Field '__metadata__' appears more than once"); + } + metadataSeen = true; + 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; + boolean dtypeSeen = false; + boolean shapeSeen = false; + boolean dataOffsetsSeen = false; + cursor.skipWhitespace(); + while (cursor.peek() != '}') { + cursor.skipWhitespace(); + final String field = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + switch (field) { + case "dtype" -> { + if (dtypeSeen) { + throw cursor.malformed("Tensor '" + name + "' field 'dtype' appears more than once"); + } + dtypeSeen = true; + dtype = cursor.parseString(); + } + case "shape" -> { + if (shapeSeen) { + throw cursor.malformed("Tensor '" + name + "' field 'shape' appears more than once"); + } + shapeSeen = true; + shape = parseIntArray(); + } + case "data_offsets" -> { + if (dataOffsetsSeen) { + throw cursor.malformed( + "Tensor '" + name + "' field 'data_offsets' appears more than once"); + } + dataOffsetsSeen = true; + 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 || dataOffsetEnd < 0) { + throw cursor.malformed("Tensor '" + name + + "' is missing dtype, shape, or data_offsets"); + } + if (dataOffsetEnd < dataOffsetBegin) { + throw cursor.malformed("Tensor '" + name + + "' has data_offsets whose end precedes their beginning"); + } + 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(); + final String value = cursor.parseString(); + if (map.putIfAbsent(key, value) != null) { + throw cursor.malformed("Metadata field '" + key + "' appears more than once"); + } + 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..5ad5cfba9a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/SafetensorsWriter.java @@ -0,0 +1,132 @@ +/* + * 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}, a dimension is less + * than 1, 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) { + throw new IllegalArgumentException("rows must be at least 1, got " + rows); + } + if (cols < 1) { + throw new IllegalArgumentException("cols must be at least 1, got " + cols); + } + if (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..4562675209 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java @@ -0,0 +1,1381 @@ +/* + * 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. To embed text, the model retrieves each tokenized piece's row, applies optional + * weights, mean-pools the rows, and optionally L2-normalizes the result. + * + *

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. Separate-file SentencePiece layouts + * carry a trained {@code .model} file in addition to the JSON vocabulary. In every layout 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. Loading rejects a poolable piece with no matrix row.

+ * + *

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 store its matrix in a {@code model.quantized} file (written by + * the {@code QuantizeModel} tool), which stores the matrix and any per-token pooling weights + * itself. A directory contains exactly one matrix file: the quantized file or the safetensors, + * not both. A directory containing 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. Without the file, all text uses subword tokenization. + * 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 copies 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]; + // Conventional WordPiece special tokens, excluded from neighbor results when present. + 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 stored in the file, 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. 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 separate-file SentencePiece directory carries a trained model + * ({@code sentencepiece.bpe.model}, {@code spiece.model}, or {@code tokenizer.model}) next to + * the JSON vocabulary and uses it for normalization and segmentation.

+ * + *

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(modelDirectory, tokenizerJsonFile, + 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 modelDirectory, Path tokenizerJsonFile, 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 Path quantizedFile = quantizedMatrixFileOrNull(modelDirectory); + final TableAndWeights tableAndWeights; + if (quantizedFile != null) { + tableAndWeights = readQuantizedTable(quantizedFile, vocabulary, terms.size(), + tokenizerJsonFile.toString()); + } else { + final Matrix matrix = readMatrix(vocabulary, terms.size(), + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), tokenizerJsonFile.toString()); + 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(), tokenizer), + 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 making " + + "that choice explicitly"); + } + 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 contains and rejects an ambiguous choice. + * The quantizer writes {@code model.quantized} next to the {@code model.safetensors} it read, + * so deleting one selects the matrix the loader should use. + * + * @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 contains 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. A table distilled with another pooling operation is rejected. + * + * @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"); + } + 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 specify the normalization"); + } + 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 floating-point 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 contains the model's 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 floating-point 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(), tokenizer), + 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 either tensor 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]; + if (dimension < 1) { + throw new InvalidFormatException("Embedding matrix '" + matrixName + "' in " + + safetensorsFile + " has dimension 0"); + } + final float[] embeddings = tensors.readFloats(matrixName); + // Distillation replaces non-finite teacher values with zero before writing. A non-finite + // value therefore marks a corrupt or incompatible file and would contaminate similarities. + 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)) { + final TensorInfo weightsInfo = tensors.tensorInfo(WEIGHTS_TENSOR_NAME); + if (weightsInfo.shape().length != 1) { + throw new InvalidFormatException("Tensor '" + WEIGHTS_TENSOR_NAME + "' in " + + safetensorsFile + " must be 1-D, but its shape is " + + java.util.Arrays.toString(weightsInfo.shape())); + } + 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"); + } + for (int row = 0; row < weights.length; row++) { + if (!Float.isFinite(weights[row])) { + throw new InvalidFormatException("Tensor '" + WEIGHTS_TENSOR_NAME + "' in " + + safetensorsFile + " holds the non-finite value " + weights[row] + " in row " + + row + "; the tensor is corrupt"); + } + } + } + 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; absent tokens do + * not set a row in the mask. + * @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 int row : vocabulary.specialRows()) { + specialRows[row] = true; + } + for (final String special : specialTokens) { + final int row = vocabulary.id(special); + if (row >= 0) { + specialRows[row] = true; + } + } + return specialRows; + } + + /** Marks special rows declared by a self-contained Unigram tokenizer. */ + private static boolean[] specialRows(EmbeddingVocabulary vocabulary, + Set specialTokens, int totalRows, + Model2VecUnigramTokenizer tokenizer) { + final boolean[] specialRows = specialRows(vocabulary, specialTokens, totalRows); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (tokenizer.isUnknown(id) || tokenizer.isControl(id)) { + markSpecialRow(specialRows, vocabulary, tokenizer.idToPiece(id)); + } + } + return specialRows; + } + + /** Marks special rows declared by a separate SentencePiece tokenizer. */ + private static boolean[] specialRows(EmbeddingVocabulary vocabulary, + Set specialTokens, int totalRows, + SentencePieceTokenizer tokenizer) { + final boolean[] specialRows = specialRows(vocabulary, specialTokens, totalRows); + for (int id = 0; id < tokenizer.vocabularySize(); id++) { + if (tokenizer.isUnknown(id) || tokenizer.isControl(id)) { + markSpecialRow(specialRows, vocabulary, tokenizer.idToPiece(id)); + } + } + return specialRows; + } + + /** Marks the row for {@code piece}, if the matrix vocabulary contains it. */ + private static void markSpecialRow(boolean[] specialRows, EmbeddingVocabulary vocabulary, + String piece) { + final int row = vocabulary.id(piece); + if (row >= 0) { + specialRows[row] = true; + } + } + + /** + * {@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 double[] sum = new double[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 double[] pooled = table.finishPooling(sum); + final float[] result = new float[dimension]; + if (normalize) { + double sumOfSquares = 0; + for (final double value : pooled) { + sumOfSquares += value * value; + } + final double norm = Math.max(Math.sqrt(sumOfSquares), NORMALIZE_EPSILON); + for (int d = 0; d < dimension; d++) { + result[d] = (float) (pooled[d] / norm); + } + } else { + for (int d = 0; d < dimension; d++) { + result[d] = finiteFloat(pooled[d]); + } + } + return result; + } + + /** + * Converts a finite double to the nearest finite float value. + * + * @param value The value to convert. + * @return The converted value. + */ + private float finiteFloat(double value) { + if (value > Float.MAX_VALUE) { + return Float.MAX_VALUE; + } + if (value < -Float.MAX_VALUE) { + return -Float.MAX_VALUE; + } + return (float) value; + } + + /** + * 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, the + * whole text follows the subword path. + * + * @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() + + "' without a matrix row after load-time vocabulary validation"); + } + 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. Equal scores retain matrix row + * order. + * + * @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 ?". Equal scores retain matrix row + * order. + * + * @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 double[] target = new double[dimension]; + for (int d = 0; d < dimension; d++) { + target[d] = (double) 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 queryTerms 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[] widened = new double[query.length]; + for (int d = 0; d < query.length; d++) { + widened[d] = query[d]; + } + return nearestNeighbors(widened, topK, sortedExcludedRows); + } + + /** + * 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. + * @return Up to {@code topK} neighbors, most similar first. + */ + private List nearestNeighbors(double[] 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 double[] 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, boundedCosine(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 : boundedCosine(dot / denominator); + } + + /** + * {@return the L2 norm of a vector} + * + * @param vector The vector to measure. + */ + private double norm(double[] vector) { + double sumOfSquares = 0; + for (final double value : vector) { + sumOfSquares += value * value; + } + return Math.sqrt(sumOfSquares); + } + + /** {@return a computed cosine bounded to its mathematical range} */ + private double boundedCosine(double similarity) { + return Math.max(-1.0, Math.min(1.0, similarity)); + } + + /** + * A bounded selection of the {@code k} highest-similarity rows, kept as a min-heap over + * primitive parallel arrays. The root is the lowest-ranked retained candidate, which permits + * one comparison for most scanned rows and avoids allocation 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 (!isWeaker(i, parent)) { + break; + } + swap(parent, i); + i = parent; + } + } else if (isStronger(similarity, row, 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 lowest-ranked retained candidate, the heap root} */ + int minRow() { + return rows[0]; + } + + /** {@return the similarity of the lowest-ranked retained candidate, the heap root} */ + double minSimilarity() { + return similarities[0]; + } + + /** Removes the lowest-ranked retained 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 && isWeaker(left, smallest)) { + smallest = left; + } + if (right < size && isWeaker(right, smallest)) { + smallest = right; + } + if (smallest == i) { + return; + } + swap(i, smallest); + i = smallest; + } + } + + /** + * Tests whether one heap entry ranks below another. For equal scores, the later matrix row is + * lower-ranked so that output ties retain row order. + * + * @param candidate The candidate heap index. + * @param other The heap index to compare against. + * @return {@code true} when {@code candidate} ranks below {@code other}. + */ + private boolean isWeaker(int candidate, int other) { + final int scoreOrder = Double.compare(similarities[candidate], similarities[other]); + return scoreOrder < 0 || scoreOrder == 0 && rows[candidate] > rows[other]; + } + + /** + * Tests whether a candidate ranks above a heap entry. + * + * @param similarity The candidate similarity. + * @param row The candidate matrix row. + * @param other The heap index to compare against. + * @return {@code true} when the candidate ranks above {@code other}. + */ + private boolean isStronger(double similarity, int row, int other) { + final int scoreOrder = Double.compare(similarity, similarities[other]); + return scoreOrder > 0 || scoreOrder == 0 && row < rows[other]; + } + + /** + * 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..619ca69c4a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TeacherTokenizer.java @@ -0,0 +1,1366 @@ +/* + * 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, including the normalizer, + * pre-tokenizer, and Unigram scores. The cleaned {@code tokenizer.json} continues to describe the + * distilled table.

+ */ +final class TeacherTokenizer { + + private static final String HEX_DIGITS = "0123456789abcdef"; + + /** 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, a vocabulary token appears more than once, 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()); + boolean seenModel = false; + boolean seenAddedTokens = false; + boolean seenPostProcessor = false; + boolean seenNormalizer = 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 (seenModel) { + throw cursor.malformed("Field 'model' appears more than once"); + } + seenModel = true; + final ModelSection model = parseModel(cursor); + modelType = model.type(); + tokensById = model.tokensById(); + unkToken = model.unkToken(); + unkId = model.unkId(); + } + case "added_tokens" -> { + if (seenAddedTokens) { + throw cursor.malformed("Field 'added_tokens' appears more than once"); + } + seenAddedTokens = true; + addedContents = parseAddedTokenContents(cursor); + } + case "post_processor" -> { + if (seenPostProcessor) { + throw cursor.malformed("Field 'post_processor' appears more than once"); + } + seenPostProcessor = true; + postProcessor = parsePostProcessor(cursor); + } + case "normalizer" -> { + if (seenNormalizer) { + throw cursor.malformed("Field 'normalizer' appears more than once"); + } + seenNormalizer = true; + 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++) { + final String token = tokensById.get(id); + final Integer previousId = idByToken.putIfAbsent(token, id); + if (previousId != null) { + throw new InvalidFormatException("Vocabulary " + tokenizerJsonFile + " declares token '" + + token + "' more than once, at ids " + previousId + " and " + 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[] {checkedTokenId(postProcessor.clsId(), "cls", tokenizerJsonFile)} + : resolveNames(postProcessor.bosNames(), postProcessor.specialTokenIds(), idByToken, + tokenizerJsonFile); + final int[] eosIds = postProcessor.sepId() != null + ? new int[] {checkedTokenId(postProcessor.sepId(), "sep", tokenizerJsonFile)} + : 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) && !keepSpecial.contains(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; + boolean seenLowerCase = false; + 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 (seenLowerCase) { + throw cursor.malformed("Field 'normalizer.lowercase' appears more than once"); + } + seenLowerCase = true; + if (cursor.consumeLiteral("true")) { + lowerCase = Boolean.TRUE; + } else if (cursor.consumeLiteral("false")) { + lowerCase = Boolean.FALSE; + } else { + throw cursor.malformed("Field 'normalizer.lowercase' must be a boolean"); + } + } 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 has no token id. + */ + 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] = checkedTokenId(specialId, names.get(i), file); + } 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; + } + + /** + * Converts a post-processor token id to the integer representation used by the tokenizer. + * + * @param id The parsed token id. + * @param description The field or token name that supplied the id. + * @param file The source file, for error messages. + * @return The token id as an integer. + * @throws InvalidFormatException Thrown if the id is negative or exceeds the integer range. + */ + private static int checkedTokenId(long id, String description, Path file) + throws InvalidFormatException { + if (id < 0 || id > Integer.MAX_VALUE) { + throw new InvalidFormatException(file + " assigns " + description + " token id " + id + + " outside the supported integer range"); + } + return (int) id; + } + + /** {@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} or one of its elements 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 (int pieceIndex = 0; pieceIndex < pieces.size(); pieceIndex++) { + final String piece = pieces.get(pieceIndex); + if (piece == null) { + throw new IllegalArgumentException("pieces[" + pieceIndex + "] must not be null"); + } + 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 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("\\u00") + .append(HEX_DIGITS.charAt(c >>> 4)) + .append(HEX_DIGITS.charAt(c & 0x0f)); + } 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; + boolean seenType = false; + boolean seenUnkToken = false; + boolean seenUnkId = false; + boolean seenVocabulary = false; + 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" -> { + if (seenType) { + throw cursor.malformed("Field 'model.type' appears more than once"); + } + seenType = true; + type = cursor.parseString(); + } + case "unk_token" -> { + if (seenUnkToken) { + throw cursor.malformed("Field 'model.unk_token' appears more than once"); + } + seenUnkToken = true; + unkToken = cursor.parseString(); + } + case "unk_id" -> { + if (seenUnkId) { + throw cursor.malformed("Field 'model.unk_id' appears more than once"); + } + seenUnkId = true; + if (!cursor.consumeLiteral("null")) { + unkId = cursor.parseLong(); + } + } + case "vocab" -> { + if (seenVocabulary) { + throw cursor.malformed("Field 'model.vocab' appears more than once"); + } + seenVocabulary = true; + 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; + boolean seenContent = false; + 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)) { + if (seenContent) { + throw cursor.malformed("Field 'added_tokens[].content' appears more than once"); + } + seenContent = true; + 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; + boolean seenType = false; + boolean seenSingle = false; + boolean seenSpecialTokens = false; + boolean seenCls = false; + boolean seenSep = 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 "type" -> { + if (seenType) { + throw cursor.malformed("Field 'post_processor.type' appears more than once"); + } + seenType = true; + type = cursor.parseString(); + } + case "single" -> { + if (seenSingle) { + throw cursor.malformed("Field 'post_processor.single' appears more than once"); + } + seenSingle = true; + final List> wrapper = parseTemplate(cursor); + bosNames = wrapper.get(0); + eosNames = wrapper.get(1); + } + case "special_tokens" -> { + if (seenSpecialTokens) { + throw cursor.malformed( + "Field 'post_processor.special_tokens' appears more than once"); + } + seenSpecialTokens = true; + specialTokenIds = parseSpecialTokenIds(cursor); + } + case "cls" -> { + if (seenCls) { + throw cursor.malformed("Field 'post_processor.cls' appears more than once"); + } + seenCls = true; + clsId = parseTokenIdPair(cursor); + } + case "sep" -> { + if (seenSep) { + throw cursor.malformed("Field 'post_processor.sep' appears more than once"); + } + seenSep = true; + 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) { + throw cursor.malformed("post_processor.type is required"); + } + return switch (type) { + case "TemplateProcessing" -> { + if (!seenSingle) { + throw cursor.malformed("post_processor.single is required for TemplateProcessing"); + } + yield new PostProcessor(bosNames, eosNames, null, null, specialTokenIds); + } + case "BertProcessing", "RobertaProcessing" -> { + if (clsId == null) { + throw cursor.malformed("post_processor.cls is required for " + type); + } + if (sepId == null) { + throw cursor.malformed("post_processor.sep is required for " + type); + } + yield 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); + int sequenceCount = 0; + 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)) { + sequenceCount++; + current = eos; + } else { + current.add(part); + } + } + requireSingleSequence(sequenceCount, cursor); + return List.of(bos, eos); + } + cursor.expect('['); + cursor.skipWhitespace(); + List current = bos; + if (cursor.peek() == ']') { + cursor.consume(); + throw cursor.malformed("A single template needs exactly one sequence placeholder; found 0"); + } + 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; + boolean seenId = false; + 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)) { + if (seenId) { + throw cursor.malformed("Field 'post_processor.single[].id' appears more than once"); + } + seenId = true; + 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)) { + if (id == null) { + throw cursor.malformed("SpecialToken template item needs an id"); + } + current.add(id); + } else if ("Sequence".equals(itemType)) { + if (id == null) { + throw cursor.malformed("Sequence template item needs an id"); + } + sequenceCount++; + 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 + "'"); + } + requireSingleSequence(sequenceCount, cursor); + return List.of(bos, eos); + } + + /** + * Verifies that a single-sequence template inserts its input once. + * + * @param count The number of sequence placeholders read. + * @param cursor The cursor used to report the source location. + * @throws InvalidFormatException Thrown unless {@code count} is one. + */ + private static void requireSingleSequence(int count, JsonCursor cursor) + throws InvalidFormatException { + if (count != 1) { + throw cursor.malformed("A single template needs exactly one sequence placeholder; found " + + count); + } + } + + /** + * {@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<>(); + final Set names = new HashSet<>(); + if (cursor.peek() == '}') { + cursor.consume(); + return ids; + } + while (true) { + cursor.skipWhitespace(); + final String name = cursor.parseString(); + cursor.skipWhitespace(); + cursor.expect(':'); + cursor.skipWhitespace(); + if (!names.add(name)) { + throw cursor.malformed("Special token '" + name + "' appears more than once"); + } + cursor.expect('{'); + cursor.skipWhitespace(); + Long id = null; + boolean seenIds = false; + 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)) { + if (seenIds) { + throw cursor.malformed("Field 'post_processor.special_tokens." + + name + ".ids' appears more than once"); + } + seenIds = true; + 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..885e96b8cf --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TensorInfo.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.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}, a dimension or the starting offset is negative, or the ending offset is + * before the starting offset. + */ + 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"); + } + for (int i = 0; i < shape.length; i++) { + if (shape[i] < 0) { + throw new IllegalArgumentException("shape[" + i + "] must not be negative"); + } + } + if (dataOffsetBegin < 0) { + throw new IllegalArgumentException("dataOffsetBegin must not be negative"); + } + if (dataOffsetEnd < dataOffsetBegin) { + throw new IllegalArgumentException( + "dataOffsetEnd must not be less than dataOffsetBegin"); + } + 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..eb6efc88ea --- /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..0dbb2b7913 --- /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 stored after the subword rows. Matching text against + * the table finds the greedily longest term at each word position, so "writ of habeas corpus" + * is preferred to "habeas corpus", which is preferred to 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..a663666622 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/TokenizerJsonVocab.java @@ -0,0 +1,389 @@ +/* + * 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.HashSet; +import java.util.List; +import java.util.Set; + +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. It uses {@link JsonCursor} and rejects input outside the expected structure. + */ +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, boolean special) { + } + + /** The matrix rows and the rows declared as special added tokens. */ + record Result(List rows, Set specialRows) { + } + + /** + * 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 { + return read(file).rows(); + } + + /** + * Reads the pieces and special-token rows of a Unigram {@code tokenizer.json}. + * + * @param file The {@code tokenizer.json} file. Must not be {@code null} and must exist. + * @return The parsed vocabulary information. + * @throws IllegalArgumentException Thrown if {@code file} is {@code null} or missing. + * @throws InvalidFormatException Thrown if the vocabulary layout is malformed. + * @throws IOException Thrown if reading the file fails. + */ + static Result read(Path file) throws IOException { + if (file == null) { + throw new IllegalArgumentException("file must not be null"); + } + if (!Files.isRegularFile(file)) { + throw new IllegalArgumentException("File does not exist or is not a regular file: " + file); + } + final 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 (vocab == null) { + throw new InvalidFormatException(file + " has no model.vocab list; it does not name " + + "the matrix rows"); + } + if (modelType == null) { + throw new InvalidFormatException(file + " has no model.type; only a Unigram tokenizer " + + "maps pieces to matrix rows here"); + } + if (!"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"); + } + 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<>(); + final Set tokenIds = new HashSet<>(); + if (cursor.peek() == ']') { + cursor.consume(); + return tokens; + } + while (true) { + cursor.skipWhitespace(); + final AddedToken token = parseAddedToken(cursor); + if (!tokenIds.add(token.id())) { + throw cursor.malformed("added token id " + token.id() + " occurs more than once"); + } + tokens.add(token); + 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; + Boolean special = 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(); + } + case "special" -> { + if (special != null) { + throw cursor.malformed("Field 'special' appears more than once in an added token"); + } + special = cursor.parseBoolean(); + } + 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, Boolean.TRUE.equals(special)); + } + + /** + * 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 gaps are rejected. + * + * @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 Result overlayAddedTokens(List vocab, + List addedTokens, Path file) + throws InvalidFormatException { + final List byId = new ArrayList<>(addedTokens); + byId.sort(Comparator.comparingLong(AddedToken::id)); + final Set specialRows = new HashSet<>(); + 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"); + } + if (token.special()) { + specialRows.add((int) token.id()); + } + } + return new Result(List.copyOf(vocab), Set.copyOf(specialRows)); + } +} 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..5a5936bc29 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/AssembleModelTool.java @@ -0,0 +1,78 @@ +/* + * 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 and validates a distilled static embedding model directory. For WordPiece models, the + * tool derives {@code vocab.txt} and {@code tokenizer_config.json} from {@code tokenizer.json}. + */ +public class AssembleModelTool extends BasicCmdLineTool { + + /** Command-line parameters accepted by this tool. */ + interface Params extends AssembleModelParams { + } + + /** {@inheritDoc} */ + @Override + public String getShortDescription() { + return "Completes and verifies a distilled static embedding model directory"; + } + + /** {@inheritDoc} */ + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + /** {@inheritDoc} */ + @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..0fd1461897 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/DistillModelTool.java @@ -0,0 +1,111 @@ +/* + * 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.math.BigDecimal; +import java.math.RoundingMode; +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 local or Hugging Face sentence-transformer into a static embedding model. The + * pipeline applies the teacher model, PCA, and Zipf weighting through {@link ModelDistiller}. + */ +public class DistillModelTool extends BasicCmdLineTool { + + /** Command-line parameters accepted by this tool. */ + interface Params extends DistillModelParams { + } + + /** {@inheritDoc} */ + @Override + public String getShortDescription() { + return "Distills a sentence-transformer teacher into a static embedding model"; + } + + /** {@inheritDoc} */ + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + /** {@inheritDoc} */ + @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 " + + formatPercentage(result.explainedVarianceRatio()) + "% of the variance"); + } + + /** + * Formats a variance ratio as a percentage with one decimal place. + * + * @param ratio The variance ratio. + * @return The percentage using a decimal point independently of the default locale. + */ + private String formatPercentage(double ratio) { + return BigDecimal.valueOf(ratio).movePointRight(2) + .setScale(1, RoundingMode.HALF_UP).toPlainString(); + } + + /** + * 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..25e83c2c78 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/QuantizeModelTool.java @@ -0,0 +1,76 @@ +/* + * 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; +import opennlp.tools.util.InvalidFormatException; + +/** + * 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. Delete {@code model.safetensors} before loading the quantized + * deployment; a model directory containing both matrix files is rejected. + */ +public class QuantizeModelTool extends BasicCmdLineTool { + + /** Command-line parameters accepted by this tool. */ + interface Params extends QuantizeModelParams { + } + + /** {@inheritDoc} */ + @Override + public String getShortDescription() { + return "Quantizes a static embedding model directory's matrix to 2-4 bits per dimension"; + } + + /** {@inheritDoc} */ + @Override + public String getHelp() { + return getBasicHelp(Params.class); + } + + /** {@inheritDoc} */ + @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 | InvalidFormatException e) { + throw new TerminateToolException(1, e.getMessage(), e); + } 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() ? ", including 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/test/java/opennlp/embeddings/EmbeddingTestFixtures.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingTestFixtures.java new file mode 100644 index 0000000000..96e7990e6d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/EmbeddingTestFixtures.java @@ -0,0 +1,484 @@ +/* + * 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.Base64; +import java.util.List; + +import opennlp.embeddings.StaticEmbeddingModel.Casing; +import opennlp.embeddings.StaticEmbeddingModel.Normalization; +import opennlp.subword.sentencepiece.SentencePieceTokenizer; + +/** + * Fixtures shared by tests in this module: deterministic WordPiece tables and JSON string + * quoting for {@code tokenizer.json} fixtures. + */ +final class EmbeddingTestFixtures { + + /** Lookup graph generated by {@code dev/embeddings/generate_test_teacher.py}. */ + private static final String LOOKUP_TEACHER_ONNX = + "CAg6wAIKOgoFdGFibGUKCWlucHV0X2lkcxIRbGFzdF9oaWRkZW5fc3RhdGUiBkdhdGhlcioLCgRh" + + "eGlzGACgAQISE0xPT0tVUF9URUFDSEVSX09OTlgqkAEICAgEEAEigAEAAAAAAAAAAAAAAAAAAAAA" + + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgD8AAAAAAAAAAAAAAAAAAIC/AABAQAAAAAAA" + + "AAAAAAAAAAAAAEAAAIA/AAAAAAAAAAAAAIC/AAAAQAAAAAAAAAAAAACAvwAAAMAAAAAAAAAAAEIF" + + "dGFibGVaJgoJaW5wdXRfaWRzEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9rZW5zYjIKEWxhc3RfaGlk" + + "ZGVuX3N0YXRlEh0KGwgBEhcKBxIFYmF0Y2gKCBIGdG9rZW5zCgIIBEIECgAQDQ=="; + + /** Graph with batch-dependent vector length, generated by the same script. */ + private static final String VARIABLE_DIMENSION_ONNX = + "CAg6twMKJgoJaW5wdXRfaWRzEghhc19mbG9hdCIEQ2FzdCoJCgJ0bxgBoAECCiMKCGFzX2Zsb2F0" + + "CgRheGVzEgZzdGF0ZXMiCVVuc3F1ZWV6ZQofCglpbnB1dF9pZHMSC2lucHV0X3NoYXBlIgVTaGFw" + + "ZQo6CgtpbnB1dF9zaGFwZQoKYmF0Y2hfYXhpcxIKYmF0Y2hfc2l6ZSIGR2F0aGVyKgsKBGF4aXMY" + + "AKABAgowCgRvbmVzCgpiYXRjaF9zaXplEgdyZXBlYXRzIgZDb25jYXQqCwoEYXhpcxgAoAECCioK" + + "BnN0YXRlcwoHcmVwZWF0cxIRbGFzdF9oaWRkZW5fc3RhdGUiBFRpbGUSF1ZBUklBQkxFX0RJTUVO" + + "U0lPTl9PTk5YKg0IARAHOgECQgRheGVzKhMIARAHOgEAQgpiYXRjaF9heGlzKg4IAhAHOgIBAUIE" + + "b25lc1omCglpbnB1dF9pZHMSGQoXCAcSEwoHEgViYXRjaAoIEgZ0b2tlbnNiOAoRbGFzdF9oaWRk" + + "ZW5fc3RhdGUSIwohCAESHQoHEgViYXRjaAoIEgZ0b2tlbnMKCBIGaGlkZGVuQgQKABAN"; + + /** A deterministic ONNX graph mapping token ids to three-dimensional hidden states. */ + private static final String TINY_TEACHER_ONNX = + "CAg66gIKJwoJaW5wdXRfaWRzEglpZHNfZmxvYXQiBENhc3QqCQoCdG8YAaABAgokCglpZHNf" + + "ZmxvYXQKBGF4ZXMSBmlkc18zZCIJVW5zcXVlZXplCiYKBmlkc18zZAoBdxIRbGFzdF9oaWRk" + + "ZW5fc3RhdGUiBk1hdE11bBIMdGlueS12ZWN0b3JzKhQIARAHQgRheGVzSggCAAAAAAAAACoX" + + "CAEIAxABQgF3SgwAAAA/AACAvwAAAEBaJgoJaW5wdXRfaWRzEhkKFwgHEhMKBxIFYmF0Y2gK" + + "CBIGdG9rZW5zWisKDmF0dGVudGlvbl9tYXNrEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9rZW5z" + + "WisKDnRva2VuX3R5cGVfaWRzEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9rZW5zYjIKEWxhc3Rf" + + "aGlkZGVuX3N0YXRlEh0KGwgBEhcKBxIFYmF0Y2gKCBIGdG9rZW5zCgIIA0IECgAQDQ=="; + + /** A tiny graph whose only input is {@code input_ids}. */ + private static final String INPUT_IDS_ONLY_ONNX = + "CAg6lwIKJwoJaW5wdXRfaWRzEglpZHNfZmxvYXQiBENhc3QqCQoCdG8YAaABAgokCglpZHNf" + + "ZmxvYXQKBGF4ZXMSBmlkc18zZCIJVW5zcXVlZXplCiYKBmlkc18zZAoBdxIRbGFzdF9oaWRk" + + "ZW5fc3RhdGUiBk1hdE11bBITdGlueS1pbnB1dC1jb250cmFjdCoUCAEQB0IEYXhlc0oIAgAA" + + "AAAAAAAqFwgBCAMQAUIBd0oMAAAAPwAAgL8AAABAWiYKCWlucHV0X2lkcxIZChcIBxITCgcS" + + "BWJhdGNoCggSBnRva2Vuc2IyChFsYXN0X2hpZGRlbl9zdGF0ZRIdChsIARIXCgcSBWJhdGNo" + + "CggSBnRva2VucwoCCANCBAoAEA0="; + + /** A tiny graph with an input the encoder cannot supply. */ + private static final String UNSUPPORTED_INPUT_ONNX = + "CAg6wgIKJwoJaW5wdXRfaWRzEglpZHNfZmxvYXQiBENhc3QqCQoCdG8YAaABAgokCglpZHNf" + + "ZmxvYXQKBGF4ZXMSBmlkc18zZCIJVW5zcXVlZXplCiYKBmlkc18zZAoBdxIRbGFzdF9oaWRk" + + "ZW5fc3RhdGUiBk1hdE11bBITdGlueS1pbnB1dC1jb250cmFjdCoUCAEQB0IEYXhlc0oIAgAA" + + "AAAAAAAqFwgBCAMQAUIBd0oMAAAAPwAAgL8AAABAWiYKCWlucHV0X2lkcxIZChcIBxITCgcS" + + "BWJhdGNoCggSBnRva2Vuc1opCgxwb3NpdGlvbl9pZHMSGQoXCAcSEwoHEgViYXRjaAoIEgZ0" + + "b2tlbnNiMgoRbGFzdF9oaWRkZW5fc3RhdGUSHQobCAESFwoHEgViYXRjaAoIEgZ0b2tlbnMK" + + "AggDQgQKABAN"; + + /** A tiny graph that returns {@link Float#MAX_VALUE} at every token position. */ + private static final String MAX_FLOAT_ONNX = + "CAg6igIKJwoJaW5wdXRfaWRzEglpZHNfZmxvYXQiBENhc3QqCQoCdG8YAaABAgokCglpZHNf" + + "ZmxvYXQKBGF4ZXMSBmlkc18zZCIJVW5zcXVlZXplCiYKBmlkc18zZAoBdxIRbGFzdF9oaWRk" + + "ZW5fc3RhdGUiBk1hdE11bBIOdGlueS1tYXgtZmxvYXQqFAgBEAdCBGF4ZXNKCAIAAAAAAAAA" + + "Kg8IAQgBEAFCAXdKBP//f39aJgoJaW5wdXRfaWRzEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9r" + + "ZW5zYjIKEWxhc3RfaGlkZGVuX3N0YXRlEh0KGwgBEhcKBxIFYmF0Y2gKCBIGdG9rZW5zCgII" + + "AUIECgAQDQ=="; + + /** A graph with INT32 {@code input_ids}. */ + private static final String INT32_INPUT_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3OsMBCiYKCWlucHV0X2lkcxIIYXNfZmxvYXQiBENhc3Qq" + + "CQoCdG8YAaABAgovCghhc19mbG9hdAoFYXhlczISEWxhc3RfaGlkZGVuX3N0YXRlIglV" + + "bnNxdWVlemUSC2ludDMyLWlucHV0Kg4IARAHOgECQgVheGVzMlomCglpbnB1dF9pZHMS" + + "GQoXCAYSEwoHEgViYXRjaAoIEgZ0b2tlbnNiIwoRbGFzdF9oaWRkZW5fc3RhdGUSDgoM" + + "CAESCAoACgAKAggBQgQKABAN"; + + /** A graph whose {@code input_ids} is rank one instead of batch by position. */ + private static final String RANK_ONE_INPUT_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3Or0BCiYKCWlucHV0X2lkcxIIYXNfZmxvYXQiBENhc3Qq" + + "CQoCdG8YAaABAgowCghhc19mbG9hdAoGYXhlczEyEhFsYXN0X2hpZGRlbl9zdGF0ZSIJ" + + "VW5zcXVlZXplEgtyYW5rMS1pbnB1dCoQCAIQBzoCAQJCBmF4ZXMxMlodCglpbnB1dF9p" + + "ZHMSEAoOCAcSCgoIEgZ0b2tlbnNiIwoRbGFzdF9oaWRkZW5fc3RhdGUSDgoMCAESCAoA" + + "CgAKAggBQgQKABAN"; + + /** A graph with an INT32 {@code attention_mask}. */ + private static final String INT32_ATTENTION_MASK_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3Ou8BCiYKCWlucHV0X2lkcxIIYXNfZmxvYXQiBENhc3Qq" + + "CQoCdG8YAaABAgovCghhc19mbG9hdAoFYXhlczISEWxhc3RfaGlkZGVuX3N0YXRlIglV" + + "bnNxdWVlemUSCmludDMyLW1hc2sqDggBEAc6AQJCBWF4ZXMyWiYKCWlucHV0X2lkcxIZ" + + "ChcIBxITCgcSBWJhdGNoCggSBnRva2Vuc1orCg5hdHRlbnRpb25fbWFzaxIZChcIBhIT" + + "CgcSBWJhdGNoCggSBnRva2Vuc2IjChFsYXN0X2hpZGRlbl9zdGF0ZRIOCgwIARIICgAK" + + "AAoCCAFCBAoAEA0="; + + /** A graph that declares a decoy rank-three float output before {@code last_hidden_state}. */ + private static final String MULTIPLE_OUTPUTS_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3OpYCCiYKCWlucHV0X2lkcxIIYXNfZmxvYXQiBENhc3Qq" + + "CQoCdG8YAaABAgovCghhc19mbG9hdAoFYXhlczISEWxhc3RfaGlkZGVuX3N0YXRlIglV" + + "bnNxdWVlemUKJAoRbGFzdF9oaWRkZW5fc3RhdGUKA3RlbhIFZGVjb3kiA011bBIQbXVs" + + "dGlwbGUtb3V0cHV0cyoOCAEQBzoBAkIFYXhlczIqDRABIgQAACBBQgN0ZW5aJgoJaW5w" + + "dXRfaWRzEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9rZW5zYhcKBWRlY295Eg4KDAgBEggKAAoA" + + "CgIIAWIjChFsYXN0X2hpZGRlbl9zdGF0ZRIOCgwIARIICgAKAAoCCAFCBAoAEA0="; + + /** A graph whose {@code input_ids} has the unsupported FLOAT element type. */ + private static final String FLOAT_INPUT_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3OpwBCjAKCWlucHV0X2lkcwoFYXhlczISEWxhc3RfaGlk" + + "ZGVuX3N0YXRlIglVbnNxdWVlemUSC2Zsb2F0LWlucHV0Kg4IARAHOgECQgVheGVzMlom" + + "CglpbnB1dF9pZHMSGQoXCAESEwoHEgViYXRjaAoIEgZ0b2tlbnNiIwoRbGFzdF9oaWRk" + + "ZW5fc3RhdGUSDgoMCAESCAoACgAKAggBQgQKABAN"; + + /** A graph whose fixed output does not follow the input batch and sequence dimensions. */ + private static final String FIXED_OUTPUT_ONNX = + "CAgSDm9wZW5ubHAtcmV2aWV3OqEBCkASEWxhc3RfaGlkZGVuX3N0YXRlIghDb25zdGFu" + + "dCohCgV2YWx1ZSoVCAEIAQgBEAEiBAAA4EBCBWZpeGVkoAEEEgxmaXhlZC1vdXRwdXRa" + + "JgoJaW5wdXRfaWRzEhkKFwgHEhMKBxIFYmF0Y2gKCBIGdG9rZW5zYicKEWxhc3RfaGlk" + + "ZGVuX3N0YXRlEhIKEAgBEgwKAggBCgIIAQoCCAFCBAoAEA0="; + + /** 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 differ, + * so pairwise cosine similarities are not all 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 + }; + + /** Tokens used by the semantic-search example; the list index is the matrix row. */ + private static final List SEARCH_VOCABULARY = List.of( + "[CLS]", "[SEP]", "[UNK]", + "home", "espresso", "machine", "how", "do", "i", "brew", "at", + "the", "history", "of", "tea", "in", "east", "asia", + "best", "grinders", "for", "pour", "over", "coffee"); + + /** + * Search rows with three directions: espresso brewing, tea history, and coffee equipment. + * The query uses the first direction, so the example has a deterministic ranking. + */ + private static final float[][] SEARCH_ROWS = { + {0f, 0f}, {0f, 0f}, {0f, 0f}, + {1f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f}, + {0f, 1f}, {0f, 1f}, {0f, 1f}, {0f, 1f}, {0f, 1f}, {0f, 1f}, {0f, 1f}, + {0.6f, 0.8f}, {0.6f, 0.8f}, {0.6f, 0.8f}, {0.6f, 0.8f}, {0.6f, 0.8f}, + {0.6f, 0.8f} + }; + + /** Not instantiable. */ + private EmbeddingTestFixtures() { + } + + /** + * Writes the lookup graph used by the distillation example. + * + * @param directory The destination directory. + * @return The graph file. + * @throws IOException Thrown if writing fails. + */ + static Path writeLookupTeacherOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, LOOKUP_TEACHER_ONNX); + } + + /** + * Writes a graph with a vector length equal to the input batch size. + * + * @param directory The destination directory. + * @return The graph file. + * @throws IOException Thrown if writing fails. + */ + static Path writeVariableDimensionOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, VARIABLE_DIMENSION_ONNX); + } + + /** + * Writes the deterministic test ONNX graph. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeTinyOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, TINY_TEACHER_ONNX); + } + + /** + * Writes a graph that declares only {@code input_ids}. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeInputIdsOnlyOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, INPUT_IDS_ONLY_ONNX); + } + + /** + * Writes a graph that also requires {@code position_ids}. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeUnsupportedInputOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, UNSUPPORTED_INPUT_ONNX); + } + + /** + * Writes a graph whose finite hidden states expose overflow in float accumulation. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeMaxFloatOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, MAX_FLOAT_ONNX); + } + + /** + * Writes a graph whose {@code input_ids} element type is INT32. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeInt32InputOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, INT32_INPUT_ONNX); + } + + /** + * Writes a graph whose {@code input_ids} is rank one. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeRankOneInputOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, RANK_ONE_INPUT_ONNX); + } + + /** + * Writes a graph whose {@code attention_mask} element type is INT32. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeInt32AttentionMaskOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, INT32_ATTENTION_MASK_ONNX); + } + + /** + * Writes a graph with two rank-three float outputs, including {@code last_hidden_state}. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeMultipleOutputsOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, MULTIPLE_OUTPUTS_ONNX); + } + + /** + * Writes a graph whose {@code input_ids} element type is FLOAT. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeFloatInputOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, FLOAT_INPUT_ONNX); + } + + /** + * Writes a graph whose output is always shaped {@code [1][1][1]}. + * + * @param directory The directory in which to create {@code model.onnx}. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + static Path writeFixedOutputOnnxModel(Path directory) throws IOException { + return writeOnnxModel(directory, FIXED_OUTPUT_ONNX); + } + + /** + * Decodes an ONNX fixture into {@code model.onnx}. + * + * @param directory The destination directory. + * @param encodedModel The base64-encoded graph. + * @return The created graph file. + * @throws IOException Thrown if the graph cannot be written. + */ + private static Path writeOnnxModel(Path directory, String encodedModel) throws IOException { + final Path file = directory.resolve("model.onnx"); + Files.write(file, Base64.getDecoder().decode(encodedModel)); + return file; + } + + /** + * 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 complete WordPiece model used by the semantic-search example. + * + * @param dir The directory to write the model files into. + * @throws IOException Thrown if writing a fixture file fails. + */ + static void writeSearchDirectory(Path dir) throws IOException { + Files.write(dir.resolve("vocab.txt"), SEARCH_VOCABULARY); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", SEARCH_ROWS)); + Files.writeString(dir.resolve("config.json"), + "{\"model_type\":\"model2vec\",\"normalize\":true}"); + 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..b260ac2444 --- /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, duplicate entries are rejected, and lookup uses + * {@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 testDuplicateTokenReportsBothLines() { + 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..7d1b3118c5 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/GaussianQuantizerTest.java @@ -0,0 +1,178 @@ +/* + * 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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests the Gaussian Lloyd-Max grids and nearest-level encoding. + */ +class GaussianQuantizerTest { + + /** + * Requires adjacent float levels to encode to their corresponding indices. + * + * @param lowerBits The float bits of the lower level. + */ + @ParameterizedTest + @ValueSource(ints = {0x3f800001, 0xbf800001, 0x00000001, 0x80000001}) + void testAdjacentLevelsEncodeToTheirIndices(int lowerBits) { + final float lower = Float.intBitsToFloat(lowerBits); + final float upper = Math.nextUp(lower); + final GaussianQuantizer quantizer = GaussianQuantizer.fromLevels( + new float[] {-2f, lower, upper, 2f}); + + assertEquals(1, quantizer.encode(lower)); + assertEquals(2, quantizer.encode(upper)); + } + + /** + * Tests float values near the exact midpoint between computed grid levels. + * + * @param bits The quantization width. + */ + @ParameterizedTest + @ValueSource(ints = {2, 3, 4}) + void testComputedGridMidpointUsesNearestLevel(int bits) { + final GaussianQuantizer quantizer = GaussianQuantizer.forBits(bits); + for (int upper = 1; upper < quantizer.levelCount(); upper++) { + final double midpoint = ((double) quantizer.level(upper - 1) + quantizer.level(upper)) / 2; + final float value = (float) midpoint; + final int expected = value <= midpoint ? upper - 1 : upper; + assertEquals(expected, quantizer.encode(value), "upper level " + upper); + } + } + + /** Selects the lower level at a representable midpoint. */ + @Test + void testExactMidpointSelectsLowerLevel() { + final GaussianQuantizer quantizer = GaussianQuantizer.fromLevels( + new float[] {-3f, -1f, 1f, 3f}); + + assertEquals(0, quantizer.encode(-2f)); + assertEquals(1, quantizer.encode(0f)); + assertEquals(2, quantizer.encode(2f)); + } + + @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 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 testFiniteGridMidpointsDoNotOverflow() { + final float maximum = Float.MAX_VALUE; + final GaussianQuantizer quantizer = GaussianQuantizer.fromLevels(new float[] { + maximum / 4f, maximum / 2f, maximum * 0.75f, maximum + }); + + assertEquals(3, quantizer.encode(maximum)); + } + + @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 testUnsupportedBitWidthsAreRejected() { + 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..2a6d4c0267 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HadamardRotationTest.java @@ -0,0 +1,142 @@ +/* + * 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; + +/** + * Tests the randomized Hadamard rotation and its inverse. + */ +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 double[] a = randomPadded(rotation, random); + final double[] 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 double[] vector = randomPadded(rotation, random); + final double[] original = vector.clone(); + rotation.rotate(vector); + rotation.inverse(vector); + assertArrayEquals(original, vector, 1e-12); + } + + @Test + void testSameSeedSameRotationDifferentSeedDifferentRotation() { + final Random random = new Random(44); + final double[] input = randomPadded(new HadamardRotation(64, 5L), random); + final double[] first = input.clone(); + final double[] second = input.clone(); + final double[] 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 equal 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 double[] oneHot = new double[rotation.paddedDimension()]; + oneHot[3] = 1.0; + rotation.rotate(oneHot); + final double expectedMagnitude = 1.0 / Math.sqrt(rotation.paddedDimension()); + for (final double 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 double[300])); + assertThrows(IllegalArgumentException.class, () -> rotation.inverse(new double[511])); + assertThrows(IllegalArgumentException.class, () -> new HadamardRotation(0, 1L)); + } + + @Test + void testDimensionOneIsTheIdentityUpToSign() { + final HadamardRotation rotation = new HadamardRotation(1, 123L); + final double[] vector = new double[] {2.5}; + rotation.rotate(vector); + assertEquals(2.5, Math.abs(vector[0]), 1e-12); + rotation.inverse(vector); + assertEquals(2.5, vector[0], 1e-12); + } + + private double[] randomPadded(HadamardRotation rotation, Random random) { + final double[] vector = new double[rotation.paddedDimension()]; + for (int i = 0; i < vector.length; i++) { + vector[i] = random.nextGaussian(); + } + return vector; + } + + private double norm(double[] vector) { + return Math.sqrt(dot(vector, vector)); + } + + private double dot(double[] a, double[] b) { + double dot = 0; + for (int i = 0; i < a.length; i++) { + dot += 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..b8a622c7de --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/HuggingFaceModelCacheTest.java @@ -0,0 +1,1005 @@ +/* + * 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.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +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 accepted only when 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 testRejectsNullTeacher() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(null, null)); + assertTrue(e.getMessage().contains("must not be null"), e.getMessage()); + } + + @Test + void testRejectsNullHubBase(@TempDir Path cacheRoot) { + assertEquals("hubBase must not be null", assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, null, cacheRoot, null)).getMessage()); + } + + @Test + void testRejectsNullCacheRoot() { + 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@/main", + "BAAI/bge-m3@main/", "BAAI/bge-m3@refs//1", "BAAI/bge-m3@a b", + "BAAI/bge-m3@main@main", "../bge-m3", "BAAI/..", "BAAI/bge-m3@..", + "BAAI/bge-m3@refs/../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()); + } + + /** + * A relative path with a {@code ..} segment is ambiguous as a teacher reference: Windows + * collapses {@code ..} lexically without checking that the segment before it exists, so a + * misspelled hub id such as {@code BAAI/..} would silently name the working directory there + * while POSIX reports it as nonexistent. The reference must be rejected on every platform, + * even when it resolves to an existing directory ({@code src} always exists below the module + * the tests run from). + */ + @Test + void testRelativePathWithParentSegmentIsRejected(@TempDir Path cacheRoot) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> HuggingFaceModelCache.resolve("src/..", hub.base(), cacheRoot, null)); + assertTrue(e.getMessage().contains("org/model"), e.getMessage()); + assertTrue(hub.requests.isEmpty(), hub.requests.toString()); + } + + /** + * 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 testAMissingEtagIsRejected(@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 testAMalformedEtagIsRejected(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()); + } + + @Test + void testAnEtagWithAnEmbeddedQuoteIsRejected(@TempDir Path cacheRoot) { + serveTeacher(); + final int middle = TOKENIZER_BLOB_SHA1.length() / 2; + final String malformed = TOKENIZER_BLOB_SHA1.substring(0, middle) + '"' + + TOKENIZER_BLOB_SHA1.substring(middle); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, malformed); + + 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; another error status must remain visible to the caller. */ + @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)); + } + + @Test + void testRejectsCommitHeaderSuppliedOnlyByRedirectTarget(@TempDir Path cacheRoot) { + hub.reply(resolvePath(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON), + new Reply(302, null, null, "/cdn/tokenizer", null)); + hub.reply("/cdn/tokenizer", + new Reply(200, COMMIT, quoted(TOKENIZER_BLOB_SHA1), null, TOKENIZER)); + hub.serve(COMMIT, ModelFileNames.TOKENIZER_JSON, TOKENIZER, quoted(TOKENIZER_BLOB_SHA1)); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256)); + + final IOException error = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(error.getMessage().contains("x-repo-commit"), error.getMessage()); + } + + @Test + void testRejectsChecksumHeaderSuppliedOnlyByRedirectTarget(@TempDir Path cacheRoot) { + hub.serve(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, TOKENIZER, + quoted(TOKENIZER_BLOB_SHA1)); + hub.reply(resolvePath(COMMIT, ModelFileNames.TOKENIZER_JSON), + new Reply(302, COMMIT, null, "/cdn/tokenizer", null)); + hub.reply("/cdn/tokenizer", + new Reply(200, null, quoted(TOKENIZER_BLOB_SHA1), null, TOKENIZER)); + hub.serve(COMMIT, ModelFileNames.ONNX_MODEL, ONNX, quoted(ONNX_SHA256)); + + final IOException error = assertThrows(IOException.class, + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + + assertTrue(error.getMessage().contains("x-linked-etag"), error.getMessage()); + } + + /** 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))); + } + + /** Concurrent writers for one teacher must not publish files from different revisions. */ + @Test + void testConcurrentResolutionsOfOneTeacherAreSerialized(@TempDir Path cacheRoot) + throws Exception { + serveTeacher(); + final CountDownLatch firstAtGraph = new CountDownLatch(1); + final CountDownLatch releaseFirst = new CountDownLatch(1); + hub.gate(COMMIT, ModelFileNames.ONNX_MODEL, firstAtGraph, releaseFirst); + + final ExecutorService executor = Executors.newFixedThreadPool(2); + try (Hub movedHub = new Hub()) { + movedHub.serve(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, BIG_TOKENIZER, + quoted(BIG_TOKENIZER_BLOB_SHA1), OTHER_COMMIT); + movedHub.serve(OTHER_COMMIT, ModelFileNames.TOKENIZER_JSON, BIG_TOKENIZER, + quoted(BIG_TOKENIZER_BLOB_SHA1), OTHER_COMMIT); + movedHub.serve(OTHER_COMMIT, ModelFileNames.ONNX_MODEL, BIG_ONNX, + quoted(BIG_ONNX_SHA256), OTHER_COMMIT); + final CountDownLatch secondReachedHub = new CountDownLatch(1); + movedHub.signalOnRequest(DEFAULT_REF, ModelFileNames.TOKENIZER_JSON, secondReachedHub); + + final Future first = executor.submit( + () -> HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null)); + assertTrue(firstAtGraph.await(5, TimeUnit.SECONDS), "first download did not reach the graph"); + final Future second = executor.submit( + () -> HuggingFaceModelCache.resolve(MODEL_ID, movedHub.base(), cacheRoot, null)); + try { + assertFalse(secondReachedHub.await(1, TimeUnit.SECONDS), + "a second writer contacted the hub while the first held the cache"); + } finally { + releaseFirst.countDown(); + } + + assertEquals(first.get(5, TimeUnit.SECONDS), second.get(5, TimeUnit.SECONDS)); + } finally { + releaseFirst.countDown(); + executor.shutdownNow(); + } + } + + @Test + void testAMarkedCacheMissingAnOptionalFileIsDownloadedAgain(@TempDir Path cacheRoot) + throws IOException { + serveTeacher(); + hub.serve(COMMIT, SENTENCEPIECE_MODEL, SENTENCEPIECE, quoted(SENTENCEPIECE_BLOB_SHA1)); + final Path cache = HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + Files.delete(cache.resolve(SENTENCEPIECE_MODEL)); + hub.requests.clear(); + + HuggingFaceModelCache.resolve(MODEL_ID, hub.base(), cacheRoot, null); + + assertArrayEquals(SENTENCEPIECE, Files.readAllBytes(cache.resolve(SENTENCEPIECE_MODEL))); + assertFalse(hub.requests.isEmpty(), "the incomplete snapshot must be downloaded again"); + } + + /** 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), + "an incomplete directory must not be returned"); + } + + /** 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)); + } + + /** + * An incomplete cache has no revision marker. Each cached file is therefore checked against the + * requested revision and reused only when its digest matches. + */ + @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 testARevisionWithSlashesIsEncodedAsOnePathSegment(@TempDir Path cacheRoot) + throws IOException { + final String requestPath = "/" + MODEL_ID + "/resolve/refs%2Fpr%2F1/" + + ModelFileNames.TOKENIZER_JSON; + hub.reply(requestPath, + new Reply(200, COMMIT, quoted(TOKENIZER_BLOB_SHA1), null, TOKENIZER)); + 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)); + assertTrue(hub.requests.contains(requestPath), hub.requests.toString()); + } + + @Test + void testARequestedCommitTheHubResolvesElsewhereIsRejected(@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 testARevisionThatCannotBePinnedIsRejected(@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 testAModelTheHubDoesNotHaveIsRejected(@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) { + } + + /** Coordinates one response with a concurrent test. */ + private record Gate(CountDownLatch entered, CountDownLatch release) { + } + + /** + * 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 Map arrivals = new ConcurrentHashMap<>(); + private final Map gates = 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); + } + + /** Records when the requested file reaches this hub. */ + private void signalOnRequest(String revision, String file, CountDownLatch arrival) { + arrivals.put(resolvePath(revision, file), arrival); + } + + /** Pauses the requested file until {@code release} is opened. */ + private void gate(String revision, String file, CountDownLatch entered, + CountDownLatch release) { + gates.put(resolvePath(revision, file), new Gate(entered, release)); + } + + /** + * 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().getRawPath(); + requests.add(path); + final CountDownLatch arrival = arrivals.get(path); + if (arrival != null) { + arrival.countDown(); + } + final Gate gate = gates.get(path); + if (gate != null) { + gate.entered().countDown(); + try { + if (!gate.release().await(5, TimeUnit.SECONDS)) { + throw new IOException("Timed out waiting to release " + path); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while waiting to release " + path, e); + } + } + 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 flattening reference characters does not create cache-name collisions. */ + @Test + void testDistinctTeachersDoNotShareACacheDirectory() { + final Set names = new HashSet<>(); + for (final String teacher : 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/JsonCursorTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/JsonCursorTest.java new file mode 100644 index 0000000000..15440cde8f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/JsonCursorTest.java @@ -0,0 +1,38 @@ +/* + * 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 org.junit.jupiter.api.Test; + +import opennlp.tools.util.InvalidFormatException; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class JsonCursorTest { + + @Test + void testRejectsExcessiveNesting() { + final String json = "[".repeat(129) + "null" + "]".repeat(129); + final JsonCursor cursor = new JsonCursor(json, "test input"); + + final InvalidFormatException error = + assertThrows(InvalidFormatException.class, cursor::skipValue); + + assertTrue(error.getMessage().contains("nesting depth"), error.getMessage()); + } +} 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..614eb52b4d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/Model2VecUnigramTokenizerTest.java @@ -0,0 +1,231 @@ +/* + * 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 java.util.stream.Stream; + +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.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class Model2VecUnigramTokenizerTest { + + private static final String NORMALIZER = + "\"normalizer\":{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}"; + private static final String PRE_TOKENIZER = + "\"pre_tokenizer\":{\"type\":\"Metaspace\",\"replacement\":\"▁\"," + + "\"prepend_scheme\":\"always\",\"split\":false}"; + private static final String MODEL = + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0,\"byte_fallback\":false," + + "\"vocab\":[[\"\",0.0],[\"▁a\",-1.0]]}"; + private static final String TOKENIZER = "{" + NORMALIZER + "," + PRE_TOKENIZER + "," + + MODEL + "}"; + + /** Verifies that normalization-added characters do not extend a piece past the source text. */ + @Test + void testPieceOffsetsReferToTheOriginalText(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, TOKENIZER); + + final Model2VecUnigramTokenizer loaded = Model2VecUnigramTokenizer.load(tokenizer); + + assertEquals(1, loaded.encode("a").getFirst().end()); + } + + /** Verifies alignment through literal insertion, marker collapse, and marker stripping. */ + @Test + void testPieceOffsetsSurvivePostNormalization(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + final String sequence = "\"normalizer\":{\"type\":\"Sequence\",\"normalizers\":[" + + "{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}," + + "{\"type\":\"Replace\",\"pattern\":{\"String\":\"▁\"}," + + "\"content\":\" ▁ \"}," + + "{\"type\":\"Replace\",\"pattern\":{\"Regex\":\"\\\\s+\"}," + + "\"content\":\" \"}," + + "{\"type\":\"Replace\",\"pattern\":{\"String\":\"a\"}," + + "\"content\":\" a \"}," + + "{\"type\":\"Strip\",\"strip_left\":false,\"strip_right\":true}]}"; + Files.writeString(tokenizer, "{" + sequence + "," + PRE_TOKENIZER + "," + MODEL + "}"); + + final Model2VecUnigramTokenizer loaded = Model2VecUnigramTokenizer.load(tokenizer); + + assertEquals(List.of(new SubwordPiece("▁a", 1, 0, 1)), loaded.encode("a")); + } + + /** Verifies that two top-level model definitions are rejected as ambiguous. */ + @Test + void testRejectsADuplicateTopLevelField(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, "{" + NORMALIZER + "," + PRE_TOKENIZER + "," + MODEL + "," + + MODEL + "}"); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer)); + + assertTrue(error.getMessage().contains("more than once"), error.getMessage()); + } + + /** + * Supplies duplicate fields in each nested object parsed by the adapter. + * + * @return The object name and tokenizer JSON for each case. + */ + private static Stream nestedDuplicateFields() { + return Stream.of( + Arguments.of("model", TOKENIZER.replace("\"unk_id\":0", + "\"unk_id\":0,\"unk_id\":0")), + Arguments.of("normalizer", TOKENIZER.replace("\"precompiled_charsmap\":\"\"", + "\"precompiled_charsmap\":\"\",\"precompiled_charsmap\":\"\"")), + Arguments.of("pre-tokenizer", TOKENIZER.replace("\"split\":false", + "\"split\":false,\"split\":false")), + Arguments.of("added token", "{\"added_tokens\":[{\"id\":1,\"id\":1," + + "\"content\":\"▁a\",\"special\":false}]," + NORMALIZER + "," + + PRE_TOKENIZER + "," + MODEL + "}")); + } + + /** Verifies duplicate fields are rejected in every nested tokenizer object. */ + @ParameterizedTest(name = "{0}") + @MethodSource("nestedDuplicateFields") + void testRejectsDuplicateNestedFields(String object, String json, @TempDir Path dir) + throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, json); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer), object); + + assertTrue(error.getMessage().contains("more than once"), error.getMessage()); + } + + /** Verifies that added-token metadata cannot assign one row more than once. */ + @Test + void testRejectsDuplicateAddedTokenIds(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, "{\"added_tokens\":[" + + "{\"id\":1,\"content\":\"▁a\",\"special\":false}," + + "{\"id\":1,\"content\":\"▁a\",\"special\":true}]," + + NORMALIZER + "," + PRE_TOKENIZER + "," + MODEL + "}"); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer)); + + assertTrue(error.getMessage().contains("added token id 1 occurs more than once"), + error.getMessage()); + } + + /** + * Supplies a trailing comma in each tokenizer structure parsed by the adapter. + * + * @return The structure name and malformed tokenizer JSON for each case. + */ + private static Stream trailingCommaJson() { + final String addedToken = "\"added_tokens\":[{\"id\":1,\"content\":\"▁a\"," + + "\"special\":false}]"; + return Stream.of( + Arguments.of("model", TOKENIZER.replace("]]}", "]],}")), + Arguments.of("vocabulary", TOKENIZER.replace("[\"▁a\",-1.0]]", + "[\"▁a\",-1.0],]")), + Arguments.of("normalizer", TOKENIZER.replace("charsmap\":\"\"}", + "charsmap\":\"\",}")), + Arguments.of("normalizer array", TOKENIZER.replace( + "\"normalizer\":{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}", + "\"normalizer\":{\"type\":\"Sequence\",\"normalizers\":[" + + "{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"},]}")), + Arguments.of("pre-tokenizer", TOKENIZER.replace("split\":false}", + "split\":false,}")), + Arguments.of("added-token object", "{" + addedToken.replace("false}", "false,}") + + "," + NORMALIZER + "," + PRE_TOKENIZER + "," + MODEL + "}"), + Arguments.of("added-token array", "{" + addedToken.replace("]", ",]") + + "," + NORMALIZER + "," + PRE_TOKENIZER + "," + MODEL + "}")); + } + + /** Verifies that the tokenizer adapter accepts only standard JSON array and object syntax. */ + @ParameterizedTest(name = "{0}") + @MethodSource("trailingCommaJson") + void testRejectsTrailingCommas(String structure, String json, @TempDir Path dir) + throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, json); + + assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer), structure); + } + + @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()); + } + + @Test + void testRejectsANormalizationStepBeforeThePrecompiledMap(@TempDir Path dir) + throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + Files.writeString(tokenizer, + "{\"normalizer\":{\"type\":\"Sequence\",\"normalizers\":[" + + "{\"type\":\"Replace\",\"pattern\":{\"Regex\":\"\\\\s+\"}," + + "\"content\":\" \"}," + + "{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}]}," + + "\"pre_tokenizer\":{\"type\":\"Metaspace\",\"replacement\":\"▁\"," + + "\"prepend_scheme\":\"always\",\"split\":false}," + + "\"model\":{\"type\":\"Unigram\",\"unk_id\":0,\"byte_fallback\":false," + + "\"vocab\":[[\"\",0.0],[\"▁a\",-1.0]]}}"); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer)); + + assertTrue(error.getMessage().contains("Precompiled normalizer must precede"), + error.getMessage()); + } + + @Test + void testRejectsAnEmptyLiteralReplacement(@TempDir Path dir) throws IOException { + final Path tokenizer = dir.resolve("tokenizer.json"); + final String sequence = "\"normalizer\":{\"type\":\"Sequence\",\"normalizers\":[" + + "{\"type\":\"Precompiled\",\"precompiled_charsmap\":\"\"}," + + "{\"type\":\"Replace\",\"pattern\":{\"String\":\"\"}," + + "\"content\":\" \"}]}"; + Files.writeString(tokenizer, "{" + sequence + "," + PRE_TOKENIZER + "," + MODEL + "}"); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> Model2VecUnigramTokenizer.load(tokenizer)); + + assertTrue(error.getMessage().contains("empty literal"), 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..e112fc6156 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelAssemblerTest.java @@ -0,0 +1,281 @@ +/* + * 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 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; + +/** + * 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 testRejectsDuplicateTopLevelModel(@TempDir Path dir) throws IOException { + final String duplicate = WORDPIECE_TOKENIZER_JSON.substring(0, + WORDPIECE_TOKENIZER_JSON.length() - 1) + ",\"model\":{\"type\":\"Unigram\"}}"; + Files.writeString(dir.resolve("tokenizer.json"), duplicate); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> ModelAssembler.assemble(dir)); + assertTrue(exception.getMessage().contains("model") + && exception.getMessage().contains("more than once"), exception.getMessage()); + } + + @Test + void testRejectsDuplicateModelType(@TempDir Path dir) throws IOException { + final String duplicate = WORDPIECE_TOKENIZER_JSON.replace( + "\"type\":\"WordPiece\"", "\"type\":\"WordPiece\",\"type\":\"Unigram\""); + Files.writeString(dir.resolve("tokenizer.json"), duplicate); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> ModelAssembler.assemble(dir)); + assertTrue(exception.getMessage().contains("model.type") + && exception.getMessage().contains("more than once"), exception.getMessage()); + } + + @Test + void testRejectsDuplicateNormalizerLowercase(@TempDir Path dir) throws IOException { + final String duplicate = WORDPIECE_TOKENIZER_JSON.replace( + "\"lowercase\":true", "\"lowercase\":true,\"lowercase\":false"); + Files.writeString(dir.resolve("tokenizer.json"), duplicate); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> ModelAssembler.assemble(dir)); + assertTrue(exception.getMessage().contains("normalizer.lowercase") + && exception.getMessage().contains("more than once"), exception.getMessage()); + } + + @Test + void testRejectsNonBooleanNormalizerLowercase(@TempDir Path dir) throws IOException { + final String malformed = WORDPIECE_TOKENIZER_JSON.replace( + "\"lowercase\":true", "\"lowercase\":\"true\""); + Files.writeString(dir.resolve("tokenizer.json"), malformed); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", ROWS)); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> ModelAssembler.assemble(dir)); + assertTrue(exception.getMessage().contains("normalizer.lowercase") + && exception.getMessage().contains("boolean"), exception.getMessage()); + } + + @Test + void testRejectsDuplicateVocabularyTokenBeforeWritingFiles(@TempDir Path dir) + throws IOException { + final String duplicate = "{\"model\":{\"type\":\"WordPiece\",\"vocab\":" + + "{\"[UNK]\":0,\"[UNK]\":1}}}"; + Files.writeString(dir.resolve("tokenizer.json"), duplicate); + Files.writeString(dir.resolve("config.json"), "{\"normalize\":false}"); + SafetensorsTestFiles.write(dir.resolve("model.safetensors"), + SafetensorsTestFiles.matrix("embeddings", new float[][] {{0f}, {1f}})); + + assertThrows(InvalidFormatException.class, () -> ModelAssembler.assemble(dir)); + assertFalse(Files.exists(dir.resolve("vocab.txt"))); + assertFalse(Files.exists(dir.resolve("tokenizer_config.json"))); + } + + @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 testToolPrintsASummaryAndRejectsABadDirectory(@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/ModelDistillerExampleTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerExampleTest.java new file mode 100644 index 0000000000..7a7fef5625 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerExampleTest.java @@ -0,0 +1,257 @@ +/* + * 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 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.assertTrue; + +/** + * Distills an original ONNX lookup table, loads the output and ranks short documents. + * The table tests data flow, not the quality of a trained language model. + */ +class ModelDistillerExampleTest { + + /** Matrix order after removal of CLS and SEP. */ + private static final List TOKENS = + List.of("[PAD]", "[UNK]", "coffee", "espresso", "tea", "history"); + + /** Longer input first to exercise term batching by sequence length. */ + private static final List REQUESTED_TERMS = + List.of("Coffee espresso tea", "coffee", "TEA HISTORY", "coffee ESPRESSO", " tea history "); + + /** Expected term order after normalization and duplicate removal. */ + private static final List TERMS = + List.of("coffee espresso tea", "tea history", "coffee espresso"); + + /** The initial coordinates after pooling CLS, content and SEP in the lookup graph. */ + private static final double[][] POOLED = { + {0, 0}, {0, 0}, {1, 0}, {2.0 / 3, 1.0 / 3}, {-1.0 / 3, 2.0 / 3}, {-1.0 / 3, -2.0 / 3}, + {4.0 / 5, 3.0 / 5}, {-0.5, 0}, {1.25, 0.25} + }; + + /** Floating-point tolerance for PCA projection and normalization. */ + private static final double TOLERANCE = 1e-5; + + /** A document and cosine similarity to the query. */ + private record Scored(String document, double score) { + } + + /** + * Tests ONNX inference, PCA, weighting, serialization and optional term batching. + * + * @param includeTerms Whether to distill additional phrases. + * @param directory The test directory. + * @throws IOException Thrown if a model file cannot be read or written. + */ + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testDistillReloadAndSearch(boolean includeTerms, @TempDir Path directory) + throws IOException { + final Path teacher = writeTeacher(directory.resolve("teacher")); + final Path output = directory.resolve("static-model"); + final List terms = includeTerms ? REQUESTED_TERMS : List.of(); + + final ModelDistiller.Result result = ModelDistiller.distill(teacher, output, 2, terms, null); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(output); + + assertEquals("WordPiece", result.family()); + assertEquals(TOKENS.size(), result.vocabularySize()); + assertEquals(includeTerms ? TERMS.size() : 0, result.termCount()); + assertEquals(4, result.teacherDimension()); + assertEquals(2, result.dimension()); + assertEquals(1, result.explainedVarianceRatio(), TOLERANCE); + assertEquals(result.dimension(), model.dimension()); + assertEquals(result.vocabularySize(), model.vocabularySize()); + assertEquals(result.termCount(), model.termCount()); + assertEquals(TOKENS, Files.readAllLines(output.resolve(ModelFileNames.VOCABULARY))); + if (includeTerms) { + assertEquals(TERMS, Files.readAllLines(output.resolve(ModelFileNames.TERMS))); + } else { + assertTrue(Files.notExists(output.resolve(ModelFileNames.TERMS))); + } + + final double[][] expected = centeredRows(TOKENS.size() + result.termCount()); + checkStoredMatrix(output, expected); + final List texts = new ArrayList<>(TOKENS); + if (includeTerms) { + texts.addAll(TERMS); + } + for (int index = 2; index < texts.size(); index++) { + final float[] vector = model.embed(texts.get(index)); + assertEquals(1, vector[0] * vector[0] + vector[1] * vector[1], TOLERANCE); + assertEquals(cosine(expected[2], expected[index]), + model.similarity("coffee", texts.get(index)), TOLERANCE, texts.get(index)); + } + assertArrayEquals(model.embed("coffee"), model.embed("COFFEE")); + assertArrayEquals(new float[2], model.embed("unlisted")); + + final List results = new ArrayList<>(); + for (final String document : List.of("history", "tea", "espresso")) { + results.add(new Scored(document, model.similarity("coffee", document))); + } + results.sort(Comparator.comparingDouble(Scored::score).reversed()); + assertEquals("espresso", results.get(0).document()); + assertTrue(results.get(0).score() > results.get(1).score()); + assertEquals("coffee", model.mostSimilar("coffee", 1).get(0).token()); + + final Path repeatedOutput = directory.resolve("repeated-model"); + assertEquals(result, ModelDistiller.distill(teacher, repeatedOutput, 2, terms, null)); + for (final String file : List.of(ModelFileNames.SAFETENSORS, ModelFileNames.TOKENIZER_JSON, + ModelFileNames.CONFIG, ModelFileNames.VOCABULARY, ModelFileNames.TOKENIZER_CONFIG)) { + assertEquals(-1L, Files.mismatch(output.resolve(file), repeatedOutput.resolve(file)), file); + } + assertArrayEquals(model.embed("coffee espresso"), + StaticEmbeddingModel.load(repeatedOutput).embed("coffee espresso")); + } + + /** + * Runs the Java distillation listing in the manual with an original test graph. + * + * @param directory The test directory. + * @throws IOException Thrown if a model file cannot be read or written. + */ + @Test + void testManualDistillationExample(@TempDir Path directory) throws IOException { + final Path teacher = writeTeacher(directory.resolve("teacher")); + final Path output = directory.resolve("static-model"); + + ModelDistiller.distill(teacher, output, 2, List.of("coffee espresso"), null); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(output); + final float[] vector = model.embed("coffee espresso"); + final List related = model.mostSimilar("coffee espresso", 3); + + assertEquals(2, vector.length); + assertEquals("coffee espresso", related.get(0).token()); + assertEquals(3, related.size()); + } + + /** + * Writes the tokenizer configuration and original ONNX lookup graph. + * + * @param directory The destination directory. + * @return The teacher directory. + * @throws IOException Thrown if writing fails. + */ + private Path writeTeacher(Path directory) throws IOException { + EmbeddingTestFixtures.writeLookupTeacherOnnxModel( + Files.createDirectories(directory.resolve("onnx"))); + Files.writeString(directory.resolve(ModelFileNames.TOKENIZER_JSON), """ + {"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}], + "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, + "coffee":4,"espresso":5,"tea":6,"history":7}}} + """); + Files.writeString(directory.resolve(ModelFileNames.TOKENIZER_CONFIG), + "{\"pad_token\":\"[PAD]\"}"); + return directory; + } + + /** + * Subtracts the mean from the analytically calculated teacher vectors. + * + * @param count The number of model entries. + * @return The expected PCA inputs. + */ + private double[][] centeredRows(int count) { + final double[] mean = new double[2]; + for (int index = 0; index < count; index++) { + for (int component = 0; component < mean.length; component++) { + mean[component] += POOLED[index][component] / count; + } + } + final double[][] result = new double[count][2]; + for (int index = 0; index < count; index++) { + for (int component = 0; component < mean.length; component++) { + result[index][component] = POOLED[index][component] - mean[component]; + } + } + return result; + } + + /** + * Checks PCA distances and Zipf scaling without depending on component signs. + * + * @param directory The saved model directory. + * @param expected The centered teacher vectors before Zipf scaling. + * @throws IOException Thrown if reading fails. + */ + private void checkStoredMatrix(Path directory, double[][] expected) throws IOException { + final SafetensorsFile file = SafetensorsFile.read(directory.resolve(ModelFileNames.SAFETENSORS)); + assertEquals(1, file.tensorNames().size()); + final String tensor = file.tensorNames().iterator().next(); + assertArrayEquals(new int[] {expected.length, 2}, file.tensorInfo(tensor).shape()); + final float[] matrix = file.readFloats(tensor); + double harmonicSum = 0; + for (int rank = 2; rank <= expected.length + 1; rank++) { + harmonicSum += 1.0 / rank; + } + final double[][] scaled = new double[expected.length][2]; + for (int index = 0; index < expected.length; index++) { + final double weight = 1e-4 / (1e-4 + 1.0 / (index + 2) / harmonicSum); + scaled[index][0] = matrix[index * 2] / weight; + scaled[index][1] = matrix[index * 2 + 1] / weight; + } + for (int left = 0; left < expected.length; left++) { + for (int right = 0; right < expected.length; right++) { + assertEquals(dot(expected[left], expected[right]), dot(scaled[left], scaled[right]), + TOLERANCE, "matrix entries " + left + ", " + right); + } + } + } + + /** + * Calculates the dot product in the original test coordinate system. + * + * @param left The initial vector. + * @param right The other vector. + * @return The dot product. + */ + private double dot(double[] left, double[] right) { + return left[0] * right[0] + left[1] * right[1]; + } + + /** + * Calculates cosine similarity for the expected test vectors. + * + * @param left The initial vector. + * @param right The other vector. + * @return Cosine similarity. + */ + private double cosine(double[] left, double[] right) { + return dot(left, right) / Math.sqrt(dot(left, left) * dot(right, right)); + } +} 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..d0219e58a3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/ModelDistillerTest.java @@ -0,0 +1,302 @@ +/* + * 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; + +/** + * Tests argument validation, Zipf weighting and output files with a small ONNX teacher. + */ +class ModelDistillerTest { + + /** A six-row WordPiece tokenizer accepted by {@link TeacherTokenizer}. */ + private static final String TINY_TEACHER_TOKENIZER = + "{\"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}]," + + "\"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,\"world\":5}}}"; + + /** Model2Vec's SIF coefficient, the value the distiller uses. */ + private static final double SIF = 1e-4; + + /** + * Writes a small teacher that can run a complete distillation without external files. + * + * @param directory The teacher directory to create. + * @return The created teacher directory. + * @throws IOException Thrown if a fixture file cannot be written. + */ + private static Path writeTinyTeacher(Path directory) throws IOException { + EmbeddingTestFixtures.writeTinyOnnxModel(Files.createDirectories(directory.resolve("onnx"))); + Files.writeString(directory.resolve("tokenizer.json"), TINY_TEACHER_TOKENIZER); + Files.writeString(directory.resolve("tokenizer_config.json"), + "{\"pad_token\":\"[PAD]\"}"); + return directory; + } + + /** + * Rejects changes in teacher vector length before creating output files. + * + * @param directory The test directory. + * @throws IOException Thrown if a fixture file cannot be written. + */ + @Test + void testRejectsChangingTeacherDimension(@TempDir Path directory) throws IOException { + final Path teacher = writeTinyTeacher(Files.createDirectory(directory.resolve("teacher"))); + EmbeddingTestFixtures.writeVariableDimensionOnnxModel(teacher.resolve("onnx")); + final Path output = directory.resolve("output"); + + final IllegalArgumentException error = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, output, 1, null)); + + assertTrue(error.getMessage().contains("hidden dimension"), error.getMessage()); + assertTrue(Files.notExists(output)); + } + + @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 testRejectsAnOutputPathThatIsAFileBeforeResolvingTheTeacher(@TempDir Path dir) + throws IOException { + final Path output = Files.writeString(dir.resolve("model.bin"), "keep me"); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill("not a model id", output, 256, null)); + + assertEquals("outputDirectory must be a directory or not exist: " + output, e.getMessage()); + assertEquals("keep me", Files.readString(output)); + } + + @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"); + } + + @Test + void testReplacesFilesDerivedByAnEarlierDistillation(@TempDir Path dir) throws IOException { + final Path teacher = writeTinyTeacher(Files.createDirectory(dir.resolve("teacher"))); + final Path output = Files.createDirectory(dir.resolve("output")); + Files.write(output.resolve(ModelFileNames.VOCABULARY), + List.of("[PAD]", "[UNK]", "stale", "rows")); + Files.writeString(output.resolve(ModelFileNames.TOKENIZER_CONFIG), + "{\"do_lower_case\":false}"); + for (final String name : ModelFileNames.SENTENCEPIECE_MODELS) { + Files.writeString(output.resolve(name), "stale model"); + } + + final ModelDistiller.Result result = ModelDistiller.distill(teacher, output, 1, null); + + assertEquals("WordPiece", result.family()); + assertEquals(List.of("[PAD]", "[UNK]", "hello", "world"), + Files.readAllLines(output.resolve(ModelFileNames.VOCABULARY))); + assertTrue(Files.readString(output.resolve(ModelFileNames.TOKENIZER_CONFIG)) + .contains("\"do_lower_case\": true")); + for (final String name : ModelFileNames.SENTENCEPIECE_MODELS) { + assertTrue(Files.notExists(output.resolve(name)), name + " must not survive the new run"); + } + final StaticEmbeddingModel model = StaticEmbeddingModel.load(output); + assertTrue(model.embed("hello")[0] != 0f); + } + + @Test + void testRejectsTheTeacherDirectoryAsItsOwnOutput(@TempDir Path dir) throws IOException { + final Path teacher = writeTinyTeacher(Files.createDirectory(dir.resolve("teacher"))); + final String tokenizer = Files.readString(teacher.resolve(ModelFileNames.TOKENIZER_JSON)); + + final IllegalArgumentException error = assertThrows(IllegalArgumentException.class, + () -> ModelDistiller.distill(teacher, teacher, 1, null)); + + assertEquals("outputDirectory must differ from teacherDirectory", error.getMessage()); + assertEquals(tokenizer, Files.readString(teacher.resolve(ModelFileNames.TOKENIZER_JSON))); + assertTrue(Files.isRegularFile(teacher.resolve(ModelFileNames.ONNX_MODEL))); + } + + @Test + void testWritesTheTeacherNameIntoTheGeneratedConfiguration(@TempDir Path dir) + throws IOException { + final Path teacher = writeTinyTeacher(Files.createDirectory(dir.resolve("teacher"))); + final Path output = dir.resolve("output"); + + final ModelDistiller.Result result = ModelDistiller.distill(teacher, output, 1, null); + + assertEquals("WordPiece", result.family()); + assertTrue(Files.readString(output.resolve(ModelFileNames.CONFIG)) + .contains("\"tokenizer_name\": \"teacher\"")); + assertEquals(1, StaticEmbeddingModel.load(output).dimension()); + } + + /** + * The JSON escape behind {@code tokenizer_name} is exercised directly, because every + * character that needs escaping (quote, backslash, control characters) is illegal in file + * names on Windows, so a real teacher directory cannot carry such a name there. + */ + @Test + void testJsonStringEscapesQuotesBackslashesAndControlCharacters() { + assertEquals("\"teacher\\\"quoted\"", ModelDistiller.jsonString("teacher\"quoted")); + assertEquals("\"back\\\\slash\"", ModelDistiller.jsonString("back\\slash")); + assertEquals("\"line\\nbreak\"", ModelDistiller.jsonString("line\nbreak")); + assertEquals("\"\\u0001\"", ModelDistiller.jsonString("\u0001")); + assertEquals("\"plain\"", ModelDistiller.jsonString("plain")); + } + + /** + * 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[0] must not be 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/NeighborTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/NeighborTest.java new file mode 100644 index 0000000000..4fef47b3d1 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/NeighborTest.java @@ -0,0 +1,45 @@ +/* + * 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 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.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; + +class NeighborTest { + + @Test + void testRejectsNullToken() { + assertThrows(IllegalArgumentException.class, () -> new Neighbor(null, 0.5)); + } + + @ParameterizedTest + @ValueSource(doubles = {Double.NaN, Double.NEGATIVE_INFINITY, -1.00001, + 1.00001, Double.POSITIVE_INFINITY}) + void testRejectsInvalidSimilarity(double similarity) { + assertThrows(IllegalArgumentException.class, () -> new Neighbor("token", similarity)); + } + + @ParameterizedTest + @ValueSource(doubles = {-1.0, 0.0, 1.0}) + void testAcceptsCosineRangeBoundaries(double similarity) { + assertDoesNotThrow(() -> new Neighbor("token", similarity)); + } +} 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..9b2581f650 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/OnnxTeacherEncoderTest.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.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; + +/** + * Validates graph inputs and pooled output using small ONNX models. + */ +class OnnxTeacherEncoderTest { + + /** + * Requires a stable vector length across batches in one session. + * + * @param initialSize The initial batch size and vector length. + * @param changedSize The next batch size and vector length. + * @param directory The test directory. + * @throws Exception Thrown if the test graph cannot be written or loaded. + */ + @ParameterizedTest + @CsvSource({"1, 3", "3, 1"}) + void testRejectsChangingHiddenDimension(int initialSize, int changedSize, + @TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeVariableDimensionOnnxModel(directory); + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertEquals(initialSize, encoder.encodeBatch(new long[initialSize][1])[0].length); + + final IllegalArgumentException error = assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[changedSize][1])); + assertTrue(error.getMessage().contains("hidden dimension"), error.getMessage()); + assertEquals(initialSize, encoder.encodeBatch(new long[initialSize][1])[0].length); + } + } + + @Test + void testRejectsNullFile() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(null)); + assertTrue(e.getMessage().contains("must not be null"), e.getMessage()); + } + + @Test + void testRejectsMissingFile(@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()); + } + + @Test + void testRejectsNullAndEmptySequences(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeTinyOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertEquals("batch[0] must not be null", assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {null})).getMessage()); + assertEquals("batch[1] must not be null", assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {{1}, null})).getMessage()); + assertEquals("batch[0] must not be empty", assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {new long[0]})).getMessage()); + } + } + + @Test + void testSupportsGraphWithoutAttentionMask(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeInputIdsOnlyOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertArrayEquals(new float[] {1f, -2f, 4f}, encoder.encodeBatch(new long[][] {{2}})[0]); + } + } + + @Test + void testRejectsUnsupportedInput(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeUnsupportedInputOnnxModel(directory); + + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(model)); + assertTrue(exception.getMessage().contains("position_ids"), exception.getMessage()); + } + + @Test + void testSupportsInt32InputIds(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeInt32InputOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertArrayEquals(new float[] {3f}, encoder.encodeBatch(new long[][] {{2, 4}})[0]); + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {{(long) Integer.MAX_VALUE + 1}})); + assertTrue(exception.getMessage().contains("input_ids[0][0]"), exception.getMessage()); + assertTrue(exception.getMessage().contains("INT32"), exception.getMessage()); + } + } + + @Test + void testRejectsRankOneInputIdsAtLoadTime(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeRankOneInputOnnxModel(directory); + + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(model)); + assertTrue(exception.getMessage().contains("input_ids"), exception.getMessage()); + assertTrue(exception.getMessage().contains("rank 2"), exception.getMessage()); + } + + @Test + void testSupportsInt32AttentionMask(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeInt32AttentionMaskOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertArrayEquals(new float[] {3f}, encoder.encodeBatch(new long[][] {{2, 4}})[0]); + } + } + + @Test + void testRejectsFloatInputIdsAtLoadTime(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeFloatInputOnnxModel(directory); + + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> OnnxTeacherEncoder.load(model)); + assertTrue(exception.getMessage().contains("input_ids"), exception.getMessage()); + assertTrue(exception.getMessage().contains("INT32 or INT64"), exception.getMessage()); + } + + @Test + void testPrefersTheNamedLastHiddenStateOutput(@TempDir Path directory) throws Exception { + final Path model = EmbeddingTestFixtures.writeMultipleOutputsOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertArrayEquals(new float[] {3f}, encoder.encodeBatch(new long[][] {{2, 4}})[0]); + } + } + + @Test + void testRejectsOutputDimensionsThatDoNotMatchTheInput(@TempDir Path directory) + throws Exception { + final Path model = EmbeddingTestFixtures.writeFixedOutputOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + final IllegalArgumentException sequenceError = assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {{2, 4}})); + assertTrue(sequenceError.getMessage().contains("sequence dimension"), + sequenceError.getMessage()); + + final IllegalArgumentException batchError = assertThrows(IllegalArgumentException.class, + () -> encoder.encodeBatch(new long[][] {{2}, {4}})); + assertTrue(batchError.getMessage().contains("batch dimension"), batchError.getMessage()); + } + } + + @Test + void testMeanPoolingDoesNotOverflowFiniteHiddenStates(@TempDir Path directory) + throws Exception { + final Path model = EmbeddingTestFixtures.writeMaxFloatOnnxModel(directory); + + try (OnnxTeacherEncoder encoder = OnnxTeacherEncoder.load(model)) { + assertArrayEquals(new float[] {Float.MAX_VALUE}, + encoder.encodeBatch(new long[][] {{1, 1}})[0]); + } + } +} 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..5f9265ec45 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixEdgeCaseTest.java @@ -0,0 +1,190 @@ +/* + * 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.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests padded dimensions and finite extreme values in a quantized matrix. + */ +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 double[] 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 stresses the transform because 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"); + } + } + + @Test + void testQuantizesEveryFiniteTwoDimensionalSignPattern() { + final float maximum = Float.MAX_VALUE; + final float[] matrix = { + maximum, maximum, + maximum, -maximum, + -maximum, maximum, + -maximum, -maximum + }; + + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, 4, 2, 4, SEED); + + for (int row = 0; row < quantized.rowCount(); row++) { + assertTrue(Double.isFinite(quantized.rowNorm(row)), "non-finite norm at row " + row); + for (final float value : quantized.decodeRow(row)) { + assertTrue(Float.isFinite(value), "non-finite decoded value at row " + row); + } + } + } + + @Test + void testExtremeFiniteRowsScoreLikeTheirDecodedValues() { + final float maximum = Float.MAX_VALUE; + final float[] matrix = { + maximum, maximum, + maximum, -maximum, + -maximum, maximum, + -maximum, -maximum + }; + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(matrix, 4, 2, 4, SEED); + + for (int row = 0; row < quantized.rowCount(); row++) { + final float[] decoded = quantized.decodeRow(row); + for (int dimension = 0; dimension < 2; dimension++) { + final float[] query = new float[2]; + query[dimension] = 1f; + final double actual = quantized.dotRotated(row, quantized.rotate(query)); + assertEquals(decoded[dimension], actual, + Math.abs((double) decoded[dimension]) * 1e-6, + "rotated scoring must match decoded coordinate " + dimension + " of row " + row); + } + } + } + + @ParameterizedTest + @CsvSource({"1, 1", "1, -1", "-1, 1", "-1, -1"}) + void testFiniteQueriesDoNotOverflowRotatedSpace(float firstSign, float secondSign) { + final QuantizedEmbeddingMatrix quantized = + QuantizedEmbeddingMatrix.quantize(new float[] {1f, 1f}, 1, 2, 4, SEED); + final float maximum = Float.MAX_VALUE; + final float[] query = {firstSign * maximum, secondSign * maximum}; + + for (final double value : quantized.rotate(query)) { + assertTrue(Double.isFinite(value), "a finite query must stay finite after rotation"); + } + } + + /** + * {@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 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..dea42aae67 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedEmbeddingMatrixTest.java @@ -0,0 +1,328 @@ +/* + * 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; + +/** + * Tests quantized matrix encoding, storage, pooling, and scoring. + */ +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 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 to exercise per-row scaling. + 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; + } + + // Minimum reconstruction quality for the deterministic fixture at each bit width. + @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 double[] 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 double[] rotatedSum = new double[quantized.paddedDimension()]; + quantized.addRowRotated(0, 1f, rotatedSum); + quantized.addRowRotated(1, 2.5f, rotatedSum); + quantized.addRowRotated(2, -0.5f, rotatedSum); + final double[] pooled = quantized.toOriginal(rotatedSum); + final float[] row0 = quantized.decodeRow(0); + final float[] row1 = quantized.decodeRow(1); + final float[] row2 = quantized.decodeRow(2); + final double[] expected = new double[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + expected[d] = row0[d] + 2.5 * row1[d] - 0.5 * row2[d]; + } + assertArrayEquals(expected, pooled, 1e-3, + "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 decoded matrix again must not change the file bytes. + 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 scale and norm metadata per row: + // 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); + // Writing the decoded matrix again must preserve the weight bytes. + 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 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 double[7])); + assertThrows(IllegalArgumentException.class, + () -> quantized.addRowRotated(0, 1f, new double[7])); + assertThrows(IllegalArgumentException.class, + () -> quantized.dotRotated(0, new double[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 malformed 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/QuantizedMatrixCompatibilityTest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixCompatibilityTest.java new file mode 100644 index 0000000000..8905d8b14b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixCompatibilityTest.java @@ -0,0 +1,178 @@ +/* + * 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.HexFormat; + +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.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Tests fixed ONQ2 files defined independently of the production writer. + * + *

Each file contains a 3-by-3 matrix padded to dimension 4. Grid levels are + * ascending odd integers, with scales [0.5, 1, 0]. Initial codes are [0,1,2,3], + * [0,2,5,7], and [0,5,10,15] for widths 2, 3, and 4. Reversing those codes at + * scale 1 gives vector 1; scale 0 gives vector 2. Expected vectors use direct + * multiplication by the normalized 4-by-4 Hadamard matrix and seed signs.

+ * + *

Header fields, grid levels, scales, norms, and optional pooling weights + * use big-endian encoding. Packed codes start at the low bits of each byte. + * The 3-bit fixture tests codes spanning byte boundaries. The stored norms + * are sqrt(5), sqrt(29), and sqrt(125), respectively, multiplied by [1,2,0].

+ */ +class QuantizedMatrixCompatibilityTest { + + private static final String ONQ2_BITS_2 = + "4f4e5132000000030000000300000002000000000000001100000004" + + "c0400000bf8000003f80000040400000" + + "3fe00000000000003ff00000000000000000000000000000" + + "4001e3779b97f4a84011e3779b97f4a80000000000000000" + + "00e41b00"; + + private static final String ONQ2_BITS_3 = + "4f4e5132000000030000000300000003ffffffffffffffff00000008" + + "c0e00000c0a00000c0400000bf8000003f8000004040000040a0000040e00000" + + "3fe00000000000003ff00000000000000000000000000000" + + "40158a68a4a8d9f340258a68a4a8d9f30000000000000000" + + "013fc00000bf00000000000000500faf000000"; + + private static final String ONQ2_BITS_4 = + "4f4e5132000000030000000300000004800000000000000000000010" + + "c1700000c1500000c1300000c1100000c0e00000c0a00000c0400000bf800000" + + "3f8000004040000040a0000040e0000041100000413000004150000041700000" + + "3fe00000000000003ff00000000000000000000000000000" + + "40265c55827df1d240365c55827df1d20000000000000000" + + "013fc00000bf0000000000000050faaf050000"; + + /** + * Loads known coordinates and checks scoring, pooling, and file output. + * + * @param bits The quantization width. + * @param seed The stored rotation seed. + * @param y The expected coordinate at index 1. + * @param z The expected coordinate at index 2. + * @param squaredNorm The expected squared norm of vector 0. + * @param dir The temporary directory. + * @throws IOException Thrown if file access fails. + */ + @ParameterizedTest + @CsvSource({"2,17,1,-2,5", "3,-1,-2,-5,29", "4,-9223372036854775808,5,-10,125"}) + void testKnownVersion2Vectors(int bits, long seed, float y, float z, + int squaredNorm, @TempDir Path dir) throws IOException { + final byte[] bytes = fixture(bits); + final Path file = dir.resolve("model.quantized"); + Files.write(file, bytes); + final QuantizedEmbeddingMatrix matrix = QuantizedEmbeddingMatrix.read(file); + + assertEquals(3, matrix.rowCount()); + assertEquals(3, matrix.dimension()); + assertEquals(4, matrix.paddedDimension()); + assertEquals(bits, matrix.bits()); + assertEquals(seed, matrix.seed()); + assertArrayEquals(new float[] {0f, y, z}, matrix.decodeRow(0), 0f); + assertArrayEquals(new float[] {0f, -2 * y, -2 * z}, matrix.decodeRow(1), 0f); + assertArrayEquals(new float[3], matrix.decodeRow(2), 0f); + assertEquals(Math.sqrt(squaredNorm), matrix.rowNorm(0), 1e-14); + assertEquals(2 * Math.sqrt(squaredNorm), matrix.rowNorm(1), 1e-14); + assertEquals(0.0, matrix.rowNorm(2)); + + final double[] query = matrix.rotate(new float[] {1f, 2f, 3f}); + final double dot = 2 * y + 3 * z; + assertEquals(dot, matrix.dotRotated(0, query), 0.0); + assertEquals(-2 * dot, matrix.dotRotated(1, query), 0.0); + assertEquals(0.0, matrix.dotRotated(2, query), 0.0); + + final double[] sum = new double[matrix.paddedDimension()]; + matrix.addRowRotated(0, 1.5f, sum); + matrix.addRowRotated(1, -0.5f, sum); + assertArrayEquals(new double[] {0, 2.5 * y, 2.5 * z}, matrix.toOriginal(sum), 0.0); + if (bits == 2) { + assertNull(matrix.poolingWeights()); + } else { + assertArrayEquals(new float[] {1.5f, -0.5f, 0f}, matrix.poolingWeights(), 0f); + } + + final Path written = dir.resolve("written.quantized"); + matrix.write(written); + assertArrayEquals(bytes, Files.readAllBytes(written)); + } + + /** + * Rejects truncated fixtures at all byte offsets with a checked exception. + * + * @param bits The quantization width. + * @param dir The temporary directory. + * @throws IOException Thrown if writing fails. + */ + @ParameterizedTest + @ValueSource(ints = {2, 3, 4}) + void testTruncatedVersion2Fixture(int bits, @TempDir Path dir) throws IOException { + final byte[] bytes = fixture(bits); + final Path file = dir.resolve("truncated.quantized"); + for (int length = 0; length < bytes.length; length++) { + Files.write(file, Arrays.copyOf(bytes, length)); + assertThrows(IOException.class, () -> QuantizedEmbeddingMatrix.read(file), + "file length " + length); + } + } + + /** + * Rejects content appended to a complete fixture. + * + * @param bits The quantization width. + * @param dir The temporary directory. + * @throws IOException Thrown if writing fails. + */ + @ParameterizedTest + @ValueSource(ints = {2, 3, 4}) + void testVersion2FixtureWithTrailingByte(int bits, @TempDir Path dir) throws IOException { + final byte[] bytes = fixture(bits); + final Path file = dir.resolve("trailing.quantized"); + Files.write(file, Arrays.copyOf(bytes, bytes.length + 1)); + + assertThrows(InvalidFormatException.class, () -> QuantizedEmbeddingMatrix.read(file)); + } + + /** + * {@return the fixed file bytes for a quantization width} + * + * @param bits The width: 2, 3, or 4. + * @throws IllegalArgumentException Thrown for an unsupported width. + */ + private byte[] fixture(int bits) { + return HexFormat.of().parseHex(switch (bits) { + case 2 -> ONQ2_BITS_2; + case 3 -> ONQ2_BITS_3; + case 4 -> ONQ2_BITS_4; + default -> throw new IllegalArgumentException("Unsupported fixture width: " + bits); + }); + } +} 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..e2959635c3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/QuantizedMatrixFormatTest.java @@ -0,0 +1,313 @@ +/* + * 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; + +/** + * Verifies the loader contract of {@link QuantizedEmbeddingMatrix#read(Path)}: malformed + * file content produces a checked {@link InvalidFormatException} before 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 = 0x4F4E5132; + + @Test + void testEarlierFormatVersionIsRejectedByMagic(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("old-version.quantized"); + try (DataOutputStream out = out(file)) { + out.writeInt(0x4F4E5131); + } + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(file)); + + assertTrue(error.getMessage().contains("magic"), error.getMessage()); + } + + @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 16 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. + writeFloatNaN(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 double at offset 100: 28 header bytes, 64 bytes of + // grid levels, and 8 bytes for row 0's scale. + writeDoubleNaN(bytes, 100); + 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 testNegativeScaleIsRejectedAsFormatError(@TempDir Path dir) throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + writeDouble(bytes, 92, -1.0); + final Path patched = dir.resolve("negative-scale.quantized"); + Files.write(patched, bytes); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + + assertTrue(error.getMessage().contains("scale"), error.getMessage()); + } + + @Test + void testScaleThatWouldOverflowDecodingIsRejectedAsFormatError(@TempDir Path dir) + throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + writeDouble(bytes, 92, Double.MAX_VALUE); + final Path patched = dir.resolve("overflowing-scale.quantized"); + Files.write(patched, bytes); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + + assertTrue(error.getMessage().contains("scale"), error.getMessage()); + } + + @Test + void testDecodedNormBeyondFloatVectorRangeIsRejectedAsFormatError(@TempDir Path dir) + throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + writeDouble(bytes, 100, Double.MAX_VALUE); + final Path patched = dir.resolve("overflowing-norm.quantized"); + Files.write(patched, bytes); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + + assertTrue(error.getMessage().contains("norm"), error.getMessage()); + } + + @Test + void testZeroScaleWithPositiveNormIsRejectedAsFormatError(@TempDir Path dir) + throws IOException { + final Path file = validFile(dir); + final byte[] bytes = Files.readAllBytes(file); + writeDouble(bytes, 92, 0.0); + final Path patched = dir.resolve("inconsistent-zero-scale.quantized"); + Files.write(patched, bytes); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + + assertTrue(error.getMessage().contains("scale"), error.getMessage()); + assertTrue(error.getMessage().contains("norm"), error.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 109, right after the weight + // presence flag at offset 108. + writeFloatNaN(bytes, 109); + 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 108 declares per-row pooling weights the file does + // not contain, so the declared total exceeds the file size. + bytes[108] = 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()); + } + + @Test + void testInvalidPoolingWeightFlagIsRejectedAsFormatError(@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); + bytes[108] = 2; + final Path patched = dir.resolve("bad-weights-flag.quantized"); + Files.write(patched, bytes); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> QuantizedEmbeddingMatrix.read(patched)); + + assertTrue(error.getMessage().contains("flag"), error.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 double scale at offset + * 92, one double decoded norm at offset 100, the weight presence flag at offset 108, and two + * packed code bytes, 111 bytes in total. + * + * @param dir The directory to write into. + * @throws IOException Thrown if writing fails. + */ + private 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 void writeFloatNaN(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; + } + + /** Overwrites eight bytes with the big-endian bits of {@link Double#NaN}. */ + private void writeDoubleNaN(byte[] bytes, int offset) { + writeDouble(bytes, offset, Double.NaN); + } + + private void writeDouble(byte[] bytes, int offset, double value) { + final long bits = Double.doubleToLongBits(value); + for (int i = 0; i < Long.BYTES; i++) { + bytes[offset + i] = (byte) (bits >>> (56 - 8 * i)); + } + } + + private 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..4887eca30a --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/RandomizedPcaTest.java @@ -0,0 +1,284 @@ +/* + * 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 testWideMatrixUsesNoMoreSampleDimensionsThanCenteredRank() { + final int rows = 3; + final int cols = 20; + final int components = 2; + final float[] data = new float[rows * cols]; + for (int c = 0; c < cols; c++) { + data[c] = c + 1; + data[cols + c] = (c + 1) * (c + 1); + data[2 * cols + c] = c % 3 - 1; + } + + final RandomizedPca.Result result = + RandomizedPca.fitTransform(data, rows, cols, components, 42); + + assertEquals(rows * components, result.transformed().length); + assertTrue(result.explainedVarianceRatio() > 0.999999, + "two components must retain all variance of three centered rows"); + for (final float value : result.transformed()) { + assertTrue(Float.isFinite(value)); + } + } + + @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. + */ + @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. The calculation must reject this case. + */ + @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..7f8d2667b0 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsFileTest.java @@ -0,0 +1,458 @@ +/* + * 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.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.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 create headers with + // invalid 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 validates its input; malformed fixtures are written directly here. + 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 testZeroLengthTensorMayPrecedeDataAtTheSameOffset(@TempDir Path dir) throws IOException { + final String header = "{\"values\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4]},\"empty\":{\"dtype\":\"F32\",\"shape\":[0]," + + "\"data_offsets\":[0,0]}}"; + final Path file = writeFile(dir, MODEL_FILE_NAME, header, floatsToLittleEndianBytes(3f)); + + final SafetensorsFile parsed = SafetensorsFile.read(file); + + assertEquals(List.of("values", "empty"), List.copyOf(parsed.tensorNames())); + assertArrayEquals(new float[0], parsed.readFloats("empty")); + assertArrayEquals(new float[] {3f}, parsed.readFloats("values")); + } + + @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 testRejectsHeaderLargerThanTheSafetensorsLimit(@TempDir Path dir) throws IOException { + final Path file = dir.resolve("oversized-header.safetensors"); + final byte[] prefix = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(100_000_001L).array(); + Files.write(file, prefix); + + final InvalidFormatException exception = + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + assertTrue(exception.getMessage().contains("100000000 bytes"), exception.getMessage()); + } + + @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 testRejectsOverlappingTensorRanges(@TempDir Path dir) throws IOException { + final String header = "{\"first\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4]},\"second\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4]}}"; + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[4]); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsGapsBetweenTensorRanges(@TempDir Path dir) throws IOException { + final String header = "{\"first\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[0,4]},\"second\":{\"dtype\":\"F32\",\"shape\":[1]," + + "\"data_offsets\":[8,12]}}"; + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[12]); + + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + } + + @Test + void testRejectsTrailingTensorData(@TempDir Path dir) throws IOException { + final String header = singleTensorHeader("w", "F32", "[1]", 0, 4); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[8]); + + 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 testRejectsInvalidUtf8InHeader(@TempDir Path dir) throws IOException { + final byte[] invalidUtf8 = {(byte) 0xC3, 0x28}; + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(invalidUtf8.length).array()); + out.write(invalidUtf8); + final Path file = dir.resolve("invalid-utf8.safetensors"); + Files.write(file, out.toByteArray()); + + final InvalidFormatException exception = + assertThrows(InvalidFormatException.class, () -> SafetensorsFile.read(file)); + assertTrue(exception.getMessage().contains("valid UTF-8"), exception.getMessage()); + } + + @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 testRejectsAFileTruncatedAfterRead(@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 be rejected rather than 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()); + } + + @ParameterizedTest + @CsvSource({ + "-1, 0, 0, shape[0]", + "1, -1, 0, dataOffsetBegin", + "1, 2, 1, dataOffsetEnd" + }) + void testTensorInfoRejectsInvalidDimensionsAndOffsets(int dimension, long begin, long end, + String messagePart) { + final IllegalArgumentException error = assertThrows(IllegalArgumentException.class, + () -> new TensorInfo("t", "F32", new int[] {dimension}, begin, end)); + + assertTrue(error.getMessage().contains(messagePart), error.getMessage()); + } + + @Test + void testTensorInfoRejectsElementCountOverflow() { + 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()); + } + + @Test + void testReadFloatsReportsElementCountOverflowAsInvalidFormat(@TempDir Path dir) + throws IOException { + final String header = singleTensorHeader("huge", "F32", + "[2147483647,2147483647,2147483647]", 0, 0); + final Path file = writeFile(dir, MODEL_FILE_NAME, header, new byte[0]); + final SafetensorsFile parsed = SafetensorsFile.read(file); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> parsed.readFloats("huge")); + + assertTrue(error.getMessage().contains("overflows"), error.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..833c835200 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsHeaderParserTest.java @@ -0,0 +1,216 @@ +/* + * 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\":{}}", + // JSON hexadecimal digits are ASCII + "{\"__metadata__\":{\"note\":\"\\" + "uＦＦＦＦ\"}}", + // unescaped control character in a string + "{\"__metadata__\":{\"note\":\"line\nbreak\"}}", + // form feed is not JSON whitespace + "{\f}", + // 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]}}", + // duplicate dtype + "{\"w\":{\"dtype\":\"F32\",\"dtype\":\"F16\",\"shape\":[1]," + + "\"data_offsets\":[0,4]}}", + // duplicate shape + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"shape\":[2]," + + "\"data_offsets\":[0,4]}}", + // duplicate data_offsets + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]," + + "\"data_offsets\":[4,8]}}", + // duplicate metadata section + "{\"__metadata__\":{},\"__metadata__\":{}}", + // duplicate metadata key + "{\"__metadata__\":{\"format\":\"pt\",\"format\":\"tf\"}}", + // 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 data offset + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[-1,4]}}", + // reversed data offsets + "{\"w\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[4,0]}}", + // 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]}}", + // leading zero in an integer + "{\"w\":{\"dtype\":\"F32\",\"shape\":[01],\"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 testRejectsSignedUnicodeEscape() { + // 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 testRejectsMalformedNumberInSkippedField() { + // 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 testRejectsLoneMinusInSkippedField() { + // 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..e2c1d74d90 --- /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 malformed bytes construct those bytes directly. + */ +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..68cc52fa5c --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SafetensorsWriterTest.java @@ -0,0 +1,179 @@ +/* + * 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(delimiter = ';', value = { + "0;2;2;rows must be at least 1, got 0", + "2;0;2;cols must be at least 1, got 0", + "-1;2;2;rows must be at least 1, got -1", + "2;3;5;values has 5 elements, not 2 x 3", + "2;3;7;values has 7 elements, not 2 x 3", + "1;1;0;values has 0 elements, not 1 x 1" + }) + void testRejectsAShapeThatDoesNotMatchTheValues(int rows, int cols, int valueCount, + String expectedMessage, @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(expectedMessage, 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..24a9b1e89f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/SentencePieceModelFixture.java @@ -0,0 +1,169 @@ +/* + * 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; + +/** + * Builds a SentencePiece-layout embedding model around the bundled tiny Unigram model. Its + * tokenizer ids and matrix rows differ, which tests row lookup by piece text. + */ +final class SentencePieceModelFixture { + + static final String MODEL_RESOURCE = "/opennlp/embeddings/tiny-unigram.model"; + private static final char[] HEX_DIGITS = "0123456789abcdef".toCharArray(); + + 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 = modelResource()) { + modelBytes = in.readAllBytes(); + } + try (InputStream in = modelResource()) { + tokenizer = SentencePieceTokenizer.load(in); + } + 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 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 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("\\u") + .append(HEX_DIGITS[(c >>> 12) & 0xF]) + .append(HEX_DIGITS[(c >>> 8) & 0xF]) + .append(HEX_DIGITS[(c >>> 4) & 0xF]) + .append(HEX_DIGITS[c & 0xF]); + } else { + quoted.append(c); + } + } + } + } + return quoted.append('"').toString(); + } + + /** + * Opens the bundled SentencePiece model. + * + * @return A new stream for the model resource. + * @throws IOException Thrown if the resource is missing. + */ + private InputStream modelResource() throws IOException { + final InputStream in = SentencePieceModelFixture.class.getResourceAsStream(MODEL_RESOURCE); + if (in == null) { + throw new IOException("Missing test resource " + MODEL_RESOURCE); + } + return in; + } +} 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..3951080572 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelQuantizedTest.java @@ -0,0 +1,395 @@ +/* + * 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; + +/** + * Tests quantized WordPiece model loading, embedding, search, and file selection. + */ +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", + "xyzzy plugh" + }; + + /** + * 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 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 The quantization result. + * @throws IOException Thrown if quantizing or deleting fails. + */ + private 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 can + // be lower than a single-row reconstruction. The threshold decreases with the bit width and + // detects errors in the rotation, grid, or scale. + final double threshold = switch (bits) { + case 2 -> 0.88; + case 3 -> 0.95; + default -> 0.98; + }; + int compared = 0; + for (final String text : SENTENCES) { + final double cosine = cosine(floatModel.embed(text), quantizedModel.embed(text)); + if (Double.isNaN(cosine)) { + // Both models produced a zero vector for out-of-vocabulary text. + continue; + } + compared++; + assertTrue(cosine >= threshold, + bits + "-bit embedding of '" + text + "' has cosine " + cosine); + } + assertEquals(SENTENCES.length - 1, compared, + "only the out-of-vocabulary sentence should produce a zero vector"); + } + + @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 { + final Path weightedDirectory = Files.createDirectory(directory.resolve("weighted")); + final Path unweightedDirectory = Files.createDirectory(directory.resolve("unweighted")); + final List vocabulary = List.of("[UNK]", "[CLS]", "[SEP]", "hello", "world"); + final float[][] rows = { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {1f, 0f}, {0f, 1f} + }; + for (final Path modelDirectory : List.of(weightedDirectory, unweightedDirectory)) { + Files.write(modelDirectory.resolve(ModelFileNames.VOCABULARY), vocabulary); + Files.writeString(modelDirectory.resolve(ModelFileNames.CONFIG), "{\"normalize\": true}"); + Files.writeString(modelDirectory.resolve(ModelFileNames.TOKENIZER_CONFIG), + "{\"do_lower_case\": true}"); + } + SafetensorsTestFiles.write(weightedDirectory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", rows), + SafetensorsTestFiles.vector(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME, + new float[] {1f, 1f, 1f, 10f, 1f})); + SafetensorsTestFiles.write(unweightedDirectory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", rows)); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(weightedDirectory); + final StaticEmbeddingModel unweightedModel = StaticEmbeddingModel.load(unweightedDirectory); + final float[] expected = floatModel.embed("hello world"); + assertTrue(cosine(expected, unweightedModel.embed("hello world")) < 0.9, + "the fixture must distinguish weighted from unweighted pooling"); + + final ModelQuantizer.Result result = deployQuantized(weightedDirectory, 4); + assertTrue(result.hasWeights(), "the weights tensor must be included"); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(weightedDirectory); + final double cosine = cosine(expected, quantizedModel.embed("hello world")); + assertTrue(cosine > 0.98, "weighted quantized embedding has cosine " + cosine); + } + + @Test + void testQuantizedDirectoryIncludesTermRows(@TempDir Path directory) throws IOException { + Files.write(directory.resolve(ModelFileNames.VOCABULARY), + List.of("[CLS]", "[SEP]", "[UNK]", "habeas", "corpus")); + Files.write(directory.resolve(ModelFileNames.TERMS), List.of("habeas corpus")); + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {1f, 0f}, {0f, 1f}, {10f, 10f} + })); + Files.writeString(directory.resolve(ModelFileNames.CONFIG), "{\"normalize\": false}"); + Files.writeString(directory.resolve(ModelFileNames.TOKENIZER_CONFIG), + "{\"do_lower_case\": true}"); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + + deployQuantized(directory, 4); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + + assertEquals(1, quantizedModel.termCount()); + assertTrue(cosine(floatModel.embed("habeas corpus"), + quantizedModel.embed("habeas corpus")) > 0.98); + assertEquals("habeas corpus", + quantizedModel.mostSimilar("habeas corpus", 1).get(0).token()); + } + + @Test + void testQuantizedModel2VecUnigramDirectoryLoadsWithoutSafetensors(@TempDir Path directory) + throws IOException { + Files.writeString(directory.resolve(ModelFileNames.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(directory.resolve(ModelFileNames.CONFIG), "{\"normalize\":false}"); + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {2f}, {4f}, {8f}, {16f} + })); + final StaticEmbeddingModel floatModel = StaticEmbeddingModel.load(directory); + + deployQuantized(directory, 4); + final StaticEmbeddingModel quantizedModel = StaticEmbeddingModel.load(directory); + + assertEquals(6, quantizedModel.vocabularySize()); + assertTrue(cosine(floatModel.embed("hello world."), + quantizedModel.embed("hello world.")) > 0.98); + } + + @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 remain first after 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()); + assertFalse(List.of("hello", "world", "apple").contains(neighbors.get(0).token()), + "the analogy result must exclude every query term"); + } + + @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 testVerificationSampleDoesNotExceedItsCap(@TempDir Path directory) throws IOException { + final float[][] rows = new float[1025][1]; + for (int row = 0; row < rows.length; row++) { + rows[row][0] = row + 1f; + } + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", rows)); + + final ModelQuantizer.Result result = ModelQuantizer.quantize(directory, 4, SEED); + + assertEquals(1024, result.sampledRows()); + } + + @Test + void testVerificationCosineStaysWithinItsMathematicalRange() { + final float[] row = {6.0943845e19f, 2.0969745e19f}; + + assertEquals(1.0, ModelQuantizer.cosine(row, 0, row.length, row)); + } + + @Test + void testRowCountMismatchIsRejected(@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 testQuantizerRejectsBadBitWidthBeforeReadingTheMatrix(@TempDir Path directory) + throws IOException { + Files.write(directory.resolve(ModelFileNames.SAFETENSORS), new byte[] {1, 2, 3, 4}); + + final IllegalArgumentException error = assertThrows(IllegalArgumentException.class, + () -> ModelQuantizer.quantize(directory, 1, SEED)); + + assertTrue(error.getMessage().contains("Bits"), error.getMessage()); + } + + @Test + void testQuantizerRejectsMatrixShapedPoolingWeights(@TempDir Path directory) + throws IOException { + SafetensorsTestFiles.write(directory.resolve(ModelFileNames.SAFETENSORS), + SafetensorsTestFiles.matrix("embeddings", new float[][] {{1f, 2f}, {3f, 4f}}), + SafetensorsTestFiles.matrix(StaticEmbeddingModel.WEIGHTS_TENSOR_NAME, + new float[][] {{1f}, {1f}})); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> ModelQuantizer.quantize(directory, 4, SEED)); + + assertTrue(error.getMessage().contains("1-D"), error.getMessage()); + } + + @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 produce equal embeddings"); + } + } + + /** + * {@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 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..b5689622cc --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceQuantizedTest.java @@ -0,0 +1,101 @@ +/* + * 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 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests quantized SentencePiece model loading, embedding, and search. + */ +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(); + } + + @ParameterizedTest + @ValueSource(strings = {"a", "the model", "hello there world"}) + void testQuantizedSentencePieceEmbedsLikeTheFloatModel(String text, @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()); + final double cosine = cosine(floatModel.embed(text), quantizedModel.embed(text)); + assertFalse(Double.isNaN(cosine), "fixture text must produce a nonzero vector: " + text); + assertTrue(cosine > 0.97, + "quantized SentencePiece embedding of '" + text + "' has 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 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..88a0a4935b --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSentencePieceTest.java @@ -0,0 +1,313 @@ +/* + * 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", model.vocabularySize())) { + 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 testDirectoryLoadExplainsAnIncompleteSeparateFileTokenizer(@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..e1f037bd5c --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelSimilarityTest.java @@ -0,0 +1,299 @@ +/* + * 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 different directions (unlike {@link StaticEmbeddingModelTest}'s + * collinear rows, which are ideal for pooling-math assertions but would make every pairwise cosine + * similarity 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 { + // Unlike gensim's convention of excluding the query word, mostSimilar excludes only special + // tokens. A single-word query's own vocabulary row is therefore its 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 testMostSimilarOrdersEqualScoresByMatrixRow(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("ties-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "query", "first", "second")); + final Path tensors = dir.resolve("ties-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {1f, 0f}, {1f, 0f}, {1f, 0f} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.CASED, Normalization.NONE); + + assertEquals(List.of("query", "first"), + model.mostSimilar("query", 2).stream().map(Neighbor::token).toList()); + assertEquals(List.of("query", "first", "second"), + model.mostSimilar("query", 3).stream().map(Neighbor::token).toList()); + } + + @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 testMostSimilarDoesNotOverflowFiniteRows(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("large-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "large")); + final Path tensors = dir.resolve("large-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {0f}, {Float.MAX_VALUE} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.NONE); + + final Neighbor neighbor = model.mostSimilar("large", 1).get(0); + + assertEquals("large", neighbor.token()); + assertTrue(Double.isFinite(neighbor.similarity())); + assertEquals(1.0, neighbor.similarity(), 1e-12); + } + + @Test + void testCosineSimilarityStaysWithinItsDocumentedRange(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("rounding-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "rounding")); + final Path tensors = dir.resolve("rounding-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {6.0943845e19f, 2.0969745e19f} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.NONE); + + assertEquals(1.0, model.similarity("rounding", "rounding")); + assertEquals(1.0, model.mostSimilar("rounding", 1).get(0).similarity()); + } + + @Test + void testAnalogyDoesNotOverflowFiniteInputVectors(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("analogy-overflow-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "a", "b", "c", "answer")); + final Path tensors = dir.resolve("analogy-overflow-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {0f}, {-Float.MAX_VALUE}, {Float.MAX_VALUE}, {Float.MAX_VALUE}, {1f} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.CASED, Normalization.NONE); + + final Neighbor neighbor = model.analogy("a", "b", "c", 1).get(0); + + assertEquals("answer", neighbor.token()); + assertTrue(Double.isFinite(neighbor.similarity())); + assertEquals(1.0, 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..5e54561397 --- /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 is rejected during loading. + */ +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 testRejectsARowCountMismatchWithTerms(@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 testRejectsAMalformedTermsFile(@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..fdb38a3bd8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingModelTest.java @@ -0,0 +1,607 @@ +/* + * 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 testMeanPoolingDoesNotOverflowFiniteRows(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("large-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "large", "value")); + final Path tensors = dir.resolve("large-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {0f}, {Float.MAX_VALUE}, {Float.MAX_VALUE} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.NONE); + + final float[] result = model.embed("large value"); + + assertArrayEquals(new float[] {Float.MAX_VALUE}, result); + } + + @Test + void testL2NormalizationHandlesFiniteVectorsWhoseNormExceedsFloatRange(@TempDir Path dir) + throws IOException { + final Path vocabulary = dir.resolve("large-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "large")); + final Path tensors = dir.resolve("large-model.safetensors"); + SafetensorsTestFiles.write(tensors, SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {Float.MAX_VALUE, Float.MAX_VALUE} + })); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.L2); + + final float[] result = model.embed("large"); + + final float normalizedCoordinate = (float) (1.0 / Math.sqrt(2.0)); + assertArrayEquals(new float[] {normalizedCoordinate, normalizedCoordinate}, result, 1e-6f); + } + + @Test + void testFinitePoolingWeightsDoNotProduceInfiniteCoordinates(@TempDir Path dir) + throws IOException { + final Path vocabulary = dir.resolve("weighted-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "large")); + final Path tensors = dir.resolve("weighted-model.safetensors"); + SafetensorsTestFiles.write(tensors, + SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f}, {0f}, {0f}, {Float.MAX_VALUE} + }), + SafetensorsTestFiles.vector("weights", new float[] {1f, 1f, 1f, Float.MAX_VALUE})); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.NONE); + + final float[] result = model.embed("large"); + + assertArrayEquals(new float[] {Float.MAX_VALUE}, result); + } + + @Test + void testL2NormalizationPrecedesFloatNarrowing(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("weighted-vocab.txt"); + Files.write(vocabulary, List.of("[CLS]", "[SEP]", "[UNK]", "large")); + final Path tensors = dir.resolve("weighted-model.safetensors"); + SafetensorsTestFiles.write(tensors, + SafetensorsTestFiles.matrix("embeddings", new float[][] { + {0f, 0f}, {0f, 0f}, {0f, 0f}, {Float.MAX_VALUE, 1f} + }), + SafetensorsTestFiles.vector("weights", new float[] {1f, 1f, 1f, Float.MAX_VALUE})); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.L2); + + final float[] result = model.embed("large"); + + assertEquals(1f, result[0]); + assertEquals((float) (1.0 / Float.MAX_VALUE), result[1], Float.MIN_VALUE); + } + + @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 testLoadRejectsAZeroDimensionMatrix(@TempDir Path dir) throws IOException { + final float[][] rows = new float[VOCAB_TOKENS.size()][0]; + final Path tensors = dir.resolve("zero-dimension.safetensors"); + SafetensorsTestFiles.write(tensors, + SafetensorsTestFiles.matrix("embeddings", rows)); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(writeVocab(dir), tensors, + Casing.UNCASED, Normalization.NONE)); + + assertTrue(error.getMessage().contains("dimension"), error.getMessage()); + } + + @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. Loading must reject it because 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 testLoadRejectsANonFiniteWeight(@TempDir Path dir) throws IOException { + final Path vocab = writeVocab(dir); + final float[] weights = {1f, 1f, 1f, 1f, Float.NaN, 1f}; + final Path nanTensors = dir.resolve("nan-weight.safetensors"); + SafetensorsTestFiles.write(nanTensors, + SafetensorsTestFiles.matrix("embeddings", ROWS), + SafetensorsTestFiles.vector("weights", weights)); + + final InvalidFormatException nan = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocab, nanTensors, + Casing.UNCASED, Normalization.NONE)); + assertTrue(nan.getMessage().contains("weights"), nan.getMessage()); + assertTrue(nan.getMessage().contains("row 4"), nan.getMessage()); + + weights[4] = Float.NEGATIVE_INFINITY; + final Path infiniteTensors = dir.resolve("infinite-weight.safetensors"); + SafetensorsTestFiles.write(infiniteTensors, + SafetensorsTestFiles.matrix("embeddings", ROWS), + SafetensorsTestFiles.vector("weights", weights)); + + final InvalidFormatException infinite = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocab, infiniteTensors, + Casing.UNCASED, Normalization.NONE)); + assertTrue(infinite.getMessage().contains("weights"), infinite.getMessage()); + 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")); + } + + @Test + void testLoadRejectsWeightsThatAreNotOneDimensional(@TempDir Path dir) throws IOException { + final Path vocabulary = dir.resolve("scalar-weight-vocab.txt"); + Files.write(vocabulary, List.of("[UNK]")); + final Path tensors = dir.resolve("scalar-weight.safetensors"); + SafetensorsTestFiles.write(tensors, + SafetensorsTestFiles.matrix("embeddings", new float[][] {{1f}}), + new SafetensorsTestFiles.Tensor("weights", new int[0], new float[] {1f})); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> StaticEmbeddingModel.load(vocabulary, tensors, + Casing.UNCASED, Normalization.NONE)); + + assertTrue(exception.getMessage().contains("weights"), exception.getMessage()); + assertTrue(exception.getMessage().contains("1-D"), exception.getMessage()); + } + + // 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, so a config declaring another operation is invalid. + 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 testTextEmbedderInterfaceMatchesDirectUseAndBatches(@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..3cb56e5cd8 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingSearchExampleTest.java @@ -0,0 +1,65 @@ +/* + * 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; + +/** + * Demonstrates semantic search with the query and documents shown in {@code embeddings.xml}. + */ +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.writeSearchDirectory(dir); + final StaticEmbeddingModel model = StaticEmbeddingModel.load(dir); + + final String query = "home espresso machine"; + final List documents = List.of( + "How do I brew espresso at home?", + "The history of tea in East Asia", + "Best grinders for pour-over coffee"); + + 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()); + + assertEquals(List.of( + "How do I brew espresso at home?", + "Best grinders for pour-over coffee", + "The history of tea in East Asia"), + results.stream().map(Scored::document).toList()); + assertTrue(results.get(0).score() > results.get(1).score()); + assertTrue(results.get(1).score() > results.get(2).score()); + } +} 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..69a404472f --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/StaticEmbeddingUsageExampleTest.java @@ -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. + */ +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; + +/** + * Demonstrates loading, embedding, similarity, nearest-neighbor search, and analogy operations. + */ +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..c3280b6173 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TeacherTokenizerTest.java @@ -0,0 +1,642 @@ +/* + * 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.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 { + + private static final String MINIMAL_WORDPIECE_MODEL = + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}"; + + // 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()); + } + + @ParameterizedTest + @ValueSource(strings = { + "{" + MINIMAL_WORDPIECE_MODEL + "," + MINIMAL_WORDPIECE_MODEL + "}", + "{\"added_tokens\":[],\"added_tokens\":[]," + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":null,\"post_processor\":null," + MINIMAL_WORDPIECE_MODEL + "}", + "{\"normalizer\":{},\"normalizer\":{}," + MINIMAL_WORDPIECE_MODEL + "}", + "{\"model\":{\"type\":\"WordPiece\",\"type\":\"WordPiece\"," + + "\"unk_token\":\"a\",\"vocab\":{\"a\":0}}}", + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\"," + + "\"unk_token\":\"a\",\"vocab\":{\"a\":0}}}", + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\"," + + "\"vocab\":{\"a\":0},\"vocab\":{\"a\":0}}}" + }) + void testRejectsDuplicateTokenizerFields(String json, @TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", json); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(exception.getMessage().contains("more than once"), exception.getMessage()); + } + + @ParameterizedTest + @ValueSource(strings = { + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"type\":\"TemplateProcessing\",\"single\":\"$A\"," + + "\"special_tokens\":{ }}," + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\"$A\",\"single\":\"$A\",\"special_tokens\":{ }}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\"$A\",\"special_tokens\":{ },\"special_tokens\":{ }}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"BertProcessing\"," + + "\"cls\":[\"a\",0],\"cls\":[\"a\",0],\"sep\":[\"a\",0]}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"BertProcessing\"," + + "\"cls\":[\"a\",0],\"sep\":[\"a\",0],\"sep\":[\"a\",0]}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":[{\"Sequence\":{\"id\":\"A\",\"id\":\"A\"," + + "\"type_id\":0}}],\"special_tokens\":{ }}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\"[CLS] $A\",\"special_tokens\":{" + + "\"[CLS]\":{\"ids\":[0]},\"[CLS]\":{\"ids\":[0]}}}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\"[CLS] $A\",\"special_tokens\":{" + + "\"[CLS]\":{\"ids\":[0],\"ids\":[0]}}}," + + MINIMAL_WORDPIECE_MODEL + "}", + "{\"added_tokens\":[{\"content\":\"a\",\"content\":\"a\"}]," + + MINIMAL_WORDPIECE_MODEL + "}" + }) + void testRejectsDuplicateNestedTokenizerFields(String json, @TempDir Path dir) + throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", json); + + final InvalidFormatException exception = assertThrows(InvalidFormatException.class, + () -> TeacherTokenizer.read(tokenizerJson, null)); + assertTrue(exception.getMessage().contains("more than once"), exception.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 an expected message fragment. + */ + @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", + "a duplicate vocabulary token;{\"model\":{\"type\":\"Unigram\",\"unk_id\":0," + + "\"vocab\":[[\"\",0.0],[\"piece\",-1.0],[\"piece\",-2.0]]}};" + + "declares token 'piece' more than once", + "an unsupported post-processor;{\"post_processor\":{\"type\":\"ByteLevel\"}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "is not supported", + "a post-processor object without a type;{\"post_processor\":{}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "post_processor.type is required", + "a TemplateProcessing post-processor without a single template;{\"post_processor\":{" + + "\"type\":\"TemplateProcessing\",\"special_tokens\":{}}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "post_processor.single is required", + "a string template without a sequence;{\"post_processor\":{" + + "\"type\":\"TemplateProcessing\",\"single\":\"[CLS] [SEP]\"," + + "\"special_tokens\":{\"[CLS]\":{\"ids\":[0]},\"[SEP]\":{\"ids\":[0]}}}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "exactly one sequence placeholder", + "a structured template with two sequences;{\"post_processor\":{" + + "\"type\":\"TemplateProcessing\",\"single\":[" + + "{\"Sequence\":{\"id\":\"A\",\"type_id\":0}}," + + "{\"Sequence\":{\"id\":\"A\",\"type_id\":0}}],\"special_tokens\":{}}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "exactly one sequence placeholder", + "a BertProcessing post-processor without cls;{\"post_processor\":{" + + "\"type\":\"BertProcessing\",\"sep\":[\"a\",0]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "post_processor.cls is required", + "a RobertaProcessing post-processor without sep;{\"post_processor\":{" + + "\"type\":\"RobertaProcessing\",\"cls\":[\"a\",0]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"a\",\"vocab\":{\"a\":0}}};" + + "post_processor.sep is required", + "a negative post-processor id;{\"post_processor\":{\"type\":\"BertProcessing\"," + + "\"cls\":[\"[CLS]\",-1],\"sep\":[\"[SEP]\",1]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0,\"[SEP]\":1}}};outside the supported integer range", + "an overflowing post-processor id;{\"post_processor\":{\"type\":\"TemplateProcessing\"," + + "\"single\":\"[CLS] $A\",\"special_tokens\":{\"[CLS]\":{" + + "\"ids\":[4294967296]}}},\"model\":{\"type\":\"WordPiece\"," + + "\"unk_token\":\"[UNK]\",\"vocab\":{\"[UNK]\":0}}};" + + "outside the supported integer range", + "a structured special token without an id;{\"post_processor\":{" + + "\"type\":\"TemplateProcessing\",\"single\":[{\"SpecialToken\":{}}]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0}}};SpecialToken template item needs an id", + "a structured sequence without an id;{\"post_processor\":{" + + "\"type\":\"TemplateProcessing\",\"single\":[{\"Sequence\":{}}]}," + + "\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[UNK]\"," + + "\"vocab\":{\"[UNK]\":0}}};Sequence template item needs an id", + "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()); + } + + @Test + void testUnusedPatternDoesNotRemoveTheUnknownToken(@TempDir Path dir) throws IOException { + final Path tokenizerJson = write(dir, "tokenizer.json", + "{\"model\":{\"type\":\"WordPiece\",\"unk_token\":\"[unused0]\"," + + "\"vocab\":{\"[unused0]\":0,\"hello\":1}}}"); + + final TeacherTokenizer tokenizer = TeacherTokenizer.read(tokenizerJson, null); + + assertEquals("[unused0]", tokenizer.unkToken()); + assertArrayEquals(new int[] {0, 1}, tokenizer.keptOriginalIds()); + } + + /** + * 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()); + assertEquals("pieces[0] must not be null", assertThrows(IllegalArgumentException.class, + () -> tokenizer.inputSequence(Collections.singletonList(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..d849cedddf --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/TokenizerJsonVocabTest.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.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 + * input outside the expected structure is rejected. + */ +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 testRejectsDuplicateAddedTokenIds() throws IOException { + final Path file = write("{\"added_tokens\":[" + + "{\"id\":0,\"content\":\"a\",\"special\":false}," + + "{\"id\":0,\"content\":\"a\",\"special\":true}]," + + "\"model\":{\"type\":\"Unigram\",\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.read(file)); + + assertTrue(error.getMessage().contains("id 0 occurs more than once"), error.getMessage()); + } + + @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 testRejectsAModelWithoutAType() throws IOException { + final Path file = write("{\"model\":{\"vocab\":[[\"a\",0.0]]}}"); + + final InvalidFormatException e = assertThrows(InvalidFormatException.class, + () -> TokenizerJsonVocab.rows(file)); + assertTrue(e.getMessage().contains("model.type"), 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 the parser must reject this form. + 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/cmdline/CLITest.java b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.java new file mode 100644 index 0000000000..9c641bf55d --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.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.cmdline; + +import java.nio.file.Path; +import java.util.Set; +import java.util.stream.Stream; + +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.MethodSource; + +import opennlp.tools.cmdline.BasicCmdLineTool; +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 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); + } + + @Test + void testQuantizeHelpNamesEveryParameter() { + final String help = new QuantizeModelTool().getHelp(); + + assertTrue(help.contains("-modelDir dir"), help); + assertTrue(help.contains("[-bits bits]"), help); + assertTrue(help.contains("[-seed seed]"), help); + } + + @Test + void testQuantizeReportsInvalidModelContentAsAUserError(@TempDir Path directory) { + final TerminateToolException error = assertThrows(TerminateToolException.class, + () -> new QuantizeModelTool().run(new String[] {"-modelDir", directory.toString()})); + + assertEquals(1, error.getCode()); + } +} 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..78a434ebd8 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/BpeEncoder.java @@ -0,0 +1,234 @@ +/* + * 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 kept intact) 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 = 112585252536688886L; + + private static final int MAX_RESEGMENT_DEPTH = 100; + + private final HashMap 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(HashMap 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. Protected-symbol flags travel as 0/1 bytes parallel to the ranges. + final IntBuilder fromB = new IntBuilder(size); + final IntBuilder toB = new IntBuilder(size); + final ByteBuilder protectedB = new ByteBuilder(size); + int position = 0; + while (position < size) { + int matched = 0; + if (userDefinedMatcher != null) { + matched = userDefinedMatcher.longestMatch(normalized, size, position); + } + final boolean protectedSymbol = matched > 0; + final int length = protectedSymbol ? matched + : Math.min(SentencePieceNormalizer.utf8Length(normalized[position]), + size - position); + fromB.append(position); + toB.append(position + length); + protectedB.append(protectedSymbol ? (byte) 1 : (byte) 0); + position += length; + } + final int symbolCount = protectedB.length(); + final int[] from = fromB.toArray(); + final int[] to = toB.toArray(); + final byte[] protectedFlags = protectedB.array(); + final int[] prev = new int[symbolCount]; + final int[] next = new int[symbolCount]; + final boolean[] protectedSymbols = new boolean[symbolCount]; + for (int i = 0; i < symbolCount; i++) { + prev[i] = i - 1; + next[i] = i + 1 < symbolCount ? i + 1 : -1; + protectedSymbols[i] = protectedFlags[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, protectedSymbols, 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, protectedSymbols, + prev[top.left()], top.left(), agenda, revMerge); + maybeAddPair(normalized, from, to, protectedSymbols, + 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 protectedSymbols 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[] protectedSymbols, + int left, int right, PriorityQueue agenda, + Map revMerge) { + if (left == -1 || right == -1 || protectedSymbols[left] || protectedSymbols[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..4cbd594482 --- /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 valid trie never produces, are rejected + * 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..64b54aca1e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ModelProtoReader.java @@ -0,0 +1,458 @@ +/* + * 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.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CodingErrorAction; +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 is rejected.

+ * + * @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.readTag(data.length); + switch (fieldOf(tag)) { + case FIELD_MODEL_PIECES -> reader.piece(model, reader.lenPayload(tag, data.length)); + case FIELD_MODEL_TRAINER_SPEC -> + reader.trainerSpec(model, reader.lenPayload(tag, data.length)); + case FIELD_MODEL_NORMALIZER_SPEC -> + reader.normalizerSpec(model, reader.lenPayload(tag, data.length)); + case FIELD_MODEL_SELF_TEST_DATA -> + reader.selfTestData(model, reader.lenPayload(tag, data.length)); + default -> reader.skip(tag, data.length); + } + } + 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 = readTag(end); + switch (fieldOf(tag)) { + case FIELD_PIECE_PIECE -> piece = utf8(lenPayload(tag, end)); + case FIELD_PIECE_SCORE -> score = fixed32Float(tag, end); + case FIELD_PIECE_TYPE -> type = (int) varintOf(tag, end); + default -> skip(tag, end); + } + } + 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 = readTag(end); + switch (fieldOf(tag)) { + case FIELD_TRAINER_MODEL_TYPE -> model.modelType = (int) varintOf(tag, end); + case FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX -> + model.treatWhitespaceAsSuffix = varintOf(tag, end) != 0; + case FIELD_TRAINER_BYTE_FALLBACK -> model.byteFallback = varintOf(tag, end) != 0; + default -> skip(tag, end); + } + } + } + + /** + * 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 = readTag(end); + switch (fieldOf(tag)) { + case FIELD_NORMALIZER_PRECOMPILED_CHARSMAP -> + model.precompiledCharsMap = bytes(lenPayload(tag, end)); + case FIELD_NORMALIZER_ADD_DUMMY_PREFIX -> + model.addDummyPrefix = varintOf(tag, end) != 0; + case FIELD_NORMALIZER_REMOVE_EXTRA_WHITESPACES -> + model.removeExtraWhitespaces = varintOf(tag, end) != 0; + case FIELD_NORMALIZER_ESCAPE_WHITESPACES -> + model.escapeWhitespaces = varintOf(tag, end) != 0; + default -> skip(tag, end); + } + } + } + + /** + * 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 = readTag(end); + if (fieldOf(tag) == FIELD_SELF_TEST_SAMPLES) { + final int sampleEnd = lenPayload(tag, end); + String input = null; + String expected = null; + while (pos < sampleEnd) { + final long sampleTag = readTag(sampleEnd); + switch (fieldOf(sampleTag)) { + case FIELD_SAMPLE_INPUT -> input = utf8(lenPayload(sampleTag, sampleEnd)); + case FIELD_SAMPLE_EXPECTED -> expected = utf8(lenPayload(sampleTag, sampleEnd)); + default -> skip(sampleTag, sampleEnd); + } + } + if (input != null && expected != null) { + model.selfTestInputs.add(input); + model.selfTestExpected.add(expected); + } + } else { + skip(tag, end); + } + } + } + + /** + * 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. + * @param limit The exclusive end offset of the enclosing message. + * @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, int limit) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_LEN) { + throw malformed("field " + fieldOf(tag) + " is not length-delimited"); + } + final long length = varint(limit); + if (length < 0 || length > limit - pos) { + throw malformed("length " + length + " crosses its message boundary"); + } + 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. + * @param limit The exclusive end offset of the enclosing message. + * @return The decoded value. + * @throws InvalidFormatException Thrown if the wire type is wrong or the varint is malformed. + */ + private long varintOf(long tag, int limit) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_VARINT) { + throw malformed("field " + fieldOf(tag) + " is not a varint"); + } + return varint(limit); + } + + /** + * 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. + * @param limit The exclusive end offset of the enclosing message. + * @return The decoded float. + * @throws InvalidFormatException Thrown if the wire type is wrong or the input is truncated. + */ + private float fixed32Float(long tag, int limit) throws InvalidFormatException { + if (wireTypeOf(tag) != WIRE_FIXED32) { + throw malformed("field " + fieldOf(tag) + " is not a 32-bit value"); + } + if (limit - pos < 4) { + 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. + * @throws InvalidFormatException Thrown if the payload is not valid UTF-8. + */ + private String utf8(int end) throws InvalidFormatException { + try { + final String decoded = StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(data, pos, end - pos)) + .toString(); + pos = end; + return decoded; + } catch (CharacterCodingException e) { + throw malformed("string is not valid UTF-8"); + } + } + + /** + * 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 and validates a field tag. + * + * @param limit The exclusive end offset of the enclosing message. + * @return The decoded tag. + * @throws InvalidFormatException Thrown if the tag is malformed or has an invalid field number. + */ + private long readTag(int limit) throws InvalidFormatException { + final long tag = varint(limit); + final long field = tag >>> 3; + if (field == 0 || field > 0x1FFFFFFFL) { + throw malformed("invalid field number " + field); + } + return tag; + } + + /** + * Reads a base-128 varint without crossing its enclosing message. + * + * @param limit The exclusive end offset of the enclosing message. + * @return The decoded value. + * @throws InvalidFormatException Thrown if the input ends mid-varint or the varint exceeds 64 + * bits. + */ + private long varint(int limit) throws InvalidFormatException { + long value = 0; + for (int shift = 0; shift < 64; shift += 7) { + if (pos >= limit) { + throw malformed("truncated varint"); + } + final byte b = data[pos++]; + if (shift == 63 && (b & 0x7E) != 0) { + throw malformed("varint exceeds 64 bits"); + } + 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. + * @param limit The exclusive end offset of the enclosing message. + * @throws InvalidFormatException Thrown if the wire type is unsupported or the value runs past + * the input. + */ + private void skip(long tag, int limit) throws InvalidFormatException { + switch (wireTypeOf(tag)) { + case WIRE_VARINT -> varint(limit); + case WIRE_FIXED64 -> advance(8, limit); + case WIRE_LEN -> pos = lenPayload(tag, limit); + case WIRE_FIXED32 -> advance(4, limit); + 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. + * @param limit The exclusive end offset of the enclosing message. + * @throws InvalidFormatException Thrown if fewer than {@code count} bytes remain. + */ + private void advance(int count, int limit) throws InvalidFormatException { + if (limit - pos < count) { + 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..5b58d225b0 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/PieceTrie.java @@ -0,0 +1,310 @@ +/* + * 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)); + + // The counting and filling passes use the same traversal over the sorted keys. + 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..cf2df762db --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceTokenizer.java @@ -0,0 +1,780 @@ +/* + * 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.InvalidClassException; +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.

+ * + *

A tokenizer can also be persisted with {@link #serialize(OutputStream)} and restored with + * {@link #deserialize(InputStream)}. Deserialization accepts only the classes used by this + * tokenizer and applies limits to 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 = 7751888381608757475L; + + /** 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 HashMap mainPieces; + private final HashMap reservedPieces; + private final int[] byteToId; + private final SentencePieceNormalizer normalizer; + private final UnigramEncoder unigramEncoder; + private final BpeEncoder bpeEncoder; + private final String[] selfTestInputs; + private final String[] 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 = model.selfTestInputs.toArray(String[]::new); + selfTestExpected = model.selfTestExpected.toArray(String[]::new); + } + + /** + * 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.

+ * + *

Only deserialize tokenizer streams from trusted sources. If the default limits reject a + * 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)); + final Object value = ois.readObject(); + if (!(value instanceof SentencePieceTokenizer tokenizer)) { + final String type = value == null ? "null" : value.getClass().getName(); + throw new InvalidClassException( + "Expected a SentencePieceTokenizer, found " + type + "."); + } + return tokenizer; + } + } + + /** + * Resource limits applied by the {@link ObjectInputFilter} used by + * {@link SentencePieceTokenizer#deserialize(InputStream, DeserializationLimits)}. + * + *

The limits bound graph traversal independently of the class allow-list. Raise the + * {@linkplain #DEFAULT default values} only when they reject a valid model.

+ * + * @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" + ); + + /** + * 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 List.of(selfTestInputs); + } + + /** {@return the embedded self-test expected segmentations} */ + List selfTestExpected() { + return List.of(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 = asciiHexValue(piece.charAt(3)); + final int low = asciiHexValue(piece.charAt(4)); + return high < 0 || low < 0 ? -1 : (high << 4) | low; + } + + /** {@return the value of an ASCII hexadecimal digit, or {@code -1} for another character} */ + private static int asciiHexValue(char c) { + if (c >= '0' && c <= '9') { + return c - '0'; + } + if (c >= 'A' && c <= 'F') { + return c - 'A' + 10; + } + if (c >= 'a' && c <= 'f') { + return c - 'a' + 10; + } + return -1; + } +} 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..189f3fb25f --- /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 testRejectsDuplicatePieces() { + 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..16c40e9e2d --- /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 testRejectsNullInputs() { + 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..f2c859eb7e --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.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.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; + +/** Validates malformed-model rejection and concurrent tokenizer use. */ +class SentencePieceModelValidationTest { + + private static final String HEX_DIGITS = "0123456789ABCDEF"; + + @Test + void testRejectsNullAndEmptyInput() { + 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 testRejectsGarbageBytes() { + 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 testRejectsTruncatedModel() 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 testRejectsAPieceFieldThatCrossesItsMessageBoundary() { + final byte[] model = { + 0x0A, 0x02, // pieces sub-message with a two-byte payload + 0x0A, 0x04, // piece string claiming four bytes outside that payload + 'a', 'b', 'c', 'd'}; + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> ModelProtoReader.read(model)); + + assertTrue(error.getMessage().contains("message boundary"), error.getMessage()); + } + + @Test + void testRejectsMalformedUtf8InAPiece() { + final byte[] model = { + 0x0A, 0x03, // pieces sub-message + 0x0A, 0x01, (byte) 0xFF}; + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> ModelProtoReader.read(model)); + + assertTrue(error.getMessage().contains("UTF-8"), error.getMessage()); + } + + /** Verifies that the tenth byte of a 64-bit varint cannot carry more than one value bit. */ + @Test + void testRejectsVarintLargerThan64Bits() { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + out.writeBytes(minimalModel(ModelProtoReader.RawModel.MODEL_TYPE_UNIGRAM)); + out.write(0x28); // Unknown field 5 with the varint wire type. + for (int i = 0; i < 9; i++) { + out.write(0x80); + } + out.write(0x02); + + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> ModelProtoReader.read(out.toByteArray())); + + assertTrue(error.getMessage().contains("64 bits"), error.getMessage()); + } + + @Test + void testRejectsUnsupportedModelType() { + // 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 testRejectsMissingUnknownPiece() { + 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 testRejectsMalformedPrecompiledCharsMap() { + // 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)); + } + + @Test + void testRejectsNonAsciiHexadecimalBytePiece() { + final InvalidFormatException error = assertThrows(InvalidFormatException.class, + () -> SentencePieceTokenizer.load(new ByteArrayInputStream( + byteFallbackModel("<0x\uff26F>")))); + + assertTrue(error.getMessage().contains("invalid"), error.getMessage()); + } + + /** + * Builds a byte-fallback model, replacing the {@code <0xFF>} piece with the supplied text. + * + * @param lastBytePiece The text of the piece assigned to byte {@code 0xFF}. + * @return The encoded model. + */ + private static byte[] byteFallbackModel(String lastBytePiece) { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + writePiece(out, "", 2); + for (int b = 0; b < 256; b++) { + final String piece = "<0x" + HEX_DIGITS.charAt(b >>> 4) + + HEX_DIGITS.charAt(b & 0x0f) + ">"; + writePiece(out, b == 255 ? lastBytePiece : piece, 6); + } + // trainer_spec { byte_fallback = true } + out.write(0x12); + out.write(3); + out.write(0x98); + out.write(0x02); + out.write(1); + return out.toByteArray(); + } + + 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..eacd8a94f4 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java @@ -0,0 +1,180 @@ +/* + * 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 an allow-listed leaf type cannot be returned as the top-level object. + */ + @Test + void testAllowListedLeafPayloadIsRejected() throws IOException { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + out.writeObject("not a tokenizer"); + } + assertThrows(InvalidClassException.class, () -> + SentencePieceTokenizer.deserialize(new ByteArrayInputStream(bytes.toByteArray()))); + } + + /** Verifies that a serialized null cannot be returned as a tokenizer. */ + @Test + void testNullPayloadIsRejected() throws IOException { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + out.writeObject(null); + } + 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..bcd16eb53c --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceUsageExampleTest.java @@ -0,0 +1,68 @@ +/* + * 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; + +/** + * Demonstrates loading a {@link SentencePieceTokenizer}, retaining original-text offsets, and + * obtaining token ids. + */ +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..faa795ccd6 --- /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 093d7cd503..dec360a7c1 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