OPENNLP-1895: Add quantized vector storage and search for static embeddings

dev/embeddings/README.md

@@ -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.
+-->
+
+# 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.

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 <output-dir>
+
+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)

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.
+ *
+ * <p>Args: modelDir sentencesFile vectorsOut warmupSeconds measureSeconds</p>
+ */
+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<String> 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);
+  }
+}

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 <model-dir> [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()

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"

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.