diff --git a/.github/workflows/allowlist-check.yml b/.github/workflows/allowlist-check.yml index 9fec5505b2..463a92ca63 100644 --- a/.github/workflows/allowlist-check.yml +++ b/.github/workflows/allowlist-check.yml @@ -32,6 +32,13 @@ on: paths: - ".github/**" +concurrency: + # A new commit on a pull request makes the run in flight obsolete, so cancel it + # instead of paying for a result nobody will read. Pushes to a branch group by ref + # and are never cancelled, so every commit on main keeps a full build. + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + permissions: contents: read diff --git a/.github/workflows/license.yml b/.github/workflows/license.yml index 1832049149..c6fbcd1c72 100644 --- a/.github/workflows/license.yml +++ b/.github/workflows/license.yml @@ -28,7 +28,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up JDK 21 - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: java-version: '21' distribution: 'temurin' diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 2ff6080ee9..1c3a9d6cb0 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -24,6 +24,13 @@ on: branches: - main +concurrency: + # A new commit on a pull request makes the run in flight obsolete, so cancel it + # instead of paying for a result nobody will read. Pushes to a branch group by ref + # and are never cancelled, so every commit on main keeps a full build. + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: build: runs-on: ${{ matrix.os }} @@ -49,7 +56,7 @@ jobs: restore-keys: | ${{ runner.os }}-maven- - name: Set up JDK ${{ matrix.java }} - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: distribution: temurin java-version: ${{ matrix.java }} @@ -89,7 +96,7 @@ jobs: sudo apt-get install -y bats - name: Set up JDK 21 - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: distribution: temurin java-version: 21 @@ -148,7 +155,7 @@ jobs: brew install bats-core - name: Set up JDK 21 - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: distribution: temurin java-version: 21 @@ -199,7 +206,7 @@ jobs: shell: pwsh - name: Set up JDK 21 - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: distribution: temurin java-version: 21 diff --git a/.github/workflows/publish-snapshots.yml b/.github/workflows/publish-snapshots.yml index bf0d6f115e..2b7ad75d3b 100644 --- a/.github/workflows/publish-snapshots.yml +++ b/.github/workflows/publish-snapshots.yml @@ -41,7 +41,7 @@ jobs: restore-keys: | ${{ runner.os }}-maven- - name: Setup Java - uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5.7.0 + uses: actions/setup-java@dd06d9cba3e5552c54d9f8ea23572deb30010f7c # v6.0.0 with: distribution: temurin java-version: 21 diff --git a/.mvn/wrapper/maven-wrapper.properties b/.mvn/wrapper/maven-wrapper.properties index 48a56c99aa..216df05897 100644 --- a/.mvn/wrapper/maven-wrapper.properties +++ b/.mvn/wrapper/maven-wrapper.properties @@ -1,19 +1,3 @@ -# 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. -wrapperVersion=3.3.2 +wrapperVersion=3.3.4 distributionType=only-script -distributionUrl=https://repo.maven.apache.org/maven2/org/apache/maven/apache-maven/3.9.6/apache-maven-3.9.6-bin.zip +distributionUrl=https://repo.maven.apache.org/maven2/org/apache/maven/apache-maven/3.9.16/apache-maven-3.9.16-bin.zip diff --git a/NOTICE b/NOTICE index 4ae212e882..5de2119011 100644 --- a/NOTICE +++ b/NOTICE @@ -162,10 +162,10 @@ List of third-party dependencies grouped by their license type. MIT License - * onnx-runtime (com.microsoft.onnxruntime:onnxruntime_gpu:1.28.0 - https://microsoft.github.io/onnxruntime/) - * onnx-runtime (com.microsoft.onnxruntime:onnxruntime:1.28.0 - https://microsoft.github.io/onnxruntime/) + * onnx-runtime (com.microsoft.onnxruntime:onnxruntime_gpu:1.29.0 - https://microsoft.github.io/onnxruntime/) + * onnx-runtime (com.microsoft.onnxruntime:onnxruntime:1.29.0 - https://microsoft.github.io/onnxruntime/) * SLF4J API Module (org.slf4j:slf4j-api:2.0.18 - http://www.slf4j.org) The MIT License (MIT) - * ClassGraph (io.github.classgraph:classgraph:4.8.186 - https://github.com/classgraph/classgraph) + * ClassGraph (io.github.classgraph:classgraph:4.8.194 - https://github.com/classgraph/classgraph) diff --git a/SECURITY.md b/SECURITY.md index 769559a582..f39f6a6d2a 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -148,9 +148,14 @@ JVM arguments, or the `OPENNLP_MAX_ENTRIES` setting. CDN (`https://dlcdn.apache.org/opennlp/`) and verifies each download against the published SHA-512 checksum before use. A mismatch fails the load. -Note that verification happens at **download** time. A model already present in -the local download cache is currently not re-verified when it is loaded again, so -treat the cache directory (`~/.opennlp` by default) as trusted storage and protect +The published checksum is stored next to the model in the download cache, and a +model already present in that cache is re-verified against it on every load. This +check is performed locally and does not contact the CDN again. A cache entry +written by an OpenNLP version that predates this behaviour has no stored checksum; +for those the published checksum is fetched once and then stored, and if it cannot +be retrieved the model is loaded and a warning is logged. + +Treat the cache directory (`~/.opennlp` by default) as trusted storage and protect it accordingly. The base URL can be overridden with the `OPENNLP_DOWNLOAD_BASE_URL` system diff --git a/dev/README-hunspell-dictionaries.md b/dev/README-hunspell-dictionaries.md new file mode 100644 index 0000000000..711435d948 --- /dev/null +++ b/dev/README-hunspell-dictionaries.md @@ -0,0 +1,108 @@ + + +# Hunspell dictionaries for the affix stemmer + +The Hunspell stemmer (`opennlp.tools.stemmer.hunspell`) reads a user-supplied +`.dic` word list and its `.aff` affix file. Apache OpenNLP bundles no dictionary +data. The dictionary's readme states its license. + +## Where dictionaries come from + +The LibreOffice project maintains Hunspell dictionaries by language at +`github.com/LibreOffice/dictionaries`. Each dictionary has a separate license. +For example, SCOWL is the source for the `en_US` dictionary, with terms in +`README_en_US.txt`. Other sources can be used when the `.aff` and `.dic` files +follow the Hunspell format. + +OpenNLP does not ship a URL catalog. Applications that manage downloads can keep a +properties file with an entry id followed by `.url`, `.sha512`, and optionally +`.filename` keys. Use a URL for a stable release or commit. + +## Option A: application catalog + +Catalog downloads stay inactive until you set `-Dopennlp.download.remote=true`. + +```java +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import opennlp.tools.stemmer.hunspell.HunspellDictionaryDownload; +import opennlp.tools.util.DictionaryCatalog; + +// JVM flag: -Dopennlp.download.remote=true +try (InputStream in = Files.newInputStream(Path.of("dictionary-catalog.properties"))) { + DictionaryCatalog catalog = DictionaryCatalog.load(in); + HunspellDictionaryDownload.downloadFromCatalog( + catalog, "en_US", Path.of("/tmp/hunspell-en_US")); +} +``` + +For `en_US`, the catalog ids are `hunspell.en_US.aff`, `hunspell.en_US.dic`, and +optionally `hunspell.en_US.readme`. A complete catalog example lives at +`opennlp-core/opennlp-runtime/src/test/resources/opennlp/tools/util/dictionary-catalog.properties`. +The download test uses local file URLs to exercise this flow without network access. + +## Option B: your own files + +Fetch `.aff` / `.dic` (and the license readme) with any tool, or with +`ResourceInstaller.install(uri, directory, sha512)`, then load them: + +```java +import java.nio.file.Path; +import opennlp.tools.stemmer.Stemmer; +import opennlp.tools.stemmer.hunspell.HunspellDictionary; +import opennlp.tools.stemmer.hunspell.HunspellStemmerFactory; + +HunspellDictionary dictionary = HunspellDictionary.load( + Path.of("/tmp/hunspell-en_US/en_US.aff"), + Path.of("/tmp/hunspell-en_US/en_US.dic")); +HunspellStemmerFactory factory = new HunspellStemmerFactory(dictionary); + +Stemmer stemmer = factory.newStemmer(); +CharSequence stem = stemmer.stem("workers"); +``` + +The result depends on the loaded dictionary. The in-tree manual example uses a +small dictionary and checks that `workers` stems to `worker`. + +The dictionary is immutable and safe to share between threads. The factory creates a +new stemmer for each call, so each thread can use its own instance. A dictionary that +declares a non-UTF-8 encoding through the `SET` directive in its `.aff` file is decoded +accordingly; no conversion is required. + +## Testing against real dictionaries + +The in-tree tests use project-authored fixtures only. An opt-in test class, `HunspellRealDictionaryTest`, also checks everyday morphology with the LibreOffice `en_US`, `de_DE_frami`, and `hu_HU` dictionaries. Point it at one directory containing all listed `.aff` and `.dic` files. A missing dictionary skips the associated test; a dictionary that cannot be loaded fails it. + +``` +./mvnw test -pl opennlp-core/opennlp-runtime -am \ + -Dtest=HunspellRealDictionaryTest -Dsurefire.failIfNoSpecifiedTests=false \ + -Dopennlp.hunspell.dict.dir=/tmp/hunspell-dicts +``` + +## What the engine supports + +The engine applies `PFX` and `SFX` rules with strip strings and character-class conditions. It supports a prefix and suffix cross-product, a double suffix sequence connected by continuation classes, identity rules in continuation paths, file-wide `FLAG` modes, file-wide `AF` aliases, and the `SET` encoding declaration. Numeric flags range from 1 through 65000. + +Compound decomposition supports `COMPOUNDFLAG`, `COMPOUNDBEGIN`, `COMPOUNDMIDDLE`, `COMPOUNDEND`, `COMPOUNDMIN`, `COMPOUNDWORDMAX`, `COMPOUNDPERMITFLAG`, `COMPOUNDFORBIDFLAG`, `CHECKCOMPOUNDDUP`, `CHECKCOMPOUNDCASE`, and `CHECKCOMPOUNDTRIPLE`. Compound boundaries and minimum lengths use Unicode code points. `NEEDAFFIX` (also named `PSEUDOROOT`), `ONLYINCOMPOUND`, `FORBIDDENWORD`, `CIRCUMFIX`, and `FULLSTRIP` control whether an analysis is accepted. + +Other directives are skipped. Their conversion, suggestion, or advanced compound behavior is not applied by this affix stemmer. Comments and unused metadata may contain legacy-encoded bytes even when the file uses UTF-8. Parsed rules and dictionary text are decoded strictly. Default and `long` flag modes preserve raw one-byte flag values used by published UTF-8 dictionaries. Invalid rule counts, aliases, flags, and compound limits fail during loading. Each affix or dictionary stream is rejected when it exceeds `HunspellDictionary.MAX_STREAM_BYTES` (64 MiB). + +Skipped directives include `ICONV`, `OCONV`, `COMPLEXPREFIXES`, `COMPOUNDRULE`, +`IGNORE`, and `KEEPCASE`. Loading a dictionary does not apply these rules; +results can differ from Hunspell for words that need them. diff --git a/dev/README-mecab-dictionaries.md b/dev/README-mecab-dictionaries.md new file mode 100644 index 0000000000..66aa183927 --- /dev/null +++ b/dev/README-mecab-dictionaries.md @@ -0,0 +1,126 @@ + + +# CJK dictionaries for the lattice tokenizer + +The lattice tokenizer (`opennlp.tools.tokenize.lattice`) segments Japanese and Korean over a MeCab-format dictionary, and the unigram segmenter handles Chinese over a plain word-frequency lexicon. Apache OpenNLP bundles no dictionary data: you download a dictionary from the project of your choice, and each dictionary contains its own license. Read the license file inside the archive before use. + +## Known MeCab-format dictionary projects + +| Catalog id | Dictionary | Language | Encoding | +|---|---|---|---| +| `mecab.ipadic` | IPADIC 2.7.0 | Japanese | EUC-JP | +| `mecab.ko-dic` | mecab-ko-dic 2.1.1 | Korean | UTF-8 | + +Example download URLs and SHA-512 digests for those ids live in the test resource +`opennlp-core/opennlp-runtime/src/test/resources/opennlp/tools/util/dictionary-catalog.properties`. Both archives are +gzip-compressed tars; `MecabDictionaryInstaller` reads the ustar, pax, and GNU +formats through `ResourceInstaller`. + +The installer extracts only the dictionary payload: the `*.csv` and `*.def` files a +`MecabDictionary` reads, plus the `dicrc` configuration file the distributions ship +alongside them. `ResourceInstaller` rejects paths outside the staging directory, +then `MecabDictionaryInstaller` flattens the selected files into the target. The +returned value is the number of dictionary files installed. + +## Option A: opt-in catalog install + +Applications supply the catalog. Catalog URLs are inactive until you set +`-Dopennlp.download.remote=true` or the equivalent system property in code. + +```java +import java.nio.file.Path; +import opennlp.tools.tokenize.lattice.MecabDictionaryInstaller; +import opennlp.tools.util.DictionaryCatalog; + +// JVM flag: -Dopennlp.download.remote=true +DictionaryCatalog catalog = DictionaryCatalog.load(catalogProperties); +int files = MecabDictionaryInstaller.installFromCatalog( + catalog, "mecab.ipadic", Path.of("ipadic")); +``` + +## Option B: your own URL and digest + +```java +import java.net.URI; +import java.nio.file.Path; +import opennlp.tools.tokenize.lattice.MecabDictionaryInstaller; + +String expectedSha512 = "..."; // the 128-hex SHA-512 of the archive +int files = MecabDictionaryInstaller.install( + URI.create("https://example.example/dict.tar.gz"), + Path.of("dict"), + expectedSha512); +``` + +A local `file:` URI may omit the digest: +`MecabDictionaryInstaller.install(localArchive.toUri(), targetDirectory)`. +HTTP and HTTPS sources require a digest. Other URI schemes are rejected. + +## Size budgets for larger dictionaries + +Fetching and unpacking go through `ResourceInstaller` and are bounded so a crafted +archive cannot fill the disk: by default one download is capped at 1 GiB, the +unpacked payload at 4 GiB, and the archive at 100000 entries. IPADIC and +mecab-ko-dic fit comfortably. For larger dictionaries, such as UniDic, raise the +limits at JVM startup: + +```bash +-Dopennlp.download.max.bytes=4294967296 \ +-Dopennlp.install.max.total.bytes=8589934592 \ +-Dopennlp.install.max.entries=200000 +``` + +Missing, invalid, and nonpositive property values use the default limits. + +## Load and tokenize + +`MecabDictionary.load(Path)` assumes UTF-8. IPADIC needs the two-argument overload: + +```java +import java.nio.charset.Charset; +import java.nio.file.Path; +import opennlp.tools.tokenize.lattice.LatticeTokenizer; +import opennlp.tools.tokenize.lattice.MecabDictionary; + +MecabDictionary dictionary = + MecabDictionary.load(Path.of("ipadic"), Charset.forName("EUC-JP")); +LatticeTokenizer tokenizer = new LatticeTokenizer(dictionary); +// "Tokyo-to ni iku" (go to the Tokyo metropolis), escaped to keep this file ASCII +String[] tokens = tokenizer.tokenize("\u6771\u4EAC\u90FD\u306B\u884C\u304F"); +``` + +For a UTF-8 dictionary such as mecab-ko-dic, `MecabDictionary.load(Path.of("ko-dic"))` +is enough. Loaded dictionaries and tokenizers are immutable and safe to share between +threads, so load once and reuse. + +## Chinese: the unigram segmenter needs only a frequency lexicon + +`opennlp.tools.tokenize.lattice.UnigramSegmenter` does not use MeCab dictionaries. It +loads a plain text lexicon, one entry per line: the word, its count, and optionally a +tag, separated by whitespace. Any word-frequency list you have the rights to use works: + +```java +import java.nio.file.Path; +import opennlp.tools.tokenize.lattice.UnigramSegmenter; + +UnigramSegmenter segmenter = UnigramSegmenter.load(Path.of("words.txt")); +// "wo laidao Beijing Tian'anmen" (I arrive at Beijing Tiananmen), escaped as above +String[] tokens = segmenter.tokenize("\u6211\u6765\u5230\u5317\u4EAC\u5929\u5B89\u95E8"); +``` + +The lexicon archive includes its license; OpenNLP bundles no data. 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/mvnw b/mvnw index 19529ddf8c..bd8896bf22 100755 --- a/mvnw +++ b/mvnw @@ -19,7 +19,7 @@ # ---------------------------------------------------------------------------- # ---------------------------------------------------------------------------- -# Apache Maven Wrapper startup batch script, version 3.3.2 +# Apache Maven Wrapper startup batch script, version 3.3.4 # # Optional ENV vars # ----------------- @@ -105,14 +105,17 @@ trim() { printf "%s" "${1}" | tr -d '[:space:]' } +scriptDir="$(dirname "$0")" +scriptName="$(basename "$0")" + # parse distributionUrl and optional distributionSha256Sum, requires .mvn/wrapper/maven-wrapper.properties while IFS="=" read -r key value; do case "${key-}" in distributionUrl) distributionUrl=$(trim "${value-}") ;; distributionSha256Sum) distributionSha256Sum=$(trim "${value-}") ;; esac -done <"${0%/*}/.mvn/wrapper/maven-wrapper.properties" -[ -n "${distributionUrl-}" ] || die "cannot read distributionUrl property in ${0%/*}/.mvn/wrapper/maven-wrapper.properties" +done <"$scriptDir/.mvn/wrapper/maven-wrapper.properties" +[ -n "${distributionUrl-}" ] || die "cannot read distributionUrl property in $scriptDir/.mvn/wrapper/maven-wrapper.properties" case "${distributionUrl##*/}" in maven-mvnd-*bin.*) @@ -130,7 +133,7 @@ maven-mvnd-*bin.*) distributionUrl="${distributionUrl%-bin.*}-$distributionPlatform.zip" ;; maven-mvnd-*) MVN_CMD=mvnd.sh _MVNW_REPO_PATTERN=/maven/mvnd/ ;; -*) MVN_CMD="mvn${0##*/mvnw}" _MVNW_REPO_PATTERN=/org/apache/maven/ ;; +*) MVN_CMD="mvn${scriptName#mvnw}" _MVNW_REPO_PATTERN=/org/apache/maven/ ;; esac # apply MVNW_REPOURL and calculate MAVEN_HOME @@ -227,7 +230,7 @@ if [ -n "${distributionSha256Sum-}" ]; then echo "Please disable validation by removing 'distributionSha256Sum' from your maven-wrapper.properties." >&2 exit 1 elif command -v sha256sum >/dev/null; then - if echo "$distributionSha256Sum $TMP_DOWNLOAD_DIR/$distributionUrlName" | sha256sum -c >/dev/null 2>&1; then + if echo "$distributionSha256Sum $TMP_DOWNLOAD_DIR/$distributionUrlName" | sha256sum -c - >/dev/null 2>&1; then distributionSha256Result=true fi elif command -v shasum >/dev/null; then @@ -252,8 +255,41 @@ if command -v unzip >/dev/null; then else tar xzf${__MVNW_QUIET_TAR:+"$__MVNW_QUIET_TAR"} "$TMP_DOWNLOAD_DIR/$distributionUrlName" -C "$TMP_DOWNLOAD_DIR" || die "failed to untar" fi -printf %s\\n "$distributionUrl" >"$TMP_DOWNLOAD_DIR/$distributionUrlNameMain/mvnw.url" -mv -- "$TMP_DOWNLOAD_DIR/$distributionUrlNameMain" "$MAVEN_HOME" || [ -d "$MAVEN_HOME" ] || die "fail to move MAVEN_HOME" + +# Find the actual extracted directory name (handles snapshots where filename != directory name) +actualDistributionDir="" + +# First try the expected directory name (for regular distributions) +if [ -d "$TMP_DOWNLOAD_DIR/$distributionUrlNameMain" ]; then + if [ -f "$TMP_DOWNLOAD_DIR/$distributionUrlNameMain/bin/$MVN_CMD" ]; then + actualDistributionDir="$distributionUrlNameMain" + fi +fi + +# If not found, search for any directory with the Maven executable (for snapshots) +if [ -z "$actualDistributionDir" ]; then + # enable globbing to iterate over items + set +f + for dir in "$TMP_DOWNLOAD_DIR"/*; do + if [ -d "$dir" ]; then + if [ -f "$dir/bin/$MVN_CMD" ]; then + actualDistributionDir="$(basename "$dir")" + break + fi + fi + done + set -f +fi + +if [ -z "$actualDistributionDir" ]; then + verbose "Contents of $TMP_DOWNLOAD_DIR:" + verbose "$(ls -la "$TMP_DOWNLOAD_DIR")" + die "Could not find Maven distribution directory in extracted archive" +fi + +verbose "Found extracted Maven distribution directory: $actualDistributionDir" +printf %s\\n "$distributionUrl" >"$TMP_DOWNLOAD_DIR/$actualDistributionDir/mvnw.url" +mv -- "$TMP_DOWNLOAD_DIR/$actualDistributionDir" "$MAVEN_HOME" || [ -d "$MAVEN_HOME" ] || die "fail to move MAVEN_HOME" clean || : exec_maven "$@" diff --git a/mvnw.cmd b/mvnw.cmd index 249bdf3822..92450f9327 100644 --- a/mvnw.cmd +++ b/mvnw.cmd @@ -19,7 +19,7 @@ @REM ---------------------------------------------------------------------------- @REM ---------------------------------------------------------------------------- -@REM Apache Maven Wrapper startup batch script, version 3.3.2 +@REM Apache Maven Wrapper startup batch script, version 3.3.4 @REM @REM Optional ENV vars @REM MVNW_REPOURL - repo url base for downloading maven distribution @@ -40,7 +40,7 @@ @SET __MVNW_ARG0_NAME__= @SET MVNW_USERNAME= @SET MVNW_PASSWORD= -@IF NOT "%__MVNW_CMD__%"=="" (%__MVNW_CMD__% %*) +@IF NOT "%__MVNW_CMD__%"=="" ("%__MVNW_CMD__%" %*) @echo Cannot start maven from wrapper >&2 && exit /b 1 @GOTO :EOF : end batch / begin powershell #> @@ -73,16 +73,30 @@ switch -wildcard -casesensitive ( $($distributionUrl -replace '^.*/','') ) { # apply MVNW_REPOURL and calculate MAVEN_HOME # maven home pattern: ~/.m2/wrapper/dists/{apache-maven-,maven-mvnd--}/ if ($env:MVNW_REPOURL) { - $MVNW_REPO_PATTERN = if ($USE_MVND) { "/org/apache/maven/" } else { "/maven/mvnd/" } - $distributionUrl = "$env:MVNW_REPOURL$MVNW_REPO_PATTERN$($distributionUrl -replace '^.*'+$MVNW_REPO_PATTERN,'')" + $MVNW_REPO_PATTERN = if ($USE_MVND -eq $False) { "/org/apache/maven/" } else { "/maven/mvnd/" } + $distributionUrl = "$env:MVNW_REPOURL$MVNW_REPO_PATTERN$($distributionUrl -replace "^.*$MVNW_REPO_PATTERN",'')" } $distributionUrlName = $distributionUrl -replace '^.*/','' $distributionUrlNameMain = $distributionUrlName -replace '\.[^.]*$','' -replace '-bin$','' -$MAVEN_HOME_PARENT = "$HOME/.m2/wrapper/dists/$distributionUrlNameMain" + +$MAVEN_M2_PATH = "$HOME/.m2" if ($env:MAVEN_USER_HOME) { - $MAVEN_HOME_PARENT = "$env:MAVEN_USER_HOME/wrapper/dists/$distributionUrlNameMain" + $MAVEN_M2_PATH = "$env:MAVEN_USER_HOME" +} + +if (-not (Test-Path -Path $MAVEN_M2_PATH)) { + New-Item -Path $MAVEN_M2_PATH -ItemType Directory | Out-Null +} + +$MAVEN_WRAPPER_DISTS = $null +if ((Get-Item $MAVEN_M2_PATH).Target[0] -eq $null) { + $MAVEN_WRAPPER_DISTS = "$MAVEN_M2_PATH/wrapper/dists" +} else { + $MAVEN_WRAPPER_DISTS = (Get-Item $MAVEN_M2_PATH).Target[0] + "/wrapper/dists" } -$MAVEN_HOME_NAME = ([System.Security.Cryptography.MD5]::Create().ComputeHash([byte[]][char[]]$distributionUrl) | ForEach-Object {$_.ToString("x2")}) -join '' + +$MAVEN_HOME_PARENT = "$MAVEN_WRAPPER_DISTS/$distributionUrlNameMain" +$MAVEN_HOME_NAME = ([System.Security.Cryptography.SHA256]::Create().ComputeHash([byte[]][char[]]$distributionUrl) | ForEach-Object {$_.ToString("x2")}) -join '' $MAVEN_HOME = "$MAVEN_HOME_PARENT/$MAVEN_HOME_NAME" if (Test-Path -Path "$MAVEN_HOME" -PathType Container) { @@ -134,7 +148,33 @@ if ($distributionSha256Sum) { # unzip and move Expand-Archive "$TMP_DOWNLOAD_DIR/$distributionUrlName" -DestinationPath "$TMP_DOWNLOAD_DIR" | Out-Null -Rename-Item -Path "$TMP_DOWNLOAD_DIR/$distributionUrlNameMain" -NewName $MAVEN_HOME_NAME | Out-Null + +# Find the actual extracted directory name (handles snapshots where filename != directory name) +$actualDistributionDir = "" + +# First try the expected directory name (for regular distributions) +$expectedPath = Join-Path "$TMP_DOWNLOAD_DIR" "$distributionUrlNameMain" +$expectedMvnPath = Join-Path "$expectedPath" "bin/$MVN_CMD" +if ((Test-Path -Path $expectedPath -PathType Container) -and (Test-Path -Path $expectedMvnPath -PathType Leaf)) { + $actualDistributionDir = $distributionUrlNameMain +} + +# If not found, search for any directory with the Maven executable (for snapshots) +if (!$actualDistributionDir) { + Get-ChildItem -Path "$TMP_DOWNLOAD_DIR" -Directory | ForEach-Object { + $testPath = Join-Path $_.FullName "bin/$MVN_CMD" + if (Test-Path -Path $testPath -PathType Leaf) { + $actualDistributionDir = $_.Name + } + } +} + +if (!$actualDistributionDir) { + Write-Error "Could not find Maven distribution directory in extracted archive" +} + +Write-Verbose "Found extracted Maven distribution directory: $actualDistributionDir" +Rename-Item -Path "$TMP_DOWNLOAD_DIR/$actualDistributionDir" -NewName $MAVEN_HOME_NAME | Out-Null try { Move-Item -Path "$TMP_DOWNLOAD_DIR/$MAVEN_HOME_NAME" -Destination $MAVEN_HOME_PARENT | Out-Null } catch { diff --git a/opennlp-api/src/main/java/opennlp/tools/document/Annotation.java b/opennlp-api/src/main/java/opennlp/tools/document/Annotation.java new file mode 100644 index 0000000000..3d3aaa1fd0 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/Annotation.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.tools.document; + +import opennlp.tools.util.Span; + +/** + * One annotation of a {@link Document}: a typed value anchored to a {@link Span} of the + * document's original text, or a span-less value under a + * {@link LayerKey.Scope#DOCUMENT document-scoped} key. + * + *

The span always refers to the text the document was created with, never to a + * normalized or otherwise derived form, so any annotation can be highlighted in what the + * caller supplied. Whether a span is present is decided by the layer key's scope, not + * per annotation: the container rejects a span-less annotation under a positional key + * and a spanned annotation under a document-scoped key. Annotations that need to + * reference other annotations, for example a dependency arc naming its head token, do + * so by the index of the target annotation within its layer, never by object + * identity.

+ * + * @param span The location of the annotation in the original text, or {@code null} for + * a value under a document-scoped key. + * @param value The annotation value. Must not be {@code null}. + * @param The type of the annotation value. + * + * @since 3.0.0 + */ +public record Annotation(Span span, T value) { + + /** + * Validates the annotation. + * + * @throws IllegalArgumentException Thrown if {@code value} is {@code null}. + */ + public Annotation { + if (value == null) { + throw new IllegalArgumentException("value must not be null"); + } + } + + /** + * Creates a span-less annotation for a {@link LayerKey.Scope#DOCUMENT + * document-scoped} layer. + * + * @param value The annotation value. Must not be {@code null}. + * @param The type of the annotation value. + * @return An {@link Annotation} without a span. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code value} is {@code null}. + */ + public static Annotation of(T value) { + return new Annotation<>(null, value); + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/Document.java b/opennlp-api/src/main/java/opennlp/tools/document/Document.java new file mode 100644 index 0000000000..4d22d3a18b --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/Document.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.tools.document; + +import java.util.List; +import java.util.Set; + +/** + * An offset-anchored annotation container: the original text of one document plus any + * number of typed annotation layers over it. + * + *

A layer is a list of {@link Annotation annotations} identified by a + * {@link LayerKey}. The container itself knows nothing about specific layers; every + * analysis capability contributes its results as one more layer without any change to + * this interface, which is what keeps new capabilities additive. All spans refer to + * {@link #text()} as supplied, never to a derived form. A + * {@link LayerKey.Scope#DOCUMENT document-scoped} layer carries whole-document values + * without spans, for example a language id.

+ * + *

A document is never modified in place: {@link #with(LayerKey, List)} leaves its + * receiver untouched and returns a new document. Thread safety is implementation + * specific.

+ * + *

Three invariants make index-based references sound. A layer preserves its + * insertion order, and the container never sorts or reorders it. A layer is immutable + * once added: the returned lists reject modification and are detached from the + * caller's input list. Providing a layer that already exists is rejected loudly: the + * add is once-only, and the exception names the offending key. An annotation that + * references another annotation by its index within a layer, for example a dependency + * arc naming its head token, therefore stays valid for the lifetime of the + * document.

+ * + * @since 3.0.0 + */ +public interface Document { + + /** + * Creates an empty {@link Document} over a text. The returned document is immutable + * and safe to share between threads: it captures the text's content at construction, + * so later changes to a mutable {@code CharSequence} do not reach the document. + * + * @param text The original document text. Must not be {@code null}. + * @return A {@link Document} without any layers. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + static Document of(CharSequence text) { + return ImmutableDocument.empty(text); + } + + /** + * @return The original text of the document. Never {@code null}. + */ + CharSequence text(); + + /** + * Retrieves the annotations of one layer. + * + * @param layer The layer to read. Must not be {@code null}. + * @param The type of the layer's annotation values. + * @return The layer's annotations in their layer order, or an empty list when the + * layer is absent. Never {@code null}; the list is unmodifiable. + * @throws IllegalArgumentException Thrown if {@code layer} is {@code null}. + */ + List> get(LayerKey layer); + + /** + * @return The keys of all layers present on the document. Never {@code null}; the set + * is unmodifiable. + */ + Set> layers(); + + /** + * Returns a new document with one layer added. + * + * @param layer The key of the layer to add. Must not be {@code null} and must not + * already be present. + * @param annotations The annotations of the layer. Must not be {@code null}, must not + * contain {@code null}, and every value must be assignable to the + * layer's type. Under a positional key every annotation must carry + * a span within the text bounds; under a document-scoped key no + * annotation may carry a span. + * @param The type of the layer's annotation values. + * @return A new {@link Document} sharing this document's text and existing layers. + * Never {@code null}. + * @throws IllegalArgumentException Thrown if any of the above constraints is violated. + */ + Document with(LayerKey layer, List> annotations); + + /** + * How {@link #merge(Document, DuplicateLayerPolicy)} treats a layer key that is + * present on both documents. + */ + enum DuplicateLayerPolicy { + + /** Reject any layer key present on both documents. */ + REJECT, + + /** + * Keep one copy of a layer key present on both documents when the two layers are + * structurally equal, for example when two parallel branches ran the same + * tokenizer. Layers whose contents differ are rejected as with {@link #REJECT}. + * Equality is {@link Annotation} equality: spans compare by offsets and type, + * never by probability, and values by their own {@code equals}. + */ + KEEP_EQUAL + } + + /** + * Returns a new document combining this document's layers with another document's + * layers over the same text, joining documents grown independently, for example by + * pipelines that ran in parallel. + * + * @param other The document whose layers are added on top of this document's layers. + * Must not be {@code null}, must carry the same text content, and must + * not provide a layer this document already has. + * @return A new {@link Document} carrying the layers of both documents. Never + * {@code null}; both source documents are left untouched. + * @throws IllegalArgumentException Thrown if {@code other} is {@code null}, if its + * text content differs, or if a layer key is present on both documents; the + * exception names the offending key. + */ + default Document merge(Document other) { + return merge(other, DuplicateLayerPolicy.REJECT); + } + + /** + * Returns a new document combining this document's layers with another document's + * layers over the same text, resolving duplicate layer keys with + * {@code duplicateLayers}. + * + * @param other The document whose layers are added on top of this document's layers. + * Must not be {@code null} and must carry the same text content. + * @param duplicateLayers How to treat a layer key present on both documents. Must + * not be {@code null}. + * @return A new {@link Document} carrying the layers of both documents. Never + * {@code null}; both source documents are left untouched. + * @throws IllegalArgumentException Thrown if either argument is {@code null}, if the + * text content differs, or if a layer key is present on both documents and + * the policy does not keep it; the exception names the offending key. + */ + default Document merge(Document other, DuplicateLayerPolicy duplicateLayers) { + if (other == null) { + throw new IllegalArgumentException("other must not be null"); + } + if (duplicateLayers == null) { + throw new IllegalArgumentException("duplicateLayers must not be null"); + } + if (!text().toString().contentEquals(other.text())) { + throw new IllegalArgumentException( + "merge requires both documents to carry the same text"); + } + Document merged = this; + for (final LayerKey layer : other.layers()) { + if (duplicateLayers == DuplicateLayerPolicy.KEEP_EQUAL + && merged.layers().contains(layer)) { + if (layersEqual(merged, layer, other)) { + continue; + } + throw new IllegalArgumentException( + "layer is present on both documents with differing contents: " + layer); + } + merged = addLayer(merged, layer, other); + } + return merged; + } + + /** + * @return Whether the two documents carry structurally equal contents for the layer. + */ + private static boolean layersEqual(Document first, LayerKey layer, Document second) { + return first.get(layer).equals(second.get(layer)); + } + + /** + * Adds one layer of {@code from} to {@code base} through {@link #with(LayerKey, List)}, + * capturing the key's value type. + */ + private static Document addLayer(Document base, LayerKey layer, Document from) { + return base.with(layer, from.get(layer)); + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnalyzer.java b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnalyzer.java new file mode 100644 index 0000000000..e12a468b5a --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnalyzer.java @@ -0,0 +1,125 @@ +/* + * 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.document; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Runs a fixed sequence of {@link DocumentAnnotator annotators} over a text, producing + * one {@link Document} that carries every step's layers. + * + *

The pipeline is validated at build time: every annotator's required layers must be + * provided by an earlier annotator, and no two annotators may provide the same layer, so + * a misordered or conflicting pipeline fails when it is assembled rather than midway + * through a document. The analyzer holds no per-call state; it is as thread-safe as the + * annotators it is built from.

+ * + * @since 3.0.0 + */ +public final class DocumentAnalyzer { + + private final List annotators; + + private DocumentAnalyzer(List annotators) { + this.annotators = annotators; + } + + /** + * @return A new {@link Builder}. Never {@code null}. + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Analyzes a text by running every annotator in order. + * + * @param text The original document text. Must not be {@code null}. + * @return The annotated {@link Document}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + public Document analyze(CharSequence text) { + Document document = Document.of(text); + for (final DocumentAnnotator annotator : annotators) { + document = annotator.annotate(document); + } + return document; + } + + /** + * Assembles a {@link DocumentAnalyzer} from annotators in execution order. + */ + public static final class Builder { + + private final List annotators = new ArrayList<>(); + + private Builder() { + } + + /** + * Appends an annotator to the pipeline. + * + * @param annotator The annotator to run after the ones already added. Must not be + * {@code null}. + * @return This {@link Builder}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code annotator} is {@code null}. + */ + public Builder add(DocumentAnnotator annotator) { + if (annotator == null) { + throw new IllegalArgumentException("annotator must not be null"); + } + annotators.add(annotator); + return this; + } + + /** + * Validates the pipeline and builds the analyzer. + * + * @return A {@link DocumentAnalyzer}. Never {@code null}. + * @throws IllegalArgumentException Thrown if the pipeline is empty, an annotator + * requires a layer no earlier annotator provides, or two annotators provide + * the same layer. + */ + public DocumentAnalyzer build() { + if (annotators.isEmpty()) { + throw new IllegalArgumentException("a pipeline needs at least one annotator"); + } + final Map, Integer> providers = new HashMap<>(); + for (int position = 0; position < annotators.size(); position++) { + final DocumentAnnotator annotator = annotators.get(position); + for (final LayerKey required : annotator.requires()) { + if (!providers.containsKey(required)) { + throw new IllegalArgumentException("annotator " + annotator + + " requires layer " + required + ", which no earlier annotator provides"); + } + } + for (final LayerKey provided : annotator.provides()) { + final Integer earlier = providers.putIfAbsent(provided, position); + if (earlier != null) { + throw new IllegalArgumentException("annotators at positions " + earlier + + " and " + position + " both provide layer " + provided); + } + } + } + return new DocumentAnalyzer(List.copyOf(annotators)); + } + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotator.java b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotator.java new file mode 100644 index 0000000000..6bc8417e6c --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotator.java @@ -0,0 +1,62 @@ +/* + * 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.document; + +import java.util.Set; + +/** + * A pipeline step that reads layers from a {@link Document} and returns a new document + * with its own layers added. + * + *

An annotator declares the layers it {@link #requires()} and {@link #provides()}, so + * a {@link DocumentAnalyzer} can validate a pipeline before running it. Annotators are + * usually thin adapters over an existing analysis component. Thread safety is + * implementation specific.

+ * + * @since 3.0.0 + */ +public interface DocumentAnnotator { + + /** + * Annotates a document. + * + *

A required layer must be present on the document, but it may be empty: an empty + * required layer is valid input and yields the annotator's provided layers present but + * empty, so a pipeline degrades gracefully on documents without content.

+ * + * @param document The document to annotate. Must not be {@code null} and must contain + * every layer named by {@link #requires()}. + * @return A new {@link Document} carrying the layers named by {@link #provides()} in + * addition to the input layers. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null} or lacks + * a required layer. + */ + Document annotate(Document document); + + /** + * @return The keys of the layers this annotator reads. Never {@code null}. + */ + default Set> requires() { + return Set.of(); + } + + /** + * @return The keys of the layers this annotator adds. Never {@code null}. + */ + Set> provides(); +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotators.java b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotators.java new file mode 100644 index 0000000000..fd5f81b858 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/DocumentAnnotators.java @@ -0,0 +1,118 @@ +/* + * 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.document; + +import java.util.List; +import java.util.Set; + +/** + * Support methods shared by {@link DocumentAnnotator} implementations: the + * required-layer check and the per-sentence walk over the token layer. + * + *

These helpers keep the annotators' shared behavior identical across + * implementations: an absent required layer is always rejected with the same message + * naming the layer, and every per-sentence adapter applies the same sentence-to-token + * mapping, including the loud rejection of a token lying outside every sentence.

+ * + * @since 3.0.0 + */ +public final class DocumentAnnotators { + + /** + * Receives one sentence's contiguous token run during + * {@link #forEachSentence(List, List, SentenceTokenConsumer)}. + */ + @FunctionalInterface + public interface SentenceTokenConsumer { + + /** + * Consumes one sentence's tokens. + * + * @param first The position of the sentence's first token in the token layer. + * @param words The sentence's token values in layer order. Never {@code null} or + * empty; the run covers the token layer positions + * {@code [first, first + words.length)}. + */ + void accept(int first, String[] words); + } + + /** + * Verifies that a document is present and carries every given layer. + * + * @param document The document to check. + * @param layers The required layers, in the order they are to be reported. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, or if + * a layer is absent; the message names the first absent layer. + */ + public static void requireLayers(Document document, LayerKey... layers) { + if (document == null) { + throw new IllegalArgumentException("document must not be null"); + } + final Set> present = document.layers(); + for (final LayerKey layer : layers) { + if (!present.contains(layer)) { + throw new IllegalArgumentException("document lacks the required layer " + layer); + } + } + } + + /** + * Walks the token layer sentence by sentence and hands each sentence's contiguous + * token run to the consumer. + * + *

Both layers must be in text order. Each sentence consumes the contiguous run of + * tokens whose spans it encloses; a sentence without tokens is skipped. Every token + * must belong to a sentence: a token lying outside every sentence is rejected loudly + * after the walk, so it can never be silently dropped.

+ * + * @param sentences The sentence layer, in text order. Must not be {@code null}. + * @param tokens The token layer, in text order. Must not be {@code null}. + * @param consumer Receives each token-carrying sentence's run. Must not be + * {@code null}. + * @throws IllegalArgumentException Thrown if a token lies outside every sentence. + */ + public static void forEachSentence(List> sentences, + List> tokens, SentenceTokenConsumer consumer) { + int next = 0; + for (final Annotation sentence : sentences) { + final int first = next; + while (next < tokens.size() + && tokens.get(next).span().getStart() >= sentence.span().getStart() + && tokens.get(next).span().getEnd() <= sentence.span().getEnd()) { + next++; + } + final int count = next - first; + if (count == 0) { + continue; + } + final String[] words = new String[count]; + for (int i = 0; i < count; i++) { + words[i] = tokens.get(first + i).value(); + } + consumer.accept(first, words); + } + if (next != tokens.size()) { + throw new IllegalArgumentException("token at " + tokens.get(next).span() + + " lies outside every sentence"); + } + } + + private DocumentAnnotators() { + // Not instantiated; this class provides static support methods only. + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/ImmutableDocument.java b/opennlp-api/src/main/java/opennlp/tools/document/ImmutableDocument.java new file mode 100644 index 0000000000..e175c05f6c --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/ImmutableDocument.java @@ -0,0 +1,193 @@ +/* + * 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.document; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.util.Span; + +/** + * The default {@link Document} implementation: an unmodifiable map from layer key to an + * unmodifiable annotation list. + * + *

Instances are immutable: {@link #with(LayerKey, List)} returns a new document that + * shares the unchanged layers with its ancestor, copying the map but not the layers, so + * documents grown from a common ancestor share their layer lists. The text is captured + * as a {@link String} at construction, so a mutable {@link CharSequence} handed to + * {@link #empty(CharSequence)} cannot change the document afterwards. That immutability + * makes instances safe to share between threads.

+ */ +final class ImmutableDocument implements Document { + + private final String text; + private final Map, List>> layers; + + private ImmutableDocument(String text, Map, List>> layers) { + this.text = text; + this.layers = Collections.unmodifiableMap(layers); + } + + /** + * Creates a document without any layers. + * + * @param text The original document text, captured as its content at this moment. + * Must not be {@code null}. + * @return An empty {@link ImmutableDocument}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + static ImmutableDocument empty(CharSequence text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + return new ImmutableDocument(text.toString(), Collections.emptyMap()); + } + + /** {@inheritDoc} */ + @Override + public CharSequence text() { + return text; + } + + /** {@inheritDoc} */ + @Override + @SuppressWarnings("unchecked") + public List> get(LayerKey layer) { + if (layer == null) { + throw new IllegalArgumentException("layer must not be null"); + } + final List> annotations = layers.get(layer); + if (annotations == null) { + return List.of(); + } + // This cast is safe because with(LayerKey, List) verified every value against the + // key's type when the layer was inserted. + return (List>) (List) annotations; + } + + /** {@inheritDoc} */ + @Override + public Set> layers() { + // The unmodifiable map exposes an unmodifiable key set and caches it, so this + // accessor allocates no wrapper per call. + return layers.keySet(); + } + + /** {@inheritDoc} */ + @Override + public Document with(LayerKey layer, List> annotations) { + if (layer == null) { + throw new IllegalArgumentException("layer must not be null"); + } + if (annotations == null) { + throw new IllegalArgumentException("annotations must not be null"); + } + if (layers.containsKey(layer)) { + throw new IllegalArgumentException("layer is already present: " + layer); + } + validate(layer, annotations); + final Map, List>> grown = new LinkedHashMap<>(layers); + grown.put(layer, List.copyOf(annotations)); + return new ImmutableDocument(text, grown); + } + + /** + * {@inheritDoc} + * This implementation copies the layer map once, not once per added layer. + */ + @Override + public Document merge(Document other, DuplicateLayerPolicy duplicateLayers) { + if (other == null) { + throw new IllegalArgumentException("other must not be null"); + } + if (duplicateLayers == null) { + throw new IllegalArgumentException("duplicateLayers must not be null"); + } + if (!text.contentEquals(other.text())) { + throw new IllegalArgumentException( + "merge requires both documents to carry the same text"); + } + final Map, List>> combined = new LinkedHashMap<>(layers); + for (final LayerKey layer : other.layers()) { + if (combined.containsKey(layer)) { + if (duplicateLayers == DuplicateLayerPolicy.KEEP_EQUAL + && get(layer).equals(other.get(layer))) { + continue; + } + throw new IllegalArgumentException(duplicateLayers == DuplicateLayerPolicy.KEEP_EQUAL + ? "layer is present on both documents with differing contents: " + layer + : "layer is already present: " + layer); + } + combined.put(layer, copyValidated(layer, other)); + } + if (combined.size() == layers.size()) { + return this; + } + return new ImmutableDocument(text, combined); + } + + /** + * {@return a validated immutable copy of one of {@code from}'s layers, capturing the + * key's value type} + */ + private List> copyValidated(LayerKey layer, Document from) { + final List> annotations = from.get(layer); + if (annotations == null) { + throw new IllegalArgumentException("annotations must not be null"); + } + validate(layer, annotations); + return List.copyOf(annotations); + } + + /** + * Checks one layer's annotations against the key's contract: no null elements, values + * assignable to the key's type, spans present and within the text bounds under a + * positional key, absent under a document-scoped key. + * + * @throws IllegalArgumentException Thrown if any check fails; the message names the + * layer. + */ + private void validate(LayerKey layer, List> annotations) { + for (final Annotation annotation : annotations) { + if (annotation == null) { + throw new IllegalArgumentException("annotations must not contain null: " + layer); + } + if (!layer.type().isInstance(annotation.value())) { + throw new IllegalArgumentException("value of type " + + annotation.value().getClass().getName() + " does not match layer " + layer); + } + final Span span = annotation.span(); + if (layer.scope() == LayerKey.Scope.POSITIONAL) { + if (span == null) { + throw new IllegalArgumentException( + "positional layer " + layer + " requires a span on every annotation"); + } + if (span.getEnd() > text.length()) { + throw new IllegalArgumentException("span " + span + " exceeds the text length " + + text.length() + " in layer " + layer); + } + } else if (span != null) { + throw new IllegalArgumentException( + "document-scoped layer " + layer + " must not carry spans"); + } + } + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/LayerKey.java b/opennlp-api/src/main/java/opennlp/tools/document/LayerKey.java new file mode 100644 index 0000000000..e29870a06c --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/LayerKey.java @@ -0,0 +1,165 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.document; + +import java.util.Objects; + +import opennlp.tools.util.StringUtil; + +/** + * Identifies one annotation layer of a {@link Document} and carries the type of that + * layer's annotation values, so reading a layer back is statically typed. + * + *

The key space is deliberately open: any producer may define new keys in its own + * package, and the container never enumerates them. Two keys are equal when their id, + * their value type, and their {@link Scope} are equal, so independently created + * constants for the same layer interoperate. Standard keys for the toolkit's own + * results live in {@link Layers}.

+ * + *

A key declares its {@link Scope}: a {@link Scope#POSITIONAL positional} key + * guarantees a span on every annotation, and a {@link Scope#DOCUMENT document-scoped} + * key carries whole-document values without spans, for example a language id or a + * category distribution. The scope is declared per key, never per annotation, so + * consumers of a positional layer never null-check a span.

+ * + * @param The type of the annotation values stored under this key. + * + * @since 3.0.0 + */ +public final class LayerKey { + + /** How the annotations of a layer relate to the document text. */ + public enum Scope { + /** Every annotation of the layer is anchored to a span of the text. */ + POSITIONAL, + /** The layer's values describe the document as a whole and carry no spans. */ + DOCUMENT + } + + private final String id; + private final Class type; + private final Scope scope; + + private LayerKey(String id, Class type, Scope scope) { + this.id = id; + this.type = type; + this.scope = scope; + } + + /** + * Creates a {@link LayerKey}. + * + * @param id The layer identifier, for example {@code opennlp:tokens}. Keys defined by + * the toolkit carry the {@code opennlp:} prefix, an extension uses its own + * prefix, and a bare id is legal for an application-local layer. Must not + * be {@code null} or blank. + * @param type The class of the annotation values stored under the key. Must not be + * {@code null}. + * @param The type of the annotation values. + * @return A {@link Scope#POSITIONAL positional} {@link LayerKey}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code id} is {@code null} or blank, or + * {@code type} is {@code null}. + */ + public static LayerKey of(String id, Class type) { + return key(id, type, Scope.POSITIONAL); + } + + /** + * Creates a {@link Scope#DOCUMENT document-scoped} {@link LayerKey} for values that + * describe the document as a whole, for example a language id, a category + * distribution, or provenance. Annotations under such a key carry no span. + * + * @param id The layer identifier, following the same prefix rules as + * {@link #of(String, Class)}. Must not be {@code null} or blank. + * @param type The class of the annotation values stored under the key. Must not be + * {@code null}. + * @param The type of the annotation values. + * @return A document-scoped {@link LayerKey}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code id} is {@code null} or blank, or + * {@code type} is {@code null}. + */ + public static LayerKey document(String id, Class type) { + return key(id, type, Scope.DOCUMENT); + } + + /** + * Validates the components and creates the key. + * + * @param id The layer identifier. + * @param type The value class. + * @param scope The declared scope. + * @param The type of the annotation values. + * @return The {@link LayerKey}. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code id} is {@code null} or blank, or + * {@code type} is {@code null}. + */ + private static LayerKey key(String id, Class type, Scope scope) { + if (id == null || StringUtil.isBlank(id)) { + throw new IllegalArgumentException("id must not be null or blank"); + } + if (type == null) { + throw new IllegalArgumentException("type must not be null"); + } + return new LayerKey<>(id, type, scope); + } + + /** + * @return The layer identifier. Never {@code null}. + */ + public String id() { + return id; + } + + /** + * @return The class of the annotation values stored under this key. Never {@code null}. + */ + public Class type() { + return type; + } + + /** + * @return The declared scope of the layer. Never {@code null}. + */ + public Scope scope() { + return scope; + } + + /** {@inheritDoc} */ + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof LayerKey other)) { + return false; + } + return id.equals(other.id) && type.equals(other.type) && scope == other.scope; + } + + /** {@inheritDoc} */ + @Override + public int hashCode() { + return Objects.hash(id, type, scope); + } + + /** {@inheritDoc} */ + @Override + public String toString() { + return id + '<' + type.getSimpleName() + '>'; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/document/Layers.java b/opennlp-api/src/main/java/opennlp/tools/document/Layers.java new file mode 100644 index 0000000000..0257afc384 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/document/Layers.java @@ -0,0 +1,134 @@ +/* + * 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.document; + +import opennlp.tools.util.StringUtil; + +/** + * The standard {@link LayerKey layer keys} for the results the toolkit produces itself. + * + *

This class is a convenience, not a registry: the key space stays open, and any + * producer may define further keys in its own package. New capabilities must never + * require an addition here to function.

+ * + *

Namespace rule: every key the toolkit itself defines carries the + * {@code opennlp:} id prefix. An extension defines its keys under its own prefix, and + * a bare id without a prefix is legal for an application-local layer, so ids from + * independent producers cannot collide. Toolkit keys are created through + * {@link #key(String, Class)} and {@link #documentKey(String, Class)}, which apply the + * prefix, so no producer spells it.

+ * + *

Gold versus predicted: a corpus may carry a hand-annotated version of a layer + * beside a produced one. The convention is a {@code gold:} id prefix on the same key + * scheme, for example {@code gold:opennlp:tokens} beside {@code opennlp:tokens}. + * Because adding a layer is once-only, competing versions of a layer always live under + * distinct keys and never replace each other.

+ * + *

Placement rule: this class holds only the keys of the core linguistic layers + * every pipeline shares (sentences, tokens, tags, entities). A capability-specific + * layer's key lives on the annotator that provides it, for example the lemma layer's + * key on its adapter, so adding a capability never touches this class.

+ * + * @since 3.0.0 + */ +public final class Layers { + + /** The id prefix of every key the toolkit defines. */ + private static final String NAMESPACE = "opennlp:"; + + /** + * Sentence boundaries; each annotation covers one sentence and carries its text. + */ + public static final LayerKey SENTENCES = key("sentences", String.class); + + /** + * Token boundaries; each annotation covers one token and carries its text. + */ + public static final LayerKey TOKENS = key("tokens", String.class); + + /** + * Part-of-speech tags; one annotation per token, aligned with {@link #TOKENS} by + * position, carrying the tag. + */ + public static final LayerKey POS_TAGS = key("pos", String.class); + + /** + * Named entities; each annotation covers one mention and carries the entity type as + * its value. The annotation's span carries offsets only. + */ + public static final LayerKey ENTITIES = key("entities", String.class); + + /** + * Creates a {@link LayerKey.Scope#POSITIONAL positional} key in the toolkit's + * {@code opennlp:} namespace, for a layer the toolkit itself produces. + * + * @param name The layer name without a namespace, for example {@code tokens}. Must + * not be {@code null}, blank, or contain {@code ':'}. + * @param type The class of the annotation values stored under the key. Must not be + * {@code null}. + * @param The type of the annotation values. + * @return A key whose id is the name under the {@code opennlp:} prefix. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null}, blank, or + * contains {@code ':'}, or {@code type} is {@code null}. + */ + public static LayerKey key(String name, Class type) { + return LayerKey.of(NAMESPACE + validName(name), type); + } + + /** + * Creates a {@link LayerKey.Scope#DOCUMENT document-scoped} key in the toolkit's + * {@code opennlp:} namespace, for a whole-document value the toolkit itself produces. + * + * @param name The layer name without a namespace, for example {@code language}. Must + * not be {@code null}, blank, or contain {@code ':'}. + * @param type The class of the annotation values stored under the key. Must not be + * {@code null}. + * @param The type of the annotation values. + * @return A document-scoped key whose id is the name under the {@code opennlp:} + * prefix. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null}, blank, or + * contains {@code ':'}, or {@code type} is {@code null}. + */ + public static LayerKey documentKey(String name, Class type) { + return LayerKey.document(NAMESPACE + validName(name), type); + } + + /** + * Validates a namespace-free layer name. + * + * @param name The name to validate. + * @return The validated name. + * @throws IllegalArgumentException Thrown if {@code name} is {@code null}, blank, or + * contains {@code ':'}. + */ + private static String validName(String name) { + if (name == null || StringUtil.isBlank(name)) { + throw new IllegalArgumentException("name must not be null or blank"); + } + if (name.indexOf(':') >= 0) { + throw new IllegalArgumentException( + "name must not contain ':', the namespace is applied by this factory: " + name); + } + return name; + } + + private Layers() { + // Not instantiated; this class provides constants and static key factories only. + } +} 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/termvector/TermVector.java b/opennlp-api/src/main/java/opennlp/tools/termvector/TermVector.java new file mode 100644 index 0000000000..42c654257d --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/termvector/TermVector.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.tools.termvector; + +import java.util.List; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.Span; + +/** + * One entry of a term vector layer: a term, how often it occurs in the document, and + * where. + * + *

The {@link #term()} is the term's identity as the producing annotator determined + * it, typically a normalized form; the {@link #spans()} are the occurrence offsets and + * always point into the document's original text, never into a normalized form, + * so an index consumer can highlight every occurrence in what the caller supplied.

+ * + *

A term vector comes in one of two shapes, told apart by whether {@link #spans()} is + * empty. A full vector carries one span per occurrence, so + * {@code spans().size() == frequency()}. A scoring-only vector carries no spans + * at all, so consumers that only need term frequencies do not pay for offset storage. + * There is no third shape: a non-empty span list must match the frequency exactly.

+ * + *

Instances are immutable: the span list is copied on construction and the copy is + * unmodifiable.

+ * + * @param term The term string. Must not be {@code null}. + * @param frequency The number of occurrences in the document. Must be at least one. + * @param spans The occurrence spans in original text coordinates, one per occurrence, + * or an empty list for a scoring-only vector. Must not be {@code null} or + * contain {@code null} and, when non-empty, must hold exactly + * {@code frequency} spans. + * + * @since 3.0.0 + */ +@ThreadSafe +public record TermVector(String term, int frequency, List spans) { + + /** + * Validates the term vector and detaches the span list from the caller's input. + * + * @throws IllegalArgumentException Thrown if {@code term} is {@code null}, + * {@code frequency} is below one, {@code spans} is or contains {@code null}, + * or a non-empty {@code spans} list does not hold exactly {@code frequency} + * spans. + */ + public TermVector { + if (term == null) { + throw new IllegalArgumentException("term must not be null"); + } + if (frequency < 1) { + throw new IllegalArgumentException("frequency must be at least one: " + frequency); + } + if (spans == null) { + throw new IllegalArgumentException("spans must not be null"); + } + for (final Span span : spans) { + if (span == null) { + throw new IllegalArgumentException("spans must not contain null"); + } + } + if (!spans.isEmpty() && spans.size() != frequency) { + throw new IllegalArgumentException("a full term vector holds one span per " + + "occurrence: frequency is " + frequency + " but spans holds " + spans.size()); + } + spans = List.copyOf(spans); + } + + /** + * Creates a full {@link TermVector} whose frequency is derived from the occurrence + * spans. + * + * @param term The term string. Must not be {@code null}. + * @param spans The occurrence spans in original text coordinates. Must not be + * {@code null} or empty. + * @return A {@link TermVector} with {@code frequency() == spans.size()}. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code term} is {@code null} or + * {@code spans} is {@code null} or empty. + */ + public static TermVector withSpans(String term, List spans) { + if (spans == null || spans.isEmpty()) { + throw new IllegalArgumentException("spans must not be null or empty"); + } + return new TermVector(term, spans.size(), spans); + } + + /** + * Creates a scoring-only {@link TermVector} that carries the occurrence count without + * any offsets. + * + * @param term The term string. Must not be {@code null}. + * @param frequency The number of occurrences in the document. Must be at least one. + * @return A {@link TermVector} whose {@link #spans()} is empty. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code term} is {@code null} or + * {@code frequency} is below one. + */ + public static TermVector count(String term, int frequency) { + return new TermVector(term, frequency, List.of()); + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java index 1cbfffe312..850f794d8a 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java @@ -24,6 +24,9 @@ */ final class BertNormalization { + /** Default maximum word length used by BERT wordpiece tokenizers. */ + static final int DEFAULT_MAX_WORD_CODE_POINTS = 100; + private BertNormalization() { } @@ -70,6 +73,12 @@ static boolean isWhitespace(int codePoint) { return Character.getType(codePoint) == Character.SPACE_SEPARATOR; } + /** Tests whether a code point is a Unicode line or paragraph separator. */ + static boolean isLineOrParagraphSeparator(int codePoint) { + final int type = Character.getType(codePoint); + return type == Character.LINE_SEPARATOR || type == Character.PARAGRAPH_SEPARATOR; + } + /** * A punctuation character in the BERT sense: any non-alphanumeric ASCII * character that is not whitespace, or any Unicode punctuation category. diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java deleted file mode 100644 index cea1fd8f53..0000000000 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java +++ /dev/null @@ -1,125 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package opennlp.tools.tokenize; - -import java.util.ArrayList; -import java.util.Set; - -import opennlp.tools.util.Span; - -/** - * A {@link Tokenizer} implementation of the full BERT tokenization pipeline: - * basic tokenization (text normalization) followed by wordpiece tokenization, - * with the classification and separator tokens framing every result. - * - * @deprecated Use {@link WordpieceEncoder} instead: - * {@link WordpieceEncoder#encodeToPieces(CharSequence)} returns the same - * {@code String[]} as {@link #tokenize(String)}, and {@code encode} additionally - * carries vocabulary ids and original-text spans. This class is scheduled for - * removal after one stable release. - * - * @see WordpieceEncoder - */ -@Deprecated(since = "3.0.0", forRemoval = true) -public class BertTokenizer implements Tokenizer { - - private final WordpieceEncoder encoder; - - /** - * Initializes a {@link BertTokenizer} for an uncased BERT model, - * with lower casing and accent stripping enabled. - * - * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. - * - * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, - * contains {@code null}, or is missing a BERT special token. - */ - public BertTokenizer(Set vocabulary) { - this(vocabulary, true); - } - - /** - * Initializes a {@link BertTokenizer} with BERT special tokens. - * - * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. - * @param lowerCase {@code true} for uncased models (lower casing and accent - * stripping), {@code false} for cased models. - * - * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, - * contains {@code null}, or is missing a BERT special token. - */ - public BertTokenizer(Set vocabulary, boolean lowerCase) { - this(vocabulary, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, - WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); - } - - /** - * Initializes a {@link BertTokenizer} with custom special tokens, for models - * like RoBERTa that do not use the BERT defaults. - * - * @param vocabulary The wordpiece vocabulary. Must not be {@code null}. - * @param lowerCase {@code true} for uncased models (lower casing and - * accent stripping), {@code false} for cased models. - * @param classificationToken The CLS token; must be in the vocabulary. - * @param separatorToken The SEP token; must be in the vocabulary. - * @param unknownToken The UNK token; must be in the vocabulary. - * - * @throws IllegalArgumentException Thrown if any argument is {@code null}, - * the vocabulary contains {@code null}, or a special token is missing - * from the vocabulary. - */ - public BertTokenizer(Set vocabulary, boolean lowerCase, - String classificationToken, String separatorToken, String unknownToken) { - if (vocabulary == null) { - throw new IllegalArgumentException("vocabulary must not be null"); - } - // The encoder assigns each piece its list index as the id. Ids are unused on the - // tokenize() path, so synthesizing them from an arbitrary set order is fine. - this.encoder = new WordpieceEncoder(new ArrayList<>(vocabulary), lowerCase, - classificationToken, separatorToken, unknownToken); - } - - /** - * Tokenizes the given text into wordpieces, surrounded by the classification - * and separator tokens. - * - * @param text The text to tokenize. Must not be {@code null}. - * - * @return The wordpiece tokens. - * - * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. - */ - @Override - public String[] tokenize(String text) { - return encoder.encodeToPieces(text); - } - - /** - * Not supported: wordpiece tokens (subwords, {@code ##} continuations and - * special tokens) have no faithful character spans in the original text. - * Use {@link WordpieceEncoder#encode(CharSequence)} for pieces with - * original-text spans. - * - * @throws UnsupportedOperationException Always. - */ - @Override - public Span[] tokenizePos(String text) { - throw new UnsupportedOperationException( - "Wordpiece tokens cannot be mapped to character spans of the original text"); - } - -} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java index 62260b7929..dca3d56aee 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java @@ -16,43 +16,58 @@ */ package opennlp.tools.tokenize; -import opennlp.tools.util.Span; - /** - * One subword unit produced by a {@link SubwordTokenizer}, carrying both the vocabulary view - * (the piece string and its id) and the exact place in the caller's text it came from. + * One subword unit produced by a {@link SubwordTokenizer}, including the model representation + * and source range. + * + *

The piece string is in the tokenizer's normalized form and need not equal the input. + * {@code start} and {@code end} are UTF-16 offsets into the original text, so the + * surface associated with this piece is {@code text.subSequence(start, end)}. A span can include + * adjacent source characters when normalization reorders characters. Pieces without source text, + * such as control symbols, report an empty span with {@code start == end}.

* - *

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

+ *

The piece, id and span triple follows the representation established by SentencePiece + * (Kudo and Richardson, 2018), whose {@code SentencePieceText.SentencePiece} message carries + * {@code piece}, {@code id} and a {@code [begin, end)} range into the input, and gives control + * symbols an empty surface with {@code begin == end}. The {@code surface} field of that message + * has no counterpart here because it is {@code text.subSequence(start, end)}. The offsets are + * UTF-16 indices rather than the UTF-8 byte offsets used there.

* - * @param piece The piece in the vocabulary's normalized form; never null or empty. - * @param id The vocabulary id of the piece. + * @param piece The piece in the vocabulary's normalized form; must not be {@code null} or empty. + * @param id The non-negative vocabulary id of the piece. * @param start The inclusive start offset in the original text. - * @param end The exclusive end offset in the original text; not less than {@code start}. + * @param end The exclusive end offset in the original text; at least {@code start}. + * + * @see SubwordTokenizer + * @see Taku Kudo, John Richardson (2018): + * SentencePiece: A simple and language independent subword tokenizer and detokenizer for + * Neural Text Processing. EMNLP 2018 (system demonstrations), pages 66-71 + * @since 3.0.0 */ public record SubwordPiece(String piece, int id, int start, int end) { /** * Instantiates a {@link SubwordPiece}. * - * @throws IllegalArgumentException Thrown if {@code piece} is null or empty, or the span is - * negative or inverted. + * @throws IllegalArgumentException Thrown if {@code piece} is {@code null} or empty, + * {@code id} is negative, or the span is negative or inverted. */ public SubwordPiece { - if (piece == null || piece.isEmpty()) { - throw new IllegalArgumentException("piece must not be null or empty"); + if (piece == null) { + throw new IllegalArgumentException("piece must not be null"); + } + if (piece.isEmpty()) { + throw new IllegalArgumentException("piece must not be empty"); } - if (start < 0 || end < start) { - throw new IllegalArgumentException( - "The span [" + start + ", " + end + ") must not be negative or inverted."); + if (id < 0) { + throw new IllegalArgumentException("id must not be negative"); + } + if (start < 0) { + throw new IllegalArgumentException("start must not be negative"); + } + if (end < start) { + throw new IllegalArgumentException("end must be at least start"); } } - /** {@return the original-text span of this piece as a {@link Span}} */ - public Span span() { - return new Span(start, end); - } } diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java index 957ff62229..056bb59429 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java @@ -19,32 +19,55 @@ import java.util.List; /** - * Splits text into subword units against a fixed vocabulary, reporting for every unit its - * vocabulary id and the exact span of the original text it covers. + * Splits text into subword units from a fixed model vocabulary, reporting the model id and + * original-text span for each unit. * - *

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

+ *

Segmentation follows model entries, not linguistic token boundaries. Each piece is in the + * model's normalized form and need not equal the input. Offsets in each {@link SubwordPiece} + * refer to the original input text.

+ * + *

Subword units were introduced to neural models to keep the vocabulary closed while still + * covering rare and unseen words: byte pair encoding (Sennrich et al., 2016), the WordPiece + * inventory (Schuster and Nakajima, 2012; Wu et al., 2016) and the unigram language model + * (Kudo, 2018). This interface abstracts over the segmentation method and describes only what + * the trained inventories have in common, following the tokenizer-independent encoding contract + * of SentencePiece (Kudo and Richardson, 2018): every unit is a vocabulary entry with an id and + * a range over the input.

+ * + *

An implementation may include model control pieces with empty source spans. Their presence + * and placement are part of that tokenizer's contract, not this interface.

* *

Thread safety is implementation specific.

+ * + * @see SubwordPiece + * @see Rico Sennrich, Barry Haddow, Alexandra Birch + * (2016): Neural Machine Translation of Rare Words with Subword Units. ACL 2016, + * pages 1715-1725 + * @see Taku Kudo (2018): Subword Regularization: + * Improving Neural Network Translation Models with Multiple Subword Candidates. ACL 2018, + * pages 66-75 + * @see Taku Kudo, John Richardson (2018): + * SentencePiece: A simple and language independent subword tokenizer and detokenizer for + * Neural Text Processing. EMNLP 2018 (system demonstrations), pages 66-71 + * @since 3.0.0 */ public interface SubwordTokenizer { /** * Encodes text into subword pieces. * - * @param text The text to encode; must not be null. - * @return The pieces in text order; empty when the text contains nothing encodable. - * @throws IllegalArgumentException Thrown if {@code text} is null. + * @param text The text to encode; must not be {@code null}. + * @return The pieces in model order; may be empty. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. */ List encode(CharSequence text); /** * Encodes text into vocabulary ids. * - * @param text The text to encode; must not be null. - * @return The ids in text order; empty when the text contains nothing encodable. - * @throws IllegalArgumentException Thrown if {@code text} is null. + * @param text The text to encode; must not be {@code null}. + * @return The ids from {@link #encode(CharSequence)}, in the same order. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. */ default int[] encodeToIds(CharSequence text) { final List pieces = encode(text); @@ -58,9 +81,9 @@ default int[] encodeToIds(CharSequence text) { /** * Encodes text into piece strings in the vocabulary's normalized form. * - * @param text The text to encode; must not be null. - * @return The pieces in text order; empty when the text contains nothing encodable. - * @throws IllegalArgumentException Thrown if {@code text} is null. + * @param text The text to encode; must not be {@code null}. + * @return The piece strings from {@link #encode(CharSequence)}, in the same order. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. */ default String[] encodeToPieces(CharSequence text) { final List pieces = encode(text); 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 8a8ed7adbc..9e620c3f5c 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java @@ -20,34 +20,59 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; -import java.util.HashSet; import java.util.List; -import java.util.Locale; import java.util.Map; -import java.util.Set; import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.StringUtil; /** - * A {@link SubwordTokenizer} running the full BERT tokenization pipeline of the reference - * implementation: basic tokenization (control removal, whitespace normalization, CJK - * isolation, optional lower casing with accent stripping, punctuation isolation) followed by - * greedy longest-match wordpiece segmentation. + * A {@link SubwordTokenizer} implementing the BERT tokenization stages: basic tokenization + * (control removal, whitespace normalization, CJK isolation, optional lower casing with accent + * stripping, punctuation isolation) followed by greedy longest-match wordpiece segmentation. * - *

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

+ *

Each result includes a vocabulary id and range in the original text. The range + * refers to the input before normalization. Classification and separator entries use empty + * ranges at the text boundaries, so {@link #encode(CharSequence)} includes both control + * entries.

* - *

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

+ *

The wordpiece inventory was introduced by Schuster and Nakajima (2012) as the + * {@code WordPieceModel}: word units learned greedily from unsegmented text to maximize the + * language-model likelihood, so that no input is out of vocabulary. Wu et al. (2016), section + * 4.1, adopt it for neural machine translation, and Devlin et al. (2019), section 3, build BERT + * on a 30,000 entry WordPiece vocabulary with a leading classification token. Those papers + * describe how an inventory is trained; the greedy longest-match-first segmentation and + * the {@code ##} continuation marker applied here are the inference conventions of the + * BERT reference implementation. Wu et al. instead mark word starts with {@code _}.

* - *

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

+ *

Ids follow the line-number convention of BERT {@code vocab.txt} files. List constructors use + * the list index, while the map constructor uses the supplied ids. The classification, separator, + * and unknown tokens must all be present in the + * 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} + * context, so a word-final Greek capital sigma becomes U+03C2 as in the reference + * implementation.

+ * + *

A word exceeding the configured maximum number of normalized Unicode code points becomes + * the unknown piece. The default is 100, the value used by the Hugging Face {@code transformers} + * BERT tokenizer; the original {@code google-research/bert} code uses 200. Both count code + * points, and a constructor parameter selects another limit.

* * @see WordpieceTokenizer + * @see + * BERT tokenization reference + * @see Mike Schuster, Kaisuke Nakajima (2012): + * Japanese and Korean Voice Search. ICASSP 2012, pages 5149-5152 + * @see Yonghui Wu et al. (2016): Google's Neural + * Machine Translation System: Bridging the Gap between Human and Machine Translation. + * arXiv:1609.08144 + * @see Jacob Devlin, Ming-Wei Chang, Kenton Lee, + * Kristina Toutanova (2019): BERT: Pre-training of Deep Bidirectional Transformers for + * Language Understanding. NAACL-HLT 2019, pages 4171-4186 + * @since 3.0.0 */ @ThreadSafe public final class WordpieceEncoder implements SubwordTokenizer { @@ -56,15 +81,17 @@ public final class WordpieceEncoder implements SubwordTokenizer { // so it can only match after the word's first piece. private static final String CONTINUATION_PREFIX = "##"; - // The reference implementation's limit: longer words become the unknown piece. - private static final int MAX_WORD_CHARACTERS = 100; + private static final int GREEK_CAPITAL_SIGMA = 0x03A3; + private static final int GREEK_SMALL_FINAL_SIGMA = 0x03C2; - private final Set vocabulary; private final Map ids; + private final VocabularyTrie initialPieces; + private final VocabularyTrie continuationPieces; private final boolean lowerCase; private final String classificationToken; private final String separatorToken; private final String unknownToken; + private final int maxWordCodePoints; private final int classificationId; private final int separatorId; private final int unknownId; @@ -72,10 +99,10 @@ public final class WordpieceEncoder implements SubwordTokenizer { /** * Instantiates an encoder for an uncased BERT model with the BERT special tokens. * - * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be null, - * must not contain nulls or duplicates. - * @throws IllegalArgumentException Thrown if the vocabulary is null, contains a null or - * duplicate entry, or a BERT special token is missing from it. + * @param vocabulary The ordered vocabulary; the list index becomes the id. Must not be {@code null} + * or contain {@code null}, empty, or duplicate entries. + * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, contains a + * {@code null}, empty, or duplicate entry, or is missing a BERT special token. */ public WordpieceEncoder(List vocabulary) { this(vocabulary, true); @@ -84,55 +111,119 @@ public WordpieceEncoder(List vocabulary) { /** * Instantiates an encoder with the BERT special tokens. * - * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be null, - * must not contain nulls or duplicates. - * @param lowerCase True for uncased models (lower casing and accent stripping), false for - * cased models. - * @throws IllegalArgumentException Thrown if the vocabulary is null, contains a null or - * duplicate entry, or a BERT special token is missing from it. + * @param vocabulary The ordered vocabulary; the list index becomes the id. Must not be {@code null} + * or contain {@code null}, empty, or duplicate entries. + * @param lowerCase {@code true} for uncased models (lower casing and accent stripping), + * {@code false} for cased models. + * @throws IllegalArgumentException Thrown if the vocabulary is {@code null}, contains a + * {@code null}, empty, or duplicate entry, or is missing a BERT special token. */ public WordpieceEncoder(List vocabulary, boolean lowerCase) { + this(vocabulary, lowerCase, BertNormalization.DEFAULT_MAX_WORD_CODE_POINTS); + } + + /** + * Instantiates an encoder with the BERT special tokens and a custom word-length limit. + * + * @param vocabulary The ordered vocabulary; the list index becomes the id. Must not be + * {@code null} or contain {@code null}, empty, or duplicate entries. + * @param lowerCase {@code true} for uncased models, {@code false} for cased models. + * @param maxWordCodePoints The non-negative maximum number of normalized Unicode code points + * in one word. + * @throws IllegalArgumentException Thrown if an argument is invalid or a BERT special token is + * missing. + */ + public WordpieceEncoder(List vocabulary, boolean lowerCase, int maxWordCodePoints) { this(vocabulary, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, - WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN, maxWordCodePoints); } /** - * Instantiates an encoder with custom special tokens, for models that do not use the BERT - * defaults. + * Instantiates an encoder with custom special tokens. * - * @param vocabulary The ordered vocabulary; a piece's id is its index. Must not be - * null, must not contain nulls or duplicates. - * @param lowerCase True for uncased models (lower casing and accent stripping), - * false for cased models. - * @param classificationToken The CLS token; must be in the vocabulary. - * @param separatorToken The SEP token; must be in the vocabulary. - * @param unknownToken The UNK token; must be in the vocabulary. - * @throws IllegalArgumentException Thrown if any argument is null, the vocabulary contains - * a null or duplicate entry, or a special token is missing from the vocabulary. + * @param vocabulary The ordered vocabulary; the list index becomes the id. Must not be + * {@code null} or contain {@code null}, empty, or duplicate entries. + * @param lowerCase {@code true} for uncased models (lower casing and accent stripping), + * {@code false} for cased models. + * @param classificationToken The CLS token; must not be {@code null} or empty and must be in the + * vocabulary. + * @param separatorToken The SEP token; must not be {@code null} or empty and must be in the + * vocabulary. + * @param unknownToken The UNK token; must not be {@code null} or empty and must be in the + * vocabulary. + * @throws IllegalArgumentException Thrown if any argument is {@code null}, the vocabulary + * contains a {@code null}, empty, or duplicate entry, or a special token is empty or missing. */ public WordpieceEncoder(List vocabulary, boolean lowerCase, String classificationToken, String separatorToken, String unknownToken) { - this(byPiece(vocabulary), lowerCase, classificationToken, separatorToken, unknownToken); + this(vocabulary, lowerCase, classificationToken, separatorToken, unknownToken, + BertNormalization.DEFAULT_MAX_WORD_CODE_POINTS); } /** - * Instantiates an encoder from an explicit piece-to-id mapping, for vocabularies whose ids - * are not contiguous line numbers. + * Instantiates an encoder with custom special tokens and a custom word-length limit. * - * @param vocabularyIds The piece-to-id mapping. Must not be null, must not contain - * null keys or values. - * @param lowerCase True for uncased models (lower casing and accent stripping), - * false for cased models. - * @param classificationToken The CLS token; must be in the vocabulary. - * @param separatorToken The SEP token; must be in the vocabulary. - * @param unknownToken The UNK token; must be in the vocabulary. - * @throws IllegalArgumentException Thrown if any argument is null, the mapping contains a - * null key or value, or a special token is missing from the vocabulary. + * @param vocabulary The ordered vocabulary; the list index becomes the id. Must not be + * {@code null} or contain {@code null}, empty, or duplicate entries. + * @param lowerCase {@code true} for uncased models, {@code false} for cased models. + * @param classificationToken The CLS token; must be present in the vocabulary. + * @param separatorToken The SEP token; must be present in the vocabulary. + * @param unknownToken The UNK token; must be present in the vocabulary. + * @param maxWordCodePoints The non-negative maximum number of normalized Unicode code points + * in one word. + * @throws IllegalArgumentException Thrown if an argument is invalid or a special token is + * missing. + */ + public WordpieceEncoder(List vocabulary, boolean lowerCase, + String classificationToken, String separatorToken, + String unknownToken, int maxWordCodePoints) { + this(byPiece(vocabulary), lowerCase, classificationToken, separatorToken, unknownToken, + maxWordCodePoints); + } + + /** + * Instantiates an encoder from an explicit piece-to-id mapping for vocabularies with + * noncontiguous ids. + * + * @param vocabularyIds The piece-to-id mapping. Must not be {@code null} or contain + * {@code null} or empty keys, {@code null} values, or negative ids. + * @param lowerCase {@code true} for uncased models (lower casing and accent stripping), + * {@code false} for cased models. + * @param classificationToken The CLS token; must not be {@code null} or empty and must be in the + * vocabulary. + * @param separatorToken The SEP token; must not be {@code null} or empty and must be in the + * vocabulary. + * @param unknownToken The UNK token; must not be {@code null} or empty and must be in the + * vocabulary. + * @throws IllegalArgumentException Thrown if any argument is {@code null}, the mapping contains + * a {@code null} or empty key, {@code null} value, or negative id, or a special token is empty + * or missing. */ public WordpieceEncoder(Map vocabularyIds, boolean lowerCase, String classificationToken, String separatorToken, String unknownToken) { + this(vocabularyIds, lowerCase, classificationToken, separatorToken, unknownToken, + BertNormalization.DEFAULT_MAX_WORD_CODE_POINTS); + } + + /** + * Instantiates an encoder from a piece-to-id mapping with a custom word-length limit. + * + * @param vocabularyIds The piece-to-id mapping. Must not be {@code null} or contain invalid + * entries. + * @param lowerCase {@code true} for uncased models, {@code false} for cased models. + * @param classificationToken The CLS token; must be present in the vocabulary. + * @param separatorToken The SEP token; must be present in the vocabulary. + * @param unknownToken The UNK token; must be present in the vocabulary. + * @param maxWordCodePoints The non-negative maximum number of normalized Unicode code points + * in one word. + * @throws IllegalArgumentException Thrown if an argument is invalid or a special token is + * missing. + */ + public WordpieceEncoder(Map vocabularyIds, boolean lowerCase, + String classificationToken, String separatorToken, + String unknownToken, int maxWordCodePoints) { if (vocabularyIds == null) { throw new IllegalArgumentException("vocabularyIds must not be null"); } @@ -145,65 +236,74 @@ public WordpieceEncoder(Map vocabularyIds, boolean lowerCase, if (unknownToken == null) { throw new IllegalArgumentException("unknownToken must not be null"); } - final Map byPiece = new HashMap<>(vocabularyIds.size() * 2); + if (classificationToken.isEmpty()) { + throw new IllegalArgumentException("classificationToken must not be empty"); + } + if (separatorToken.isEmpty()) { + throw new IllegalArgumentException("separatorToken must not be empty"); + } + if (unknownToken.isEmpty()) { + throw new IllegalArgumentException("unknownToken must not be empty"); + } + if (maxWordCodePoints < 0) { + throw new IllegalArgumentException("maxWordCodePoints must not be negative"); + } + final Map byPiece = HashMap.newHashMap(vocabularyIds.size()); for (final Map.Entry entry : vocabularyIds.entrySet()) { if (entry.getKey() == null || entry.getValue() == null) { - throw new IllegalArgumentException( - "The vocabulary must not contain null pieces or ids: " + entry); + throw new IllegalArgumentException("vocabularyIds must not contain null pieces or ids"); + } + if (entry.getKey().isEmpty()) { + throw new IllegalArgumentException("vocabularyIds must not contain an empty piece"); + } + if (entry.getValue() < 0) { + throw new IllegalArgumentException("vocabularyIds must not contain a negative id for piece '" + + entry.getKey() + "'"); } byPiece.put(entry.getKey(), entry.getValue()); } - this.vocabulary = new HashSet<>(byPiece.keySet()); - this.ids = byPiece; + this.ids = Map.copyOf(byPiece); + this.initialPieces = new VocabularyTrie(this.ids, false); + this.continuationPieces = new VocabularyTrie(this.ids, true); this.lowerCase = lowerCase; this.classificationToken = classificationToken; this.separatorToken = separatorToken; this.unknownToken = unknownToken; - this.classificationId = requiredId(byPiece, classificationToken); - this.separatorId = requiredId(byPiece, separatorToken); - this.unknownId = requiredId(byPiece, unknownToken); + this.maxWordCodePoints = maxWordCodePoints; + this.classificationId = requiredId(classificationToken); + this.separatorId = requiredId(separatorToken); + this.unknownId = requiredId(unknownToken); } - /** - * Converts an ordered vocabulary list into the piece-to-id mapping, assigning each piece its - * index as the id. - * - * @param vocabulary The ordered vocabulary. - * @return The piece-to-id mapping. - * @throws IllegalArgumentException Thrown if the list is null or contains a null or duplicate - * entry. - */ + /** Converts an ordered vocabulary to a piece-to-id mapping. */ private static Map byPiece(List vocabulary) { if (vocabulary == null) { throw new IllegalArgumentException("vocabulary must not be null"); } - final Map byPiece = new HashMap<>(vocabulary.size() * 2); + final Map byPiece = HashMap.newHashMap(vocabulary.size()); for (int id = 0; id < vocabulary.size(); id++) { final String piece = vocabulary.get(id); if (piece == null) { - throw new IllegalArgumentException("The vocabulary contains null at index " + id + "."); + throw new IllegalArgumentException("vocabulary must not contain null at index " + id); + } + if (piece.isEmpty()) { + throw new IllegalArgumentException( + "vocabulary must not contain an empty piece at index " + id); } if (byPiece.putIfAbsent(piece, id) != null) { - throw new IllegalArgumentException("The vocabulary contains '" + piece - + "' more than once; ids would be ambiguous."); + throw new IllegalArgumentException( + "vocabulary must not contain duplicate piece '" + piece + "'"); } } return byPiece; } - /** - * Looks up the id of a special token that must be present in the vocabulary. - * - * @param ids The piece-to-id mapping. - * @param specialToken The token to look up. - * @return The token's id. - * @throws IllegalArgumentException Thrown if the token is not in the vocabulary. - */ - private static int requiredId(Map ids, String specialToken) { + /** Returns the id of a required special token. */ + private int requiredId(String specialToken) { final Integer id = ids.get(specialToken); if (id == null) { - throw new IllegalArgumentException("The special token '" + specialToken - + "' is not in the vocabulary; every emitted piece must have an id."); + throw new IllegalArgumentException( + "vocabulary must contain special token '" + specialToken + "'"); } return id; } @@ -216,8 +316,7 @@ public List encode(CharSequence text) { } final String original = text.toString(); - // The normalized text, one original-text range per char, built through the reference - // pipeline's transformations in the reference order. + // Stores an original-text range for each normalized char. MappedText mapped = cleanAndIsolateCjk(original); if (lowerCase) { mapped = lowerCaseAndStripAccents(mapped); @@ -244,19 +343,11 @@ public List encode(CharSequence text) { return pieces; } - /** - * Greedily longest-match segments one whitespace-delimited word; the pieces are emitted only if - * the whole word is representable, otherwise the word becomes a single unknown piece. - * - * @param mapped The normalized text with per-character original-text ranges. - * @param from The inclusive start of the word in {@code mapped}. - * @param to The exclusive end of the word in {@code mapped}. - * @param pieces The output list to append to. - */ + /** Appends the longest vocabulary segmentation of one normalized word. */ private void encodeWord(MappedText mapped, int from, int to, List pieces) { final int wordStart = mapped.starts[from]; final int wordEnd = mapped.ends[to - 1]; - if (to - from > MAX_WORD_CHARACTERS) { + if (Character.codePointCount(mapped.chars, from, to - from) > maxWordCodePoints) { pieces.add(new SubwordPiece(unknownToken, unknownId, wordStart, wordEnd)); return; } @@ -264,25 +355,15 @@ private void encodeWord(MappedText mapped, int from, int to, List int start = from; boolean found = true; while (start < to) { - int end = to; - found = false; - while (start < end) { - String substring = new String(mapped.chars, start, end - start); - if (start > from) { - substring = CONTINUATION_PREFIX + substring; - } - if (vocabulary.contains(substring)) { - wordPieces.add(new SubwordPiece(substring, ids.get(substring), - mapped.starts[start], mapped.ends[end - 1])); - start = end; - found = true; - break; - } - end--; - } - if (!found) { + final VocabularyMatch match = (start == from ? initialPieces : continuationPieces) + .longestMatch(mapped.chars, start, to); + if (match == null) { + found = false; break; } + wordPieces.add(new SubwordPiece(match.piece(), match.id(), + mapped.starts[start], mapped.ends[match.end() - 1])); + start = match.end(); } if (found) { pieces.addAll(wordPieces); @@ -291,9 +372,75 @@ private void encodeWord(MappedText mapped, int from, int to, List } } + /** A vocabulary match and the exclusive end offset in normalized text. */ + private record VocabularyMatch(String piece, int id, int end) { + } + + /** Finds vocabulary entries without creating strings for candidate lengths. */ + private static final class VocabularyTrie { + + private final TrieNode root = new TrieNode(); + + /** + * Indexes the initial or the continuation pieces of a vocabulary. + * + * @param vocabulary The piece-to-id mapping. + * @param continuation {@code true} to index pieces starting with the continuation prefix, + * stored without it; {@code false} to index all other pieces. + */ + private VocabularyTrie(Map vocabulary, boolean continuation) { + for (final Map.Entry entry : vocabulary.entrySet()) { + final String piece = entry.getKey(); + final boolean continuationPiece = piece.startsWith(CONTINUATION_PREFIX); + if (continuationPiece != continuation) { + continue; + } + final int offset = continuationPiece ? CONTINUATION_PREFIX.length() : 0; + if (offset == piece.length()) { + continue; + } + TrieNode node = root; + for (int i = offset; i < piece.length(); i++) { + node = node.children.computeIfAbsent(piece.charAt(i), key -> new TrieNode()); + } + node.piece = piece; + node.id = entry.getValue(); + } + } + + /** Returns the longest entry that starts at {@code from}. */ + private VocabularyMatch longestMatch(char[] text, int from, int to) { + TrieNode node = root; + VocabularyMatch longest = null; + int index = from; + while (index < to) { + final int codePoint = Character.codePointAt(text, index, to); + final int width = Character.charCount(codePoint); + for (int i = 0; i < width; i++) { + node = node.children.get(text[index + i]); + if (node == null) { + return longest; + } + } + index += width; + if (node.piece != null) { + longest = new VocabularyMatch(node.piece, node.id, index); + } + } + return longest; + } + } + + /** One node of the vocabulary trie. */ + private static final class TrieNode { + private final Map children = new HashMap<>(); + private String piece; + private int id; + } + /** - * The normalized text with, for every char, the original-text range it came from. Characters - * inserted by the pipeline (isolation spaces) carry an empty range at the insertion point. + * The normalized text with the original-text range for each character. Characters inserted by + * the pipeline (isolation spaces) use an empty range at the insertion point. */ private static final class MappedText { private char[] chars; @@ -301,24 +448,14 @@ private static final class MappedText { private int[] ends; private int length; - /** - * Instantiates an empty mapped text. - * - * @param capacity The initial capacity hint in chars. - */ + /** Creates an empty mapping with space for the expected number of UTF-16 code units. */ private MappedText(int capacity) { chars = new char[capacity]; starts = new int[capacity]; ends = new int[capacity]; } - /** - * Appends one char with the original-text range it came from. - * - * @param c The char to append. - * @param originalStart The inclusive original-text start of the char. - * @param originalEnd The exclusive original-text end of the char. - */ + /** Adds one UTF-16 code unit with a source range. */ private void add(char c, int originalStart, int originalEnd) { if (length == chars.length) { final int capacity = Math.max(16, length * 2); @@ -332,28 +469,31 @@ private void add(char c, int originalStart, int originalEnd) { length++; } - /** - * Appends every char of a string, all sharing one original-text range. - * - * @param s The string to append. - * @param originalStart The inclusive original-text start shared by all chars. - * @param originalEnd The exclusive original-text end shared by all chars. - */ + /** Adds a string with one source range shared by all code units. */ private void add(String s, int originalStart, int originalEnd) { for (int i = 0; i < s.length(); i++) { add(s.charAt(i), originalStart, originalEnd); } } + + /** Adds one code point with a source range shared by all code units. */ + private void addCodePoint(int codePoint, int originalStart, int originalEnd) { + if (Character.isBmpCodePoint(codePoint)) { + add((char) codePoint, originalStart, originalEnd); + } else { + add(Character.highSurrogate(codePoint), originalStart, originalEnd); + add(Character.lowSurrogate(codePoint), originalStart, originalEnd); + } + } + + /** Returns the mapped character content. */ + private String text() { + return new String(chars, 0, length); + } } - /** - * Cleans the text (control and whitespace normalization) and isolates CJK code points in one - * pass, recording the original-text range of every output character. - * - * @param original The original input text. - * @return The cleaned, CJK-isolated text with per-character ranges. - */ - private static MappedText cleanAndIsolateCjk(String original) { + /** Removes BERT control characters, normalizes whitespace, and isolates CJK ideographs. */ + private MappedText cleanAndIsolateCjk(String original) { final MappedText out = new MappedText(original.length() + 16); int i = 0; while (i < original.length()) { @@ -363,7 +503,8 @@ private static MappedText cleanAndIsolateCjk(String original) { i += width; continue; } - if (BertNormalization.isWhitespace(codePoint) || isLineOrParagraphSeparator(codePoint)) { + if (BertNormalization.isWhitespace(codePoint) + || BertNormalization.isLineOrParagraphSeparator(codePoint)) { out.add(' ', i, i + width); } else if (BertNormalization.isCjk(codePoint)) { out.add(' ', i, i); @@ -382,27 +523,10 @@ private static MappedText cleanAndIsolateCjk(String original) { } /** - * A Unicode line or paragraph separator ({@code Zl}, {@code Zp}). These are not whitespace in - * the BERT {@code _is_whitespace} sense, but the reference pipeline's - * {@code whitespace_tokenize} (Python's {@code str.split()}) still breaks words on them, so - * they must become word boundaries here as well. - * - * @param codePoint The code point to classify. - * @return True if the code point is a line or paragraph separator. - */ - private static boolean isLineOrParagraphSeparator(int codePoint) { - final int type = Character.getType(codePoint); - return type == Character.LINE_SEPARATOR || type == Character.PARAGRAPH_SEPARATOR; - } - - /** - * Isolates punctuation, surrounding each punctuation code point with spaces, preserving the - * original-text range of every character. - * - * @param in The input text with per-character ranges. - * @return The punctuation-isolated text with per-character ranges. + * Surrounds each BERT punctuation character with mapped spaces. Keep the character + * classification in sync with {@link BertNormalization#isolatePunctuation(String)}. */ - private static MappedText isolatePunctuation(MappedText in) { + private MappedText isolatePunctuation(MappedText in) { final MappedText out = new MappedText(in.length + 16); int i = 0; while (i < in.length) { @@ -424,15 +548,8 @@ private static MappedText isolatePunctuation(MappedText in) { return out; } - /** - * Lower cases and strips accents, preserving the original-text range of every character. When a - * contextual case mapping prevents a per-character range from being recovered, the whole - * whitespace run falls back to its full range, which widens spans but never misplaces them. - * - * @param in The input text with per-character ranges. - * @return The lower-cased, accent-stripped text with per-character ranges. - */ - private static MappedText lowerCaseAndStripAccents(MappedText in) { + /** Applies BERT lower casing and accent removal while retaining source ranges. */ + private MappedText lowerCaseAndStripAccents(MappedText in) { final MappedText out = new MappedText(in.length + 16); int from = 0; while (from < in.length) { @@ -451,78 +568,153 @@ private static MappedText lowerCaseAndStripAccents(MappedText in) { return out; } + /** Normalizes one non-whitespace run and appends it to the output mapping. */ + private void transformRun(MappedText in, int from, int to, MappedText out) { + final MappedText lower = lowerCaseRun(in, from, to); + if (Normalizer.isNormalized(lower.text(), Normalizer.Form.NFD)) { + appendWithoutAccents(lower, out); + return; + } + + final MappedText decomposedByCodePoint = decomposeByCodePoint(lower); + final String decomposedRun = Normalizer.normalize(lower.text(), Normalizer.Form.NFD); + if (decomposedRun.equals(decomposedByCodePoint.text())) { + appendWithoutAccents(decomposedByCodePoint, out); + } else { + // Canonical reordering can cross character ranges, so each result covers the source run. + appendWithoutAccents(decomposedRun, in.starts[from], in.ends[to - 1], out); + } + } + + /** Lowercases one run and retains a source range for each result code point. */ + private MappedText lowerCaseRun(MappedText in, int from, int to) { + final String source = new String(in.chars, from, to - from); + final String lower = StringUtil.toLowerCase(source); + final MappedText mapped = new MappedText(lower.length()); + int sourceIndex = from; + int lowerIndex = 0; + while (sourceIndex < to) { + final int sourceCodePoint = codePointAt(in, sourceIndex); + final int sourceWidth = Character.charCount(sourceCodePoint); + int lowerCodePoint = lower.codePointAt(lowerIndex); + final int lowerWidth = Character.charCount(lowerCodePoint); + if (sourceCodePoint == GREEK_CAPITAL_SIGMA + && isFinalSigma(in, from, to, sourceIndex)) { + lowerCodePoint = GREEK_SMALL_FINAL_SIGMA; + } + mapped.addCodePoint(lowerCodePoint, in.starts[sourceIndex], + in.ends[sourceIndex + sourceWidth - 1]); + sourceIndex += sourceWidth; + lowerIndex += lowerWidth; + } + return mapped; + } + + /** Tests whether the capital sigma at {@code index} is word-final. */ + private boolean isFinalSigma(MappedText in, int from, int to, int index) { + int before = index; + boolean followsCased = false; + while (before > from) { + final int codePoint = Character.codePointBefore(in.chars, before); + before -= Character.charCount(codePoint); + if (!isCaseIgnorable(codePoint)) { + followsCased = isCased(codePoint); + break; + } + } + if (!followsCased) { + return false; + } + + int after = index + Character.charCount(GREEK_CAPITAL_SIGMA); + while (after < to) { + final int codePoint = codePointAt(in, after); + if (!isCaseIgnorable(codePoint)) { + return !isCased(codePoint); + } + after += Character.charCount(codePoint); + } + return true; + } + + /** Tests whether a code point participates in Unicode casing. */ + private boolean isCased(int codePoint) { + return Character.isUpperCase(codePoint) + || Character.isLowerCase(codePoint) + || Character.isTitleCase(codePoint); + } + /** - * Lower cases and accent-strips one non-space run, emitting per-character ranges when the - * transformation is reproducible per code point and the run's full range otherwise. + * Tests the general categories and Unicode 15.0 word-break values used by the + * {@code Case_Ignorable} property. * - * @param in The input text with per-character ranges. - * @param from The inclusive start of the run in {@code in}. - * @param to The exclusive end of the run in {@code in}. - * @param out The output text to append to. + * @see + * Unicode 15.0 word-break properties */ - private static void transformRun(MappedText in, int from, int to, MappedText out) { - final String run = new String(in.chars, from, to - from); - // Locale.ROOT lower casing is the reference behavior of BERT's do_lower_case: the reference - // pipeline applies the full locale-independent Unicode case mappings (including one-to-many - // ones like the dotted capital I), which a per-code-point mapping cannot reproduce. - final String content = stripAccents(run.toLowerCase(Locale.ROOT)); - - // Rerun per code point to learn how many output chars each input code point produces. - final StringBuilder rerun = new StringBuilder(content.length()); - final int[] produced = new int[to - from]; - int i = from; - while (i < to) { - final int codePoint = codePointAt(in, i); + private boolean isCaseIgnorable(int codePoint) { + return switch (Character.getType(codePoint)) { + case Character.NON_SPACING_MARK, Character.ENCLOSING_MARK, + Character.FORMAT, Character.MODIFIER_LETTER, + Character.MODIFIER_SYMBOL -> true; + default -> switch (codePoint) { + case 0x0027, 0x002E, 0x003A, 0x00B7, 0x0387, 0x055F, 0x05F4, 0x2018, + 0x2019, 0x2024, 0x2027, 0xFE13, 0xFE52, 0xFE55, 0xFF07, 0xFF0E, + 0xFF1A -> true; + default -> false; + }; + }; + } + + /** Canonically decomposes each code point while preserving the source range. */ + private MappedText decomposeByCodePoint(MappedText text) { + final MappedText decomposed = new MappedText(text.length + 16); + int index = 0; + while (index < text.length) { + final int codePoint = codePointAt(text, index); final int width = Character.charCount(codePoint); - final String transformed = stripAccents( - new String(Character.toChars(codePoint)).toLowerCase(Locale.ROOT)); - rerun.append(transformed); - produced[i - from] = transformed.length(); - i += width; + final String source = new String(text.chars, index, width); + final String normalized = Normalizer.isNormalized(source, Normalizer.Form.NFD) + ? source : Normalizer.normalize(source, Normalizer.Form.NFD); + decomposed.add(normalized, text.starts[index], text.ends[index + width - 1]); + index += width; } + return decomposed; + } - if (rerun.toString().equals(content)) { - int at = from; - int emitted = 0; - while (at < to) { - final int width = Character.charCount(codePointAt(in, at)); - for (int c = 0; c < produced[at - from]; c++) { - out.add(content.charAt(emitted++), in.starts[at], in.ends[at + width - 1]); + /** Appends mapped characters excluding non-spacing marks. */ + private void appendWithoutAccents(MappedText text, MappedText out) { + final int outputStart = out.length; + int index = 0; + while (index < text.length) { + final int codePoint = codePointAt(text, index); + final int width = Character.charCount(codePoint); + if (Character.getType(codePoint) == Character.NON_SPACING_MARK) { + if (out.length > outputStart) { + out.ends[out.length - 1] = text.ends[index + width - 1]; + } + } else { + for (int i = 0; i < width; i++) { + out.add(text.chars[index + i], text.starts[index + i], text.ends[index + i]); } - at += width; } - } else { - // Contextual case mapping changed the content; the run's chars share the run's range. - out.add(content, in.starts[from], in.ends[to - 1]); + index += width; } } - /** - * Removes combining marks after NFD decomposition, the accent stripping of BERT's - * {@code do_lower_case} mode. - * - * @param text The text to strip. - * @return The text without non-spacing marks. - */ - private static String stripAccents(String text) { - final String decomposed = Normalizer.normalize(text, Normalizer.Form.NFD); - final StringBuilder stripped = new StringBuilder(decomposed.length()); - decomposed.codePoints().forEach(codePoint -> { + /** Appends characters excluding non-spacing marks with one source range. */ + private void appendWithoutAccents(String text, int start, int end, MappedText out) { + int index = 0; + while (index < text.length()) { + final int codePoint = text.codePointAt(index); if (Character.getType(codePoint) != Character.NON_SPACING_MARK) { - stripped.appendCodePoint(codePoint); + out.addCodePoint(codePoint, start, end); } - }); - return stripped.toString(); + index += Character.charCount(codePoint); + } } - /** - * Reads the code point at an index, joining a surrogate pair when one starts there. - * - * @param text The text to read from. - * @param index The char index to read at. - * @return The code point at {@code index}. - */ - private static int codePointAt(MappedText text, int index) { + /** Returns a code point from the mapped UTF-16 buffer. */ + private int codePointAt(MappedText text, int index) { final char c = text.chars[index]; if (Character.isHighSurrogate(c) && index + 1 < text.length && Character.isLowSurrogate(text.chars[index + 1])) { diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java index fa1014ab2c..d2520a1e06 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceTokenizer.java @@ -17,6 +17,7 @@ package opennlp.tools.tokenize; +import java.util.HashSet; import java.util.LinkedList; import java.util.List; import java.util.Set; @@ -35,16 +36,11 @@ * no accent stripping, no control character removal. Text that does not match * the vocabulary's casing - for uncased models that includes every capitalized * word - is mapped to the unknown token. Use {@link WordpieceEncoder} for the - * full BERT tokenization pipeline. + * BERT normalization and wordpiece segmentation stages. *

- * As of OpenNLP 3.0.0 the behavior matches the reference BERT wordpiece - * implementation in three respects that differ from earlier releases: - * runs of punctuation (and non-ASCII punctuation) are split into individual - * single-character tokens, words that cannot be fully represented by - * vocabulary pieces become a single unknown token instead of the matched - * prefix pieces followed by the unknown token, and {@link #tokenizePos(String)} - * throws {@link UnsupportedOperationException} instead of returning - * {@code null}. + * Runs of punctuation are split into individual tokens. A word that cannot be + * fully represented by vocabulary pieces becomes one unknown token. + * {@link #tokenizePos(String)} is not supported. *

* For reference see: *

    @@ -62,6 +58,8 @@ */ public class WordpieceTokenizer implements Tokenizer { + private static final String CONTINUATION_PREFIX = "##"; + /** BERT classification token: {@code [CLS]}. */ public static final String BERT_CLS_TOKEN = "[CLS]"; /** BERT separator token: {@code [SEP]}. */ @@ -80,28 +78,33 @@ public class WordpieceTokenizer implements Tokenizer { private final String classificationToken; private final String separatorToken; private final String unknownToken; - private int maxTokenLength = 50; + private final int maxTokenLength; /** * Initializes a {@link WordpieceTokenizer} with a {@code vocabulary} and a default - * {@code maxTokenLength} of {@code 50}. + * maximum token length of 100 Unicode code points. * - * @param vocabulary A set of tokens considered the vocabulary. + * @param vocabulary A set of tokens considered the vocabulary; must not be {@code null} + * or contain {@code null} or empty entries. + * @throws IllegalArgumentException Thrown if {@code vocabulary} is invalid. */ public WordpieceTokenizer(Set vocabulary) { - this(vocabulary, BERT_CLS_TOKEN, BERT_SEP_TOKEN, BERT_UNK_TOKEN); + this(vocabulary, BERT_CLS_TOKEN, BERT_SEP_TOKEN, BERT_UNK_TOKEN, + BertNormalization.DEFAULT_MAX_WORD_CODE_POINTS); } /** * Initializes a {@link WordpieceTokenizer} with a {@code vocabulary} and a custom * {@code maxTokenLength}. * - * @param vocabulary A set of tokens considered the vocabulary. - * @param maxTokenLength A non-negative number that is used as maximum token length. + * @param vocabulary A set of tokens considered the vocabulary; must not be {@code null} + * or contain {@code null} or empty entries. + * @param maxTokenLength The non-negative maximum number of Unicode code points in one token. + * @throws IllegalArgumentException Thrown if {@code vocabulary} is invalid or + * {@code maxTokenLength} is negative. */ public WordpieceTokenizer(Set vocabulary, int maxTokenLength) { - this(vocabulary); - this.maxTokenLength = requireNonNegative(maxTokenLength); + this(vocabulary, BERT_CLS_TOKEN, BERT_SEP_TOKEN, BERT_UNK_TOKEN, maxTokenLength); } /** @@ -111,31 +114,33 @@ public WordpieceTokenizer(Set vocabulary, int maxTokenLength) { * use different special tokens instead of the BERT * defaults. * - * @param vocabulary The vocabulary. - * @param classificationToken The CLS token. - * @param separatorToken The SEP token. - * @param unknownToken The UNK token. + * @param vocabulary The vocabulary; must not be {@code null} or contain {@code null} + * or empty entries. + * @param classificationToken The CLS token; must not be {@code null} or empty. + * @param separatorToken The SEP token; must not be {@code null} or empty. + * @param unknownToken The UNK token; must not be {@code null} or empty. + * @throws IllegalArgumentException Thrown if an argument is invalid. */ public WordpieceTokenizer( final Set vocabulary, final String classificationToken, final String separatorToken, final String unknownToken) { - this.vocabulary = vocabulary; - this.classificationToken = classificationToken; - this.separatorToken = separatorToken; - this.unknownToken = unknownToken; + this(vocabulary, classificationToken, separatorToken, unknownToken, + BertNormalization.DEFAULT_MAX_WORD_CODE_POINTS); } /** * Initializes a {@link WordpieceTokenizer} with a {@code vocabulary}, * custom special tokens and a custom {@code maxTokenLength}. * - * @param vocabulary The vocabulary. - * @param classificationToken The CLS token. - * @param separatorToken The SEP token. - * @param unknownToken The UNK token. - * @param maxTokenLength A non-negative number that is used as maximum token length. + * @param vocabulary The vocabulary; must not be {@code null} or contain {@code null} + * or empty entries. + * @param classificationToken The CLS token; must not be {@code null} or empty. + * @param separatorToken The SEP token; must not be {@code null} or empty. + * @param unknownToken The UNK token; must not be {@code null} or empty. + * @param maxTokenLength The non-negative maximum number of Unicode code points in one token. + * @throws IllegalArgumentException Thrown if an argument is invalid. */ public WordpieceTokenizer( final Set vocabulary, @@ -143,11 +148,44 @@ public WordpieceTokenizer( final String separatorToken, final String unknownToken, final int maxTokenLength) { - this(vocabulary, classificationToken, separatorToken, unknownToken); + this.vocabulary = copyVocabulary(vocabulary); + this.classificationToken = requireToken(classificationToken, "classificationToken"); + this.separatorToken = requireToken(separatorToken, "separatorToken"); + this.unknownToken = requireToken(unknownToken, "unknownToken"); this.maxTokenLength = requireNonNegative(maxTokenLength); } - private static int requireNonNegative(final int maxTokenLength) { + /** Validates and copies a vocabulary. */ + private Set copyVocabulary(Set vocabulary) { + if (vocabulary == null) { + throw new IllegalArgumentException("vocabulary must not be null"); + } + final Set copy = new HashSet<>(vocabulary.size()); + for (final String piece : vocabulary) { + if (piece == null) { + throw new IllegalArgumentException("vocabulary must not contain null"); + } + if (piece.isEmpty()) { + throw new IllegalArgumentException("vocabulary must not contain an empty piece"); + } + copy.add(piece); + } + return Set.copyOf(copy); + } + + /** Validates a special token. */ + private String requireToken(String token, String name) { + if (token == null) { + throw new IllegalArgumentException(name + " must not be null"); + } + if (token.isEmpty()) { + throw new IllegalArgumentException(name + " must not be empty"); + } + return token; + } + + /** Validates the maximum token length. */ + private int requireNonNegative(final int maxTokenLength) { if (maxTokenLength < 0) { throw new IllegalArgumentException( "maxTokenLength must be non-negative: " + maxTokenLength); @@ -167,9 +205,18 @@ public Span[] tokenizePos(final String text) { "Wordpiece tokens cannot be mapped to character spans of the original text"); } + /** + * {@inheritDoc} + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ @Override public String[] tokenize(final String text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final List tokens = new LinkedList<>(); tokens.add(classificationToken); @@ -189,7 +236,7 @@ public String[] tokenize(final String text) { final char[] characters = token.toCharArray(); - if (characters.length <= maxTokenLength) { + if (Character.codePointCount(characters, 0, characters.length) <= maxTokenLength) { // The pieces of this word. Only added to the result if the whole word matches. final List wordPieces = new LinkedList<>(); @@ -213,7 +260,7 @@ public String[] tokenize(final String text) { // This is a substring so prefix it with ##. if (start > 0) { - substring = "##" + substring; + substring = CONTINUATION_PREFIX + substring; } // See if the substring is in the vocabulary. @@ -231,7 +278,7 @@ public String[] tokenize(final String text) { } // Subtract 1 from the end to find the next longest piece in the vocabulary. - end--; + end -= Character.charCount(Character.codePointBefore(characters, end)); } diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/CategoryTable.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/CategoryTable.java new file mode 100644 index 0000000000..48229674d5 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/CategoryTable.java @@ -0,0 +1,301 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +import opennlp.tools.tokenize.lattice.MecabDictionary.Category; + +/** + * The {@code char.def} code point to category mappings over the Unicode code point + * range. The first category on a mapping supplies the unknown-word settings. Each + * listed category can keep a group running while following characters also list it. + * + *

    The Basic Multilingual Plane is stored in a directly indexed array. The + * supplementary planes are stored as a sorted, non-overlapping range table searched by + * binary search, because dictionaries map them in a handful of large blocks.

    + */ +final class CategoryTable { + + private final CategoryAssignment[] bmp; + private final int[] rangeStart; + private final int[] rangeEnd; + private final CategoryAssignment[] rangeCategory; + + /** + * Creates a table from resolved BMP entries and supplementary ranges. + * + * @param bmp The directly indexed BMP assignments. + * @param rangeStart The inclusive starts of the supplementary ranges. + * @param rangeEnd The inclusive upper bounds of the supplementary ranges. + * @param rangeCategory The assignment for each supplementary range. + */ + private CategoryTable(CategoryAssignment[] bmp, int[] rangeStart, int[] rangeEnd, + CategoryAssignment[] rangeCategory) { + this.bmp = bmp; + this.rangeStart = rangeStart; + this.rangeEnd = rangeEnd; + this.rangeCategory = rangeCategory; + } + + /** + * Looks up the categories a {@code char.def} mapping gives a code point. + * + * @param codePoint The code point to classify. + * @return The assignment, or {@code null} when no mapping covers the code point. + */ + CategoryAssignment categoriesOf(int codePoint) { + if (codePoint <= Character.MAX_VALUE) { + return bmp[codePoint]; + } + int low = 0; + int high = rangeStart.length - 1; + while (low <= high) { + final int middle = (low + high) >>> 1; + if (codePoint < rangeStart[middle]) { + high = middle - 1; + } else if (codePoint > rangeEnd[middle]) { + low = middle + 1; + } else { + return rangeCategory[middle]; + } + } + return null; + } + + /** + * The categories assigned to one code point, stored both in mapping order and as a + * mask over the dictionary's dense category ids. + */ + static final class CategoryAssignment { + + private final Category[] categories; + private final int categoryMask; + + /** + * Creates an assignment with the first entry as the primary category. + * + * @param categories The categories in mapping order. Must not be empty. + * @throws IllegalArgumentException Thrown if {@code categories} is {@code null} or + * empty. + */ + CategoryAssignment(Category[] categories) { + if (categories == null || categories.length == 0) { + throw new IllegalArgumentException("categories must not be null or empty"); + } + this.categories = categories.clone(); + int mask = 0; + for (final Category category : categories) { + mask |= 1 << category.id(); + } + this.categoryMask = mask; + } + + /** + * Returns the first category on the mapping. + * + * @return The category that supplies unknown-word settings. + */ + Category primary() { + return categories[0]; + } + + /** + * Computes the run end after comparing this assignment with the next character. + * MeCab's + * + * {@code seekToOtherType} replaces the current mask after each accepted + * character, so successive assignments must overlap. + * + * @param next The next character's assignment, or {@code null} at the end of text. + * @param nextRunEnd The run end calculated at the next character. + * @param characterEnd The exclusive end of the current character. + * @return {@code nextRunEnd} when the assignments intersect; + * {@code characterEnd} otherwise. + */ + int continuedRunEnd(CategoryAssignment next, int nextRunEnd, + int characterEnd) { + return next != null && (categoryMask & next.categoryMask) != 0 + ? nextRunEnd : characterEnd; + } + } + + /** + * Collects {@code char.def} mappings in file order and builds a + * {@link CategoryTable}, giving a later mapping precedence over an earlier one that + * covers the same code point, which is what direct indexing does for the BMP. + */ + static final class Builder { + + /** + * One mapping line retained for validation after all categories have been read. + * + * @param sourceStart The first code point on the mapping line, used in error + * messages even if a later mapping replaces it. + * @param categories The category names from the mapping, primary first. + */ + private record Mapping(int sourceStart, String[] categories) { + } + + private final String[][] bmp = new String[Character.MAX_VALUE + 1][]; + private final List bounds = new ArrayList<>(); + private final List names = new ArrayList<>(); + private final List mappings = new ArrayList<>(); + + /** + * Records one inclusive code point range's categories. + * + * @param from The first code point of the range. + * @param to The last code point of the range, inclusive. + * @param categories The category names to give the range, primary first. Must not + * be {@code null} or empty. + */ + void map(int from, int to, String[] categories) { + // All positions written by this mapping store this array reference. build() + // relies on array identity to resolve one CategoryAssignment per mapping line. + mappings.add(new Mapping(from, categories)); + for (int c = from; c <= Math.min(to, Character.MAX_VALUE); c++) { + bmp[c] = categories; + } + if (to > Character.MAX_VALUE) { + bounds.add(new int[] {Math.max(from, Character.MAX_VALUE + 1), to}); + names.add(categories); + } + } + + /** + * Builds the lookup table from the recorded mappings. + * + * @param categories The categories the {@code char.def} category section defined, + * keyed by name. + * @return The table. Not {@code null}. + * @throws IOException Thrown if a mapping names a category that was not defined. + */ + CategoryTable build(Map categories) throws IOException { + // Cut the supplementary ranges at every boundary they introduce, so that each + // resulting elementary interval is covered by a single winning range and the + // table stays sorted and non-overlapping for binary search. + final int[] edges = new int[bounds.size() * 2]; + for (int i = 0; i < bounds.size(); i++) { + edges[i * 2] = bounds.get(i)[0]; + edges[i * 2 + 1] = bounds.get(i)[1] + 1; + } + Arrays.sort(edges); + final List intervals = new ArrayList<>(); + final List winners = new ArrayList<>(); + for (int i = 0; i < edges.length - 1; i++) { + if (edges[i] == edges[i + 1]) { + continue; + } + final String[] winner = lastCovering(edges[i]); + if (winner == null) { + continue; + } + final int previous = intervals.size() - 1; + if (previous >= 0 && intervals.get(previous)[1] == edges[i] - 1 + && Arrays.equals(winners.get(previous), winner)) { + intervals.get(previous)[1] = edges[i + 1] - 1; + } else { + intervals.add(new int[] {edges[i], edges[i + 1] - 1}); + winners.add(winner); + } + } + final int[] starts = new int[intervals.size()]; + final int[] ends = new int[intervals.size()]; + for (int i = 0; i < intervals.size(); i++) { + starts[i] = intervals.get(i)[0]; + ends[i] = intervals.get(i)[1]; + } + final Map resolvedAssignments = + new IdentityHashMap<>(); + // Validate mappings in file order, including mappings replaced across their + // ranges. This makes typo detection independent of range precedence. + for (final Mapping mapping : mappings) { + resolve(mapping.categories(), categories, resolvedAssignments, + mapping.sourceStart()); + } + final CategoryAssignment[] resolvedBmp = new CategoryAssignment[bmp.length]; + for (int c = 0; c < bmp.length; c++) { + if (bmp[c] != null) { + resolvedBmp[c] = resolve(bmp[c], categories, resolvedAssignments, c); + } + } + final CategoryAssignment[] resolvedRanges = new CategoryAssignment[winners.size()]; + for (int i = 0; i < winners.size(); i++) { + resolvedRanges[i] = resolve(winners.get(i), categories, resolvedAssignments, + starts[i]); + } + return new CategoryTable(resolvedBmp, starts, ends, resolvedRanges); + } + + /** + * Resolves mapped category names using the defined categories. A mapping to an + * undefined category fails at load and names the offending code point. + * + * @param names The category names on a mapping line, primary first. + * @param categories The defined categories, keyed by name. + * @param resolvedAssignments Previously resolved mapping lines, indexed by their + * shared category-name arrays. + * @param codePoint A code point the mapping covers, for the error message. + * @return The resolved assignment. Not {@code null}. + * @throws IOException Thrown if any named category was not defined. + */ + private CategoryAssignment resolve(String[] names, Map categories, + Map resolvedAssignments, int codePoint) + throws IOException { + final CategoryAssignment cached = resolvedAssignments.get(names); + if (cached != null) { + return cached; + } + final Category[] resolved = new Category[names.length]; + for (int i = 0; i < names.length; i++) { + resolved[i] = categories.get(names[i]); + if (resolved[i] == null) { + throw new IOException(String.format(Locale.ROOT, + MecabDictionary.CHAR_DEF + " declaration at U+%04X names the" + + " undefined category %s", codePoint, names[i])); + } + } + final CategoryAssignment assignment = new CategoryAssignment(resolved); + resolvedAssignments.put(names, assignment); + return assignment; + } + + /** + * Finds the category of the last recorded range covering a code point. + * + * @param codePoint The code point to look up. + * @return The category names, or {@code null} when no stored range covers it. + */ + private String[] lastCovering(int codePoint) { + for (int i = bounds.size() - 1; i >= 0; i--) { + final int[] range = bounds.get(i); + if (codePoint >= range[0] && codePoint <= range[1]) { + return names.get(i); + } + } + return null; + } + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/DoubleArrayLexicon.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/DoubleArrayLexicon.java new file mode 100644 index 0000000000..fb7b805c94 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/DoubleArrayLexicon.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.tools.tokenize.lattice; + +import java.util.ArrayDeque; +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +import opennlp.tools.tokenize.lattice.MecabDictionary.PrefixMatchConsumer; +import opennlp.tools.tokenize.lattice.MecabDictionary.WordEntry; + +/** + * The lexicon as a double-array trie: one transition is one array read and one + * comparison. Characters are recoded into dense labels ordered by descending + * frequency before the array is built, which keeps the array compact; a character + * absent from the lexicon misses in the recode table before the array is consulted. + * + *

    The layout is the classic base/check pair: from state {@code s}, label + * {@code c} leads to {@code t = base[s] + c} exactly when {@code check[t] == s}. + * Label {@code 0} terminates a surface and leads to a state. Its negative base + * encodes the index of the surface's entry list.

    + */ +final class DoubleArrayLexicon { + + private final int[] base; + private final int[] check; + private final int[] codeOf; + private final List[] values; + + /** + * Creates a lexicon from completed double-array tables and entry lists. + * + * @param base The transition offsets. + * @param check The parent state for each occupied index. + * @param codeOf The dense label for each UTF-16 character. + * @param values The entries for each surface. + */ + private DoubleArrayLexicon(int[] base, int[] check, int[] codeOf, + List[] values) { + this.base = base; + this.check = check; + this.codeOf = codeOf; + this.values = values; + } + + /** + * Builds the trie from the surface-keyed lexicon. + * + * @param lexicon The entries keyed by surface form. + * @return The built trie. Not {@code null}. + */ + @SuppressWarnings({"rawtypes", "unchecked"}) + static DoubleArrayLexicon build(Map> lexicon) { + final String[] surfaces = lexicon.keySet().toArray(new String[0]); + Arrays.sort(surfaces); + final List[] values = new List[surfaces.length]; + for (int i = 0; i < surfaces.length; i++) { + values[i] = List.copyOf(lexicon.get(surfaces[i])); + } + + // Dense recode: labels ordered by descending frequency get the small codes, so + // busy transitions cluster at the low end of the array. + final int[] frequency = new int[Character.MAX_VALUE + 1]; + for (final String surface : surfaces) { + for (int i = 0; i < surface.length(); i++) { + frequency[surface.charAt(i)]++; + } + } + final long[] rankedCharacters = new long[Character.MAX_VALUE + 1]; + int distinct = 0; + for (int c = 0; c <= Character.MAX_VALUE; c++) { + if (frequency[c] > 0) { + rankedCharacters[distinct++] = + ((long) (Integer.MAX_VALUE - frequency[c]) << 16) | c; + } + } + Arrays.sort(rankedCharacters, 0, distinct); + final int[] codeOf = new int[Character.MAX_VALUE + 1]; + Arrays.fill(codeOf, -1); + for (int rank = 0; rank < distinct; rank++) { + codeOf[(int) (rankedCharacters[rank] & Character.MAX_VALUE)] = rank + 1; + } + + final Builder builder = new Builder(surfaces, codeOf); + builder.insert(0, surfaces.length, 0, Builder.ROOT); + return new DoubleArrayLexicon(Arrays.copyOf(builder.base, builder.high + 1), + Arrays.copyOf(builder.check, builder.high + 1), codeOf, values); + } + + /** + * Reports every surface starting at a text position, walking the array once. + * + * @param text The text being segmented. + * @param from The position surfaces must start at. + * @param to The exclusive end of the searchable stretch. + * @param consumer Receives each match length with its entries. + */ + void prefixMatches(String text, int from, int to, + PrefixMatchConsumer consumer) { + int state = Builder.ROOT; + for (int i = from; i < to; i++) { + final char c = text.charAt(i); + final int code = codeOf[c]; + if (code < 0) { + return; + } + final int next = base[state] + code; + if (next >= check.length || check[next] != state) { + return; + } + state = next; + final int terminal = base[state]; + if (terminal < check.length && check[terminal] == state && base[terminal] < 0) { + consumer.accept(i - from + 1, values[-base[terminal] - 1]); + } + } + } + + /** + * Builds the trie from sorted surface ranges. Each node places children at a common + * free base, and a moving watermark makes that search near-linear over real + * lexicons. The traversal uses an explicit stack so a long surface cannot exhaust + * the thread stack. + */ + private static final class Builder { + + private static final int ROOT = 1; + private static final int EMPTY = -1; + + private final String[] surfaces; + private final int[] codeOf; + private int[] base; + private int[] check; + private int high = ROOT; + private int watermark = ROOT + 1; + private int valueIndex; + + /** + * Initializes storage for the sorted surfaces and their dense character codes. + * + * @param surfaces The sorted surface forms. + * @param codeOf The dense label for each UTF-16 character. + */ + private Builder(String[] surfaces, int[] codeOf) { + this.surfaces = surfaces; + this.codeOf = codeOf; + base = new int[1 << 16]; + check = new int[1 << 16]; + Arrays.fill(check, EMPTY); + } + + /** + * A surface range with children waiting to be placed. + * + * @param left The first surface in the range. + * @param right The exclusive end of the range. + * @param depth The character depth of the node. + * @param state The node's array index. + */ + private record PendingNode(int left, int right, int depth, int state) { + } + + /** + * Places one trie and the descendants without consuming the thread stack. + * + * @param left The first surface of the node's range. + * @param right The exclusive last surface of the node's range. + * @param depth The character depth of the node. + * @param state The node's own slot. + */ + private void insert(int left, int right, int depth, int state) { + final ArrayDeque pending = new ArrayDeque<>(); + pending.push(new PendingNode(left, right, depth, state)); + while (!pending.isEmpty()) { + final PendingNode node = pending.pop(); + final int[] labels = new int[node.right() - node.left()]; + int labelCount = 0; + int previous = -2; + for (int k = node.left(); k < node.right(); k++) { + final int label = surfaces[k].length() == node.depth() + ? 0 : codeOf[surfaces[k].charAt(node.depth())]; + if (label != previous) { + labels[labelCount++] = label; + previous = label; + } + } + final int found = findBase(labels, labelCount); + base[node.state()] = found; + for (int k = 0; k < labelCount; k++) { + final int child = found + labels[k]; + check[child] = node.state(); + if (child > high) { + high = child; + } + } + + int childEnd = node.right(); + for (int k = labelCount - 1; k >= 0; k--) { + final int label = labels[k]; + int childStart = childEnd - 1; + while (childStart > node.left() + && (surfaces[childStart - 1].length() == node.depth() + ? 0 : codeOf[surfaces[childStart - 1].charAt(node.depth())]) == label) { + childStart--; + } + final int child = found + label; + if (label == 0) { + base[child] = -(++valueIndex); + } else { + pending.push(new PendingNode( + childStart, childEnd, node.depth() + 1, child)); + } + childEnd = childStart; + } + } + } + + /** + * Finds the lowest base at which all labels use free indices. Labels are in + * surface-character order, so this method computes the numeric bounds. + * + * @param labels The child labels to place. + * @param labelCount How many leading elements of {@code labels} are in use. + * @return The base offset every label fits at. + */ + private int findBase(int[] labels, int labelCount) { + int smallest = labels[0]; + int largest = labels[0]; + for (int k = 1; k < labelCount; k++) { + smallest = Math.min(smallest, labels[k]); + largest = Math.max(largest, labels[k]); + } + int candidate = Math.max(1, watermark - smallest); + while (true) { + ensureCapacity(candidate + largest); + boolean fits = true; + for (int k = 0; fits && k < labelCount; k++) { + fits = check[candidate + labels[k]] == EMPTY; + } + if (fits) { + while (watermark < check.length && check[watermark] != EMPTY) { + watermark++; + } + return candidate; + } + candidate++; + } + } + + /** + * Grows the base and check arrays until a slot is addressable. + * + * @param slot The highest slot index that has to be writable. + */ + private void ensureCapacity(int slot) { + if (slot >= check.length) { + int capacity = check.length; + while (capacity <= slot) { + capacity += capacity >> 1; + } + base = Arrays.copyOf(base, capacity); + final int old = check.length; + check = Arrays.copyOf(check, capacity); + Arrays.fill(check, old, capacity, EMPTY); + } + } + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/LatticeTokenizer.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/LatticeTokenizer.java new file mode 100644 index 0000000000..d8ddcdcdb4 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/LatticeTokenizer.java @@ -0,0 +1,383 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.util.ArrayList; +import java.util.List; + +import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.tokenize.lattice.CategoryTable.CategoryAssignment; +import opennlp.tools.tokenize.lattice.MecabDictionary.Category; +import opennlp.tools.tokenize.lattice.MecabDictionary.WordEntry; +import opennlp.tools.util.Span; +import opennlp.tools.util.StringUtil; + +/** + * Dictionary-driven segmentation for languages written without spaces: a Viterbi + * search over the word lattice of a {@link MecabDictionary}, minimizing the sum of + * word costs and connection costs. This is the segmentation approach behind Japanese + * and Korean morphological analysis. A single decoder serves both because the supplied + * dictionary provides the language-specific data. + * + *

    Unknown text is handled through the dictionary's character categories: where the + * lexicon has no entry, or a category always invokes them, unknown-word candidates are + * generated per category template, grouping runs of same-category characters when the + * category requests it. A multi-category run continues while successive assignments + * overlap. Whitespace cannot join or appear as a morpheme. + * Every reported span is in original text coordinates.

    + * + *

    {@link #analyze(String)} returns full morphemes with their dictionary features; + * the {@link Tokenizer} view reports just the surfaces and spans.

    + * + *

    The tokenizer reads only immutable dictionary state and is safe to share between + * threads.

    + * + * @since 3.0.0 + */ +public final class LatticeTokenizer implements Tokenizer { + + /** The context id of the beginning and end of text. */ + private static final int BOUNDARY_CONTEXT = 0; + + private final MecabDictionary dictionary; + + /** + * Initializes the tokenizer. + * + * @param dictionary The dictionary to segment with. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code dictionary} is {@code null}. + */ + public LatticeTokenizer(MecabDictionary dictionary) { + if (dictionary == null) { + throw new IllegalArgumentException("dictionary must not be null"); + } + this.dictionary = dictionary; + } + + /** + * One lattice node: a candidate morpheme with its best path cost so far. Nodes + * ending at one position chain through {@link #nextEndingHere}. + */ + private static final class Node { + private final int start; + private final int end; + private final WordEntry entry; + private final boolean unknown; + private long pathCost = Long.MAX_VALUE; + private Node previous; + private Node nextEndingHere; + + /** + * Creates one lattice candidate. + * + * @param start The candidate start in original text coordinates. + * @param end The exclusive candidate end. + * @param entry The dictionary entry. + * @param unknown Whether unknown-word handling generated the candidate. + */ + private Node(int start, int end, WordEntry entry, boolean unknown) { + this.start = start; + this.end = end; + this.entry = entry; + this.unknown = unknown; + } + } + + /** + * Segments a text into morphemes with their dictionary features. + * + * @param text The text to segment. Must not be {@code null}. + * @return The morphemes in text order, spans in original coordinates, whitespace + * omitted. Never {@code null}; empty for empty or all-whitespace input. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + * @throws IllegalStateException Thrown if the dictionary offers no candidate at some + * position, which a {@code unk.def} without a {@code DEFAULT} template does. + */ + public List analyze(String text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final List morphemes = new ArrayList<>(); + int start = 0; + while (start < text.length()) { + if (StringUtil.isWhitespace(text.charAt(start))) { + start++; + continue; + } + int end = start; + while (end < text.length() && !StringUtil.isWhitespace(text.charAt(end))) { + end++; + } + decode(text, start, end, morphemes); + start = end; + } + return morphemes; + } + + /** + * {@inheritDoc} + * + *

    Reports the segmented surfaces, whitespace omitted.

    + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + * @throws IllegalStateException Thrown if the dictionary offers no candidate at some + * position; see {@link #analyze(String)}. + */ + @Override + public String[] tokenize(String text) { + final List morphemes = analyze(text); + final String[] tokens = new String[morphemes.size()]; + for (int i = 0; i < tokens.length; i++) { + tokens[i] = morphemes.get(i).surface(); + } + return tokens; + } + + /** + * {@inheritDoc} + * + *

    Reports the segmented spans in original text coordinates, whitespace omitted.

    + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + * @throws IllegalStateException Thrown if the dictionary offers no candidate at some + * position; see {@link #analyze(String)}. + */ + @Override + public Span[] tokenizePos(String text) { + final List morphemes = analyze(text); + final Span[] spans = new Span[morphemes.size()]; + for (int i = 0; i < spans.length; i++) { + spans[i] = morphemes.get(i).span(); + } + return spans; + } + + /** + * Runs the Viterbi search over one whitespace-free stretch of text. + * + * @param text The text being segmented. + * @param from The stretch start. + * @param to The exclusive stretch end. + * @param morphemes Receives the cheapest path's morphemes, in text order. + * @throws IllegalStateException Thrown if no path reaches the end of the stretch. + */ + private void decode(String text, int from, int to, List morphemes) { + final int length = to - from; + // Each element heads the chain of nodes ending at that position. + final Node[] endingAt = new Node[length + 1]; + + final CategoryAssignment[] categoryAt = new CategoryAssignment[length]; + final int[] runEndAt = new int[length]; + computeCategoryRuns(text, from, to, categoryAt, runEndAt); + + final List candidates = new ArrayList<>(); + for (int i = 0; i < length; i++) { + if (i > 0 && endingAt[i] == null) { + continue; + } + candidates.clear(); + candidates(text, from, to, i, categoryAt[i], runEndAt[i], candidates); + for (final Node candidate : candidates) { + relax(candidate, i == 0 ? null : endingAt[i]); + if (candidate.pathCost < Long.MAX_VALUE) { + final int end = candidate.end - from; + candidate.nextEndingHere = endingAt[end]; + endingAt[end] = candidate; + } + } + } + + Node best = null; + long bestTotal = Long.MAX_VALUE; + for (Node node = endingAt[length]; node != null; node = node.nextEndingHere) { + final long total = node.pathCost + + dictionary.connectionCost(node.entry.rightId(), BOUNDARY_CONTEXT); + if (best == null || total < bestTotal) { + best = node; + bestTotal = total; + } + } + if (best == null) { + throw new IllegalStateException( + "no segmentation path for \"" + text.subSequence(from, to) + "\""); + } + + final List reversed = new ArrayList<>(); + for (Node node = best; node != null; node = node.previous) { + reversed.add(new Morpheme(new Span(node.start, node.end), + text.substring(node.start, node.end), node.entry.features(), node.unknown)); + } + for (int i = reversed.size() - 1; i >= 0; i--) { + morphemes.add(reversed.get(i)); + } + } + + /** + * Connects a candidate to the cheapest predecessor ending where it starts. + * + * @param candidate The node to give a path cost and a predecessor. + * @param predecessors The head of the chain of nodes ending where the candidate + * starts, or {@code null} when it starts at the stretch start. + */ + private void relax(Node candidate, Node predecessors) { + if (predecessors == null) { + candidate.pathCost = candidate.entry.cost() + + dictionary.connectionCost(BOUNDARY_CONTEXT, candidate.entry.leftId()); + return; + } + for (Node predecessor = predecessors; predecessor != null; + predecessor = predecessor.nextEndingHere) { + final long total = predecessor.pathCost + + dictionary.connectionCost(predecessor.entry.rightId(), candidate.entry.leftId()) + + candidate.entry.cost(); + if (total < candidate.pathCost) { + candidate.pathCost = total; + candidate.previous = predecessor; + } + } + } + + /** + * Computes the per-position categories and connected-category run end for one + * stretch, in one right-to-left pass over the code points. Positions inside a + * surrogate sequence have a {@code null} assignment; candidates do not start there. + * + * @param text The text being segmented. + * @param from The stretch start. + * @param to The exclusive stretch end. + * @param categoryAt Receives each position's categories, indexed by {@code + * position - from}. + * @param runEndAt Receives each position's exclusive connected-category run end, + * indexed the same way. + */ + private void computeCategoryRuns(String text, int from, int to, + CategoryAssignment[] categoryAt, int[] runEndAt) { + CategoryAssignment next = null; + int nextRunEnd = to; + for (int position = to; position > from; ) { + final int codePoint = text.codePointBefore(position); + position -= Character.charCount(codePoint); + final int index = position - from; + categoryAt[index] = dictionary.categoriesOf(codePoint); + runEndAt[index] = categoryAt[index].continuedRunEnd(next, nextRunEnd, + position + Character.charCount(codePoint)); + next = categoryAt[index]; + nextRunEnd = runEndAt[index]; + } + } + + /** + * Gathers lexicon matches and unknown-word candidates starting at one position. + * + * @param text The text being segmented. + * @param from The stretch start. + * @param to The exclusive stretch end, which no candidate may reach past. + * @param offset The candidate start, relative to {@code from}. + * @param positionCategories The categories of that position, or {@code null} for a + * position inside a surrogate sequence. + * @param positionRunEnd The exclusive end of the overlapping-category run starting + * there. Used only when {@code positionCategories} is not + * {@code null}. + * @param candidates Receives the candidates. Must be empty on entry. + * @throws IllegalStateException Thrown if neither the lexicon, the position's + * category, nor the {@code DEFAULT} template offers a candidate. + */ + private void candidates(String text, int from, int to, int offset, + CategoryAssignment positionCategories, int positionRunEnd, List candidates) { + final int position = from + offset; + dictionary.prefixMatches(text, position, to, (length, entries) -> { + for (final WordEntry entry : entries) { + candidates.add(new Node(position, position + length, entry, false)); + } + }); + final boolean lexiconMatch = !candidates.isEmpty(); + + final int codePoint = text.codePointAt(position); + final Category category; + final int runEnd; + if (positionCategories == null) { + // Only a lexicon surface ending inside a surrogate pair can make such a + // position reachable; classify the stray code unit on the spot so the lattice + // stays connected. + category = dictionary.categoryOf(codePoint); + runEnd = position + Character.charCount(codePoint); + } else { + category = positionCategories.primary(); + runEnd = positionRunEnd; + } + if (!lexiconMatch || category.invoke()) { + final List templates = dictionary.unknownEntries(category.name()); + if (templates != null) { + addUnknown(candidates, text, position, runEnd, category, templates); + } + } + if (candidates.isEmpty()) { + // A category with no grouped or fixed-length candidate still provides one + // character. Incomplete dictionaries fall back to the DEFAULT template. + List fallback = dictionary.unknownEntries(category.name()); + if (fallback == null) { + fallback = dictionary.unknownEntries(MecabDictionary.DEFAULT_CATEGORY); + } + if (fallback != null) { + for (final WordEntry entry : fallback) { + candidates.add( + new Node(position, position + Character.charCount(codePoint), entry, true)); + } + } + } + if (candidates.isEmpty()) { + throw new IllegalStateException("dictionary provides no candidate at position " + + position + "; unk.def lacks a DEFAULT template"); + } + } + + /** + * Emits unknown-word candidates per the category's grouping and length settings. + * + *

    Candidates remain inside a run connected by overlapping category assignments. + * Lengths count code points, not UTF-16 code units.

    + * + * @param candidates Receives the candidates. + * @param text The text being segmented. + * @param position The position the candidates start at. + * @param runEnd The exclusive end of the connected-category run starting at + * {@code position}. + * @param category The category of that run. + * @param templates The category's unknown-word templates. + */ + private void addUnknown(List candidates, String text, int position, + int runEnd, Category category, List templates) { + if (category.group()) { + for (final WordEntry entry : templates) { + candidates.add(new Node(position, runEnd, entry, true)); + } + } + final int lengths = category.length(); + int end = position; + for (int length = 1; length <= lengths && end < runEnd; length++) { + end += Character.charCount(text.codePointAt(end)); + if (category.group() && end == runEnd) { + // This length coincides with the grouped run emitted above; skip the duplicate. + continue; + } + for (final WordEntry entry : templates) { + candidates.add(new Node(position, end, entry, true)); + } + } + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionary.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionary.java new file mode 100644 index 0000000000..ebd7c9f34b --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionary.java @@ -0,0 +1,683 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.BufferedReader; +import java.io.IOException; +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.nio.file.DirectoryStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.BitSet; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import opennlp.tools.tokenize.lattice.CategoryTable.CategoryAssignment; +import opennlp.tools.util.ResourceLimits; +import opennlp.tools.util.StringUtil; + +/** + * An immutable, in-memory dictionary in the + * MeCab directory format: lexicon entries + * from the {@code *.csv} files, connection costs from {@code matrix.def}, character + * categories from {@code char.def}, and unknown-word templates from {@code unk.def}, + * loaded from a user-supplied dictionary directory. No dictionary data is bundled or + * downloaded by this class. + * + *

    The same format serves multiple languages: the Japanese + * IPADIC and + * UniDic distributions and the Korean + * mecab-ko-dic all load + * through this one reader, with the feature columns passed through untouched because + * their schemas differ.

    + * + *

    Each instance uses about 0.75 MB for category tables indexed by the 16-bit + * code-unit space, so load once and share. Lexicon CSV files under the directory are + * read in sorted path order so tie-breaking is stable across file systems. Connection + * costs must cover all matrix cells; missing and duplicate entries are + * rejected instead of being treated as cost zero. Matrix dimensions and the lexicon + * entry count are bounded by {@link ResourceLimits#MAX_ENTRIES}, and the matrix cell + * count by {@link ResourceLimits#MAX_MATRIX_CELLS}. Lexicon CSV fields may be + * MeCab-quoted with {@code ""} escapes. Word and connection costs must fit in a signed + * 16-bit integer. An {@code unk.def} template must name a category defined by + * {@code char.def}.

    + * + *

    Instances are immutable and safe to share between threads.

    + * + * @see LatticeTokenizer + * @since 3.0.0 + */ +public final class MecabDictionary { + + /** + * The category name every {@code char.def} must define; unmapped code points and + * unknown-word handling fall back to it. + */ + static final String DEFAULT_CATEGORY = "DEFAULT"; + + private static final String MATRIX_DEF = "matrix.def"; + static final String CHAR_DEF = "char.def"; + private static final String UNK_DEF = "unk.def"; + + /** + * Maximum category count accepted by MeCab's + * + * character-property compiler. + */ + private static final int MAX_CATEGORY_COUNT = 17; + + /** + * Maximum value of MeCab's + * + * 4-bit category length field. + */ + private static final int MAX_CATEGORY_LENGTH = 15; + + static final String LEXICON_EXTENSION = ".csv"; + static final String DEFINITION_EXTENSION = ".def"; + static final String CONFIGURATION_FILE = "dicrc"; + private static final String LEXICON_GLOB = "*" + LEXICON_EXTENSION; + private static final char COMMENT_MARKER = '#'; + + /** The code point prefix used by {@code char.def}, in either letter case. */ + private static final String HEX_PREFIX = "0x"; + + /** The separator between the two ends of a {@code char.def} code point range. */ + private static final String RANGE_SEPARATOR = ".."; + + /** The {@code char.def} field value that turns a category flag on. */ + private static final String FLAG_ON = "1"; + + /** The {@code char.def} field value that turns a category flag off. */ + private static final String FLAG_OFF = "0"; + + /** + * One lexicon or unknown-word entry. + * + * @param leftId The left context id, an index into the connection matrix. + * @param rightId The right context id, an index into the connection matrix. + * @param cost The entry's own cost. + * @param features The entry's feature columns, in file order. + */ + record WordEntry(int leftId, int rightId, int cost, List features) { + } + + /** + * One character category's unknown-word behavior from {@code char.def}. + * + * @param id The dense zero-based category id. + * @param name The category name. + * @param invoke Whether unknown-word candidates are generated even where the lexicon + * matched. + * @param group Whether a full run connected by overlapping category assignments is + * emitted as one candidate. + * @param length How many leading characters of the run are offered as candidates. + */ + record Category(int id, String name, boolean invoke, boolean group, int length) { + } + + /** Receives one common-prefix match during {@link #prefixMatches}. */ + interface PrefixMatchConsumer { + + /** + * Accepts one match. + * + * @param length The matched surface length in characters. + * @param entries The lexicon entries for that surface. + */ + void accept(int length, List entries); + } + + private final DoubleArrayLexicon lexicon; + private final short[] connectionCosts; + private final int rightSize; + private final CategoryTable categoryTable; + private final CategoryAssignment defaultCategories; + private final Map> unknownEntries; + + /** + * Creates an immutable dictionary from parsed lexicon and definition data. + * + * @param lexicon The surface lexicon. + * @param connectionCosts The flattened connection matrix. + * @param rightSize The number of left-context columns in each matrix record. + * @param categories The character categories by name. + * @param categoryTable The code point category assignments. + * @param unknownEntries The unknown-word templates by category. + */ + private MecabDictionary(DoubleArrayLexicon lexicon, + short[] connectionCosts, int rightSize, Map categories, + CategoryTable categoryTable, Map> unknownEntries) { + this.lexicon = lexicon; + this.connectionCosts = connectionCosts; + this.rightSize = rightSize; + this.categoryTable = categoryTable; + final Category defaultCategory = Objects.requireNonNull( + categories.get(DEFAULT_CATEGORY), "DEFAULT category"); + this.defaultCategories = new CategoryAssignment( + new Category[] {defaultCategory}); + final Map> copy = new HashMap<>(unknownEntries.size()); + for (final Map.Entry> entry : unknownEntries.entrySet()) { + copy.put(entry.getKey(), List.copyOf(entry.getValue())); + } + this.unknownEntries = Map.copyOf(copy); + } + + /** + * Loads a dictionary directory encoded in UTF-8. + * + * @param directory The unpacked dictionary directory. Must not be {@code null}. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if reading fails or a file is malformed. + * @throws IllegalArgumentException Thrown if {@code directory} is {@code null}. + */ + public static MecabDictionary load(Path directory) throws IOException { + return load(directory, StandardCharsets.UTF_8); + } + + /** + * Loads a dictionary directory. + * + * @param directory The unpacked dictionary directory holding the {@code *.csv} + * lexicon files, {@code matrix.def}, {@code char.def}, and + * {@code unk.def}. Must not be {@code null}. + * @param charset The encoding the distribution uses, for example UTF-8 or EUC-JP. + * Must not be {@code null}. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if reading fails, a required file is missing, a file is + * malformed, or a lexicon entry's context ids are outside the + * {@code matrix.def} dimensions. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static MecabDictionary load(Path directory, Charset charset) throws IOException { + if (directory == null) { + throw new IllegalArgumentException("directory must not be null"); + } + if (charset == null) { + throw new IllegalArgumentException("charset must not be null"); + } + // The connection matrix is read first because its dimensions are what every + // lexicon entry's context ids have to be inside of. + final Path matrixFile = directory.resolve(MATRIX_DEF); + if (!Files.exists(matrixFile)) { + throw new IOException("required dictionary file is missing: " + matrixFile); + } + final int leftSize; + final int rightSize; + final short[] costs; + final int cellCount; + try (BufferedReader reader = Files.newBufferedReader(matrixFile, charset)) { + final String rawHeader = reader.readLine(); + if (rawHeader == null) { + throw new IOException("empty " + MATRIX_DEF + " under " + directory); + } + final String headerLine = StringUtil.trimUnicodeWhitespace(rawHeader); + if (headerLine.isEmpty()) { + throw new IOException("empty " + MATRIX_DEF + " under " + directory); + } + final String[] header = splitWhitespace(headerLine); + if (header.length != 2) { + throw new IOException("malformed " + MATRIX_DEF + " header: " + headerLine); + } + leftSize = parseInt(header[0], MATRIX_DEF, 1); + rightSize = parseInt(header[1], MATRIX_DEF, 1); + if (leftSize < 1 || rightSize < 1) { + throw new IOException(MATRIX_DEF + " dimensions must be positive, got " + + leftSize + " " + rightSize); + } + if (leftSize > ResourceLimits.MAX_ENTRIES + || rightSize > ResourceLimits.MAX_ENTRIES) { + throw new IOException(MATRIX_DEF + " dimensions " + leftSize + " x " + rightSize + + " exceed safe limit of " + ResourceLimits.MAX_ENTRIES); + } + final long cells = (long) leftSize * rightSize; + if (cells > Integer.MAX_VALUE) { + throw new IOException(MATRIX_DEF + " dimensions " + leftSize + " x " + rightSize + + " overflow the addressable connection matrix"); + } + if (cells > ResourceLimits.MAX_MATRIX_CELLS) { + throw new IOException(MATRIX_DEF + " dimensions " + leftSize + " x " + rightSize + + " exceed safe limit of " + ResourceLimits.MAX_MATRIX_CELLS); + } + cellCount = (int) cells; + costs = new short[cellCount]; + // leftSize bounds right-context ids and rightSize bounds left-context ids, matching + // MeCab's connector.h layout (the names read transposed against the id names). + final BitSet filled = new BitSet(cellCount); + int lineNumber = 1; + String raw; + while ((raw = reader.readLine()) != null) { + lineNumber++; + final String line = StringUtil.trimUnicodeWhitespace(raw); + if (line.isEmpty()) { + continue; + } + final String[] fields = splitWhitespace(line); + if (fields.length != 3) { + throw new IOException("malformed " + MATRIX_DEF + " line " + lineNumber); + } + final int right = parseInt(fields[0], MATRIX_DEF, lineNumber); + final int left = parseInt(fields[1], MATRIX_DEF, lineNumber); + if (right < 0 || right >= leftSize || left < 0 || left >= rightSize) { + throw new IOException("malformed " + MATRIX_DEF + " line " + lineNumber + + ": context ids " + right + " " + left + + " are outside the declared dimensions " + leftSize + " " + rightSize); + } + final int cost = parseInt(fields[2], MATRIX_DEF, lineNumber); + if (cost < Short.MIN_VALUE || cost > Short.MAX_VALUE) { + throw new IOException("malformed " + MATRIX_DEF + " line " + lineNumber + + ": connection cost " + cost + " is outside the 16-bit range the" + + " format defines"); + } + final int index = right * rightSize + left; + if (filled.get(index)) { + throw new IOException("duplicate " + MATRIX_DEF + " entry " + right + " " + + left + " at line " + lineNumber); + } + costs[index] = (short) cost; + filled.set(index); + } + if (filled.cardinality() != cellCount) { + throw new IOException(MATRIX_DEF + " declares " + leftSize + " x " + rightSize + + " connection costs but only " + filled.cardinality() + + " pairs are listed"); + } + } + + final Map> lexicon = new HashMap<>(); + final List csvFiles = new ArrayList<>(); + try (DirectoryStream stream = Files.newDirectoryStream(directory, LEXICON_GLOB)) { + for (final Path csv : stream) { + csvFiles.add(csv); + } + } + Collections.sort(csvFiles); + final int[] entryCount = {0}; + for (final Path csv : csvFiles) { + readLexicon(csv, charset, lexicon, leftSize, rightSize, entryCount); + } + if (lexicon.isEmpty()) { + throw new IOException("no lexicon entries found under " + directory); + } + + final Map categories = new HashMap<>(); + final CategoryTable.Builder categoryTable = new CategoryTable.Builder(); + readCharacterDefinition(directory.resolve(CHAR_DEF), charset, categories, + categoryTable); + final Map> unknown = new HashMap<>(); + final Path unkFile = directory.resolve(UNK_DEF); + readLexicon(unkFile, charset, unknown, leftSize, rightSize, new int[] {0}); + for (final String category : unknown.keySet()) { + if (!categories.containsKey(category)) { + throw new IOException( + UNK_DEF + " names the undefined category " + category + ": " + unkFile); + } + } + + return new MecabDictionary(DoubleArrayLexicon.build(lexicon), costs, + rightSize, categories, categoryTable.build(categories), unknown); + } + + /** + * Parses one lexicon-format CSV file, rejecting entries with context ids the + * connection matrix cannot be indexed with. + * + * @param file The file to read. + * @param charset The encoding to decode with. + * @param target Receives the entries, keyed by surface form. + * @param leftSize The first {@code matrix.def} dimension, which bounds right context + * ids. + * @param rightSize The second {@code matrix.def} dimension, which bounds left context + * ids. + * @param entryCount A one-element running total of entries read so far, shared across + * the lexicon files of one load. + * @throws IOException Thrown if the file is missing, an entry is malformed or has + * an empty surface, an entry's context id is outside the matrix dimensions, + * or the running entry count exceeds {@link ResourceLimits#MAX_ENTRIES}. + */ + private static void readLexicon(Path file, Charset charset, + Map> target, int leftSize, int rightSize, int[] entryCount) + throws IOException { + if (!Files.exists(file)) { + throw new IOException("required dictionary file is missing: " + file); + } + int lineNumber = 0; + try (BufferedReader reader = Files.newBufferedReader(file, charset)) { + String line; + while ((line = reader.readLine()) != null) { + lineNumber++; + if (line.isEmpty()) { + continue; + } + final List fields = splitCsv(line); + if (fields.size() < 4) { + throw new IOException("malformed entry at " + file + " line " + lineNumber); + } + final String surface = fields.get(0); + if (surface.isEmpty()) { + throw new IOException("malformed entry at " + file + " line " + lineNumber + + ": surface must not be empty"); + } + final int leftId = parseInt(fields.get(1), file.toString(), lineNumber); + final int rightId = parseInt(fields.get(2), file.toString(), lineNumber); + if (leftId < 0 || leftId >= rightSize) { + throw new IOException("malformed entry at " + file + " line " + lineNumber + + ": left context id " + leftId + " is outside the " + MATRIX_DEF + + " dimensions " + leftSize + " " + rightSize); + } + if (rightId < 0 || rightId >= leftSize) { + throw new IOException("malformed entry at " + file + " line " + lineNumber + + ": right context id " + rightId + " is outside the " + MATRIX_DEF + + " dimensions " + leftSize + " " + rightSize); + } + if (entryCount[0] >= ResourceLimits.MAX_ENTRIES) { + throw new IOException("lexicon entry count exceeds safe limit of " + + ResourceLimits.MAX_ENTRIES); + } + entryCount[0]++; + final int cost = parseInt(fields.get(3), file.toString(), lineNumber); + if (cost < Short.MIN_VALUE || cost > Short.MAX_VALUE) { + throw new IOException("malformed entry at " + file + " line " + lineNumber + + ": word cost " + cost + + " is outside the 16-bit range the format defines"); + } + final WordEntry entry = new WordEntry(leftId, rightId, cost, + List.copyOf(fields.subList(4, fields.size()))); + target.computeIfAbsent(surface, key -> new ArrayList<>(1)).add(entry); + } + } + } + + /** + * Reads {@code char.def}: the category behavior lines and the code point mapping + * lines, in file order, so that a later mapping wins over an earlier one. + * + * @param file The file to read. + * @param charset The encoding to decode with. + * @param categories Receives the defined categories, keyed by name. + * @param categoryTable Receives the code point to category name mappings. + * @throws IOException Thrown if the file is missing, a line is malformed, a code + * point is outside the Unicode range, a range descends, a category is + * duplicated, category count or length exceeds the MeCab format, or the file + * defines no {@code DEFAULT} category. + */ + private static void readCharacterDefinition(Path file, Charset charset, + Map categories, CategoryTable.Builder categoryTable) + throws IOException { + if (!Files.exists(file)) { + throw new IOException("required dictionary file is missing: " + file); + } + int lineNumber = 0; + try (BufferedReader reader = Files.newBufferedReader(file, charset)) { + String raw; + while ((raw = reader.readLine()) != null) { + lineNumber++; + final String line = StringUtil.trimUnicodeWhitespace(stripComment(raw)); + if (line.isEmpty()) { + continue; + } + final String[] fields = splitWhitespace(line); + if (fields[0].regionMatches(true, 0, HEX_PREFIX, 0, HEX_PREFIX.length())) { + final int rangeSeparator = fields[0].indexOf(RANGE_SEPARATOR); + final int from; + final int to; + if (rangeSeparator >= 0) { + from = parseCodePoint(fields[0].substring(0, rangeSeparator), file, + lineNumber); + to = parseCodePoint( + fields[0].substring(rangeSeparator + RANGE_SEPARATOR.length()), file, + lineNumber); + } else { + from = parseCodePoint(fields[0], file, lineNumber); + to = from; + } + if (fields.length < 2) { + throw new IOException( + "mapping without category at " + file + " line " + lineNumber); + } + if (from > to) { + throw new IOException("code point range descends at " + file + " line " + + lineNumber); + } + categoryTable.map(from, to, Arrays.copyOfRange(fields, 1, fields.length)); + } else { + if (fields.length < 4) { + throw new IOException( + "malformed category at " + file + " line " + lineNumber); + } + if (!isFlag(fields[1]) || !isFlag(fields[2])) { + throw new IOException( + "malformed category flag at " + file + " line " + lineNumber); + } + final int length = parseInt(fields[3], file.toString(), lineNumber); + if (length < 0 || length > MAX_CATEGORY_LENGTH) { + throw new IOException( + "category LENGTH must be between 0 and " + MAX_CATEGORY_LENGTH + " at " + + file + " line " + lineNumber); + } + if (categories.containsKey(fields[0])) { + throw new IOException("duplicate " + CHAR_DEF + " category " + + fields[0] + " at line " + lineNumber); + } + if (categories.size() >= MAX_CATEGORY_COUNT) { + throw new IOException(CHAR_DEF + " defines " + (categories.size() + 1) + + " categories; MeCab supports at most " + MAX_CATEGORY_COUNT); + } + final Category category = new Category(categories.size(), fields[0], + FLAG_ON.equals(fields[1]), FLAG_ON.equals(fields[2]), length); + categories.put(fields[0], category); + } + } + } + if (!categories.containsKey(DEFAULT_CATEGORY)) { + throw new IOException( + CHAR_DEF + " defines no " + DEFAULT_CATEGORY + " category: " + file); + } + } + + /** + * Reports every lexicon surface starting at a text position, walking the trie once + * with no substring allocation. + * + * @param text The text being segmented. + * @param from The position surfaces must start at. + * @param to The exclusive end of the searchable stretch. + * @param consumer Receives each match. + */ + void prefixMatches(String text, int from, int to, PrefixMatchConsumer consumer) { + lexicon.prefixMatches(text, from, to, consumer); + } + + /** + * Reads the connection cost between two adjacent nodes. + * + * @param rightId The right context id of the earlier node. + * @param leftId The left context id of the later node. + * @return The connection cost. + */ + int connectionCost(int rightId, int leftId) { + return connectionCosts[rightId * rightSize + leftId]; + } + + /** + * Classifies a character by code point, so a character outside the Basic Multilingual + * Plane is classified once, not as separate surrogate code units. + * + * @param codePoint The code point to classify. + * @return Its category, falling back to {@code DEFAULT} when no {@code char.def} + * mapping covers the code point. Never {@code null}. + */ + Category categoryOf(int codePoint) { + return categoriesOf(codePoint).primary(); + } + + /** + * Classifies a character into all categories assigned by {@code char.def}. The first + * category supplies unknown-word behavior and all categories participate in run + * grouping. + * + * @param codePoint The code point to classify. + * @return The category assignment, falling back to {@code DEFAULT} when no mapping + * covers the code point. + */ + CategoryAssignment categoriesOf(int codePoint) { + final CategoryAssignment categories = categoryTable.categoriesOf(codePoint); + return categories != null ? categories : defaultCategories; + } + + /** + * Looks up the unknown-word templates of a category. + * + * @param category The category name. + * @return The templates, or {@code null} when the category has none. + */ + List unknownEntries(String category) { + return unknownEntries.get(category); + } + + /** + * Removes a trailing {@code #} comment from a {@code char.def} line. + * + * @param line The raw line. + * @return The line up to but excluding the first {@code #}, or the complete line when + * there is none. + */ + private static String stripComment(String line) { + final int hash = line.indexOf(COMMENT_MARKER); + return hash < 0 ? line : line.substring(0, hash); + } + + /** + * Reports whether a {@code char.def} category flag field is exactly {@code 0} or + * {@code 1}. + * + * @param field The flag field text. + * @return {@code true} when the field is a recognized flag value. + */ + private static boolean isFlag(String field) { + return FLAG_ON.equals(field) || FLAG_OFF.equals(field); + } + + /** + * Splits a lexicon line on commas, honoring MeCab-style {@code "..."} quoting with + * {@code ""} escapes inside a quoted field. + * + * @param line The line to split. + * @return The fields in order, empty fields included. Never {@code null}. + */ + private static List splitCsv(String line) { + final List fields = new ArrayList<>(); + final StringBuilder field = new StringBuilder(); + boolean inQuotes = false; + for (int i = 0; i < line.length(); i++) { + final char c = line.charAt(i); + if (inQuotes) { + if (c == '"') { + if (i + 1 < line.length() && line.charAt(i + 1) == '"') { + field.append('"'); + i++; + } else { + inQuotes = false; + } + } else { + field.append(c); + } + } else if (c == '"') { + inQuotes = true; + } else if (c == ',') { + fields.add(field.toString()); + field.setLength(0); + } else { + field.append(c); + } + } + fields.add(field.toString()); + return fields; + } + + /** + * Splits a line into its whitespace-separated fields. + * + * @param line The line to split. + * @return The non-empty fields in order. Never {@code null}. + */ + private static String[] splitWhitespace(String line) { + final List parts = new ArrayList<>(); + int start = -1; + for (int i = 0; i <= line.length(); i++) { + if (i == line.length() || StringUtil.isWhitespace(line.charAt(i))) { + if (start >= 0) { + parts.add(line.substring(start, i)); + start = -1; + } + } else if (start < 0) { + start = i; + } + } + return parts.toArray(new String[0]); + } + + /** + * Parses a decimal integer field, reporting the file and line on failure. + * + * @param text The field text. + * @param file The file being read, for the error message. + * @param lineNumber The line being read, for the error message. + * @return The parsed value. + * @throws IOException Thrown if the field is not a valid integer. + */ + private static int parseInt(String text, String file, int lineNumber) + throws IOException { + try { + return Integer.parseInt(StringUtil.trimUnicodeWhitespace(text)); + } catch (NumberFormatException e) { + throw new IOException("malformed number in " + file + " line " + lineNumber, e); + } + } + + /** + * Parses a {@code 0x}-prefixed hexadecimal code point from {@code char.def}. + * + * @param text The field text including the {@code 0x} prefix. + * @param file The file being read, for the error message. + * @param lineNumber The line being read, for the error message. + * @return The parsed code point, which may be in a supplementary plane. + * @throws IOException Thrown if the field is shorter than the prefix, is not a valid + * hexadecimal number, or names a value no Unicode code point has. + */ + private static int parseCodePoint(String text, Path file, int lineNumber) + throws IOException { + final int codePoint; + try { + codePoint = Integer.parseInt( + StringUtil.trimUnicodeWhitespace(text).substring(HEX_PREFIX.length()), 16); + } catch (RuntimeException e) { + throw new IOException("malformed code point in " + file + " line " + lineNumber, e); + } + if (!Character.isValidCodePoint(codePoint)) { + throw new IOException("code point out of range in " + file + " line " + lineNumber); + } + return codePoint; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/Morpheme.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/Morpheme.java new file mode 100644 index 0000000000..a6694ddc58 --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/Morpheme.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.util.List; + +import opennlp.tools.util.Span; + +/** + * One morpheme from lattice segmentation: the {@link Span} it covers in the original + * text, its surface form, and the feature columns its dictionary entry carries. + * + *

    The features are the entry's columns exactly as listed in the dictionary, since + * different dictionaries carry different schemas: part of speech first by convention, + * then dictionary-specific columns such as conjugation, base form, or reading. A + * morpheme produced by unknown-word handling has the unknown entry's features and is + * marked as such.

    + * + * @param span The location of the morpheme in the original text. Must not be + * {@code null}. + * @param surface The covered text. Must not be {@code null} or empty. + * @param features The dictionary feature columns. Must not be {@code null}. + * @param unknown Whether the morpheme came from unknown-word handling rather than a + * lexicon entry. + * + * @since 3.0.0 + */ +public record Morpheme(Span span, String surface, List features, + boolean unknown) { + + /** + * Validates the morpheme. + * + * @throws IllegalArgumentException Thrown if {@code span}, {@code surface}, or + * {@code features} is {@code null}, or {@code surface} is empty. + */ + public Morpheme { + if (span == null) { + throw new IllegalArgumentException("span must not be null"); + } + if (surface == null || surface.isEmpty()) { + throw new IllegalArgumentException("surface must not be null or empty"); + } + if (features == null) { + throw new IllegalArgumentException("features must not be null"); + } + features = List.copyOf(features); + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/UnigramSegmenter.java b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/UnigramSegmenter.java new file mode 100644 index 0000000000..91f65bcafc --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/lattice/UnigramSegmenter.java @@ -0,0 +1,426 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.Charset; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.util.ResourceLimits; +import opennlp.tools.util.Span; +import opennlp.tools.util.StringUtil; + +/** + * Frequency-driven segmentation for Chinese and similar scripts: a Viterbi search that + * maximizes the summed log-probability of the words in a user-supplied frequency + * lexicon, with unlisted characters falling back to single-character words. This is the + * unigram model behind common Chinese segmenters. It omits the connection costs used + * by {@link LatticeTokenizer}, making it suitable for lexicons containing words and + * counts. + * + *

    The lexicon format is one entry per line: the word, its count, and optionally a + * tag, separated by whitespace. The lexicon file is user-supplied; no lexicon data is + * bundled. Every reported span is in original text coordinates.

    + * + *

    Instances are immutable and safe to share between threads.

    + * + * @since 3.0.0 + */ +public final class UnigramSegmenter implements Tokenizer { + + /** The log-probability charged to a character the lexicon does not know. */ + private final double unknownLogProbability; + + private final WordTrie trie; + + /** + * One immutable trie node: children are a sorted character array with a parallel + * node array, found by binary search, so a descent avoids boxing a {@link Character}. + */ + private static final class WordTrie { + + private final char[] keys; + private final WordTrie[] nodes; + private final double logProbability; + + /** + * Creates an immutable trie node. + * + * @param keys The sorted child labels. + * @param nodes The child nodes, parallel to {@code keys}. + * @param logProbability The word score, or {@link Double#NaN} for a nonterminal node. + */ + private WordTrie(char[] keys, WordTrie[] nodes, double logProbability) { + this.keys = keys; + this.nodes = nodes; + this.logProbability = logProbability; + } + + /** + * Descends one character. + * + * @param c The next surface character. + * @return The child node, or {@code null} when no surface continues with {@code c}. + */ + private WordTrie child(char c) { + final int index = Arrays.binarySearch(keys, c); + return index >= 0 ? nodes[index] : null; + } + } + + /** One mutable trie node during construction, copied into a {@link WordTrie}. */ + private static final class WordTrieBuilder { + + private final Map children = new HashMap<>(); + private double logProbability = Double.NaN; + private WordTrie built; + + /** + * One pending post-order traversal step. + * + * @param node The mutable node. + * @param childrenBuilt Whether the node's children have been copied. + */ + private record BuildStep(WordTrieBuilder node, boolean childrenBuilt) { + } + + /** + * Copies the mutable tree without consuming the thread stack. Each mutable node + * uses a step before child visits and a step afterward, preserving post-order + * construction for surfaces of any supported length. + */ + private WordTrie build() { + final ArrayDeque pending = new ArrayDeque<>(); + pending.push(new BuildStep(this, false)); + while (!pending.isEmpty()) { + final BuildStep step = pending.pop(); + final WordTrieBuilder node = step.node(); + if (!step.childrenBuilt()) { + pending.push(new BuildStep(node, true)); + for (final WordTrieBuilder child : node.children.values()) { + pending.push(new BuildStep(child, false)); + } + continue; + } + final char[] keys = new char[node.children.size()]; + int i = 0; + for (final Character key : node.children.keySet()) { + keys[i++] = key; + } + Arrays.sort(keys); + final WordTrie[] nodes = new WordTrie[keys.length]; + for (int k = 0; k < keys.length; k++) { + final WordTrieBuilder child = node.children.get(keys[k]); + nodes[k] = child.built; + child.built = null; + } + node.built = new WordTrie(keys, nodes, node.logProbability); + } + final WordTrie result = built; + built = null; + return result; + } + } + + /** + * Creates a segmenter from a word trie and the unknown-character score. + * + * @param trie The word trie. + * @param unknownLogProbability The score for an unlisted character. + */ + private UnigramSegmenter(WordTrie trie, double unknownLogProbability) { + this.trie = trie; + this.unknownLogProbability = unknownLogProbability; + } + + /** + * Loads a frequency lexicon encoded in UTF-8. + * + * @param lexicon The lexicon file. Must not be {@code null}. + * @return The segmenter. Not {@code null}. + * @throws IOException Thrown if reading fails or the lexicon is empty or malformed. + * @throws IllegalArgumentException Thrown if {@code lexicon} is {@code null}. + */ + public static UnigramSegmenter load(Path lexicon) throws IOException { + return load(lexicon, StandardCharsets.UTF_8); + } + + /** + * Loads a frequency lexicon. + * + * @param lexicon The lexicon file: one word, its count, and an optional tag per + * line. Must not be {@code null}. + * @param charset The lexicon encoding. Must not be {@code null}. + * @return The segmenter. Not {@code null}. + * @throws IOException Thrown if reading fails or the lexicon is empty or malformed. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static UnigramSegmenter load(Path lexicon, Charset charset) throws IOException { + if (lexicon == null) { + throw new IllegalArgumentException("lexicon must not be null"); + } + if (charset == null) { + throw new IllegalArgumentException("charset must not be null"); + } + try (InputStream in = Files.newInputStream(lexicon)) { + return load(in, charset); + } + } + + /** + * Loads a frequency lexicon from a stream. + * + * @param lexiconStream The lexicon content. Must not be {@code null}. Not closed. + * @param charset The lexicon encoding. Must not be {@code null}. + * @return The segmenter. Not {@code null}. + * @throws IOException Thrown if reading fails or the lexicon is empty or malformed. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static UnigramSegmenter load(InputStream lexiconStream, Charset charset) + throws IOException { + return loadInternal(lexiconStream, charset, ResourceLimits.MAX_ENTRIES); + } + + /** + * Loads a frequency lexicon under an entry limit. + * + * @param lexiconStream The lexicon content. Must not be {@code null}. Not closed. + * @param charset The lexicon encoding. Must not be {@code null}. + * @param maxEntries The inclusive limit on distinct lexicon entries. + * @return The segmenter. Not {@code null}. + * @throws IOException Thrown if reading fails or the lexicon is empty or malformed. + * @throws IllegalArgumentException Thrown if a parameter is invalid. + */ + private static UnigramSegmenter loadInternal(InputStream lexiconStream, Charset charset, + int maxEntries) throws IOException { + if (lexiconStream == null) { + throw new IllegalArgumentException("lexiconStream must not be null"); + } + if (charset == null) { + throw new IllegalArgumentException("charset must not be null"); + } + if (maxEntries < 1) { + throw new IllegalArgumentException("maxEntries must be positive"); + } + final Map counts = new HashMap<>(); + long total = 0; + final BufferedReader reader = new BufferedReader(new InputStreamReader( + lexiconStream, charset.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT))); + int lineNumber = 0; + String raw; + while ((raw = reader.readLine()) != null) { + lineNumber++; + final String line = StringUtil.trimUnicodeWhitespace(raw); + if (line.isEmpty()) { + continue; + } + final int wordEnd = whitespaceIndex(line); + if (wordEnd < 0) { + throw new IOException("lexicon line " + lineNumber + " has no count"); + } + final String word = line.substring(0, wordEnd); + int countStart = wordEnd; + while (countStart < line.length() && StringUtil.isWhitespace(line.charAt(countStart))) { + countStart++; + } + int countEnd = countStart; + while (countEnd < line.length() && !StringUtil.isWhitespace(line.charAt(countEnd))) { + countEnd++; + } + final long count; + try { + count = Long.parseLong(line.substring(countStart, countEnd)); + } catch (NumberFormatException e) { + throw new IOException("malformed count at lexicon line " + lineNumber, e); + } + if (count <= 0) { + throw new IOException("count must be positive at lexicon line " + lineNumber); + } + if (!counts.containsKey(word) && counts.size() >= maxEntries) { + throw new IOException( + "lexicon entry count exceeds safe limit of " + maxEntries); + } + counts.merge(word, count, Long::sum); + try { + total = Math.addExact(total, count); + } catch (ArithmeticException e) { + throw new IOException("lexicon count total overflows at line " + lineNumber, e); + } + } + if (counts.isEmpty()) { + throw new IOException("the lexicon lists no words"); + } + + final WordTrieBuilder root = new WordTrieBuilder(); + final double logTotal = Math.log(total); + for (final Map.Entry entry : counts.entrySet()) { + WordTrieBuilder node = root; + final String word = entry.getKey(); + for (int c = 0; c < word.length(); c++) { + node = node.children.computeIfAbsent(word.charAt(c), key -> new WordTrieBuilder()); + } + node.logProbability = Math.log(entry.getValue()) - logTotal; + } + // Charge an unlisted character half of one count out of the total, which makes it + // rarer than any listed word: every listed count is at least one. + final double unknown = Math.log(0.5) - logTotal; + return new UnigramSegmenter(root.build(), unknown); + } + + /** + * Loads a frequency lexicon under a caller-supplied entry limit. + * + * @param lexiconStream The lexicon content. Must not be {@code null}. Not closed. + * @param charset The lexicon encoding. Must not be {@code null}. + * @param maxEntries The inclusive limit on distinct lexicon entries. + * @return The segmenter. Not {@code null}. + * @throws IOException Thrown if reading fails or the lexicon is empty or malformed. + * @throws IllegalArgumentException Thrown if a parameter is invalid. + */ + static UnigramSegmenter load(InputStream lexiconStream, Charset charset, int maxEntries) + throws IOException { + return loadInternal(lexiconStream, charset, maxEntries); + } + + /** + * {@inheritDoc} + * + *

    Reports the segmented surfaces, whitespace omitted.

    + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + @Override + public String[] tokenize(String text) { + return Span.spansToStrings(tokenizePos(text), text); + } + + /** + * {@inheritDoc} + * + *

    Reports the segmented spans in original text coordinates, whitespace omitted.

    + * + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + @Override + public Span[] tokenizePos(String text) { + if (text == null) { + throw new IllegalArgumentException("text must not be null"); + } + final List spans = new ArrayList<>(); + int start = 0; + while (start < text.length()) { + if (StringUtil.isWhitespace(text.charAt(start))) { + start++; + continue; + } + int end = start; + while (end < text.length() && !StringUtil.isWhitespace(text.charAt(end))) { + end++; + } + decode(text, start, end, spans); + start = end; + } + return spans.toArray(new Span[0]); + } + + /** + * Viterbi over word log-probabilities within one whitespace-free stretch. + * + * @param text The text being segmented. + * @param from The stretch start. + * @param to The exclusive stretch end. + * @param spans Receives the best path's spans, in text order and in original text + * coordinates. + */ + private void decode(String text, int from, int to, List spans) { + final int length = to - from; + final double[] best = new double[length + 1]; + final int[] previous = new int[length + 1]; + for (int i = 1; i <= length; i++) { + best[i] = Double.NEGATIVE_INFINITY; + } + for (int i = 0; i < length; i++) { + if (best[i] == Double.NEGATIVE_INFINITY) { + continue; + } + // A single-character step at the unknown log-probability keeps every position + // reachable even where no lexicon word matches. The step advances one code + // point instead of one code unit, so an unknown supplementary character is advanced + // over as one unit and no span boundary can occur inside its surrogate pair. + final int width = Character.charCount(text.codePointAt(from + i)); + final double fallback = best[i] + unknownLogProbability; + if (i + width <= length && fallback > best[i + width]) { + best[i + width] = fallback; + previous[i + width] = i; + } + WordTrie node = trie; + for (int j = from + i; j < to; j++) { + node = node.child(text.charAt(j)); + if (node == null) { + break; + } + if (!Double.isNaN(node.logProbability)) { + final int end = j - from + 1; + final double score = best[i] + node.logProbability; + if (score > best[end]) { + best[end] = score; + previous[end] = i; + } + } + } + } + final List reversed = new ArrayList<>(); + for (int end = length; end > 0; end = previous[end]) { + reversed.add(new Span(from + previous[end], from + end)); + } + for (int i = reversed.size() - 1; i >= 0; i--) { + spans.add(reversed.get(i)); + } + } + + /** + * Finds the first whitespace character in a lexicon line. + * + * @param text The line to scan. + * @return The index of the first whitespace character, or {@code -1} when the line + * contains none. + */ + private static int whitespaceIndex(String text) { + for (int i = 0; i < text.length(); i++) { + if (StringUtil.isWhitespace(text.charAt(i))) { + return i; + } + } + return -1; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/util/ResourceLimits.java b/opennlp-api/src/main/java/opennlp/tools/util/ResourceLimits.java new file mode 100644 index 0000000000..0f6f4d846d --- /dev/null +++ b/opennlp-api/src/main/java/opennlp/tools/util/ResourceLimits.java @@ -0,0 +1,107 @@ +/* + * 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.util; + +/** + * Shared upper bounds for counts read from user-supplied resources, so a crafted + * file cannot force an outsized allocation before validation completes. + * + * @since 3.0.0 + */ +public final class ResourceLimits { + + /** + * System property for overriding {@link #MAX_ENTRIES}. + * Set at JVM startup, e.g. {@code -DOPENNLP_MAX_ENTRIES=5000000}. + * Falls back to {@code 10_000_000} if absent or invalid. + */ + public static final String MAX_ENTRIES_PROPERTY = "OPENNLP_MAX_ENTRIES"; + + /** + * Upper bound on count fields and resource sizes that drive allocations + * (matrix dimensions, lexicon entries, model outcome counts, and similar). + * Configurable via {@link #MAX_ENTRIES_PROPERTY}. + */ + public static final int MAX_ENTRIES = initLimit(MAX_ENTRIES_PROPERTY, 10_000_000); + + /** + * System property for overriding {@link #MAX_MATRIX_CELLS}. + * Set at JVM startup, e.g. {@code -Dopennlp.max.matrix.cells=20000000}. + * Falls back to {@code 134_217_728} if absent or invalid. + */ + public static final String MAX_MATRIX_CELLS_PROPERTY = "opennlp.max.matrix.cells"; + + /** + * Upper bound on the cell count of a two-dimensional cost table, whose entries + * are far smaller than the record-sized entries {@link #MAX_ENTRIES} bounds. + * The default of 2^27 cells caps a 16-bit cost matrix at 256 MiB, which admits + * every published MeCab-format distribution (mecab-ko-dic 2.1.1 alone declares + * 3822 x 2693, above {@link #MAX_ENTRIES}) while still refusing the roughly + * 4 GiB allocation a crafted {@code 46340 46340} header would force. + * Configurable via {@link #MAX_MATRIX_CELLS_PROPERTY}. + */ + public static final int MAX_MATRIX_CELLS = + initLimit(MAX_MATRIX_CELLS_PROPERTY, 134_217_728); + + private ResourceLimits() { + } + + /** + * Reads a positive integer limit from the given system property. + * + * @param property The system property name. Must not be {@code null}. + * @param defaultValue The value used when the property is absent or invalid. + * @return The configured limit, or {@code defaultValue}. + */ + static int initLimit(String property, int defaultValue) { + final String prop = System.getProperty(property, "").trim(); + if (!prop.isEmpty()) { + try { + final int val = Integer.parseInt(prop); + if (val > 0) { + return val; + } + } catch (NumberFormatException ignore) { + // Fall through to the default. + } + } + return defaultValue; + } + + /** + * Reads a positive long limit from a system property. + * + * @param property The system property name. Must not be {@code null}. + * @param defaultValue The value used when the property is absent or invalid. + * @return The configured limit, or {@code defaultValue}. + */ + static long initLimit(String property, long defaultValue) { + final String prop = System.getProperty(property, "").trim(); + if (!prop.isEmpty()) { + try { + final long val = Long.parseLong(prop); + if (val > 0) { + return val; + } + } catch (NumberFormatException ignore) { + // Use the default value. + } + } + return defaultValue; + } +} diff --git a/opennlp-api/src/main/java/opennlp/tools/util/StringUtil.java b/opennlp-api/src/main/java/opennlp/tools/util/StringUtil.java index 98cca59891..44f683d588 100644 --- a/opennlp-api/src/main/java/opennlp/tools/util/StringUtil.java +++ b/opennlp-api/src/main/java/opennlp/tools/util/StringUtil.java @@ -20,6 +20,7 @@ import java.nio.CharBuffer; import java.util.ArrayList; import java.util.List; +import java.util.Locale; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -268,6 +269,30 @@ public static boolean isEmpty(CharSequence theString) { return theString.length() == 0; } + /** + * Determines whether a {@link CharSequence} is blank: empty, or made up entirely of + * code points that {@link #isWhitespace(int)} accepts. Unlike + * {@link String#isBlank()}, this follows the toolkit's whitespace definition, which + * includes the no-break spaces the JDK predicate leaves out, so a value spelled + * entirely from them cannot pass a blank check as content. Unlike + * {@link #isUnicodeBlank(CharSequence)}, it resolves through the active + * {@link WhitespaceMode} and does not treat {@code null} as blank. + * + * @param theString The {@link CharSequence} to examine. Must not be {@code null}. + * @return {@code true} if {@code theString} is empty or all whitespace. + * @throws NullPointerException Thrown if {@code theString} is {@code null}. + */ + public static boolean isBlank(CharSequence theString) { + for (int i = 0; i < theString.length(); ) { + final int codePoint = Character.codePointAt(theString, i); + if (!isWhitespace(codePoint)) { + return false; + } + i += Character.charCount(codePoint); + } + return true; + } + /** * Get the minimum of three values. * @@ -485,8 +510,8 @@ public static String decodeShortestEditScript(String wordForm, String permutatio * @return Retrieves the Shortest Edit Script (SES) required to go from a word to a lemma. */ public static String getShortestEditScript(String wordForm, String lemma) { - String reversedWF = new StringBuffer(wordForm.toLowerCase()).reverse().toString(); - String reversedLemma = new StringBuffer(lemma.toLowerCase()).reverse().toString(); + String reversedWF = new StringBuffer(wordForm.toLowerCase(Locale.ROOT)).reverse().toString(); + String reversedLemma = new StringBuffer(lemma.toLowerCase(Locale.ROOT)).reverse().toString(); StringBuffer permutations = new StringBuffer(); String ses; if (!reversedWF.equals(reversedLemma)) { diff --git a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/CharClass.java b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/CharClass.java index 7abae2a230..7719a1eb3f 100644 --- a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/CharClass.java +++ b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/CharClass.java @@ -431,6 +431,97 @@ public AlignedText collapsePreservingAligned(CharSequence text, CodePointSet kee return new AlignedText(text, out.toString(), alignment.build(length)); } + /** + * Collapses runs of members like {@link #collapse(CharSequence)}, but emits + * {@code paragraphReplacement} when a run contains two or more logical line breaks, or the usual + * {@code replacement} when it contains at most one. A carriage return immediately followed by a + * line feed counts as one break, not two. The hard-wrap unwrap for Gutenberg-style prose uses + * this with the line-break code points as {@code lineBreaks} and {@code '\n'} as + * {@code paragraphReplacement}. + * + * @param text The text to collapse. + * @param lineBreaks The member code points that count as a line break when tallying a run. + * @param paragraphReplacement The replacement emitted for a run with two or more logical breaks. + * @return The collapsed text. + * @throws IllegalArgumentException Thrown if {@code text} or {@code lineBreaks} is {@code null}, + * or {@code paragraphReplacement} is not a valid code point. + */ + public String collapseParagraphPreserving(CharSequence text, CodePointSet lineBreaks, + int paragraphReplacement) { + requireNonNullArg(text, "text"); + requireNonNullArg(lineBreaks, "lineBreaks"); + requireValidCodePoint(paragraphReplacement); + final StringBuilder out = new StringBuilder(text.length()); + final int length = text.length(); + int i = 0; + while (i < length) { + final At cp = CodePoints.at(text, i); + if (members.contains(cp.codePoint())) { + int j = cp.nextIndex(i); + while (j < length) { + final At next = CodePoints.at(text, j); + if (!members.contains(next.codePoint())) { + break; + } + j = next.nextIndex(j); + } + final int emitted = countLogicalLineBreaks(text, i, j, lineBreaks) >= 2 + ? paragraphReplacement : replacement; + out.appendCodePoint(emitted); + i = j; + } else { + out.appendCodePoint(cp.codePoint()); + i = cp.nextIndex(i); + } + } + return out.toString(); + } + + /** + * Like {@link #collapseParagraphPreserving(CharSequence, CodePointSet, int)} but also produces + * the {@link Alignment} back to the original text. + * + * @param text The text to collapse. + * @param lineBreaks The member code points that count as a line break when tallying a run. + * @param paragraphReplacement The replacement emitted for a run with two or more logical breaks. + * @return The collapsed text and its alignment. + * @throws IllegalArgumentException Thrown if {@code text} or {@code lineBreaks} is {@code null}, + * or {@code paragraphReplacement} is not a valid code point. + */ + public AlignedText collapseParagraphPreservingAligned(CharSequence text, CodePointSet lineBreaks, + int paragraphReplacement) { + requireNonNullArg(text, "text"); + requireNonNullArg(lineBreaks, "lineBreaks"); + requireValidCodePoint(paragraphReplacement); + final StringBuilder out = new StringBuilder(text.length()); + final Alignment.Builder alignment = new Alignment.Builder(text.length()); + final int length = text.length(); + int i = 0; + while (i < length) { + final At cp = CodePoints.at(text, i); + if (members.contains(cp.codePoint())) { + int j = cp.nextIndex(i); + while (j < length) { + final At next = CodePoints.at(text, j); + if (!members.contains(next.codePoint())) { + break; + } + j = next.nextIndex(j); + } + final int emitted = countLogicalLineBreaks(text, i, j, lineBreaks) >= 2 + ? paragraphReplacement : replacement; + out.appendCodePoint(emitted); + alignment.replace(j - i, Character.charCount(emitted)); + i = j; + } else { + out.appendCodePoint(cp.codePoint()); + alignment.equal(cp.charCount()); + i = cp.nextIndex(i); + } + } + return new AlignedText(text, out.toString(), alignment.build(length)); + } + /** * Like {@link #trim(CharSequence)} but also produces the {@link Alignment} back to the original * text. The trimmed leading and trailing members appear as deletions, so a span never reports @@ -619,6 +710,41 @@ private int skipRun(CharSequence text, int runStart) { return i; } + /** + * Counts logical line breaks in a whitespace run. A carriage return immediately followed by a + * line feed counts as one break. + * + * @param text The text containing the run. + * @param runStart The index where the whitespace run starts. + * @param runEnd The index of the first code point after the run. + * @param lineBreaks The code points that count as a line break when tallying the run. + * @return The number of logical line breaks in {@code text[runStart..runEnd)}. + */ + private int countLogicalLineBreaks(CharSequence text, int runStart, int runEnd, + CodePointSet lineBreaks) { + int count = 0; + int k = runStart; + while (k < runEnd) { + final At cp = CodePoints.at(text, k); + if (lineBreaks.contains(cp.codePoint())) { + if (cp.codePoint() == 0x000D) { + final int next = cp.nextIndex(k); + if (next < runEnd && CodePoints.at(text, next).codePoint() == 0x000A) { + k = CodePoints.at(text, next).nextIndex(next); + } else { + k = next; + } + } else { + k = cp.nextIndex(k); + } + count++; + } else { + k = cp.nextIndex(k); + } + } + return count; + } + /** * Validates that {@code codePoint} is a Unicode code point. * diff --git a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/OffsetAwareNormalizer.java b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/OffsetAwareNormalizer.java index 156a1d61fc..1b562ad8ab 100644 --- a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/OffsetAwareNormalizer.java +++ b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/OffsetAwareNormalizer.java @@ -22,16 +22,16 @@ * character offsets of the original. * *

    Length-changing folds move offsets: collapsing a run of whitespace, folding a supplementary - * dash to one ASCII hyphen, or stripping invisible controls all shift every later character. A rung + * dash to one ASCII hyphen, or stripping invisible controls all shift every later character. A normalizer * that performs such a fold over the cursor-based {@link CharClass} engine can record those edits - * and expose them through {@link #normalizeAligned(CharSequence)}. A rung that delegates to + * and expose them through {@link #normalizeAligned(CharSequence)}. A normalizer that delegates to * {@link java.text.Normalizer} (NFC/NFKC) or to a stemmer cannot report its edits, so it does not * implement this interface; that is a deliberate capability split rather than an oversight.

    * *

    {@code TextNormalizer.Builder.buildAligned()} composes a chain of these into a single * offset-aware pipeline whose {@link AlignedText} maps a match all the way back to the original * input. An interface-typed caller tests for the capability - * ({@code normalizer instanceof OffsetAwareNormalizer}) instead of depending on a concrete rung, + * ({@code normalizer instanceof OffsetAwareNormalizer}) instead of depending on a concrete normalizer, * the same plain {@code instanceof} pattern used by * {@code OffsetMappingNameFinder} (in the DL layer) rather than reflection.

    */ diff --git a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/UnicodeWhitespace.java b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/UnicodeWhitespace.java index 3712f09066..59728a3083 100644 --- a/opennlp-api/src/main/java/opennlp/tools/util/normalizer/UnicodeWhitespace.java +++ b/opennlp-api/src/main/java/opennlp/tools/util/normalizer/UnicodeWhitespace.java @@ -161,8 +161,11 @@ public String toUnicodeNotation() { private static final int[] CODE_POINTS = new int[WHITESPACE.size()]; private static final List LINE_BREAKS = new ArrayList<>(); private static final List NON_BREAKING = new ArrayList<>(); + private static final CodePointSet LINE_BREAK_CODE_POINT_SET; static { + final int[] lineBreakCodePointBuffer = new int[WHITESPACE.size()]; + int lineBreakCodePointCount = 0; for (int i = 0; i < WHITESPACE.size(); i++) { final WhitespaceCharacter ws = WHITESPACE.get(i); BY_CODE_POINT.put(ws.codePoint(), ws); @@ -170,6 +173,7 @@ public String toUnicodeNotation() { CODE_POINTS[i] = ws.codePoint(); if (ws.isLineBreak()) { LINE_BREAKS.add(ws); + lineBreakCodePointBuffer[lineBreakCodePointCount++] = ws.codePoint(); } if (ws.isNonBreaking()) { NON_BREAKING.add(ws); @@ -178,6 +182,10 @@ public String toUnicodeNotation() { for (final RelatedCharacter related : LOOKALIKES) { LOOKALIKE_MEMBERSHIP.set(related.codePoint()); } + + final int[] lineBreakCodePoints = new int[lineBreakCodePointCount]; + System.arraycopy(lineBreakCodePointBuffer, 0, lineBreakCodePoints, 0, lineBreakCodePointCount); + LINE_BREAK_CODE_POINT_SET = CodePointSet.of(lineBreakCodePoints); } private UnicodeWhitespace() { @@ -230,6 +238,14 @@ public static List lineBreaks() { return List.copyOf(LINE_BREAKS); } + /** + * {@return the {@link CodePointSet} of whitespace characters that force a line or paragraph + * break} + */ + public static CodePointSet lineBreakCodePointSet() { + return LINE_BREAK_CODE_POINT_SET; + } + /** {@return the non-breaking whitespace characters} */ public static List nonBreaking() { return List.copyOf(NON_BREAKING); diff --git a/opennlp-api/src/test/java/opennlp/tools/document/DocumentAnnotatorsTest.java b/opennlp-api/src/test/java/opennlp/tools/document/DocumentAnnotatorsTest.java new file mode 100644 index 0000000000..bef5db118f --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/document/DocumentAnnotatorsTest.java @@ -0,0 +1,121 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.document; + +import java.util.ArrayList; +import java.util.List; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.util.Span; + +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 {@link DocumentAnnotators} support methods directly: the required-layer + * check's exact rejection messages and the per-sentence walk's slicing, skipping, and + * loud rejection of a token outside every sentence. The adapter tests exercise the same + * behavior through the annotators; this class pins the helpers as public API on their + * own. + */ +public class DocumentAnnotatorsTest { + + @Test + void testRequireLayersAcceptsPresentLayers() { + final Document document = Document.of("the") + .with(Layers.SENTENCES, List.of()) + .with(Layers.TOKENS, List.of()); + DocumentAnnotators.requireLayers(document, Layers.SENTENCES, Layers.TOKENS); + DocumentAnnotators.requireLayers(document); + } + + @Test + void testRequireLayersRejectsNullDocument() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> DocumentAnnotators.requireLayers(null, Layers.TOKENS)); + assertEquals("document must not be null", e.getMessage()); + } + + /** + * Verifies that an absent layer is rejected with the shared message naming the first + * absent layer in the order the caller listed them. + */ + @Test + void testRequireLayersNamesTheFirstAbsentLayer() { + final Document document = Document.of("the").with(Layers.TOKENS, List.of()); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> DocumentAnnotators.requireLayers(document, + Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS)); + assertEquals("document lacks the required layer opennlp:sentences", + e.getMessage()); + } + + /** + * Verifies the walk contract: each sentence receives the contiguous run of tokens its + * span encloses with the run's first token layer position, and a sentence without + * tokens is skipped rather than reported as an empty run. + */ + @Test + void testForEachSentenceSlicesContiguousRuns() { + final List> sentences = List.of( + new Annotation<>(new Span(0, 9), "Ana runs."), + new Annotation<>(new Span(10, 11), "!"), + new Annotation<>(new Span(12, 21), "Bob sits.")); + final List> tokens = List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(12, 15), "Bob"), + new Annotation<>(new Span(16, 21), "sits.")); + + final List firsts = new ArrayList<>(); + final List> runs = new ArrayList<>(); + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + firsts.add(first); + runs.add(List.of(words)); + }); + + assertEquals(List.of(0, 2), firsts); + assertEquals(List.of( + List.of("Ana", "runs."), + List.of("Bob", "sits.")), runs); + } + + @Test + void testForEachSentenceRejectsTokenOutsideEverySentence() { + final List> sentences = List.of( + new Annotation<>(new Span(0, 9), "Ana runs.")); + final List> tokens = List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(10, 13), "Bob")); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + })); + assertEquals("token at [10..13) lies outside every sentence", e.getMessage()); + } + + @Test + void testForEachSentenceOverEmptyLayersConsumesNothing() { + final List> runs = new ArrayList<>(); + DocumentAnnotators.forEachSentence(List.of(), List.of(), + (first, words) -> runs.add(List.of(words))); + assertTrue(runs.isEmpty()); + } +} diff --git a/opennlp-api/src/test/java/opennlp/tools/document/DocumentContractTest.java b/opennlp-api/src/test/java/opennlp/tools/document/DocumentContractTest.java new file mode 100644 index 0000000000..27a6a1948f --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/document/DocumentContractTest.java @@ -0,0 +1,644 @@ +/* + * 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.document; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Pins down the observable contract of the document container at its edges: key + * equality, layer ordering, span boundary cases, list immutability, the exact rejection + * messages, and the analyzer's build-time validation. Every expected value in this class + * is asserted exactly so any behavioral drift is caught, not just gross breakage. + */ +public class DocumentContractTest { + + private static final LayerKey WORDS = LayerKey.of("words", String.class); + + /** + * Verifies that two independently created keys with the same id and the same type are + * equal, hash alike, and therefore address the same layer, while remaining distinct + * instances. This is what lets separately compiled producers agree on a layer without + * sharing a constant. + */ + @Test + void testKeysWithSameIdAndTypeAddressTheSameLayer() { + final LayerKey first = LayerKey.of("words", String.class); + final LayerKey second = LayerKey.of("words", String.class); + assertNotSame(first, second); + assertEquals(first, second); + assertEquals(first.hashCode(), second.hashCode()); + assertEquals("words", first.toString()); + + final Document document = Document.of("the") + .with(first, List.of(new Annotation<>(new Span(0, 3), "the"))); + // Reading through the other, equal key yields the very layer added above. + assertEquals(1, document.get(second).size()); + assertEquals("the", document.get(second).get(0).value()); + // A duplicate add through the equal key is rejected like any duplicate. + assertThrows(IllegalArgumentException.class, () -> document.with(second, List.of())); + } + + /** + * Verifies that two keys sharing an id but differing in value type are unequal and + * denote two independent layers that can coexist on one document. + */ + @Test + void testKeysWithSameIdButDifferentTypesAreDifferentLayers() { + final LayerKey asString = LayerKey.of("marks", String.class); + final LayerKey asInteger = LayerKey.of("marks", Integer.class); + assertNotEquals(asString, asInteger); + + final Document document = Document.of("ab") + .with(asString, List.of(new Annotation<>(new Span(0, 1), "a"))) + .with(asInteger, List.of(new Annotation<>(new Span(1, 2), 7))); + assertEquals(Set.of(asString, asInteger), document.layers()); + assertEquals("a", document.get(asString).get(0).value()); + assertEquals(7, document.get(asInteger).get(0).value()); + } + + /** + * Verifies the per-key scope contract: a document-scoped layer carries span-less + * values and round-trips them, a spanned annotation under a document-scoped key is + * rejected with a message naming the layer, a span-less annotation under a + * positional key is rejected likewise, and two keys differing only in scope are + * unequal and never address the same layer. + */ + @Test + void testDocumentScopedLayersCarrySpanlessValues() { + final LayerKey language = LayerKey.document("language", String.class); + assertEquals(LayerKey.Scope.DOCUMENT, language.scope()); + assertEquals(LayerKey.Scope.POSITIONAL, WORDS.scope()); + + final Document document = Document.of("the dog") + .with(language, List.of(Annotation.of("eng"))); + assertEquals(1, document.get(language).size()); + assertEquals("eng", document.get(language).get(0).value()); + assertNull(document.get(language).get(0).span()); + + final IllegalArgumentException spanned = assertThrows(IllegalArgumentException.class, + () -> Document.of("the").with(language, + List.of(new Annotation<>(new Span(0, 3), "eng")))); + assertEquals("document-scoped layer language must not carry spans", + spanned.getMessage()); + + final IllegalArgumentException spanless = assertThrows(IllegalArgumentException.class, + () -> Document.of("the").with(WORDS, List.of(Annotation.of("the")))); + assertEquals("positional layer words requires a span on every annotation", + spanless.getMessage()); + + final LayerKey positionalTwin = LayerKey.of("language", String.class); + assertNotEquals(language, positionalTwin); + final Document both = Document.of("the") + .with(language, List.of(Annotation.of("eng"))) + .with(positionalTwin, List.of(new Annotation<>(new Span(0, 3), "the"))); + assertEquals(2, both.layers().size()); + } + + /** + * Verifies the toolkit-namespace factories: {@link Layers#key(String, Class)} and + * {@link Layers#documentKey(String, Class)} apply the {@code opennlp:} prefix and + * yield keys equal to independently spelled ones, and a name that is null, blank, or + * already carries a namespace is rejected. + */ + @Test + void testToolkitKeysCarryTheNamespacePrefix() { + assertEquals(LayerKey.of("opennlp:things", String.class), + Layers.key("things", String.class)); + assertEquals("opennlp:things", Layers.key("things", String.class).id()); + assertEquals(Layers.TOKENS, Layers.key("tokens", String.class)); + + final LayerKey whole = Layers.documentKey("language", String.class); + assertEquals("opennlp:language", whole.id()); + assertEquals(LayerKey.Scope.DOCUMENT, whole.scope()); + + assertThrows(IllegalArgumentException.class, () -> Layers.key(" ", String.class)); + assertThrows(IllegalArgumentException.class, () -> Layers.key(null, String.class)); + assertThrows(IllegalArgumentException.class, () -> Layers.key("things", null)); + final IllegalArgumentException nested = assertThrows(IllegalArgumentException.class, + () -> Layers.key("opennlp:things", String.class)); + assertEquals("name must not contain ':', the namespace is applied by this factory: " + + "opennlp:things", nested.getMessage()); + assertThrows(IllegalArgumentException.class, + () -> Layers.documentKey("app:x", String.class)); + } + + /** + * Verifies the gold-layer convention: a hand-annotated version of a layer lives + * under the {@code gold:} prefixed key beside the produced layer, both are readable + * independently, and the once-only rule keeps either from replacing the other. + */ + @Test + void testGoldLayerLivesBesideThePredictedLayer() { + final LayerKey gold = LayerKey.of("gold:" + Layers.TOKENS.id(), String.class); + final Document document = Document.of("the dog") + .with(Layers.TOKENS, List.of(new Annotation<>(new Span(0, 3), "the"))) + .with(gold, List.of( + new Annotation<>(new Span(0, 3), "the"), + new Annotation<>(new Span(4, 7), "dog"))); + assertEquals("gold:opennlp:tokens", gold.id()); + assertEquals(1, document.get(Layers.TOKENS).size()); + assertEquals(2, document.get(gold).size()); + assertThrows(IllegalArgumentException.class, + () -> document.with(gold, List.of())); + } + + /** + * Verifies that a layer preserves the insertion order of its annotations: the + * container does not sort by span, so a producer that wants span order must supply + * span order. + */ + @Test + void testLayerPreservesInsertionOrder() { + final Document document = Document.of("the dog") + .with(WORDS, List.of( + new Annotation<>(new Span(4, 7), "dog"), + new Annotation<>(new Span(0, 3), "the"))); + final List> words = document.get(WORDS); + assertEquals(2, words.size()); + assertEquals(new Span(4, 7), words.get(0).span()); + assertEquals("dog", words.get(0).value()); + assertEquals(new Span(0, 3), words.get(1).span()); + assertEquals("the", words.get(1).value()); + } + + /** + * Verifies that several annotations may share one span within a layer, for example + * alternative readings of the same region, and all of them are retained in order. + */ + @Test + void testAnnotationsWithIdenticalSpansAreAllRetained() { + final Document document = Document.of("bank") + .with(WORDS, List.of( + new Annotation<>(new Span(0, 4), "institution"), + new Annotation<>(new Span(0, 4), "riverside"))); + final List> words = document.get(WORDS); + assertEquals(2, words.size()); + assertEquals("institution", words.get(0).value()); + assertEquals("riverside", words.get(1).value()); + assertEquals(words.get(0).span(), words.get(1).span()); + } + + /** + * Verifies that zero-length spans are accepted anywhere within the bounds, including + * at the very end of the text where start and end equal the text length. + */ + @Test + void testZeroLengthSpansAreAccepted() { + final Document document = Document.of("ab") + .with(WORDS, List.of( + new Annotation<>(new Span(1, 1), "between"), + new Annotation<>(new Span(2, 2), "at the end"))); + final List> words = document.get(WORDS); + assertEquals(new Span(1, 1), words.get(0).span()); + assertEquals(new Span(2, 2), words.get(1).span()); + assertEquals(0, words.get(0).span().length()); + } + + /** + * Verifies that spans are indexed in {@code char} units, like {@link Span} itself: a + * supplementary-plane character counts as two, so a span over it covers the whole + * surrogate pair and the char-based text length is what bounds a span. + */ + @Test + void testSpansAreCharIndexedOverSupplementaryCharacters() { + // U+1F600, a supplementary-plane character, is two chars in the text + final String text = "\uD83D\uDE00 ok"; + final Document document = Document.of(text) + .with(WORDS, List.of( + new Annotation<>(new Span(0, 2), "emoji"), + new Annotation<>(new Span(3, 5), "ok"))); + final List> words = document.get(WORDS); + assertEquals("\uD83D\uDE00", + words.get(0).span().getCoveredText(document.text()).toString()); + assertEquals("ok", words.get(1).span().getCoveredText(document.text()).toString()); + assertThrows(IllegalArgumentException.class, () -> Document.of("\uD83D\uDE00") + .with(WORDS, List.of(new Annotation<>(new Span(0, 3), "past the end")))); + } + + /** + * Verifies that a span reaching past the end of the text is rejected on insertion + * with a message naming the span, the text length, and the layer. + */ + @Test + void testSpanBeyondTextLengthIsRejectedWithExactMessage() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> Document.of("the").with(WORDS, + List.of(new Annotation<>(new Span(0, 4), "the?")))); + assertEquals("span [0..4) exceeds the text length 3 in layer words", + e.getMessage()); + } + + /** + * Verifies that adding a layer under a key that is already present is rejected with a + * message naming the offending layer. + */ + @Test + void testDuplicateLayerIsRejectedWithExactMessage() { + final Document document = Document.of("the").with(WORDS, List.of()); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> document.with(WORDS, List.of())); + assertEquals("layer is already present: words", e.getMessage()); + } + + /** + * Verifies that reading an absent layer yields an empty, unmodifiable list rather + * than {@code null}, so callers can iterate without a presence check. + */ + @Test + void testAbsentLayerReadsAsUnmodifiableEmptyList() { + final List> absent = Document.of("the").get(WORDS); + assertTrue(absent.isEmpty()); + assertThrows(UnsupportedOperationException.class, + () -> absent.add(new Annotation<>(new Span(0, 3), "the"))); + } + + /** + * Verifies that the list returned for a present layer is unmodifiable and detached + * from the caller's input list: mutating the input after the add does not change the + * document. + */ + @Test + void testPresentLayerListIsUnmodifiableAndDetachedFromInput() { + final List> input = new ArrayList<>(); + input.add(new Annotation<>(new Span(0, 3), "the")); + final Document document = Document.of("the").with(WORDS, input); + + final List> words = document.get(WORDS); + assertThrows(UnsupportedOperationException.class, () -> words.remove(0)); + assertThrows(UnsupportedOperationException.class, + () -> words.add(new Annotation<>(new Span(0, 3), "the"))); + + input.clear(); + assertEquals(1, document.get(WORDS).size()); + } + + /** + * Verifies that the layer key set exposed by a document cannot be mutated by callers. + */ + @Test + void testLayerKeySetIsUnmodifiable() { + final Document document = Document.of("the").with(WORDS, List.of()); + assertThrows(UnsupportedOperationException.class, () -> document.layers().clear()); + } + + /** + * Verifies that an analyzer whose annotator requires a layer no earlier annotator + * provides fails at build time with a message naming the annotator and the missing + * layer. The annotator under test overrides {@code toString()} so the whole message + * can be asserted exactly. + */ + @Test + void testUnsatisfiedRequirementFailsAtBuildTimeWithExactMessage() { + final DocumentAnnotator needsTags = new DocumentAnnotator() { + + @Override + public Document annotate(Document document) { + throw new IllegalStateException("must never run; the pipeline must not build"); + } + + @Override + public Set> requires() { + return Set.of(Layers.POS_TAGS); + } + + @Override + public Set> provides() { + return Set.of(WORDS); + } + + @Override + public String toString() { + return "tag-consumer"; + } + }; + final DocumentAnalyzer.Builder builder = DocumentAnalyzer.builder().add(needsTags); + final IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, builder::build); + assertEquals("annotator tag-consumer requires layer opennlp:pos," + + " which no earlier annotator provides", e.getMessage()); + } + + /** + * Verifies that an analyzer whose annotators would provide the same layer twice fails + * at build time with a message naming the layer and the positions of both providers, + * instead of crashing midway through the first document. + */ + @Test + void testDuplicateProviderFailsAtBuildTimeWithExactMessage() { + final DocumentAnnotator provider = new DocumentAnnotator() { + + @Override + public Document annotate(Document document) { + throw new IllegalStateException("must never run; the pipeline must not build"); + } + + @Override + public Set> provides() { + return Set.of(WORDS); + } + }; + final DocumentAnalyzer.Builder builder = DocumentAnalyzer.builder() + .add(provider).add(provider); + final IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, builder::build); + assertEquals("annotators at positions 0 and 1 both provide layer words", + e.getMessage()); + } + + /** + * Verifies that building an analyzer without any annotator fails with a message + * stating that a pipeline needs at least one annotator. + */ + @Test + void testEmptyPipelineFailsWithExactMessage() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> DocumentAnalyzer.builder().build()); + assertEquals("a pipeline needs at least one annotator", e.getMessage()); + } + + /** + * Verifies that a caller can add a new layer type without changing the document + * container. + */ + @Test + void testCustomLayerNeedsNoContainerChange() { + record Sentiment(String polarity, double score) { + } + final LayerKey sentiment = LayerKey.of("sentiment", Sentiment.class); + final DocumentAnnotator annotator = new DocumentAnnotator() { + + @Override + public Document annotate(Document document) { + final Span all = new Span(0, document.text().length()); + return document.with(sentiment, + List.of(new Annotation<>(all, new Sentiment("positive", 0.9d)))); + } + + @Override + public Set> provides() { + return Set.of(sentiment); + } + }; + final Document document = DocumentAnalyzer.builder().add(annotator).build() + .analyze("good dog"); + assertEquals("positive", document.get(sentiment).get(0).value().polarity()); + } + + /** + * Verifies that the value type travels through {@link LayerKey}: a layer added under + * an {@code Integer} key reads back as {@code Annotation}, so its values + * participate in arithmetic without a cast, and a mismatched value can never enter + * the layer in the first place. + */ + @Test + void testValueTypeTravelsThroughTheKey() { + final LayerKey counts = LayerKey.of("counts", Integer.class); + final Document document = Document.of("ab cd") + .with(counts, List.of( + new Annotation<>(new Span(0, 2), 2), + new Annotation<>(new Span(3, 5), 40))); + int sum = 0; + for (final Annotation count : document.get(counts)) { + sum += count.value(); + } + assertEquals(42, sum); + + // The insertion-time check backs the typed read: a raw-typed caller cannot place a + // String under the Integer key. + @SuppressWarnings({"unchecked", "rawtypes"}) + final LayerKey raw = (LayerKey) LayerKey.of("counts2", Integer.class); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> Document.of("ab").with(raw, + List.of(new Annotation<>(new Span(0, 2), "not a number")))); + assertEquals("value of type java.lang.String does not match layer counts2", + e.getMessage()); + } + + /** + * Verifies that merge joins two documents grown independently over the same text and + * leaves both sources untouched. Text content decides equality, not the + * {@link CharSequence} implementation. + */ + @Test + void testMergeJoinsLayersOfDocumentsOverTheSameText() { + final LayerKey lengths = LayerKey.of("lengths", Integer.class); + final Document words = Document.of("the dog") + .with(WORDS, List.of( + new Annotation<>(new Span(0, 3), "the"), + new Annotation<>(new Span(4, 7), "dog"))); + final Document counted = Document.of(new StringBuilder("the dog")) + .with(lengths, List.of( + new Annotation<>(new Span(0, 3), 3), + new Annotation<>(new Span(4, 7), 3))); + + final Document merged = words.merge(counted); + + assertEquals("the dog", merged.text().toString()); + assertEquals(Set.of(WORDS, lengths), merged.layers()); + assertEquals("the", merged.get(WORDS).get(0).value()); + assertEquals(3, merged.get(lengths).get(0).value().intValue()); + assertEquals(Set.of(WORDS), words.layers()); + assertEquals(Set.of(lengths), counted.layers()); + } + + /** + * Verifies that merge rejects a null argument, a document over a different text even + * when their layers are disjoint, and a layer key present on both documents, naming + * the offending key. + */ + @Test + void testMergeRejectsNullDifferentTextAndDuplicateLayers() { + final Document words = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 3), "the"))); + + final IllegalArgumentException nullOther = assertThrows(IllegalArgumentException.class, + () -> words.merge(null)); + assertEquals("other must not be null", nullOther.getMessage()); + + final IllegalArgumentException differentText = assertThrows( + IllegalArgumentException.class, () -> words.merge(Document.of("the cat"))); + assertEquals("merge requires both documents to carry the same text", + differentText.getMessage()); + + final IllegalArgumentException duplicate = assertThrows(IllegalArgumentException.class, + () -> words.merge(words)); + assertEquals("layer is already present: words", duplicate.getMessage()); + } + + /** + * Verifies that {@link Document.DuplicateLayerPolicy#KEEP_EQUAL} keeps one copy of a + * layer both documents rebuilt identically while still joining the disjoint layers. + */ + @Test + void testMergeKeepingEqualLayersToleratesIdenticalCopies() { + final List> tokens = List.of( + new Annotation<>(new Span(0, 3), "the"), + new Annotation<>(new Span(4, 7), "dog")); + final LayerKey lengths = LayerKey.of("lengths", Integer.class); + final Document words = Document.of("the dog").with(WORDS, tokens); + final Document recounted = Document.of("the dog") + .with(WORDS, tokens) + .with(lengths, List.of( + new Annotation<>(new Span(0, 3), 3), + new Annotation<>(new Span(4, 7), 3))); + + final Document merged = words.merge(recounted, Document.DuplicateLayerPolicy.KEEP_EQUAL); + + assertEquals(Set.of(WORDS, lengths), merged.layers()); + // The shared layer is kept once, not concatenated. + assertEquals(2, merged.get(WORDS).size()); + assertEquals(2, merged.get(lengths).size()); + } + + /** + * Verifies that {@link Document.DuplicateLayerPolicy#KEEP_EQUAL} still rejects a layer + * whose two copies differ, naming the key, and rejects a null policy. + */ + @Test + void testMergeKeepingEqualLayersRejectsDifferingCopiesAndNullPolicy() { + final Document words = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 3), "the"))); + final Document retokenized = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 7), "the dog"))); + + final IllegalArgumentException differing = assertThrows(IllegalArgumentException.class, + () -> words.merge(retokenized, Document.DuplicateLayerPolicy.KEEP_EQUAL)); + assertEquals("layer is present on both documents with differing contents: words", + differing.getMessage()); + + final IllegalArgumentException nullPolicy = assertThrows(IllegalArgumentException.class, + () -> words.merge(Document.of("the dog"), null)); + assertEquals("duplicateLayers must not be null", nullPolicy.getMessage()); + } + + /** + * Verifies that a document implementation that does not override merge gets the same + * semantics from the interface default: disjoint layers join, a layer both documents + * rebuilt identically is kept once under KEEP_EQUAL, and differing copies are rejected + * with the same message the default implementation produces. + */ + @Test + void testMergeDefaultImplementationServesForeignDocuments() { + final List> tokens = List.of(new Annotation<>(new Span(0, 3), "the")); + final LayerKey lengths = LayerKey.of("lengths", Integer.class); + final Document words = new DelegatingDocument(Document.of("the dog").with(WORDS, tokens)); + final Document counted = Document.of("the dog") + .with(WORDS, tokens) + .with(lengths, List.of(new Annotation<>(new Span(0, 3), 3))); + + final Document merged = words.merge(counted, Document.DuplicateLayerPolicy.KEEP_EQUAL); + assertEquals(Set.of(WORDS, lengths), merged.layers()); + assertEquals(1, merged.get(WORDS).size()); + + final Document retokenized = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 7), "the dog"))); + final IllegalArgumentException differing = assertThrows(IllegalArgumentException.class, + () -> words.merge(retokenized, Document.DuplicateLayerPolicy.KEEP_EQUAL)); + assertEquals("layer is present on both documents with differing contents: words", + differing.getMessage()); + } + + /** + * Verifies that merge validates the layers it takes from the other document instead of + * trusting them: a foreign implementation can hand out annotations that were never + * checked, and an out-of-bounds span among them is rejected by name. + */ + @Test + void testMergeRevalidatesLayersOfForeignDocuments() { + final Document words = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 3), "the"))); + final LayerKey stale = LayerKey.of("stale", String.class); + final Document unvalidated = new UnvalidatedDocument("the dog", stale, + List.of(new Annotation<>(new Span(0, 99), "out of bounds"))); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> words.merge(unvalidated)); + assertEquals("span [0..99) exceeds the text length 7 in layer stale", + e.getMessage()); + } + + /** + * A pass-through wrapper that overrides none of the interface defaults, so calling + * merge on it runs the interface's default implementation. + */ + private record DelegatingDocument(Document delegate) implements Document { + + @Override + public CharSequence text() { + return delegate.text(); + } + + @Override + public List> get(LayerKey layer) { + return delegate.get(layer); + } + + @Override + public Set> layers() { + return delegate.layers(); + } + + @Override + public Document with(LayerKey layer, List> annotations) { + return new DelegatingDocument(delegate.with(layer, annotations)); + } + } + + /** + * A document whose single layer bypassed all validation, standing in for a foreign + * implementation that does not enforce the layer contract itself. + */ + private record UnvalidatedDocument(String rawText, LayerKey key, + List> annotations) + implements Document { + + @Override + public CharSequence text() { + return rawText; + } + + @Override + @SuppressWarnings("unchecked") + public List> get(LayerKey layer) { + return key.equals(layer) ? (List>) (List) annotations : List.of(); + } + + @Override + public Set> layers() { + return Set.of(key); + } + + @Override + public Document with(LayerKey layer, List> annotations) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/opennlp-api/src/test/java/opennlp/tools/document/DocumentTest.java b/opennlp-api/src/test/java/opennlp/tools/document/DocumentTest.java new file mode 100644 index 0000000000..787c9d8edb --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/document/DocumentTest.java @@ -0,0 +1,148 @@ +/* + * 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.document; + +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests the {@link Document} container: typed layer access, copy-on-add immutability, + * and the insertion-time validation that protects the layer invariants. + */ +public class DocumentTest { + + private static final LayerKey WORDS = LayerKey.of("words", String.class); + private static final LayerKey NUMBERS = LayerKey.of("numbers", Integer.class); + + @Test + void testEmptyDocument() { + final Document document = Document.of("the dog"); + assertEquals("the dog", document.text()); + assertTrue(document.layers().isEmpty()); + assertTrue(document.get(WORDS).isEmpty()); + } + + /** + * Verifies that the text is captured at construction: mutating a + * {@link StringBuilder} handed to {@link Document#of(CharSequence)} does not reach + * the document, so the span bounds validated on insertion stay valid for its + * lifetime. + */ + @Test + void testTextIsCapturedAtConstruction() { + final StringBuilder mutable = new StringBuilder("the dog"); + final Document document = Document.of(mutable) + .with(WORDS, List.of(new Annotation<>(new Span(4, 7), "dog"))); + mutable.setLength(0); + assertEquals("the dog", document.text().toString()); + assertEquals("dog", + document.get(WORDS).get(0).span().getCoveredText(document.text()).toString()); + } + + @Test + void testWithAddsATypedLayer() { + final Document document = Document.of("the dog") + .with(WORDS, List.of(new Annotation<>(new Span(0, 3), "the"), + new Annotation<>(new Span(4, 7), "dog"))); + assertEquals(Set.of(WORDS), document.layers()); + final List> words = document.get(WORDS); + assertEquals(2, words.size()); + assertEquals("dog", words.get(1).value()); + assertEquals(new Span(4, 7), words.get(1).span()); + } + + @Test + void testWithIsCopyOnAdd() { + final Document empty = Document.of("42"); + final Document grown = empty.with(NUMBERS, + List.of(new Annotation<>(new Span(0, 2), 42))); + assertTrue(empty.layers().isEmpty()); + assertEquals(Set.of(NUMBERS), grown.layers()); + // unchanged layers are shared, not copied + final Document both = grown.with(WORDS, List.of()); + assertSame(grown.get(NUMBERS), both.get(NUMBERS)); + } + + @Test + void testEqualKeysFromDifferentConstantsInteroperate() { + final Document document = Document.of("the") + .with(LayerKey.of("words", String.class), + List.of(new Annotation<>(new Span(0, 3), "the"))); + assertEquals(1, document.get(WORDS).size()); + assertNotEquals(WORDS, LayerKey.of("words", CharSequence.class)); + } + + @Test + void testDuplicateLayerThrows() { + final Document document = Document.of("the").with(WORDS, List.of()); + assertThrows(IllegalArgumentException.class, () -> document.with(WORDS, List.of())); + } + + @Test + void testSpanBeyondTextThrows() { + assertThrows(IllegalArgumentException.class, () -> Document.of("the") + .with(WORDS, List.of(new Annotation<>(new Span(0, 4), "the?")))); + } + + @Test + void testValueTypeIsCheckedOnInsertion() { + // a raw-typed caller cannot smuggle a mismatched value past the layer type + @SuppressWarnings({"unchecked", "rawtypes"}) + final LayerKey raw = (LayerKey) NUMBERS; + assertThrows(IllegalArgumentException.class, () -> Document.of("the") + .with(raw, List.of(new Annotation<>(new Span(0, 3), "not a number")))); + } + + @Test + void testNullArgumentsThrow() { + final Document document = Document.of("the"); + assertThrows(IllegalArgumentException.class, () -> Document.of(null)); + assertThrows(IllegalArgumentException.class, () -> document.get(null)); + assertThrows(IllegalArgumentException.class, () -> document.with(null, List.of())); + assertThrows(IllegalArgumentException.class, () -> document.with(WORDS, null)); + } + + @Test + void testAnnotationValidation() { + // a span-less annotation is legal to build; the container judges it per key scope + assertThrows(IllegalArgumentException.class, () -> new Annotation<>(new Span(0, 3), null)); + assertThrows(IllegalArgumentException.class, () -> Annotation.of(null)); + assertThrows(IllegalArgumentException.class, () -> Document.of("the") + .with(WORDS, List.of(Annotation.of("the")))); + } + + @Test + void testLayerKeyValidation() { + assertThrows(IllegalArgumentException.class, () -> LayerKey.of(" ", String.class)); + assertThrows(IllegalArgumentException.class, () -> LayerKey.of(null, String.class)); + assertThrows(IllegalArgumentException.class, () -> LayerKey.of("words", null)); + assertThrows(IllegalArgumentException.class, () -> LayerKey.document(" ", String.class)); + assertThrows(IllegalArgumentException.class, () -> LayerKey.document(null, String.class)); + assertThrows(IllegalArgumentException.class, () -> LayerKey.document("lang", null)); + } +} 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-api/src/test/java/opennlp/tools/termvector/TermVectorTest.java b/opennlp-api/src/test/java/opennlp/tools/termvector/TermVectorTest.java new file mode 100644 index 0000000000..8448f465d3 --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/termvector/TermVectorTest.java @@ -0,0 +1,105 @@ +/* + * 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.termvector; + +import java.util.ArrayList; +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.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies the {@link TermVector} invariants: the two legal shapes (full with one span + * per occurrence, scoring-only with none), the validation of everything in between, and + * the immutability of the span list. + */ +public class TermVectorTest { + + @Test + void testWithSpansDerivesTheFrequency() { + final TermVector vector = + TermVector.withSpans("dog", List.of(new Span(4, 7), new Span(19, 22))); + assertEquals("dog", vector.term()); + assertEquals(2, vector.frequency()); + assertEquals(List.of(new Span(4, 7), new Span(19, 22)), vector.spans()); + } + + @Test + void testCountCarriesNoSpans() { + final TermVector vector = TermVector.count("dog", 3); + assertEquals("dog", vector.term()); + assertEquals(3, vector.frequency()); + assertTrue(vector.spans().isEmpty()); + } + + @Test + void testSpanListIsDetachedFromTheCallersInput() { + final List spans = new ArrayList<>(List.of(new Span(0, 3))); + final TermVector vector = TermVector.withSpans("the", spans); + spans.add(new Span(4, 7)); + assertEquals(1, vector.spans().size()); + assertThrows(UnsupportedOperationException.class, + () -> vector.spans().add(new Span(8, 11))); + } + + @Test + void testNullTermIsRejected() { + assertThrows(IllegalArgumentException.class, + () -> new TermVector(null, 1, List.of(new Span(0, 1)))); + } + + @Test + void testNullSpanListIsRejected() { + assertThrows(IllegalArgumentException.class, () -> new TermVector("dog", 1, null)); + } + + @Test + void testNullSpanElementIsRejected() { + final List spans = new ArrayList<>(); + spans.add(new Span(0, 3)); + spans.add(null); + assertThrows(IllegalArgumentException.class, () -> new TermVector("dog", 2, spans)); + assertThrows(IllegalArgumentException.class, () -> TermVector.withSpans("dog", spans)); + } + + @ParameterizedTest + @ValueSource(ints = {0, -1, Integer.MIN_VALUE}) + void testFrequencyBelowOneIsRejected(int frequency) { + assertThrows(IllegalArgumentException.class, () -> TermVector.count("dog", frequency)); + } + + @Test + void testEmptySpanListCannotDeriveAFrequency() { + assertThrows(IllegalArgumentException.class, + () -> TermVector.withSpans("dog", List.of())); + } + + @Test + void testPartialSpanListIsRejected() { + // Two occurrences but only one recorded span: neither full nor scoring-only. + assertThrows(IllegalArgumentException.class, + () -> new TermVector("dog", 2, List.of(new Span(0, 3)))); + } +} diff --git a/opennlp-api/src/test/java/opennlp/tools/util/ResourceLimitsTest.java b/opennlp-api/src/test/java/opennlp/tools/util/ResourceLimitsTest.java new file mode 100644 index 0000000000..84f1e20a61 --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/util/ResourceLimitsTest.java @@ -0,0 +1,64 @@ +/* + * 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.util; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +/** Tests system-property parsing for resource limits. */ +public class ResourceLimitsTest { + + private static final String PROPERTY = "opennlp.test.resource.limit"; + + /** Clears the test property after each test. */ + @AfterEach + void clearProperty() { + System.clearProperty(PROPERTY); + } + + /** Verifies positive integer and long overrides. */ + @Test + void testPositiveOverrides() { + System.setProperty(PROPERTY, "1024"); + Assertions.assertEquals(1024, ResourceLimits.initLimit(PROPERTY, 7)); + Assertions.assertEquals(1024L, ResourceLimits.initLimit(PROPERTY, 7L)); + } + + /** Verifies that an absent property uses the supplied defaults. */ + @Test + void testAbsentPropertyUsesDefaults() { + Assertions.assertEquals(7, ResourceLimits.initLimit(PROPERTY, 7)); + Assertions.assertEquals(7L, ResourceLimits.initLimit(PROPERTY, 7L)); + } + + /** + * Verifies that an invalid property uses the supplied defaults. + * + * @param value The invalid property value. + */ + @ParameterizedTest(name = "value {0} uses the default") + @ValueSource(strings = {"", " ", "abc", "-1", "0"}) + void testInvalidPropertyUsesDefaults(String value) { + System.setProperty(PROPERTY, value); + Assertions.assertEquals(7, ResourceLimits.initLimit(PROPERTY, 7)); + Assertions.assertEquals(7L, ResourceLimits.initLimit(PROPERTY, 7L)); + } +} diff --git a/opennlp-api/src/test/java/opennlp/tools/util/StringUtilLocaleTest.java b/opennlp-api/src/test/java/opennlp/tools/util/StringUtilLocaleTest.java new file mode 100644 index 0000000000..a58db568c3 --- /dev/null +++ b/opennlp-api/src/test/java/opennlp/tools/util/StringUtilLocaleTest.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.tools.util; + +import java.util.Locale; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Tests that the {@link StringUtil} operations feeding the lemmatizer's shortest edit script + * are independent of the JVM's default {@link Locale}. + */ +public class StringUtilLocaleTest { + + /** + * Turkish folds {@code 'I'} to the dotless {@code 'ı'} (U+0131) instead of {@code 'i'}, + * which makes it the canonical probe for an unqualified {@code String#toLowerCase()}. + */ + private static final Locale TURKISH = Locale.of("tr", "TR"); + + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + + /** + * The shortest edit script becomes an outcome label inside a trained lemmatizer model. + * If it varied with the default locale, a model trained on a Turkish JVM would carry + * labels no other JVM could reproduce. + */ + @Test + void testGetShortestEditScriptIsIndependentOfDefaultLocale() { + Locale.setDefault(TURKISH); + + Assertions.assertEquals("D0s", StringUtil.getShortestEditScript("IMPORTS", "import")); + Assertions.assertEquals("R2ioR1cuI1s", StringUtil.getShortestEditScript("MICE", "mouse")); + Assertions.assertEquals("D3iD2cR0sx", StringUtil.getShortestEditScript("INDICES", "index")); + Assertions.assertEquals("O", StringUtil.getShortestEditScript("Illinois", "illinois")); + } + + /** + * Guards the choice of {@link Locale#ROOT} over the code point based + * {@link StringUtil#toLowerCase(CharSequence)}: the two disagree on context sensitive + * mappings such as the Greek final sigma. Existing models were trained with + * {@link Locale#ROOT} semantics, so only {@link Locale#ROOT} keeps them readable. + */ + @Test + void testGetShortestEditScriptUsesRootFoldingNotCodePointFolding() { + // Locale.ROOT maps the trailing capital sigma to the final form 'ς', so both inputs + // fold to the same string and no edit is required. + Assertions.assertEquals("O", StringUtil.getShortestEditScript("ΟΔΟΣ", + "οδος")); + + // Character#toLowerCase is context free and would yield the non final 'σ' instead, + // which is a different string and hence a different, model breaking outcome label. + Assertions.assertNotEquals("οδος", + StringUtil.toLowerCase("ΟΔΟΣ")); + } +} diff --git a/opennlp-api/src/test/java/opennlp/tools/util/normalizer/CharClassTest.java b/opennlp-api/src/test/java/opennlp/tools/util/normalizer/CharClassTest.java index 29a92f1c6a..9fb6e8137a 100644 --- a/opennlp-api/src/test/java/opennlp/tools/util/normalizer/CharClassTest.java +++ b/opennlp-api/src/test/java/opennlp/tools/util/normalizer/CharClassTest.java @@ -22,7 +22,6 @@ import org.junit.jupiter.api.Test; import opennlp.tools.util.Span; -import opennlp.tools.util.normalizer.UnicodeWhitespace.WhitespaceCharacter; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -52,8 +51,7 @@ private static String cp(int codePoint) { } private static CodePointSet lineBreaks() { - return CodePointSet.of(UnicodeWhitespace.lineBreaks().stream() - .mapToInt(WhitespaceCharacter::codePoint).toArray()); + return UnicodeWhitespace.lineBreakCodePointSet(); } // --- membership -------------------------------------------------------------------------- @@ -141,6 +139,33 @@ void testSquishNewlineAndTabEdgeCases() { assertEquals(" ", WS.collapsePreserving("\t", keep, '\n')); } + // --- paragraph unwrap (collapseParagraphPreserving) -------------------------------------- + + @Test + void testCollapseParagraphPreservingUnwrapsSingleBreaks() { + final CodePointSet breaks = lineBreaks(); + assertEquals("a b", WS.collapseParagraphPreserving("a \n b", breaks, '\n')); + assertEquals("a b", WS.collapseParagraphPreserving("a\r\nb", breaks, '\n')); + assertEquals("on the bank", WS.collapseParagraphPreserving("on the\nbank", breaks, '\n')); + } + + @Test + void testCollapseParagraphPreservingKeepsBlankLines() { + final CodePointSet breaks = lineBreaks(); + assertEquals("a\nb", WS.collapseParagraphPreserving("a\n\nb", breaks, '\n')); + assertEquals("a\nb", WS.collapseParagraphPreserving("a\n\n\n\nb", breaks, '\n')); + assertEquals("a\nb", WS.collapseParagraphPreserving("a\r\n\r\nb", breaks, '\n')); + } + + @Test + void testCollapseParagraphPreservingEdgeCases() { + final CodePointSet breaks = lineBreaks(); + assertEquals("a b", WS.collapseParagraphPreserving("a \t b", breaks, '\n')); + assertEquals("\nabc", WS.collapseParagraphPreserving("\n\nabc", breaks, '\n')); + assertEquals(" ", WS.collapseParagraphPreserving("\n", breaks, '\n')); + assertEquals(" ", WS.collapseParagraphPreserving("\t", breaks, '\n')); + } + // --- trim / removeAll -------------------------------------------------------------------- @Test @@ -447,6 +472,13 @@ void testCollapsePreservingAlignedRunWithoutKeepCollapsesToReplacement() { assertSpan(1, 4, at.toOriginalSpan(1, 2)); } + @Test + void testCollapseParagraphPreservingAligned() { + final AlignedText at = WS.collapseParagraphPreservingAligned("a\n\n\t\tb", lineBreaks(), '\n'); + assertEquals("a\nb", at.normalized()); + assertSpan(1, 5, at.toOriginalSpan(1, 2)); // the paragraph break covers the whole run + } + // Every aligned operation must produce exactly the same string as its plain counterpart; only the // alignment is extra. This pins that contract across a battery of inputs so the two code paths // cannot drift apart. @@ -480,6 +512,9 @@ void testAlignedOperationsAgreeWithPlainOutput() { assertEquals(charClass.collapsePreserving(input, keep, '\n'), charClass.collapsePreservingAligned(input, keep, '\n').normalized(), "collapsePreserving vs collapsePreservingAligned for [" + input + "]"); + assertEquals(charClass.collapseParagraphPreserving(input, keep, '\n'), + charClass.collapseParagraphPreservingAligned(input, keep, '\n').normalized(), + "collapseParagraphPreserving vs collapseParagraphPreservingAligned for [" + input + "]"); } } } @@ -495,6 +530,8 @@ void nullParametersAreRejectedWithIllegalArgumentException() { assertThrows(IllegalArgumentException.class, () -> ws.collapse(null)); assertThrows(IllegalArgumentException.class, () -> ws.collapsePreserving(null, nl, '\n')); assertThrows(IllegalArgumentException.class, () -> ws.collapsePreserving("x", null, '\n')); + assertThrows(IllegalArgumentException.class, () -> ws.collapseParagraphPreserving(null, nl, '\n')); + assertThrows(IllegalArgumentException.class, () -> ws.collapseParagraphPreserving("x", null, '\n')); assertThrows(IllegalArgumentException.class, () -> ws.trim(null)); assertThrows(IllegalArgumentException.class, () -> ws.removeAll(null)); assertThrows(IllegalArgumentException.class, () -> ws.normalizeAligned(null)); @@ -503,6 +540,10 @@ void nullParametersAreRejectedWithIllegalArgumentException() { () -> ws.collapsePreservingAligned(null, nl, '\n')); assertThrows(IllegalArgumentException.class, () -> ws.collapsePreservingAligned("x", null, '\n')); + assertThrows(IllegalArgumentException.class, + () -> ws.collapseParagraphPreservingAligned(null, nl, '\n')); + assertThrows(IllegalArgumentException.class, + () -> ws.collapseParagraphPreservingAligned("x", null, '\n')); assertThrows(IllegalArgumentException.class, () -> ws.trimAligned(null)); assertThrows(IllegalArgumentException.class, () -> ws.removeAllAligned(null)); assertThrows(IllegalArgumentException.class, () -> CharClass.substitute(null, cp -> null)); diff --git a/opennlp-api/src/test/java/opennlp/tools/util/normalizer/UnicodeWhitespaceTest.java b/opennlp-api/src/test/java/opennlp/tools/util/normalizer/UnicodeWhitespaceTest.java index bd040efc05..1859025ac9 100644 --- a/opennlp-api/src/test/java/opennlp/tools/util/normalizer/UnicodeWhitespaceTest.java +++ b/opennlp-api/src/test/java/opennlp/tools/util/normalizer/UnicodeWhitespaceTest.java @@ -148,6 +148,15 @@ void testLineBreaksAreExactlyTheSeven() { .map(WhitespaceCharacter::codePoint).collect(Collectors.toSet())); } + @Test + void testLineBreakCodePointSetMatchesLineBreaks() { + final CodePointSet lineBreaks = UnicodeWhitespace.lineBreakCodePointSet(); + assertEquals(UnicodeWhitespace.lineBreaks().size(), lineBreaks.size()); + for (final WhitespaceCharacter ws : UnicodeWhitespace.lineBreaks()) { + assertTrue(lineBreaks.contains(ws.codePoint())); + } + } + @Test void testNonBreakingAreExactlyTheThree() { final Set expected = Set.of(0x00A0, 0x2007, 0x202F); diff --git a/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/DetailedFMeasureListener.java b/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/DetailedFMeasureListener.java index 5f67501751..7c6e18631c 100644 --- a/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/DetailedFMeasureListener.java +++ b/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/DetailedFMeasureListener.java @@ -116,7 +116,7 @@ private Stats initStatsForOutcomeAndGet(String type) { + " [target: %3d; tp: %3d; fp: %3d]"; public String createReport() { - return createReport(Locale.getDefault()); + return createReport(Locale.ROOT); } public String createReport(Locale locale) { diff --git a/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/FineGrainedReportListener.java b/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/FineGrainedReportListener.java index b0ae3ba936..25a6481a6b 100644 --- a/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/FineGrainedReportListener.java +++ b/opennlp-core/opennlp-cli/src/main/java/opennlp/tools/cmdline/FineGrainedReportListener.java @@ -29,6 +29,7 @@ import java.util.Iterator; import java.util.LinkedList; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Map.Entry; import java.util.Set; @@ -163,6 +164,19 @@ private SortedSet getConfusionMatrixTagset(String token) { return stats.getConfusionMatrixTagset(token); } + /** + * Applies a {@link MessageFormat} number pattern using {@link Locale#ROOT}, so that the + * figures in a report stay identical regardless of the JVM's default {@link Locale}. + * + * @param pattern The {@link MessageFormat} pattern to apply. + * @param value The value to format. + * + * @return The formatted value. + */ + private static String formatNumber(String pattern, Object value) { + return new MessageFormat(pattern, Locale.ROOT).format(new Object[] {value}); + } + private double[][] getConfusionMatrix() { return stats.getConfusionMatrix(); } @@ -187,7 +201,7 @@ private String matrixToString(SortedSet tagset, double[][] data, boolean minColumnSize = matrix[i][j].length(); } } - matrix[i][j] = MessageFormat.format("{0,number,#.##%}", data[i][j]); + matrix[i][j] = formatNumber("{0,number,#.##%}", data[i][j]); if (data[i][j] == 1 && filter) { initialIndex = i + 1; } @@ -229,12 +243,12 @@ protected void printGeneralStatistics() { String.format("%21s: %6s", "Max sentence size", getMaxSentenceSize())).append("\n"); printStream.append( String.format("%21s: %6s", "Average sentence size", - MessageFormat.format("{0,number,#.##}", getAverageSentenceSize()))).append("\n"); + formatNumber("{0,number,#.##}", getAverageSentenceSize()))).append("\n"); printStream.append( String.format("%21s: %6s", "Tags count", getNumberOfTags())).append("\n"); printStream.append( String.format("%21s: %6s", "Accuracy", - MessageFormat.format("{0,number,#.##%}", getAccuracy()))).append("\n"); + formatNumber("{0,number,#.##%}", getAccuracy()))).append("\n"); printFooter("Evaluation Corpus Statistics"); } @@ -308,7 +322,7 @@ protected void printTokenErrorRank() { String tok = tokIterator.next(); int ocurrencies = getTokenFrequency(tok); int errors = getTokenErrors(tok); - String rate = MessageFormat.format("{0,number,#.##%}", (double) errors + String rate = formatNumber("{0,number,#.##%}", (double) errors / ocurrencies); printStream.append(String.format(format, tok, errors, ocurrencies, rate) @@ -347,7 +361,7 @@ protected void printTagsErrorRank() { for (String tag : tags) { int ocurrencies = getTagFrequency(tag); int errors = getTagErrors(tag); - String rate = MessageFormat.format("{0,number,#.###}", (double) errors + String rate = formatNumber("{0,number,#.###}", (double) errors / ocurrencies); double p = getTagPrecision(tag); @@ -355,9 +369,9 @@ protected void printTagsErrorRank() { double f = getTagFMeasure(tag); printStream.append(String.format(format, tag, errors, ocurrencies, rate, - MessageFormat.format("{0,number,#.###}", p > 0 ? p : 0), - MessageFormat.format("{0,number,#.###}", r > 0 ? r : 0), - MessageFormat.format("{0,number,#.###}", f > 0 ? f : 0)) + formatNumber("{0,number,#.###}", p > 0 ? p : 0), + formatNumber("{0,number,#.###}", r > 0 ? r : 0), + formatNumber("{0,number,#.###}", f > 0 ? f : 0)) ); } @@ -405,7 +419,7 @@ protected void printDetailedConfusionMatrix() { .append("]\n") .append( String.format("%12s: %-8s", "Accuracy", - MessageFormat.format("{0,number,#.##%}", acc))) + formatNumber("{0,number,#.##%}", acc))) .append("\n"); printStream.append( String.format("%12s: %-8s", "Ocurrencies", diff --git a/opennlp-core/opennlp-cli/src/test/java/opennlp/tools/cmdline/FineGrainedReportListenerTest.java b/opennlp-core/opennlp-cli/src/test/java/opennlp/tools/cmdline/FineGrainedReportListenerTest.java new file mode 100644 index 0000000000..36f6dca6d7 --- /dev/null +++ b/opennlp-core/opennlp-cli/src/test/java/opennlp/tools/cmdline/FineGrainedReportListenerTest.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.cmdline; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.Locale; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import opennlp.tools.cmdline.postag.POSTaggerFineGrainedReportListener; +import opennlp.tools.postag.POSSample; + +/** + * Tests that the figures written by a {@link FineGrainedReportListener} do not change with + * the JVM's default {@link Locale}. + */ +public class FineGrainedReportListenerTest { + + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + + @Test + void testReportIsIndependentOfDefaultLocale() throws Exception { + Locale.setDefault(Locale.GERMANY); + + final String report = createReport(); + + // Two out of three tags are correct. A locale with a comma decimal separator would + // render this as "66,67%" and change every figure in the report along with it. + Assertions.assertTrue(report.contains("66.67%"), + () -> "Report should contain the Locale.ROOT rendering of the accuracy:\n" + report); + Assertions.assertFalse(report.contains("66,67%"), + () -> "Report should not render figures with the default locale:\n" + report); + } + + private static String createReport() throws Exception { + final String[] sentence = {"He", "runs", "fast"}; + final POSSample reference = new POSSample(sentence, new String[] {"PRP", "VBZ", "RB"}); + final POSSample prediction = new POSSample(sentence, new String[] {"PRP", "VBZ", "JJ"}); + + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + final POSTaggerFineGrainedReportListener listener = + new POSTaggerFineGrainedReportListener(out); + listener.misclassified(reference, prediction); + listener.writeReport(); + return out.toString(StandardCharsets.UTF_8); + } + } +} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java index 8f59c6f844..ecbaaf5b07 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/AbstractDL.java @@ -23,6 +23,7 @@ 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; @@ -36,8 +37,9 @@ import ai.onnxruntime.OrtException; import ai.onnxruntime.OrtSession; -import opennlp.tools.tokenize.BertTokenizer; -import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.tokenize.SubwordPiece; +import opennlp.tools.tokenize.SubwordTokenizer; +import opennlp.tools.tokenize.WordpieceEncoder; import opennlp.tools.tokenize.WordpieceTokenizer; import opennlp.tools.util.Span; import opennlp.tools.util.normalizer.AlignedText; @@ -55,7 +57,7 @@ public abstract class AbstractDL implements AutoCloseable { protected final OrtEnvironment env; protected final OrtSession session; - protected final Tokenizer tokenizer; + protected final SubwordTokenizer tokenizer; protected final Map vocab; private final AtomicBoolean closed = new AtomicBoolean(); @@ -105,7 +107,7 @@ protected AbstractDL(final File model, final File vocabulary, final OrtSession createdSession = env.createSession(model.getPath(), sessionOptions); try { this.vocab = Map.copyOf(loadVocabFile(vocabulary)); - this.tokenizer = createPipelineTokenizer(vocab, lowerCase); + this.tokenizer = createWordpieceEncoder(vocab, lowerCase); } catch (IOException | RuntimeException e) { // Vocabulary/tokenizer init failed after the native session was created; close it // so a partially constructed instance never leaks the ONNX session. @@ -136,7 +138,7 @@ protected AbstractDL(final OrtEnvironment env, final OrtSession session, this.env = env; this.session = session; this.vocab = vocab; - this.tokenizer = createPipelineTokenizer(vocab, lowerCase); + this.tokenizer = createWordpieceEncoder(vocab, lowerCase); } /** @@ -238,53 +240,58 @@ static WordpieceTokenizer createWordpieceTokenizer( } /** - * Creates a {@link BertTokenizer} that performs the full BERT tokenization - * pipeline: basic tokenization (text normalization) followed by wordpiece. - * The special tokens are selected based on the vocabulary: if it contains - * RoBERTa-style tokens, those are used, otherwise the BERT defaults. - * - * @param vocab The vocabulary map. - * @param lowerCase {@code true} for uncased models (lower casing and accent - * stripping), {@code false} for cased models. - * @return A configured {@link BertTokenizer}. - * @throws IllegalArgumentException Thrown if the selected special tokens - * are not all present in the vocabulary. - */ - // The deprecated BertTokenizer stays the return type until its removal in 3.1, so that - // already-compiled subclasses overriding this method keep overriding it. - @SuppressWarnings("removal") - protected BertTokenizer createTokenizer( - final Map vocab, final boolean lowerCase) { - - return createPipelineTokenizer(vocab, lowerCase); - } - - /** - * Builds the pipeline tokenizer, selecting the RoBERTa special tokens when the vocabulary - * carries them and the BERT defaults otherwise. + * Builds the BERT encoder, selecting RoBERTa special tokens when present and BERT defaults + * otherwise. * * @param vocab The vocabulary map. * @param lowerCase {@code true} for uncased models, {@code false} for cased models. - * @return A configured {@link BertTokenizer}. - * @throws IllegalArgumentException Thrown if the selected special tokens are not all present in - * the vocabulary. + * @return A configured {@link WordpieceEncoder}. + * @throws IllegalArgumentException Thrown if {@code vocab} is {@code null} or the selected + * special tokens are not all present in it. */ - // BertTokenizer is deprecated for removal in 3.1; built here until then. - @SuppressWarnings("removal") - static BertTokenizer createPipelineTokenizer( + static WordpieceEncoder createWordpieceEncoder( final Map vocab, final boolean lowerCase) { + if (vocab == null) { + throw new IllegalArgumentException("vocab must not be null"); + } if (vocab.containsKey( WordpieceTokenizer.ROBERTA_CLS_TOKEN) && vocab.containsKey( WordpieceTokenizer.ROBERTA_SEP_TOKEN)) { - return new BertTokenizer( - vocab.keySet(), + return new WordpieceEncoder( + vocab, lowerCase, WordpieceTokenizer.ROBERTA_CLS_TOKEN, WordpieceTokenizer.ROBERTA_SEP_TOKEN, resolveUnknownToken(vocab)); } - return new BertTokenizer(vocab.keySet(), lowerCase); + return new WordpieceEncoder(vocab, lowerCase, + WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, + WordpieceTokenizer.BERT_UNK_TOKEN); + } + + /** + * Encodes text as the inputs consumed by a BERT model: the piece strings, their vocabulary + * ids, an attention mask of ones, and single-segment token types. + * + * @param text The text to encode; must not be {@code null}. + * @return The encoded pieces, ids, attention mask, and token types. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + */ + protected final Tokens encodeTokens(CharSequence text) { + final List pieces = tokenizer.encode(text); + final String[] tokens = new String[pieces.size()]; + final long[] ids = new long[pieces.size()]; + final long[] mask = new long[pieces.size()]; + final long[] types = new long[pieces.size()]; + for (int i = 0; i < pieces.size(); i++) { + final SubwordPiece piece = pieces.get(i); + tokens[i] = piece.piece(); + ids[i] = piece.id(); + } + Arrays.fill(mask, 1); + return new Tokens(tokens, ids, mask, types); } /** diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/InferenceOptions.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/InferenceOptions.java index b59aac226b..857fc12785 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/InferenceOptions.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/InferenceOptions.java @@ -88,7 +88,7 @@ public boolean isNormalizeWhitespace() { * produces still align with the input. Off by default. * *

    This is a one-for-one replacement, not the collapse-and-trim whitespace fold of the runtime - * {@code TextNormalizer.whitespace()} rung: runs of whitespace are not merged and leading or + * {@code TextNormalizer.whitespace()} normalizer: runs of whitespace are not merged and leading or * trailing whitespace is not removed, so offsets are preserved.

    * * @param normalizeWhitespace Whether to normalize whitespace. 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/doccat/DocumentCategorizerDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java index 7b17cb3028..8b3020d453 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/doccat/DocumentCategorizerDL.java @@ -23,7 +23,6 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.ArrayList; -import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -376,18 +375,7 @@ private List tokenize(final String input) { final List t = new ArrayList<>(groups.size()); for (final String group : groups) { - // Now we can tokenize the group and continue. - final String[] tokens = tokenizer.tokenize(group); - - final long[] ids = tokenIds(tokens, vocab); - - final long[] mask = new long[ids.length]; - Arrays.fill(mask, 1); - - final long[] types = new long[ids.length]; - Arrays.fill(types, 0); - - t.add(new Tokens(tokens, ids, mask, types)); + t.add(encodeTokens(group)); } @@ -395,33 +383,6 @@ private List tokenize(final String input) { } - /** - * Maps tokens to their vocabulary ids. - * - * @param tokens The tokens to map. - * @param vocab The vocabulary map. - * @return The token ids. - * - * @throws IllegalArgumentException Thrown if a token is not present in the - * vocabulary. - */ - static long[] tokenIds(final String[] tokens, final Map vocab) { - - final long[] ids = new long[tokens.length]; - - for (int x = 0; x < tokens.length; x++) { - final Integer id = vocab.get(tokens[x]); - if (id == null) { - throw new IllegalArgumentException("Token '" + tokens[x] - + "' is not present in the vocabulary; the vocabulary file does not match the model."); - } - ids[x] = id; - } - - return ids; - - } - /** * Applies softmax to an array of values. * diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java index 956f177d61..db55827b1c 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java @@ -21,7 +21,6 @@ import java.io.IOException; import java.nio.LongBuffer; import java.util.ArrayList; -import java.util.Arrays; import java.util.Comparator; import java.util.HashMap; import java.util.List; @@ -789,18 +788,7 @@ private List tokenize(final String text) { final List t = new ArrayList<>(chunks.size()); for (final TextChunk chunk : chunks) { - // Now we can tokenize the group and continue. - final String[] tokens = tokenizer.tokenize(chunk.text()); - - final long[] ids = tokenIds(tokens, vocab); - - final long[] mask = new long[ids.length]; - Arrays.fill(mask, 1); - - final long[] types = new long[ids.length]; - Arrays.fill(types, 0); - - t.add(new ChunkTokens(new Tokens(tokens, ids, mask, types), chunk.start(), chunk.end())); + t.add(new ChunkTokens(encodeTokens(chunk.text()), chunk.start(), chunk.end())); } @@ -808,31 +796,4 @@ private List tokenize(final String text) { } - /** - * Maps tokens to their vocabulary ids. - * - * @param tokens The tokens to map. - * @param vocab The vocabulary map. - * @return The token ids. - * - * @throws IllegalArgumentException Thrown if a token is not present in the - * vocabulary. - */ - static long[] tokenIds(final String[] tokens, final Map vocab) { - - final long[] ids = new long[tokens.length]; - - for (int x = 0; x < tokens.length; x++) { - final Integer id = vocab.get(tokens[x]); - if (id == null) { - throw new IllegalArgumentException("Token '" + tokens[x] - + "' is not present in the vocabulary; the vocabulary file does not match the model."); - } - ids[x] = id; - } - - return ids; - - } - } 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 f1250ea601..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,19 +20,22 @@ import java.io.File; import java.io.IOException; import java.nio.LongBuffer; -import java.util.Arrays; +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.tokenize.Tokenizer; - +import opennlp.tools.embeddings.TextEmbedder; /** * Facilitates the generation of sentence vectors using @@ -56,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 @@ -94,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); } @@ -103,11 +116,16 @@ 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 { - final Tokens tokens = tokenize(sentence, tokenizer, vocab); + if (sentence == null) { + throw new IllegalArgumentException("sentence must not be null"); + } + + final Tokens tokens = encodeTokens(sentence); final Map inputs = new HashMap<>(); @@ -133,40 +151,154 @@ public float[] getVectors(final String sentence) throws OrtException { } /** - * Encodes text as model inputs: wordpiece token ids, an attention mask of ones, - * and single-segment (all zero) token type ids. + * {@inheritDoc} * - * @param text The text to encode. - * @param tokenizer The wordpiece tokenizer matching the {@code vocab}. - * @param vocab The vocabulary map. - * @return The encoded {@link Tokens}. + *

    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 the tokenizer emits a token that is - * not present in the vocabulary. + * @throws IllegalArgumentException Thrown if {@code text} is {@code null}. + * @throws IllegalStateException Thrown if inference fails; the cause carries the + * underlying {@link OrtException}. */ - static Tokens tokenize(final String text, final Tokenizer tokenizer, - final Map vocab) { + @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); + } + } - final String[] tokens = tokenizer.tokenize(text); + /** + * {@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; + } - final long[] ids = new long[tokens.length]; + /** + * 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 { - for (int x = 0; x < tokens.length; x++) { - final Integer id = vocab.get(tokens[x]); - if (id == null) { - throw new IllegalArgumentException("Token '" + tokens[x] - + "' is not present in the vocabulary; the vocabulary file does not match the model."); - } - ids[x] = id; + 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 long[] mask = new long[ids.length]; - Arrays.fill(mask, 1); + 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)); - final long[] types = new long[ids.length]; + inputs.put(TOKEN_TYPE_IDS, OnnxTensor.createTensor(env, LongBuffer.wrap(types), shape)); - return new Tokens(tokens, ids, mask, types); + 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 5ae7e3ebee..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 @@ -21,11 +21,14 @@ import java.util.Map; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; -import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.tokenize.SubwordTokenizer; 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.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -52,34 +55,34 @@ private static Map robertaVocab() { } @Test - void testCreatesLowerCasingPipelineTokenizer() { - final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), true); + void testCreatesLowerCasingWordpieceEncoder() { + final SubwordTokenizer tokenizer = AbstractDL.createWordpieceEncoder(bertVocab(), true); // Capitalized input must be lower cased before the wordpiece lookup. assertArrayEquals(new String[] { WordpieceTokenizer.BERT_CLS_TOKEN, "hello", "world", WordpieceTokenizer.BERT_SEP_TOKEN}, - tokenizer.tokenize("Hello World")); + tokenizer.encodeToPieces("Hello World")); } @Test - void testCreatesCasePreservingPipelineTokenizer() { - final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), false); + void testCreatesCasePreservingWordpieceEncoder() { + final SubwordTokenizer tokenizer = AbstractDL.createWordpieceEncoder(bertVocab(), false); // Without lower casing, capitalized words miss the lowercase-only vocabulary. assertArrayEquals(new String[] { WordpieceTokenizer.BERT_CLS_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN, "world", WordpieceTokenizer.BERT_SEP_TOKEN}, - tokenizer.tokenize("Hello world")); + tokenizer.encodeToPieces("Hello world")); } @Test void testSelectsRobertaSpecialTokens() { - final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(robertaVocab(), false); + final SubwordTokenizer tokenizer = AbstractDL.createWordpieceEncoder(robertaVocab(), false); assertArrayEquals(new String[] { WordpieceTokenizer.ROBERTA_CLS_TOKEN, "hello", WordpieceTokenizer.ROBERTA_UNK_TOKEN, WordpieceTokenizer.ROBERTA_SEP_TOKEN}, - tokenizer.tokenize("hello missing")); + tokenizer.encodeToPieces("hello missing")); } @Test @@ -88,12 +91,12 @@ void testFallsBackToBertUnknownToken() { vocab.remove(WordpieceTokenizer.ROBERTA_UNK_TOKEN); vocab.put(WordpieceTokenizer.BERT_UNK_TOKEN, 2); - final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(vocab, false); + final SubwordTokenizer tokenizer = AbstractDL.createWordpieceEncoder(vocab, false); assertArrayEquals(new String[] { WordpieceTokenizer.ROBERTA_CLS_TOKEN, "hello", WordpieceTokenizer.BERT_UNK_TOKEN, WordpieceTokenizer.ROBERTA_SEP_TOKEN}, - tokenizer.tokenize("hello missing")); + tokenizer.encodeToPieces("hello missing")); } @Test @@ -102,23 +105,47 @@ void testRejectsRobertaVocabularyWithoutUnknownToken() { vocab.remove(WordpieceTokenizer.ROBERTA_UNK_TOKEN); assertThrows(IllegalArgumentException.class, - () -> AbstractDL.createPipelineTokenizer(vocab, false)); + () -> AbstractDL.createWordpieceEncoder(vocab, false)); assertThrows(IllegalArgumentException.class, () -> AbstractDL.createWordpieceTokenizer(vocab)); } - @Test - void testTokenizePosIsUnsupported() { - final Tokenizer tokenizer = AbstractDL.createPipelineTokenizer(bertVocab(), true); - assertThrows(UnsupportedOperationException.class, () -> tokenizer.tokenizePos("the fox")); + @ParameterizedTest + @ValueSource(strings = { + WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, + WordpieceTokenizer.BERT_UNK_TOKEN + }) + void testRejectsBertVocabularyMissingSpecialTokensAtCreation(String missingToken) { + final Map vocab = bertVocab(); + vocab.remove(missingToken); + + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> AbstractDL.createWordpieceEncoder(vocab, true)); + assertEquals("vocabulary must contain special token '" + missingToken + "'", + exception.getMessage()); } @Test - void testRejectsBertVocabularyMissingSpecialTokensAtCreation() { - final Map vocab = bertVocab(); - vocab.remove(WordpieceTokenizer.BERT_UNK_TOKEN); + void testRejectsNullVocabularyAtCreation() { + final IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> AbstractDL.createWordpieceEncoder(null, true)); + assertEquals("vocab must not be null", exception.getMessage()); + } - assertThrows(IllegalArgumentException.class, - () -> AbstractDL.createPipelineTokenizer(vocab, true)); + @Test + void testDlEncodingPreservesVocabularyIds() { + final Map vocab = Map.of( + WordpieceTokenizer.BERT_CLS_TOKEN, 101, + WordpieceTokenizer.BERT_SEP_TOKEN, 205, + WordpieceTokenizer.BERT_UNK_TOKEN, 999, + "hello", 42); + + 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()); + assertArrayEquals(new long[] {1, 1, 1}, tokens.mask()); + assertArrayEquals(new long[] {0, 0, 0}, tokens.types()); } @Test 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 197d8a3bc4..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, () -> @@ -106,27 +106,9 @@ void testConstructorRejectsNullInferenceOptions() { new AverageClassificationScoringStrategy(), null)); } - @Test - void testTokenIdsMapsTokensToVocabularyIds() { - final long[] ids = DocumentCategorizerDL.tokenIds( - new String[] {WordpieceTokenizer.BERT_CLS_TOKEN, "hello", "world", - WordpieceTokenizer.BERT_SEP_TOKEN}, vocab()); - - assertArrayEquals(new long[] {0, 3, 4, 1}, ids); - } - - @Test - void testTokenIdsRejectsTokensMissingFromVocabulary() { - final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> - DocumentCategorizerDL.tokenIds(new String[] {"hello", "missing"}, vocab())); - - assertTrue(e.getMessage().contains("missing"), - "the error message should name the missing token: " + e.getMessage()); - } - @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()); @@ -136,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"), @@ -199,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)); @@ -209,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/namefinder/NameFinderDLTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/namefinder/NameFinderDLTest.java index 23db5e05ca..931b09d03f 100644 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/namefinder/NameFinderDLTest.java +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/namefinder/NameFinderDLTest.java @@ -27,7 +27,6 @@ import opennlp.tools.tokenize.WordpieceTokenizer; import opennlp.tools.util.Span; -import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -55,24 +54,6 @@ private static Map vocab() { return vocab; } - @Test - void testTokenIdsMapsTokensToVocabularyIds() { - final long[] ids = NameFinderDL.tokenIds( - new String[] {WordpieceTokenizer.BERT_CLS_TOKEN, "hello", "world", - WordpieceTokenizer.BERT_SEP_TOKEN}, vocab()); - - assertArrayEquals(new long[] {0, 3, 4, 1}, ids); - } - - @Test - void testTokenIdsRejectsTokensMissingFromVocabulary() { - final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> - NameFinderDL.tokenIds(new String[] {"hello", "missing"}, vocab())); - - assertTrue(e.getMessage().contains("missing"), - "the error message should name the missing token: " + e.getMessage()); - } - @Test void testDecodeSpansUsesBioEntityTypesAndBoundedProbabilities() { final String text = "Alice visited New York City."; 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/java/opennlp/dl/vectors/SentenceVectorsDLTest.java b/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLTest.java deleted file mode 100644 index 422c7773a9..0000000000 --- a/opennlp-core/opennlp-ml/opennlp-dl/src/test/java/opennlp/dl/vectors/SentenceVectorsDLTest.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * 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.util.HashMap; -import java.util.Map; - -import org.junit.jupiter.api.Test; - -import opennlp.dl.Tokens; -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.assertThrows; - -public class SentenceVectorsDLTest { - - private static Map vocab() { - final Map vocab = new HashMap<>(); - vocab.put(WordpieceTokenizer.BERT_CLS_TOKEN, 0); - vocab.put(WordpieceTokenizer.BERT_SEP_TOKEN, 1); - vocab.put(WordpieceTokenizer.BERT_UNK_TOKEN, 2); - vocab.put("hello", 3); - vocab.put("world", 4); - return vocab; - } - - @Test - void testTokenizeUsesSingleSegmentBertEncoding() { - final Map vocab = vocab(); - final WordpieceTokenizer tokenizer = new WordpieceTokenizer(vocab.keySet()); - - final Tokens tokens = SentenceVectorsDL.tokenize("hello world", tokenizer, vocab); - - assertArrayEquals(new String[] { - WordpieceTokenizer.BERT_CLS_TOKEN, "hello", "world", WordpieceTokenizer.BERT_SEP_TOKEN}, - tokens.tokens()); - assertArrayEquals(new long[] {0, 3, 4, 1}, tokens.ids()); - // The attention mask must be 1 for every real token. - assertArrayEquals(new long[] {1, 1, 1, 1}, tokens.mask()); - // Single-segment input: all token type ids must be 0. - assertArrayEquals(new long[] {0, 0, 0, 0}, tokens.types()); - } - - @Test - void testTokenizeMapsOutOfVocabularyWordsToUnknownToken() { - final Map vocab = vocab(); - final WordpieceTokenizer tokenizer = new WordpieceTokenizer(vocab.keySet()); - - final Tokens tokens = SentenceVectorsDL.tokenize("hello xyz", tokenizer, vocab); - - assertArrayEquals(new long[] {0, 3, 2, 1}, tokens.ids()); - assertEquals(WordpieceTokenizer.BERT_UNK_TOKEN, tokens.tokens()[2]); - } - - @Test - void testTokenizeRejectsTokensMissingFromVocabulary() { - final Map vocab = vocab(); - vocab.remove(WordpieceTokenizer.BERT_UNK_TOKEN); - final WordpieceTokenizer tokenizer = new WordpieceTokenizer(vocab.keySet()); - - assertThrows(IllegalArgumentException.class, () -> - SentenceVectorsDL.tokenize("hello xyz", tokenizer, vocab)); - } -} diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py new file mode 100644 index 0000000000..05a67ccd5f --- /dev/null +++ b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/gen_tiny_vectors_model.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Generates tiny-vectors.onnx, the deterministic model behind SentenceVectorsDLEmbedderTest. +# +# The graph computes output[b][t][d] = float(input_ids[b][t]) * W[0][d] with +# W = [[0.5, -1.0, 2.0]], so the vector at any token position is that token's vocabulary id +# times W, hand-computable in the test. It declares the same three inputs a BERT-style +# encoder declares (input_ids, attention_mask, token_type_ids; the latter two are accepted +# and ignored) and one output of shape [batch, tokens, 3] so the hidden dimension is static +# in the model metadata. +# +# Regenerate with: python3 gen_tiny_vectors_model.py (requires the onnx package) + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +W = np.array([[0.5, -1.0, 2.0]], dtype=np.float32) + +cast = helper.make_node("Cast", ["input_ids"], ["ids_float"], to=TensorProto.FLOAT) +unsqueeze = helper.make_node("Unsqueeze", ["ids_float", "axes"], ["ids_3d"]) +matmul = helper.make_node("MatMul", ["ids_3d", "w"], ["last_hidden_state"]) + + +def encoder_input(name): + return helper.make_tensor_value_info(name, TensorProto.INT64, ["batch", "tokens"]) + + +graph = helper.make_graph( + [cast, unsqueeze, matmul], + "tiny-vectors", + [encoder_input("input_ids"), encoder_input("attention_mask"), + encoder_input("token_type_ids")], + [helper.make_tensor_value_info( + "last_hidden_state", TensorProto.FLOAT, ["batch", "tokens", 3])], + [numpy_helper.from_array(np.array([2], dtype=np.int64), name="axes"), + numpy_helper.from_array(W, name="w")], +) + +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) +model.ir_version = 8 +onnx.checker.check_model(model) +onnx.save(model, "tiny-vectors.onnx") +print("wrote tiny-vectors.onnx,", len(model.SerializeToString()), "bytes") diff --git a/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx new file mode 100644 index 0000000000..7d63c91322 Binary files /dev/null and b/opennlp-core/opennlp-ml/opennlp-dl/src/test/resources/opennlp/dl/vectors/tiny-vectors.onnx differ diff --git a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/BeamSearch.java b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/BeamSearch.java index 804d676366..923bd64cd5 100644 --- a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/BeamSearch.java +++ b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/BeamSearch.java @@ -18,7 +18,6 @@ package opennlp.tools.ml; import java.util.Arrays; -import java.util.List; import java.util.PriorityQueue; import java.util.Queue; @@ -75,6 +74,92 @@ private static final class CacheState { } } + /** + * Immutable node in a backward-linked chain of outcome candidates, used only inside + * {@link #bestSequences(int, Object[], Object[], double, BeamSearchContextGenerator, SequenceValidator)}. + * A node stores its own outcome plus a parent link, so extending a candidate is O(1); + * the full outcome array is materialized only when a candidate is expanded. + * Score accumulation ({@code parent.score + StrictMath.log(prob)}) mirrors + * {@link Sequence#Sequence(Sequence, String, double)}, and descending score order + * mirrors {@link Sequence#compareTo(Sequence)}, so non-tied search results stay + * bit-identical to a search over {@link Sequence} instances. Unlike + * {@link Sequence#compareTo(Sequence)}, {@link #compareTo(SearchNode)} additionally + * resolves exact score ties into a canonical outcome order; under ties the + * {@link Sequence}-based search has no defined order at all. + */ + private static final class SearchNode implements Comparable { + private final SearchNode parent; + private final String outcome; // null on the root + private final double prob; + private final double score; + private final int size; + + /** + * Creates the root node: the empty candidate with score {@code 0}. + */ + private SearchNode() { + this.parent = null; + this.outcome = null; + this.prob = 0d; + this.score = 0d; + this.size = 0; + } + + /** + * Creates a candidate extending {@code parent} by one outcome. + * + * @param parent The candidate to extend. Must not be {@code null}. + * @param outcome The outcome to append. + * @param prob The probability of {@code outcome}. + */ + private SearchNode(SearchNode parent, String outcome, double prob) { + this.parent = parent; + this.outcome = outcome; + this.prob = prob; + this.score = parent.score + StrictMath.log(prob); + this.size = parent.size + 1; + } + + /** + * @return The outcomes on the path from the root to this node, in sequence order. + */ + private String[] outcomes() { + final String[] outcomes = new String[size]; + SearchNode node = this; + for (int i = size - 1; i >= 0; i--) { + outcomes[i] = node.outcome; + node = node.parent; + } + return outcomes; + } + + /** + * Orders nodes by descending score, mirroring {@link Sequence#compareTo(Sequence)}. + * Exact score ties are broken lexicographically by the outcome chain from the + * root, so tied candidates expand and are returned in the same canonical order + * on every run and JVM. {@link PriorityQueue} gives no ordering guarantee for + * elements whose comparison returns {@code 0}, and ties decide both which + * candidates survive beam truncation and the order of equal-scored winners. + */ + @Override + public int compareTo(SearchNode other) { + final int byScore = Double.compare(other.score, this.score); + if (byScore != 0) { + return byScore; + } + final String[] mine = outcomes(); + final String[] theirs = other.outcomes(); + final int common = StrictMath.min(mine.length, theirs.length); + for (int i = 0; i < common; i++) { + final int byOutcome = mine[i].compareTo(theirs[i]); + if (byOutcome != 0) { + return byOutcome; + } + } + return Integer.compare(mine.length, theirs.length); + } + } + /** * Initializes a {@link BeamSearch} instance. * @@ -113,10 +198,10 @@ public Sequence[] bestSequences(final int numSequences, final T[] sequence, final CacheState state = threadState.get(); - Queue prev = new PriorityQueue<>(size); - Queue next = new PriorityQueue<>(size); - Queue tmp; - prev.add(new Sequence()); + Queue prev = new PriorityQueue<>(size); + Queue next = new PriorityQueue<>(size); + Queue tmp; + prev.add(new SearchNode()); Object[] context = additionalContext; if (context == null) { @@ -127,9 +212,8 @@ public Sequence[] bestSequences(final int numSequences, final T[] sequence, final int sz = StrictMath.min(size, prev.size()); for (int sc = 0; prev.size() > 0 && sc < sz; sc++) { - final Sequence top = prev.remove(); - final List tmpOutcomes = top.getOutcomes(); - final String[] outcomes = tmpOutcomes.toArray(new String[0]); + final SearchNode top = prev.remove(); + final String[] outcomes = top.outcomes(); final String[] contexts = cg.getContext(i, sequence, outcomes, context); final double[] scores; if (state.cache != null) { @@ -157,8 +241,8 @@ public Sequence[] bestSequences(final int numSequences, final T[] sequence, if (scores[p] >= min) { final String out = model.getOutcome(p); if (validator.validSequence(i, sequence, outcomes, out)) { - final Sequence ns = new Sequence(top, out, scores[p]); - if (ns.getScore() > minSequenceScore) { + final SearchNode ns = new SearchNode(top, out, scores[p]); + if (ns.score > minSequenceScore) { next.add(ns); } } @@ -169,8 +253,8 @@ public Sequence[] bestSequences(final int numSequences, final T[] sequence, for (int p = 0; p < scores.length; p++) { final String out = model.getOutcome(p); if (validator.validSequence(i, sequence, outcomes, out)) { - final Sequence ns = new Sequence(top, out, scores[p]); - if (ns.getScore() > minSequenceScore) { + final SearchNode ns = new SearchNode(top, out, scores[p]); + if (ns.score > minSequenceScore) { next.add(ns); } } @@ -189,7 +273,22 @@ public Sequence[] bestSequences(final int numSequences, final T[] sequence, final Sequence[] topSequences = new Sequence[numSeq]; for (int seqIndex = 0; seqIndex < numSeq; seqIndex++) { - topSequences[seqIndex] = prev.remove(); + final SearchNode winner = prev.remove(); + final String[] outs = new String[winner.size]; + final double[] probs = new double[winner.size]; + SearchNode node = winner; + for (int j = winner.size - 1; j >= 0; j--) { + outs[j] = node.outcome; + probs[j] = node.prob; + node = node.parent; + } + // Sequence.add accumulates score += StrictMath.log(p) per element, so rebuilding in + // chain order yields a score bit-identical to the node's accumulated score. + final Sequence seq = new Sequence(); + for (int j = 0; j < outs.length; j++) { + seq.add(outs[j], probs[j]); + } + topSequences[seqIndex] = seq; } return topSequences; diff --git a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModel.java b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModel.java index 33131fb532..be6d6d1e82 100644 --- a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModel.java +++ b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModel.java @@ -18,8 +18,10 @@ package opennlp.tools.ml.model; import java.text.DecimalFormat; +import java.text.DecimalFormatSymbols; import java.util.Arrays; import java.util.LinkedHashMap; +import java.util.Locale; import java.util.Map; import java.util.Objects; @@ -31,7 +33,8 @@ */ public abstract class AbstractModel implements MaxentModel { - private static final DecimalFormat DECIMAL_FORMAT = new DecimalFormat("0.0000"); + private static final DecimalFormat DECIMAL_FORMAT = + new DecimalFormat("0.0000", DecimalFormatSymbols.getInstance(Locale.ROOT)); /** Mapping between predicates/contexts and an integer representing them. */ protected Map pmap; diff --git a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModelReader.java b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModelReader.java index 8325b04f21..c24349401f 100644 --- a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModelReader.java +++ b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/ml/model/AbstractModelReader.java @@ -24,6 +24,8 @@ import java.util.StringTokenizer; import java.util.zip.GZIPInputStream; +import opennlp.tools.util.ResourceLimits; + /** * An abstract, basic implementation of a model reader. */ @@ -32,31 +34,17 @@ public abstract class AbstractModelReader { /** * System property for overriding the maximum number of entries (outcomes, predicates, * outcome patterns, chunk counts) that may be read from a model file or training data. - * Set at JVM startup, e.g. {@code -DOPENNLP_MAX_ENTRIES=5000000}. - * Falls back to {@code 10_000_000} if absent or invalid. + * Alias of {@link ResourceLimits#MAX_ENTRIES_PROPERTY}. */ - public static final String MAX_ENTRIES_PROPERTY = "OPENNLP_MAX_ENTRIES"; + public static final String MAX_ENTRIES_PROPERTY = ResourceLimits.MAX_ENTRIES_PROPERTY; /** * Upper bound on count fields read from a model file. - * Prevents OOM on crafted inputs with oversized array size declarations. - * Configurable via the {@link #MAX_ENTRIES_PROPERTY} system property. - *

    + * Alias of {@link ResourceLimits#MAX_ENTRIES}. * Public so that deserializers outside this package which implement their own binary * format can apply the same bound to their count fields. */ - public static final int MAX_ENTRIES = initMaxEntries(); - - private static int initMaxEntries() { - String prop = System.getProperty(MAX_ENTRIES_PROPERTY, "").trim(); - if (!prop.isEmpty()) { - try { - int val = Integer.parseInt(prop); - if (val > 0) return val; - } catch (NumberFormatException ignore) { } - } - return 10_000_000; - } + public static final int MAX_ENTRIES = ResourceLimits.MAX_ENTRIES; /** * The number of predicates contained in a model. diff --git a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/util/TrainingParameters.java b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/util/TrainingParameters.java index 6af97efe43..39f69de847 100644 --- a/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/util/TrainingParameters.java +++ b/opennlp-core/opennlp-ml/opennlp-ml-commons/src/main/java/opennlp/tools/util/TrainingParameters.java @@ -22,6 +22,7 @@ import java.io.OutputStream; import java.util.Collections; import java.util.HashMap; +import java.util.Locale; import java.util.Map; import java.util.Map.Entry; import java.util.Properties; @@ -329,12 +330,12 @@ public static TrainingParameters setParams(String[] params) { mlParams.put(Parameters.ALGORITHM_PARAM , "MAXENT"); mlParams.put(Parameters.TRAINER_TYPE_PARAM , EventTrainer.EVENT_VALUE); mlParams.put(Parameters.ITERATIONS_PARAM , - null != getIntParameter("-" + Parameters.ITERATIONS_PARAM.toLowerCase() , params) ? - getIntParameter("-" + Parameters.ITERATIONS_PARAM.toLowerCase() , params) : + null != getIntParameter("-" + Parameters.ITERATIONS_PARAM.toLowerCase(Locale.ROOT) , params) ? + getIntParameter("-" + Parameters.ITERATIONS_PARAM.toLowerCase(Locale.ROOT) , params) : Parameters.ITERATIONS_DEFAULT_VALUE); mlParams.put(Parameters.CUTOFF_PARAM , - null != getIntParameter("-" + Parameters.CUTOFF_PARAM.toLowerCase() , params) ? - getIntParameter("-" + Parameters.CUTOFF_PARAM.toLowerCase() , params) : + null != getIntParameter("-" + Parameters.CUTOFF_PARAM.toLowerCase(Locale.ROOT) , params) ? + getIntParameter("-" + Parameters.CUTOFF_PARAM.toLowerCase(Locale.ROOT) , params) : Parameters.CUTOFF_DEFAULT_VALUE); return mlParams; diff --git a/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/ml/BeamSearchBenchmark.java b/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/ml/BeamSearchBenchmark.java new file mode 100644 index 0000000000..18e2f8212f --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/ml/BeamSearchBenchmark.java @@ -0,0 +1,166 @@ +/* + * 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.ml; + +import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; +import org.openjdk.jmh.runner.options.TimeValue; + +import opennlp.tools.util.BeamSearchContextGenerator; +import opennlp.tools.util.SequenceValidator; + +/** + * JMH benchmark for {@link BeamSearch} on long input sequences. + *

    + * One op = one {@code bestSequence} call on a synthetic token sequence of + * {@code sequenceLength} tokens. Long sequences are what expose the quadratic + * per-candidate outcome-list copying removed with OPENNLP-1903; sentences of + * 5-10 tokens, as used by the ME benchmarks, show no measurable difference. + * Only public API that predates OPENNLP-1903 is used + * ({@code BeamSearch(int, MaxentModel, int)} and + * {@code bestSequence(T[], Object[], BeamSearchContextGenerator, SequenceValidator)}), + * so the same compiled class exercises both implementations: to produce a + * baseline, place an {@code opennlp-ml-commons} jar built before OPENNLP-1903 + * on the classpath ahead of the freshly built classes and rerun. + */ +@BenchmarkMode(Mode.Throughput) +@OutputTimeUnit(TimeUnit.SECONDS) +@Warmup(iterations = 5, time = 2) +@Measurement(iterations = 10, time = 2) +@Fork(2) +public class BeamSearchBenchmark { + + private static final int BEAM_SIZE = 3; + private static final int NUM_INPUTS = 64; + private static final int VOCAB_SIZE = 17; + /** Seed for both the input generator and the {@link SeededMaxentModel}. */ + private static final long SEED = 0x5eedL; + private static final String[] MODEL_OUTCOMES = {"start", "cont", "other"}; + + private static final SequenceValidator ACCEPT_ALL = + (i, input, outcomes, outcome) -> true; + + @State(Scope.Benchmark) + public static class SearchState { + + @Param({"8", "64", "256"}) + int sequenceLength; + + @Param({"0", "64"}) + int cacheSize; + + BeamSearch beamSearch; + TokenContextGenerator contextGenerator; + String[][] inputs; + final AtomicInteger cursor = new AtomicInteger(); + + @Setup(Level.Trial) + public void create() { + beamSearch = new BeamSearch(BEAM_SIZE, + new SeededMaxentModel(MODEL_OUTCOMES, SEED), cacheSize); + contextGenerator = new TokenContextGenerator(); + inputs = new String[NUM_INPUTS][]; + Random rnd = new Random(SEED); + for (int n = 0; n < NUM_INPUTS; n++) { + String[] input = new String[sequenceLength]; + for (int i = 0; i < sequenceLength; i++) { + input[i] = "t" + rnd.nextInt(VOCAB_SIZE); + } + inputs[n] = input; + } + } + + String[] nextInput() { + // Rotate through the input pool so repeated invocations decode different + // sequences instead of hammering one fully cache-warm input. + return inputs[Math.floorMod(cursor.getAndIncrement(), NUM_INPUTS)]; + } + } + + /** + * Derives contexts from the current token and the previous outcome, interned + * so identical context content maps to the same {@code String[]} instance + * and the identity-keyed contexts cache in {@link BeamSearch} produces hits. + * Thread-safe. + */ + static final class TokenContextGenerator implements BeamSearchContextGenerator { + + private final ConcurrentHashMap intern = new ConcurrentHashMap<>(); + + @Override + public String[] getContext(int index, String[] sequence, + String[] priorDecisions, Object[] additionalContext) { + String prev = index > 0 ? priorDecisions[index - 1] : ""; + String[] ctx = {"tok=" + sequence[index], "prev=" + prev}; + String key = ctx[0] + '|' + ctx[1]; + String[] existing = intern.putIfAbsent(key, ctx); + return existing != null ? existing : ctx; + } + } + + @Benchmark + @Threads(1) + public void bestSequenceSingle(SearchState state, Blackhole bh) { + bh.consume(state.beamSearch.bestSequence(state.nextInput(), null, + state.contextGenerator, ACCEPT_ALL)); + } + + @Benchmark + @Threads(Threads.MAX) + public void bestSequenceConcurrent(SearchState state, Blackhole bh) { + bh.consume(state.beamSearch.bestSequence(state.nextInput(), null, + state.contextGenerator, ACCEPT_ALL)); + } + + /** + * Quick local iteration only: {@code forks(0)} disables JVM fork isolation + * (unlike {@code mvn} with the {@code jmh} profile). + * Use the Maven-invoked configuration for publishable numbers. + */ + public static void main(String[] args) throws Exception { + Options opt = new OptionsBuilder() + .include(BeamSearchBenchmark.class.getSimpleName()) + .forks(0) + .warmupIterations(1) + .warmupTime(TimeValue.seconds(1)) + .measurementIterations(1) + .measurementTime(TimeValue.seconds(1)) + .build(); + new Runner(opt).run(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationBenchmark.java b/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationBenchmark.java new file mode 100644 index 0000000000..851020eae4 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/jmh/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationBenchmark.java @@ -0,0 +1,208 @@ +/* + * 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.sentdetect; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Locale; +import java.util.concurrent.TimeUnit; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; + +import opennlp.tools.dictionary.Dictionary; +import opennlp.tools.formats.ResourceAsStreamFactory; +import opennlp.tools.util.InputStreamFactory; +import opennlp.tools.util.ObjectStream; +import opennlp.tools.util.PlainTextByLineStream; +import opennlp.tools.util.StringList; +import opennlp.tools.util.TrainingParameters; + +/** + * JMH benchmark for the abbreviation veto of + * {@link SentenceDetectorME#isAcceptableBreak(CharSequence, int, int)}. + *

    + * One op is one {@link SentenceDetectorME#sentPosDetect(CharSequence)} call over a document of + * {@code documentChars} characters. Three variants run over the same input: + *

      + *
    • {@code noDictionary} is the floor, the veto is switched off entirely, so what is left is + * the shared cost of scanning and of the maxent evaluation per candidate;
      • + *
    • {@code legacyVeto} is {@link LegacyAbbreviationSentenceDetectorME}, which holds the + * previous full-text scan;
      • + *
    • {@code boundedWindowVeto} is the current implementation.
      • + *
    + * The {@code documentChars} axis is the point of the benchmark: the previous implementation + * searched the whole text once per dictionary entry per candidate, so its cost per document grows + * quadratically, while the bounded window makes it grow linearly. The {@code dictionaryEntries} + * axis exposes the second factor of that product. + */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 3, time = 2) +@Measurement(iterations = 5, time = 2) +@Fork(2) +@Threads(1) +public class SentenceDetectorMEAbbreviationBenchmark { + + /** + * The abbreviations of the shipped English dictionary. They occur in the benchmark document. + */ + private static final String[] PRESENT_ABBREVIATIONS = {"Mr.", "Mrs.", "Ms.", "tel."}; + + @State(Scope.Benchmark) + public static class DocumentState { + + /** + * The document length in characters. Doubling it is what separates a quadratic veto from a + * linear one. + */ + @Param({"12500", "25000", "50000", "100000"}) + public int documentChars; + + /** + * The number of dictionary entries. The shipped German list has 215, the shipped English one + * that the runtime module tests against has 4. + */ + @Param({"10", "200"}) + public int dictionaryEntries; + + SentenceModel model; + String document; + SentenceDetectorME noDictionary; + SentenceDetectorME legacy; + SentenceDetectorME boundedWindow; + + @Setup(Level.Trial) + public void prepare() throws IOException { + model = trainModel(); + document = buildDocument(documentChars); + final Dictionary dictionary = buildDictionary(dictionaryEntries); + noDictionary = new SentenceDetectorME(model, (Dictionary) null); + legacy = new LegacyAbbreviationSentenceDetectorME(model, dictionary); + boundedWindow = new SentenceDetectorME(model, dictionary); + } + } + + @Benchmark + public void noDictionary(DocumentState state, Blackhole bh) { + bh.consume(state.noDictionary.sentPosDetect(state.document)); + } + + @Benchmark + public void legacyVeto(DocumentState state, Blackhole bh) { + bh.consume(state.legacy.sentPosDetect(state.document)); + } + + @Benchmark + public void boundedWindowVeto(DocumentState state, Blackhole bh) { + bh.consume(state.boundedWindow.sentPosDetect(state.document)); + } + + /* + * ------------------------------------------------------------------------------------------ + * Fixtures. + * ------------------------------------------------------------------------------------------ + */ + + /** + * @return A model trained on the bundled English samples, without an abbreviation dictionary, + * so all three variants share one feature generator and differ only in the veto. + * @throws IOException Thrown if the training samples cannot be read. + */ + private static SentenceModel trainModel() throws IOException { + final InputStreamFactory in = new ResourceAsStreamFactory( + SentenceDetectorMEAbbreviationBenchmark.class, + "/opennlp/tools/sentdetect/Sentences.txt"); + final ObjectStream samples = new SentenceSampleStream( + new PlainTextByLineStream(in, StandardCharsets.UTF_8)); + return SentenceDetectorME.train("eng", samples, + new SentenceDetectorFactory("eng", true, null, null), + TrainingParameters.defaultParams()); + } + + /** + * Builds a document of about {@code chars} characters that contains abbreviations at a + * realistic rate, so the veto is reached often rather than in a corner. + * + * @param chars The minimum document length in characters. + * @return The document. + */ + private static String buildDocument(int chars) { + final String paragraph = + "Mr. Smith left the building at noon. She told me he lived in Edinburgh. " + + "Mrs. Clark called tel. 555 1234 and asked for Ms. Adams. " + + "The driver got badly injured near the old bridge. " + + "OpenNLP provides tools for natural language processing. " + + "I wrote him a letter right away and posted it the same day. "; + final StringBuilder document = new StringBuilder(chars + paragraph.length()); + while (document.length() < chars) { + document.append(paragraph); + } + return document.toString(); + } + + /** + * Builds a case-insensitive dictionary of {@code entries} abbreviations, as the shipped + * dictionaries are. The first entries occur in the document; the rest are synthetic and do not, + * which is the realistic case, since a dictionary covers a language and a document uses a + * handful of its entries. + * + * @param entries The number of entries to produce. + * @return The dictionary. + */ + private static Dictionary buildDictionary(int entries) { + final Dictionary dictionary = new Dictionary(false); + for (String abbreviation : PRESENT_ABBREVIATIONS) { + dictionary.put(new StringList(abbreviation)); + } + for (int i = PRESENT_ABBREVIATIONS.length; i < entries; i++) { + dictionary.put(new StringList(String.format(Locale.ROOT, "zq%d.", i))); + } + return dictionary; + } + + /** + * Quick local iteration only: {@code forks(0)} disables JVM fork isolation + * (unlike {@code mvn} with the {@code jmh} profile). + * Use the Maven-invoked configuration for publishable numbers. + */ + public static void main(String[] args) throws Exception { + Options opt = new OptionsBuilder() + .include(SentenceDetectorMEAbbreviationBenchmark.class.getSimpleName()) + .forks(0) + .warmupIterations(2) + .measurementIterations(3) + .build(); + new Runner(opt).run(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/chunker/ChunkerAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/chunker/ChunkerAnnotator.java new file mode 100644 index 0000000000..8bd4a17a44 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/chunker/ChunkerAnnotator.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.tools.chunker; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +/** + * Adapts a {@link Chunker} to the document pipeline: reads {@link Layers#SENTENCES}, + * {@link Layers#TOKENS}, and {@link Layers#POS_TAGS} and provides {@link #CHUNKS}, one + * annotation per phrase chunk carrying the chunk type, for example {@code NP} or + * {@code VP}, on the span from its first to its last token. + * + *

    Each sentence is chunked separately with its tokens and tags as one sequence, the + * way the chunker contract expects its input. A chunker's spans index tokens within the + * sentence; the adapter maps them onto the token spans, which already refer to the + * original text, so a chunk covers exactly the text of its tokens. Chunks are emitted in + * text order.

    + * + *

    The adapter holds no per-call state; it is as thread-safe as the chunker it + * wraps.

    + * + * @since 3.0.0 + */ +public final class ChunkerAnnotator implements DocumentAnnotator { + + /** + * Phrase chunks; each annotation covers one chunk and carries its type, ordered by + * text position. + */ + public static final LayerKey CHUNKS = Layers.key("chunks", String.class); + + private final Chunker chunker; + + /** + * Initializes the adapter. + * + * @param chunker The chunker to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code chunker} is {@code null}. + */ + public ChunkerAnnotator(Chunker chunker) { + if (chunker == null) { + throw new IllegalArgumentException("chunker must not be null"); + } + this.chunker = chunker; + } + + /** + * Chunks the document sentence by sentence and adds the {@link #CHUNKS} layer. + * + *

    The required layers must be present, but they may be empty: a document without + * sentences or tokens yields a present-but-empty chunk layer. The token and tag + * layers must be aligned one to one.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#SENTENCES}, {@link Layers#TOKENS}, and + * {@link Layers#POS_TAGS} layers, with every token lying inside a + * sentence. + * @return A new {@link Document} with the {@link #CHUNKS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, a + * required layer is absent, the token and tag layers differ in size, a token + * lies outside every sentence, or the chunker returns a span outside the + * sentence, an empty span, or a span without a type. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.SENTENCES, Layers.TOKENS, + Layers.POS_TAGS); + final List> sentences = document.get(Layers.SENTENCES); + final List> tokens = document.get(Layers.TOKENS); + final List> tags = document.get(Layers.POS_TAGS); + if (tags.size() != tokens.size()) { + throw new IllegalArgumentException("document needs aligned " + + Layers.TOKENS + " and " + Layers.POS_TAGS + " layers"); + } + final List> chunks = new ArrayList<>(); + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + final String[] sentenceTags = new String[words.length]; + for (int i = 0; i < words.length; i++) { + sentenceTags[i] = tags.get(first + i).value(); + } + for (final Span chunk : chunker.chunkAsSpans(words, sentenceTags)) { + if (chunk.getStart() < 0 || chunk.getEnd() > words.length + || chunk.getStart() >= chunk.getEnd()) { + throw new IllegalArgumentException("chunker returned chunk " + chunk + + " outside the sentence's " + words.length + " tokens"); + } + if (chunk.getType() == null) { + throw new IllegalArgumentException( + "chunker returned chunk " + chunk + " without a type"); + } + chunks.add(new Annotation<>(new Span( + tokens.get(first + chunk.getStart()).span().getStart(), + tokens.get(first + chunk.getEnd() - 1).span().getEnd()), chunk.getType())); + } + }); + return document.with(CHUNKS, chunks); + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(CHUNKS); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/DictionaryLemmatizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/DictionaryLemmatizer.java index 8b68cf6760..c6459702cb 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/DictionaryLemmatizer.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/DictionaryLemmatizer.java @@ -32,6 +32,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; /** @@ -161,7 +162,7 @@ public Map, List> getDictMap() { * @return Retrieves the dictionary keys (word and postag). */ private List getDictKeys(final String word, final String postag) { - return new ArrayList<>(Arrays.asList(word.toLowerCase(), postag)); + return new ArrayList<>(Arrays.asList(word.toLowerCase(Locale.ROOT), postag)); } diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerAnnotator.java new file mode 100644 index 0000000000..2c633798e5 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerAnnotator.java @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.lemmatizer; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; + +/** + * Adapts a {@link Lemmatizer} to the document pipeline: reads {@link Layers#SENTENCES}, + * {@link Layers#TOKENS}, and {@link Layers#POS_TAGS} and provides {@link #LEMMAS}, one + * annotation per token on the token's span. + * + *

    Each sentence is lemmatized separately, the way the lemmatizer contract expects its + * input, so lemmatization decisions never cross a sentence boundary. Token spans already + * refer to the original document text, so only the token and tag sequences handed to the + * lemmatizer are sliced per sentence; the produced lemma layer stays aligned with + * {@link Layers#TOKENS} by position.

    + * + * @since 3.0.0 + */ +public final class LemmatizerAnnotator implements DocumentAnnotator { + + /** + * The lemma layer. It is aligned with the token layer by position, and each annotation + * carries the lemma of its token on that token's span. + */ + public static final LayerKey LEMMAS = Layers.key("lemmas", String.class); + + private final Lemmatizer lemmatizer; + + /** + * Initializes the adapter. + * + * @param lemmatizer The lemmatizer to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code lemmatizer} is {@code null}. + */ + public LemmatizerAnnotator(Lemmatizer lemmatizer) { + if (lemmatizer == null) { + throw new IllegalArgumentException("lemmatizer must not be null"); + } + this.lemmatizer = lemmatizer; + } + + /** + * Lemmatizes the document sentence by sentence and adds the {@link #LEMMAS} layer. + * + *

    For every sentence, the tokens whose spans lie inside the sentence span are + * lemmatized as one sequence together with their tags, and each lemma is emitted on + * its token's span. The required layers must be present, but they may be empty: a + * document without sentences or tokens yields a present-but-empty lemma layer, and a + * sentence containing no tokens contributes nothing.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#SENTENCES} and {@link Layers#TOKENS} layers and a + * {@link Layers#POS_TAGS} layer with exactly one tag per token, with + * every token lying inside a sentence. + * @return A new {@link Document} with the {@link #LEMMAS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, the + * sentence layer, the token layer, or the tag layer is absent, the tag layer + * does not have exactly one tag per token, a token lies outside every + * sentence, or the lemmatizer does not return one lemma per token of a + * sentence. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, + Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS); + final List> sentences = document.get(Layers.SENTENCES); + final List> tokens = document.get(Layers.TOKENS); + final List> tags = document.get(Layers.POS_TAGS); + if (tags.size() != tokens.size()) { + throw new IllegalArgumentException("document needs aligned " + + Layers.TOKENS + " and " + Layers.POS_TAGS + " layers"); + } + final List> layer = new ArrayList<>(tokens.size()); + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + final String[] posTags = new String[words.length]; + for (int i = 0; i < words.length; i++) { + posTags[i] = tags.get(first + i).value(); + } + final String[] lemmas = lemmatizer.lemmatize(words, posTags); + if (lemmas.length != words.length) { + throw new IllegalArgumentException("lemmatizer returned " + lemmas.length + + " lemmas for " + words.length + " tokens"); + } + for (int i = 0; i < words.length; i++) { + layer.add(new Annotation<>(tokens.get(first + i).span(), lemmas[i])); + } + }); + return document.with(LEMMAS, layer); + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(LEMMAS); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerME.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerME.java index c5dafaf5e0..468857bae8 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerME.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/lemmatizer/LemmatizerME.java @@ -22,6 +22,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import opennlp.tools.commons.ThreadSafe; @@ -168,7 +169,7 @@ public String[][] predictLemmas(int numLemmas, String[] toks, String[] tags) { public static String[] decodeLemmas(String[] toks, String[] preds) { List lemmas = new ArrayList<>(); for (int i = 0; i < toks.length; i++) { - String lemma = StringUtil.decodeShortestEditScript(toks[i].toLowerCase(), preds[i]); + String lemma = StringUtil.decodeShortestEditScript(toks[i].toLowerCase(Locale.ROOT), preds[i]); if (lemma.length() == 0) { lemma = "_"; } diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/NameFinderAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/NameFinderAnnotator.java new file mode 100644 index 0000000000..43baa231e0 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/NameFinderAnnotator.java @@ -0,0 +1,149 @@ +/* + * 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.namefind; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +/** + * Adapts a {@link TokenNameFinder} to the document pipeline: reads + * {@link Layers#SENTENCES} and {@link Layers#TOKENS}, maps the finder's token-index + * spans to character spans on the original text, and provides {@link Layers#ENTITIES}. + * The entity type is carried as the annotation value; the annotation's span carries + * offsets only. + * + *

    Each sentence's tokens are passed to {@link TokenNameFinder#find(String[])} as one + * sequence, the way the finder contract expects its input, so no mention can straddle a + * sentence boundary. The finder's adaptive data is cleared exactly once per call, as the + * {@link TokenNameFinder#clearAdaptiveData()} contract asks, whether annotation succeeds + * or fails, so no document can leak finder state into the next one.

    + * + *

    Spans the finder returns without a type are recorded with the {@link #UNTYPED} + * entity type.

    + * + * @since 3.0.0 + */ +public final class NameFinderAnnotator implements DocumentAnnotator { + + /** + * The entity type recorded when the wrapped finder returns a span without a type. It + * is {@link NameSample#DEFAULT_TYPE}. Type-aware consumers should treat this label as + * an unknown type rather than as a distinct one, since it carries no information about + * what kind of entity was found. + */ + public static final String UNTYPED = NameSample.DEFAULT_TYPE; + + private final TokenNameFinder finder; + + /** + * Initializes the adapter. + * + * @param finder The name finder to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code finder} is {@code null}. + */ + public NameFinderAnnotator(TokenNameFinder finder) { + if (finder == null) { + throw new IllegalArgumentException("finder must not be null"); + } + this.finder = finder; + } + + /** + * Finds names sentence by sentence and adds the {@link Layers#ENTITIES} layer. + * + *

    For every sentence, the tokens whose spans lie inside the sentence span are + * passed to the finder as one sequence, and each sentence-local mention is mapped + * through the sentence's first token position onto character spans of the original + * text. The required layers must be present, but they may be empty: a document + * without sentences or tokens yields a present-but-empty entity layer, and a sentence + * containing no tokens contributes nothing. A mention without a type is recorded with + * the type {@link #UNTYPED} as the annotation value.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#SENTENCES} and {@link Layers#TOKENS} layers, with + * every token lying inside a sentence. + * @return A new {@link Document} with the {@link Layers#ENTITIES} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, the + * sentence layer or the token layer is absent, a token lies outside every + * sentence, or the finder returns a mention that is empty or whose token + * indices lie outside its sentence's tokens. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.SENTENCES, Layers.TOKENS); + final List> sentences = document.get(Layers.SENTENCES); + final List> tokens = document.get(Layers.TOKENS); + final List> entities = new ArrayList<>(); + // The adaptive data is cleared even when annotation fails, so a rejected document + // cannot leak finder state into the next one. + try { + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + // The finder indexes within the sentence; shifting by the sentence's first + // token position turns every mention boundary into a document-wide token + // index, whose token spans already refer to the original text. An empty + // mention is rejected with the out-of-bounds ones: it covers no token, so it + // has no character span. + for (final Span mention : finder.find(words)) { + if (mention.getStart() < 0 || mention.getEnd() > words.length + || mention.getStart() >= mention.getEnd()) { + throw new IllegalArgumentException("finder returned mention " + mention + + " outside the sentence's " + words.length + " tokens"); + } + final int start = tokens.get(first + mention.getStart()).span().getStart(); + final int end = tokens.get(first + mention.getEnd() - 1).span().getEnd(); + final String type = mention.getType() == null ? UNTYPED : mention.getType(); + entities.add(new Annotation<>(new Span(start, end), type)); + } + }); + } finally { + finder.clearAdaptiveData(); + } + return document.with(Layers.ENTITIES, entities); + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.SENTENCES, Layers.TOKENS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(Layers.ENTITIES); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinder.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinder.java index b8239c28ee..c62841ebd2 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinder.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinder.java @@ -29,6 +29,9 @@ /** * A {@link TokenNameFinder} implementation based on a series of regular expressions. + *

    + * Please note: RegEx is not allowed in the processing path of OpenNLP. + * However, a RegEx step is an obvious exception. */ public final class RegexNameFinder implements TokenNameFinder { diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinderFactory.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinderFactory.java index 7c2bc79ecc..003d945125 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinderFactory.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/namefind/RegexNameFinderFactory.java @@ -26,6 +26,9 @@ * * Returns a {@link RegexNameFinder} based on a selection of * defaults or a configuration and a selection of defaults. + *

    + * Please note: RegEx is not allowed in the processing path of OpenNLP. + * However, a RegEx step is an obvious exception. */ public class RegexNameFinderFactory { @@ -107,12 +110,18 @@ public String getType() { @Override public Map getRegexMap() { Pattern[] p = new Pattern[1]; - p[0] = Pattern.compile("([a-z0-9!#$%&'*+/=?^_`{|}~-]+(?:\\.[a-z0-9!#$%&'*+/=?^_`{|}~-]+)*" + - "|\"([\\x01-\\x08\\x0b\\x0c\\x0e-\\x1f\\x21\\x23-\\x5b\\x5d-\\x7f]|\\\\[\\x01-\\x09" + - "\\x0b\\x0c\\x0e-\\x7f])*\")@(?:(?:[a-z0-9]([a-z0-9-]*[a-z0-9])?\\.)+[a-z0-9](?:[a-z0-9-]" + - "*[a-z0-9])?|\\[((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\.){3}(?:25[0-5]|2[0-4][0-9]|[01]" + - "?[0-9][0-9]?|[a-z0-9-]*[a-z0-9]:(?:[\\x01-\\x08\\x0b\\x0c\\x0e-\\x1f\\x21-\\x5a\\x53-\\x7f]" + - "|\\\\[\\x01-\\x09\\x0b\\x0c\\x0e-\\x7f])+)\\])", Pattern.CASE_INSENSITIVE); + // Every quantifier is bounded by a constant, which removes both the exponential + // backtracking and the recursion depth that made the old pattern a ReDoS vector. + // Limits follow RFC 5321: local part <= 64 chars, each domain label <= 63 chars. + p[0] = Pattern.compile( + "(? regexMap = new HashMap<>(); regexMap.put(getType(), p); return regexMap; @@ -127,16 +136,19 @@ public String getType() { @Override public Map getRegexMap() { Pattern[] p = new Pattern[1]; - p[0] = Pattern.compile("\\b(((ht|f)tp(s?)\\:\\/\\/|~\\/|\\/)|www.)" - + "(\\w+:\\w+@)?(([-\\w]+\\.)+(com|org|net|gov" + // Flattened to single-level groups and every quantifier bounded by a constant. + // The bounds alone remove the nested-quantifier backtracking and the recursion + // depth that caused the StackOverflowError; they must stay bounded ({0,255} / + // {1,63} rather than * / +) or a long path segment reintroduces the overflow. + p[0] = Pattern.compile("\\b(?:(?:ht|f)tps?://|~/|/|www\\.)" + + "(?:\\w{1,63}:\\w{1,63}@)?" + + "(?:[-\\w]{1,63}\\.){1,20}(?:com|org|net|gov" + "|mil|biz|info|mobi|name|aero|jobs|museum" - + "|travel|[a-z]{2}))(:[\\d]{1,5})?" - + "(((\\/([-\\w~!$+|.,=]|%[a-f\\d]{2})+)+|\\/)+|\\?|#)?" - + "((\\?([-\\w~!$+|.,*:]|%[a-f\\d{2}])+=?" - + "([-\\w~!$+|.,*:=]|%[a-f\\d]{2})*)" - + "(&(?:[-\\w~!$+|.,*:]|%[a-f\\d{2}])+=?" - + "([-\\w~!$+|.,*:=]|%[a-f\\d]{2})*)*)*" - + "(#([-\\w~!$+|.,*:=]|%[a-f\\d]{2})*)?\\b", Pattern.CASE_INSENSITIVE); + + "|travel|[a-z]{2})(?::\\d{1,5})?" + + "(?:/(?:[-\\w~!$+|.,=]|%[a-f\\d]{2}){0,255}){0,50}" + + "(?:\\?(?:[-\\w~!$+|.,*:=]|%[a-f\\d]{2}){0,255}" + + "(?:&(?:[-\\w~!$+|.,*:=]|%[a-f\\d]{2}){0,255}){0,50})?" + + "(?:#(?:[-\\w~!$+|.,*:=]|%[a-f\\d]{2}){0,255})?\\b", Pattern.CASE_INSENSITIVE); Map regexMap = new HashMap<>(); regexMap.put(getType(), p); return regexMap; diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/parser/ParserAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/parser/ParserAnnotator.java new file mode 100644 index 0000000000..e38b2fdb22 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/parser/ParserAnnotator.java @@ -0,0 +1,208 @@ +/* + * 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.parser; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; +import opennlp.tools.util.StringUtil; + +/** + * Adapts a constituency {@link Parser} to the document pipeline: reads + * {@link Layers#SENTENCES} and {@link Layers#TOKENS} and provides {@link #PHRASES}, one + * annotation per phrase node of each sentence's parse, carrying the phrase label and the + * span of its head token. + * + *

    Each sentence is parsed from its tokens as one sequence. Every node above the + * part-of-speech level except the root becomes an annotation on the span from its first + * to its last token, in pre-order, so an enclosing phrase precedes the phrases it + * contains and phrases nest by span containment. Part-of-speech nodes are left to the + * {@link Layers#POS_TAGS} layer and token nodes to {@link Layers#TOKENS}. The head token + * is the one the parser's head rules select, so a consumer can read the head of a noun + * phrase without its own rules.

    + * + *

    The adapter holds no per-call state; it is as thread-safe as the parser it + * wraps.

    + * + * @since 3.0.0 + */ +public final class ParserAnnotator implements DocumentAnnotator { + + /** + * One phrase of a constituency parse: its label, such as {@code NP} or {@code VP}, + * and the span of the token that heads it. The phrase's own span is the annotation's + * span. + * + * @param label The phrase label. Must not be {@code null} or blank. + * @param head The span of the head token in the document text. Must not be + * {@code null}. + * + * @since 3.0.0 + */ + public record Phrase(String label, Span head) { + + /** + * Validates the phrase. + * + * @throws IllegalArgumentException Thrown if {@code label} is {@code null} or + * blank, or {@code head} is {@code null}. + */ + public Phrase { + if (label == null || StringUtil.isBlank(label)) { + throw new IllegalArgumentException("label must not be null or blank"); + } + if (head == null) { + throw new IllegalArgumentException("head must not be null"); + } + } + } + + /** + * Parse phrases; each annotation covers one phrase and carries its {@link Phrase}, + * in pre-order of the parse tree. + */ + public static final LayerKey PHRASES = Layers.key("phrases", Phrase.class); + + private final Parser parser; + + /** + * Initializes the adapter. + * + * @param parser The parser to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code parser} is {@code null}. + */ + public ParserAnnotator(Parser parser) { + if (parser == null) { + throw new IllegalArgumentException("parser must not be null"); + } + this.parser = parser; + } + + /** + * Parses the document sentence by sentence and adds the {@link #PHRASES} layer. + * + *

    The required layers must be present, but they may be empty: a document without + * sentences or tokens yields a present-but-empty phrase layer.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#SENTENCES} and {@link Layers#TOKENS} layers, with + * every token lying inside a sentence. + * @return A new {@link Document} with the {@link #PHRASES} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, a + * required layer is absent, a token lies outside every sentence, or the + * parser returns a node outside the sentence's tokens. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.SENTENCES, Layers.TOKENS); + final List> sentences = document.get(Layers.SENTENCES); + final List> tokens = document.get(Layers.TOKENS); + final List> phrases = new ArrayList<>(); + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + final Parse root = parser.parse(Parse.createFromTokens(words)); + if (root == null) { + throw new IllegalArgumentException("parser returned no parse"); + } + // The parse text is the tokens joined by single spaces, so a token's start in + // that text identifies its index. + final int[] starts = new int[words.length]; + for (int i = 1; i < words.length; i++) { + starts[i] = starts[i - 1] + words[i - 1].length() + 1; + } + final int length = starts[words.length - 1] + words[words.length - 1].length(); + for (final Parse child : root.getChildren()) { + collect(child, first, starts, length, tokens, phrases); + } + }); + return document.with(PHRASES, phrases); + } + + /** Emits a node and, in pre-order, every phrase node below it. */ + private void collect(Parse node, int first, int[] starts, int length, + List> tokens, List> phrases) { + if (node.isPosTag() || Parser.TOK_NODE.equals(node.getType())) { + return; + } + final int from = tokenIndex(starts, length, node.getSpan().getStart(), node); + final int to = tokenIndex(starts, length, node.getSpan().getEnd(), node); + final int head = node.getHeadIndex(); + if (head < 0 || head >= starts.length) { + throw new IllegalArgumentException("parser returned node " + node.getType() + + " with head " + head + " outside the sentence's " + starts.length + " tokens"); + } + phrases.add(new Annotation<>(new Span(tokens.get(first + from).span().getStart(), + tokens.get(first + to).span().getEnd()), + new Phrase(node.getType(), tokens.get(first + head).span()))); + for (final Parse child : node.getChildren()) { + collect(child, first, starts, length, tokens, phrases); + } + } + + /** + * Maps an offset in the parse text, the tokens joined by single spaces, to the index + * of the token it lies in or, for a span end, ends. + */ + private int tokenIndex(int[] starts, int length, int offset, Parse node) { + if (offset < 0 || offset > length) { + throw new IllegalArgumentException("parser returned node " + node.getType() + + " at " + node.getSpan() + " outside the sentence's " + starts.length + + " tokens"); + } + int low = 0; + int high = starts.length - 1; + while (low < high) { + final int mid = (low + high + 1) >>> 1; + if (starts[mid] <= offset) { + low = mid; + } else { + high = mid - 1; + } + } + return low; + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.SENTENCES, Layers.TOKENS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(PHRASES); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/postag/POSTaggerAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/postag/POSTaggerAnnotator.java new file mode 100644 index 0000000000..feb34b6aa4 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/postag/POSTaggerAnnotator.java @@ -0,0 +1,118 @@ +/* + * 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.postag; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; + +/** + * Adapts a {@link POSTagger} to the document pipeline: reads {@link Layers#SENTENCES} + * and {@link Layers#TOKENS} and provides {@link Layers#POS_TAGS}, one tag annotation per + * token on the token's span. + * + *

    Each sentence is tagged separately, the way the tagger contract expects its input, + * so tagging decisions never cross a sentence boundary. Token spans already refer to the + * original document text, so only the token sequence handed to the tagger is sliced per + * sentence; the produced tag layer stays aligned with {@link Layers#TOKENS} by + * position.

    + * + * @since 3.0.0 + */ +public final class POSTaggerAnnotator implements DocumentAnnotator { + + private final POSTagger tagger; + + /** + * Initializes the adapter. + * + * @param tagger The tagger to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code tagger} is {@code null}. + */ + public POSTaggerAnnotator(POSTagger tagger) { + if (tagger == null) { + throw new IllegalArgumentException("tagger must not be null"); + } + this.tagger = tagger; + } + + /** + * Tags the document sentence by sentence and adds the {@link Layers#POS_TAGS} layer. + * + *

    For every sentence, the tokens whose spans lie inside the sentence span are + * tagged as one sequence, and each tag is emitted on its token's span. The required + * layers must be present, but they may be empty: a document without sentences or + * tokens yields a present-but-empty tag layer, and a sentence containing no tokens + * contributes nothing.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#SENTENCES} and {@link Layers#TOKENS} layers, with + * every token lying inside a sentence. + * @return A new {@link Document} with the {@link Layers#POS_TAGS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, the + * sentence layer or the token layer is absent, a token lies outside every + * sentence, or the tagger does not return one tag per token of a sentence. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.SENTENCES, Layers.TOKENS); + final List> sentences = document.get(Layers.SENTENCES); + final List> tokens = document.get(Layers.TOKENS); + final List> tagAnnotations = new ArrayList<>(tokens.size()); + DocumentAnnotators.forEachSentence(sentences, tokens, (first, words) -> { + final String[] tags = tagger.tag(words); + if (tags.length != words.length) { + throw new IllegalArgumentException( + "tagger returned " + tags.length + " tags for " + words.length + " tokens"); + } + for (int i = 0; i < words.length; i++) { + tagAnnotations.add(new Annotation<>(tokens.get(first + i).span(), tags[i])); + } + }); + return document.with(Layers.POS_TAGS, tagAnnotations); + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.SENTENCES, Layers.TOKENS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(Layers.POS_TAGS); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorAnnotator.java new file mode 100644 index 0000000000..0586daedc3 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorAnnotator.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.tools.sentdetect; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +/** + * Adapts a {@link SentenceDetector} to the document pipeline: provides + * {@link Layers#SENTENCES} from the document text. + * + *

    The wrapped detector stays the primary API for single-task use; this adapter calls + * it like any other caller would.

    + * + * @since 3.0.0 + */ +public final class SentenceDetectorAnnotator implements DocumentAnnotator { + + private final SentenceDetector detector; + + /** + * Initializes the adapter. + * + * @param detector The sentence detector to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code detector} is {@code null}. + */ + public SentenceDetectorAnnotator(SentenceDetector detector) { + if (detector == null) { + throw new IllegalArgumentException("detector must not be null"); + } + this.detector = detector; + } + + /** + * Detects sentences over the document text and adds the {@link Layers#SENTENCES} + * layer, each sentence annotated with its covered text on its span. + * + * @param document The document to annotate. Must not be {@code null}. + * @return A new {@link Document} with the {@link Layers#SENTENCES} layer added. + * Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null} or + * already carries the {@link Layers#SENTENCES} layer. + */ + @Override + public Document annotate(Document document) { + if (document == null) { + throw new IllegalArgumentException("document must not be null"); + } + final CharSequence text = document.text(); + final List> sentences = new ArrayList<>(); + for (final Span span : detector.sentPosDetect(text)) { + sentences.add(new Annotation<>(span, span.getCoveredText(text).toString())); + } + return document.with(Layers.SENTENCES, sentences); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(Layers.SENTENCES); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorME.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorME.java index 7be0d10b1c..9e2109b313 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorME.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/sentdetect/SentenceDetectorME.java @@ -21,9 +21,12 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.SortedSet; +import java.util.TreeSet; import opennlp.tools.commons.ThreadSafe; import opennlp.tools.dictionary.Dictionary; @@ -94,9 +97,10 @@ private static final class SentenceDetectorState { s -> s.sentProbs = new ArrayList<>()); /** - * The {@link Dictionary abbreviation dictionary} if available (may be {@code null}). + * The abbreviation dictionary index that backs {@link #isAcceptableBreak(CharSequence, int, int)}. + * It is {@code null} if no abbreviation dictionary is available for the underlying model. */ - private final Dictionary abbDict; + private final AbbreviationIndex abbIndex; protected final boolean useTokenEnd; @@ -124,10 +128,12 @@ public SentenceDetectorME(SentenceModel model) { * * @param model The {@link SentenceModel} to be used. * @param abbDict The {@link Dictionary} to be used. It must fit the language of the {@code model}. + * Its entries are read once, here; later changes to the {@code abbDict} instance do not + * affect this detector. */ public SentenceDetectorME(SentenceModel model, Dictionary abbDict) { this.model = model.getMaxentModel(); - this.abbDict = abbDict; + this.abbIndex = AbbreviationIndex.of(abbDict); SentenceDetectorFactory sdFactory = model.getFactory(); cgen = sdFactory.getSDContextGenerator(); scanner = sdFactory.getEndOfSentenceScanner(); @@ -152,7 +158,7 @@ public SentenceDetectorME(SentenceModel model, Factory factory) { getAbbreviations(model.getAbbreviations()), customEOSCharacters); scanner = factory.createEndOfSentenceScanner(customEOSCharacters); } - abbDict = model.getAbbreviations(); + abbIndex = AbbreviationIndex.of(model.getAbbreviations()); useTokenEnd = model.useTokenEnd(); } @@ -387,47 +393,192 @@ public double[] getSentenceProbabilities() { *

    Note: The implementation always returns {@code true} if no * abbreviation dictionary is available for the underlying model.

    * + *

    Only a bounded region of {@code s} is consulted. An abbreviation occurrence starting at + * {@code p} with length {@code L} can veto the break only when it starts at or before the + * candidate and reaches it, that is when + * {@code candidateIndex - L <= p <= candidateIndex}. Positions outside that window are + * irrelevant by construction, so the decision costs time proportional to the longest + * dictionary entry rather than to the length of {@code s}.

    + * + *

    A {@code candidateIndex} past the end of {@code s} cannot carry an abbreviation, so it + * is accepted rather than raising an {@link IndexOutOfBoundsException}. At + * {@code candidateIndex == s.length()}, the decision matches the previous implementation when + * the match is not at {@code fromIndex}; only the match-at-{@code fromIndex} case is an + * accept-instead-of-throw relaxation.

    + * * @param s the {@link CharSequence} in which the break occurred. * @param fromIndex the start of the segment currently being evaluated. - * @param candidateIndex the index of the candidate sentence ending. + * @param candidateIndex the index of the candidate sentence ending. Must be greater than or + * equal to {@code fromIndex} and a valid index into {@code s}. * @return {@code true} if the break is acceptable, {@code false} otherwise. */ protected boolean isAcceptableBreak(CharSequence s, int fromIndex, int candidateIndex) { - if (abbDict == null) - return true; - - final String text = s.toString(); - final boolean caseSensitive = abbDict.isCaseSensitive(); - final String searchText = caseSensitive ? text : StringUtil.toLowerCase(text); - for (StringList abb : abbDict) { - final String abbToken = caseSensitive ? abb.getToken(0) - : StringUtil.toLowerCase(abb.getToken(0)); - final int tokenLength = abbToken.length(); - int tokenStartPos = searchText.indexOf(abbToken, fromIndex); - while (tokenStartPos != -1) { - if (tokenStartPos > candidateIndex) { - break; // past candidate position, no point searching further - } - if (tokenStartPos == fromIndex - && searchText.substring(tokenStartPos, candidateIndex + 1).equals(abbToken)) { - return false; // full abbreviation match at segment start -> no acceptable break - } - final char prevChar = s.charAt(tokenStartPos == fromIndex ? tokenStartPos : tokenStartPos - 1); - if (tokenStartPos + tokenLength >= candidateIndex + return abbIndex == null || abbIndex.allowsBreak(s, fromIndex, candidateIndex); + } + + /** + * An immutable, length-bucketed index over an abbreviation {@link Dictionary}. It answers + * {@link #isAcceptableBreak(CharSequence, int, int)} by enumerating the few text positions + * that can carry a relevant abbreviation and asking a hash set what is there, instead of + * searching the whole text once per dictionary entry. + */ + private static final class AbbreviationIndex { + + /** + * The dictionary entries, folded to lower case unless the dictionary is case-sensitive. + * Only the first token of a multi-token entry participates, as before. + */ + private final Set entries; + + /** + * The distinct entry lengths, ascending. These are the only window sizes worth probing. + */ + private final int[] entryLengths; + + /** + * The longest entry length, that is how far past a candidate a relevant match can reach. + */ + private final int maxEntryLength; + + /** + * Whether the dictionary matches case-sensitively. + */ + private final boolean caseSensitive; + + /** + * @param abbDict The {@link Dictionary} to index, may be {@code null}. + * @return An index over {@code abbDict}, or {@code null} if {@code abbDict} is {@code null}. + */ + static AbbreviationIndex of(Dictionary abbDict) { + return abbDict == null ? null : new AbbreviationIndex(abbDict); + } + + /** + * Reads {@code abbDict} once, so this instance is immutable and safe to share. + * + * @param abbDict The {@link Dictionary} to index. Must not be {@code null}. + */ + private AbbreviationIndex(Dictionary abbDict) { + caseSensitive = abbDict.isCaseSensitive(); + final Set tokens = new HashSet<>(); + final SortedSet lengths = new TreeSet<>(); + for (StringList abb : abbDict) { + final String token = caseSensitive ? abb.getToken(0) + : StringUtil.toLowerCase(abb.getToken(0)); + tokens.add(token); + lengths.add(token.length()); + } + entries = tokens; + entryLengths = lengths.stream().mapToInt(Integer::intValue).toArray(); + maxEntryLength = entryLengths.length == 0 ? 0 : entryLengths[entryLengths.length - 1]; + } + + /** + * @param s The text in which the break occurred. + * @param fromIndex The start of the segment currently being evaluated. + * @param candidateIndex The index of the candidate sentence ending. + * @return {@code true} if a break at {@code candidateIndex} is allowed. + */ + boolean allowsBreak(CharSequence s, int fromIndex, int candidateIndex) { + final int textLength = s.length(); + if (entryLengths.length == 0 || candidateIndex < fromIndex || candidateIndex < 0 + || candidateIndex > textLength) { + return true; + } + // Occurrences before the segment start do not participate. + final int scanStart = StrictMath.max(0, fromIndex); + // The window holds every position a relevant occurrence can start at, plus the longest + // entry so that the text of an occurrence starting at the candidate is covered too. + final int windowStart = codePointStart(s, + StrictMath.max(scanStart, candidateIndex - maxEntryLength)); + final int windowEnd = codePointEnd(s, + StrictMath.min(textLength, candidateIndex + maxEntryLength)); + // Case folding is applied to the window only. It is per code point, so it preserves + // indices and yields exactly the characters a folding of the whole text would. + final String window = caseSensitive + ? s.subSequence(windowStart, windowEnd).toString() + : toLowerCase(s, windowStart, windowEnd); + + for (final int tokenLength : entryLengths) { + for (int pos = StrictMath.max(scanStart, candidateIndex - tokenLength); + pos <= candidateIndex; pos++) { + final int endPos = pos + tokenLength; + if (endPos > textLength) { + break; // the entry no longer fits, and it fits even less further right + } + if (!entries.contains(window.substring(pos - windowStart, endPos - windowStart))) { + continue; + } + if (pos == fromIndex && endPos == candidateIndex + 1) { + return false; // full abbreviation match at segment start -> no acceptable break + } + final char prevChar = s.charAt(pos == fromIndex ? pos : pos - 1); /* * Note: * Skip abbreviation candidate if regular characters exist directly before it, * That is, any letter or digit except: a whitespace, an apostrophe, or an opening round bracket. * This prevents mismatches from overlaps close to an actual sentence end. */ - && (Character.isWhitespace(prevChar) || isApostrophe(prevChar) || prevChar == '(')) { - return false; // in case of a valid abbreviation: the (sentence) break is not accepted + if (Character.isWhitespace(prevChar) || isApostrophe(prevChar) || prevChar == '(') { + return false; // in case of a valid abbreviation: the (sentence) break is not accepted + } } - // Try next occurrence of this abbreviation in the text - tokenStartPos = searchText.indexOf(abbToken, tokenStartPos + 1); } + return true; // no abbreviation(s) at given positions: valid sentence boundary + } + + /** + * Lower-cases {@code [from, to)} exactly as {@link StringUtil#toLowerCase(CharSequence)} + * lower-cases a whole text: per code point via {@link Character#toLowerCase(int)}. + * + *

    Folding only the window is correct only because that mapping is 1:1 in {@code char} + * count, so {@code pos - windowStart} taken from the unfolded text still indexes the folded + * window. {@link String#toLowerCase()} must not be used here: full case mapping can expand + * a code point (for example {@code İ} / U+0130) and would corrupt every offset in the window.

    + * + * @param s The text to read from. + * @param from The first index to fold, at a code point boundary. + * @param to The index to stop at, at a code point boundary. + * @return The folded characters of {@code [from, to)}. + */ + private static String toLowerCase(CharSequence s, int from, int to) { + final StringBuilder folded = new StringBuilder(to - from); + int i = from; + while (i < to) { + final int cp = Character.codePointAt(s, i); + folded.appendCodePoint(Character.toLowerCase(cp)); + i += Character.charCount(cp); + } + return folded.toString(); + } + + /** + * @param s The text {@code index} refers to. + * @param index The index to align. + * @return {@code index}, moved one character left if it points at the trailing half of a + * surrogate pair, which is the only index a code point walk cannot start at. + */ + private static int codePointStart(CharSequence s, int index) { + if (index > 0 && index < s.length() && Character.isLowSurrogate(s.charAt(index)) + && Character.isHighSurrogate(s.charAt(index - 1))) { + return index - 1; + } + return index; + } + + /** + * @param s The text {@code index} refers to. + * @param index The index to align. + * @return {@code index}, moved one character right if it splits a surrogate pair, which is + * the only index a code point walk cannot stop at. + */ + private static int codePointEnd(CharSequence s, int index) { + if (index > 0 && index < s.length() && Character.isHighSurrogate(s.charAt(index - 1)) + && Character.isLowSurrogate(s.charAt(index))) { + return index + 1; + } + return index; } - return true; // no abbreviation(s) at given positions: valid sentence boundary } /** diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/StemmerAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/StemmerAnnotator.java new file mode 100644 index 0000000000..900fe188a6 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/StemmerAnnotator.java @@ -0,0 +1,107 @@ +/* + * 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.stemmer; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.DocumentAnnotators; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; + +/** + * Adapts a {@link Stemmer} to the document pipeline: stems the token layer and provides + * {@link #STEMS}, one annotation per token on the token's span. + * + *

    Stemming operates on the token surface alone, so this annotator requires only the + * token layer; no part-of-speech tags are involved.

    + * + * @since 3.0.0 + */ +public final class StemmerAnnotator implements DocumentAnnotator { + + /** + * The stem layer. It is aligned with the token layer by position, and each annotation + * carries the stem of its token on that token's span. + */ + public static final LayerKey STEMS = Layers.key("stems", String.class); + + private final Stemmer stemmer; + + /** + * Initializes the adapter. + * + * @param stemmer The stemmer to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code stemmer} is {@code null}. + */ + public StemmerAnnotator(Stemmer stemmer) { + if (stemmer == null) { + throw new IllegalArgumentException("stemmer must not be null"); + } + this.stemmer = stemmer; + } + + /** + * Stems the token layer and adds the {@link #STEMS} layer. + * + *

    The token layer must be present, but it may be empty: a document without tokens + * yields a present-but-empty stem layer.

    + * + * @param document The document to annotate. Must not be {@code null} and must carry + * the {@link Layers#TOKENS} layer. + * @return A new {@link Document} with the {@link #STEMS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null} or the + * token layer is absent. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.TOKENS); + final List> tokens = document.get(Layers.TOKENS); + final List> layer = new ArrayList<>(tokens.size()); + for (final Annotation token : tokens) { + layer.add(new Annotation<>(token.span(), stemmer.stem(token.value()).toString())); + } + return document.with(STEMS, layer); + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.TOKENS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(STEMS); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/AffixCondition.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/AffixCondition.java new file mode 100644 index 0000000000..19ce657ccc --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/AffixCondition.java @@ -0,0 +1,165 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.stemmer.hunspell; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** + * One parsed affix condition: a fixed-length sequence of literal code points and + * bracketed character classes, matched with a single scan and no regular expressions. + * A suffix condition anchors at the end of the candidate stem, a prefix condition at + * its start; the condition {@code .} matches everything. Positions are Unicode code + * points so supplementary characters agree with {@code FLAG UTF-8} flag reading. + */ +final class AffixCondition { + + /** The shared instance for the condition {@code .}, which accepts every stem. */ + private static final AffixCondition ANY = new AffixCondition(new int[0][], null, true); + + /** Per position: the accepted code points, or {@code null} for any code point. */ + private final int[][] accepted; + /** Per position with a class: whether the class is negated; {@code null} rows unused. */ + private final boolean[] negated; + /** Whether the owning rule is a suffix rule, which anchors the condition at the end. */ + private final boolean suffix; + + /** + * Initializes the condition. + * + * @param accepted The accepted code points per position. + * @param negated The negation marker per position. + * @param suffix Whether the owning rule is a suffix rule. + */ + private AffixCondition(int[][] accepted, boolean[] negated, boolean suffix) { + this.accepted = accepted; + this.negated = negated; + this.suffix = suffix; + } + + /** + * Parses a condition field. Each pattern position is a literal code point, a + * {@code .} matching any code point, or a bracketed class such as {@code [sx]}; a + * class starting with {@code ^} is negated and matches any code point outside it. + * + * @param pattern The condition text from the affix rule. + * @param suffix Whether the owning rule is a suffix rule. + * @param lineNumber The affix file line, for error messages. + * @return The parsed condition. Never {@code null}. + * @throws IOException Thrown if a character class is unterminated. + */ + static AffixCondition parse(String pattern, boolean suffix, int lineNumber) + throws IOException { + if (".".equals(pattern)) { + return ANY; + } + final List positions = new ArrayList<>(); + final List negations = new ArrayList<>(); + int i = 0; + while (i < pattern.length()) { + final int codePoint = pattern.codePointAt(i); + if (codePoint == '[') { + final int end = pattern.indexOf(']', i + 1); + if (end < 0) { + throw new IOException("unterminated character class at line " + lineNumber); + } + String members = pattern.substring(i + 1, end); + boolean negate = false; + if (members.startsWith("^")) { + negate = true; + members = members.substring(1); + } + positions.add(toCodePoints(members)); + negations.add(negate); + i = end + 1; + } else if (codePoint == '.') { + positions.add(null); + negations.add(false); + i++; + } else { + positions.add(new int[] {codePoint}); + negations.add(false); + i += Character.charCount(codePoint); + } + } + final int[][] accepted = positions.toArray(new int[0][]); + final boolean[] negated = new boolean[accepted.length]; + for (int p = 0; p < negated.length; p++) { + negated[p] = negations.get(p); + } + return new AffixCondition(accepted, negated, suffix); + } + + /** + * Collects the code points of a character-class body. + * + * @param members The class body text. + * @return The code points in order. Never {@code null}. + */ + private static int[] toCodePoints(String members) { + final int[] codePoints = new int[members.codePointCount(0, members.length())]; + int i = 0; + int out = 0; + while (i < members.length()) { + final int codePoint = members.codePointAt(i); + codePoints[out++] = codePoint; + i += Character.charCount(codePoint); + } + return codePoints; + } + + /** + * Tests a candidate stem against the condition at its anchored side: the last + * positions of the stem for a suffix condition, the first positions for a prefix + * condition. A stem shorter than the condition never matches. Length is in code + * points. + * + * @param stem The candidate stem after affix removal and strip restoration. + * @return {@code true} if the stem satisfies the condition. + */ + boolean matches(String stem) { + if (accepted.length == 0) { + return true; + } + final int stemPoints = stem.codePointCount(0, stem.length()); + if (stemPoints < accepted.length) { + return false; + } + int offset = suffix ? stem.offsetByCodePoints(0, stemPoints - accepted.length) : 0; + for (int p = 0; p < accepted.length; p++) { + final int[] members = accepted[p]; + final int codePoint = stem.codePointAt(offset); + offset += Character.charCount(codePoint); + if (members == null) { + continue; + } + boolean member = false; + for (final int candidate : members) { + if (candidate == codePoint) { + member = true; + break; + } + } + if (member == negated[p]) { + return false; + } + } + return true; + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionary.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionary.java new file mode 100644 index 0000000000..2bb06935d5 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionary.java @@ -0,0 +1,1675 @@ +/* + * 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.stemmer.hunspell; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.Charset; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.IllegalCharsetNameException; +import java.nio.charset.StandardCharsets; +import java.nio.charset.UnsupportedCharsetException; +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 opennlp.tools.commons.ThreadSafe; +import opennlp.tools.util.StringUtil; + +/** + * An immutable, in-memory Hunspell-format dictionary loaded from user-supplied + * {@code .aff} and {@code .dic} files. OpenNLP includes no dictionary data. + * + *

    Supported affix features are {@code PFX} and {@code SFX} rules with strip + * strings, character-class conditions, cross-product combinations, and a double suffix + * connected by continuation classes; {@code FLAG} modes {@code char}, {@code UTF-8}, + * {@code long}, and {@code num}; the {@code AF} alias table; and the {@code SET} + * encoding declaration. Compound decomposition supports {@code COMPOUNDFLAG}, + * {@code COMPOUNDBEGIN}, {@code COMPOUNDMIDDLE}, {@code COMPOUNDEND}, + * {@code COMPOUNDMIN}, {@code COMPOUNDWORDMAX}, {@code COMPOUNDPERMITFLAG}, + * {@code COMPOUNDFORBIDFLAG}, {@code CHECKCOMPOUNDDUP}, + * {@code CHECKCOMPOUNDCASE}, and {@code CHECKCOMPOUNDTRIPLE}. The blocking flags + * {@code NEEDAFFIX} (also named {@code PSEUDOROOT}), {@code ONLYINCOMPOUND}, and + * {@code FORBIDDENWORD}, plus {@code CIRCUMFIX} and {@code FULLSTRIP}, are also + * applied.

    + * + *

    Other directives are skipped. Their spelling, conversion, suggestion, or + * advanced compound behavior is not applied by this affix stemmer. Dictionary + * morphology fields are also ignored.

    + * + *

    Instances are immutable and safe to share between threads.

    + * + * @see HunspellStemmer + * @see HunspellStemmerFactory + * @since 3.0.0 + */ +@ThreadSafe +public final class HunspellDictionary { + + /** The file suffix for a Hunspell affix file. */ + public static final String AFFIX_FILE_SUFFIX = ".aff"; + + /** The file suffix for a Hunspell word-list file. */ + public static final String DICTIONARY_FILE_SUFFIX = ".dic"; + + /** + * Inclusive upper bound on bytes buffered from one affix or dictionary stream + * during {@link #load(InputStream, InputStream)}. The limit is 64 MiB. Larger + * streams fail with {@link IOException}. + */ + public static final int MAX_STREAM_BYTES = 64 * 1024 * 1024; + + /** + * One parsed affix rule of a {@code PFX} or {@code SFX} block. + * + * @param flag The flag naming the rule's block, which an entry carries to accept it. + * @param crossProduct Whether the rule may combine with an affix of the opposite kind. + * @param strip The stem material the rule replaces, restored during analysis. + * @param affix The surface material the rule adds to the stem. + * @param condition The condition the stem must satisfy for the rule to apply. + * @param continuation The flags of the further affixes that may stack on this one. + */ + record Affix(int flag, boolean crossProduct, String strip, String affix, + AffixCondition condition, int[] continuation) { + + /** + * Checks whether a further affix may stack on this one. + * + * @param otherFlag The stacking affix's flag. + * @return {@code true} if this affix's continuation classes allow it. + */ + boolean allowsContinuation(int otherFlag) { + for (final int candidate : continuation) { + if (candidate == otherFlag) { + return true; + } + } + return false; + } + } + + /** The place a part takes in a compound, deciding which positional flag admits it. */ + enum CompoundPosition { + /** The first part. */ + BEGIN, + /** Any part between the first and the last. */ + MIDDLE, + /** The last part. */ + END + } + + /** The shared empty bucket answered for characters no affix rule is keyed under. */ + private static final List NO_AFFIXES = List.of(); + + /** The line tag of a prefix block and of every rule line inside it. */ + private static final String PREFIX_TAG = "PFX"; + + /** The line tag of a suffix block and of every rule line inside it. */ + private static final String SUFFIX_TAG = "SFX"; + + /** The directive that selects the file-wide flag representation. */ + private static final String FLAG_TAG = "FLAG"; + + /** The directive that defines the file-wide flag alias table. */ + private static final String ALIAS_TAG = "AF"; + + /** Prefix used by comment lines. */ + private static final String COMMENT_PREFIX = "#"; + + /** The affix-file directive that declares the character encoding. */ + private static final String SET_TAG = "SET"; + + /** The {@code SET} directive followed by a space. */ + private static final String SET_PREFIX = SET_TAG + " "; + + /** The {@code SET} directive followed by a tab. */ + private static final String SET_TAB_PREFIX = SET_TAG + "\t"; + + /** The affix format's marker for absent strip or affix material. */ + private static final String NO_MATERIAL = "0"; + + /** + * Largest flag value permitted by {@code FLAG num}, as specified by the + * + * Hunspell format manual. + */ + private static final int MAX_NUMERIC_FLAG = 65_000; + + /** Largest {@code COMPOUNDMIN} value that can be doubled without overflow. */ + private static final int MAX_COMPOUND_MIN = Integer.MAX_VALUE / 2; + + private final Map> entries; + private final BoundaryIndex suffixesByLast; + private final List suffixesWithoutMaterial; + private final BoundaryIndex prefixesByFirst; + private final List prefixesWithoutMaterial; + private final int compoundFlag; + private final int compoundBegin; + private final int compoundEnd; + private final int compoundMin; + private final int needAffix; + private final int onlyInCompound; + private final int forbiddenWord; + private final int circumfix; + private final int compoundMiddle; + private final int compoundPermit; + private final int compoundForbid; + private final int compoundWordMax; + private final boolean checkCompoundDup; + private final boolean checkCompoundCase; + private final boolean checkCompoundTriple; + private final boolean fullStrip; + + /** + * Initializes the dictionary from the two parsed files. + * + * @param entries The words mapped to the flag sets of their entries. + * @param affix The parsed affix file. + */ + private HunspellDictionary(Map> entries, AffixFile affix) { + this.compoundFlag = affix.compoundFlag; + this.compoundBegin = affix.compoundBegin; + this.compoundEnd = affix.compoundEnd; + this.compoundMin = affix.compoundMin; + this.needAffix = affix.needAffix; + this.onlyInCompound = affix.onlyInCompound; + this.forbiddenWord = affix.forbiddenWord; + this.circumfix = affix.circumfix; + this.compoundMiddle = affix.compoundMiddle; + this.compoundPermit = affix.compoundPermit; + this.compoundForbid = affix.compoundForbid; + this.compoundWordMax = affix.compoundWordMax; + this.checkCompoundDup = affix.checkCompoundDup; + this.checkCompoundCase = affix.checkCompoundCase; + this.checkCompoundTriple = affix.checkCompoundTriple; + this.fullStrip = affix.fullStrip; + this.entries = entries; + // A material-bearing rule can only be undone from a word whose boundary + // character matches its affix material, so bucketing by that character + // narrows each scan to one bucket plus the strip-only rules. + final List suffixesWithout = new ArrayList<>(); + this.suffixesByLast = bucketByBoundary(affix.suffixes, true, suffixesWithout); + this.suffixesWithoutMaterial = List.copyOf(suffixesWithout); + final List prefixesWithout = new ArrayList<>(); + this.prefixesByFirst = bucketByBoundary(affix.prefixes, false, prefixesWithout); + this.prefixesWithoutMaterial = List.copyOf(prefixesWithout); + } + + /** + * An immutable index of affix rules keyed by the boundary code point of their affix + * material, answering each lookup by binary search so the per-word scans in + * {@link HunspellStemmer} allocate nothing. + */ + private static final class BoundaryIndex { + + /** The boundary code points, sorted ascending. */ + private final int[] boundaries; + /** The rule bucket for each boundary, aligned with {@link #boundaries}. */ + private final List> buckets; + + /** + * Initializes the index from mutable buckets, freezing each one. + * + * @param byBoundary The rule buckets keyed by boundary code point. + */ + private BoundaryIndex(Map> byBoundary) { + this.boundaries = new int[byBoundary.size()]; + int b = 0; + for (final Integer boundary : byBoundary.keySet()) { + boundaries[b++] = boundary; + } + Arrays.sort(boundaries); + this.buckets = new ArrayList<>(boundaries.length); + for (final int boundary : boundaries) { + buckets.add(List.copyOf(byBoundary.get(boundary))); + } + } + + /** + * The rules bucketed under a boundary code point. + * + * @param codePoint The boundary code point to look up. + * @return The bucket, possibly empty. Never {@code null}. + */ + List bucket(int codePoint) { + final int index = Arrays.binarySearch(boundaries, codePoint); + return index >= 0 ? buckets.get(index) : NO_AFFIXES; + } + } + + /** + * Buckets affix rules by the boundary code point of their affix material, the last + * code point for a suffix rule and the first for a prefix rule. + * + * @param rules The rules of one kind, in file order. + * @param suffix Whether the rules are suffix rules. + * @param withoutMaterial Collects the rules with empty affix material, which no + * boundary code point keys. + * @return The rules indexed by their boundary code point. Never {@code null}. + */ + private static BoundaryIndex bucketByBoundary(List rules, + boolean suffix, List withoutMaterial) { + final Map> byBoundary = new HashMap<>(); + for (final Affix rule : rules) { + final String material = rule.affix(); + if (material.isEmpty()) { + withoutMaterial.add(rule); + } else { + final int boundary = suffix + ? material.codePointBefore(material.length()) + : material.codePointAt(0); + byBoundary.computeIfAbsent(boundary, key -> new ArrayList<>()).add(rule); + } + } + return new BoundaryIndex(byBoundary); + } + + /** + * Loads a dictionary from its two files. + * + * @param affixFile The {@code .aff} affix file. Must not be {@code null}. + * @param dictionaryFile The {@code .dic} word list. Must not be {@code null}. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if reading fails or a file is malformed. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static HunspellDictionary load(Path affixFile, Path dictionaryFile) + throws IOException { + if (affixFile == null) { + throw new IllegalArgumentException("affixFile must not be null"); + } + if (dictionaryFile == null) { + throw new IllegalArgumentException("dictionaryFile must not be null"); + } + try (InputStream affix = Files.newInputStream(affixFile); + InputStream dictionary = Files.newInputStream(dictionaryFile)) { + return load(affix, dictionary); + } + } + + /** + * Loads a dictionary from its two streams. Each stream is buffered up to + * {@link #MAX_STREAM_BYTES} bytes; a larger stream fails with {@link IOException}. + * + * @param affixStream The {@code .aff} affix content. Must not be {@code null}. Not + * closed. + * @param dictionaryStream The {@code .dic} word list content. Must not be + * {@code null}. Not closed. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if reading fails, a stream exceeds + * {@link #MAX_STREAM_BYTES}, or the content is malformed. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static HunspellDictionary load(InputStream affixStream, + InputStream dictionaryStream) throws IOException { + if (affixStream == null) { + throw new IllegalArgumentException("affixStream must not be null"); + } + if (dictionaryStream == null) { + throw new IllegalArgumentException("dictionaryStream must not be null"); + } + byte[] affixBytes = readBounded(affixStream, MAX_STREAM_BYTES, "affix stream"); + final Charset charset = declaredCharset(affixBytes); + maskIgnoredAffixLines(affixBytes); + final boolean rawUtf8Flags = StandardCharsets.UTF_8.equals(charset) + && !usesUnicodeOrNumericFlags(affixBytes); + affixBytes = normalizeUtf8ByteFlags(affixBytes, charset); + final AffixFile affix = parseAffix(decode(affixBytes, charset, "affix stream")); + byte[] dictionaryBytes = readBounded(dictionaryStream, MAX_STREAM_BYTES, + "dictionary stream"); + if (rawUtf8Flags) { + dictionaryBytes = normalizeDictionaryByteFlags(dictionaryBytes); + } + final Map> entries = parseWordList( + decode(dictionaryBytes, charset, "dictionary stream"), + affix.flagMode, affix.flagAliases); + return new HunspellDictionary(entries, affix); + } + + /** + * Replaces comments and unused directive lines with ASCII spaces before strict + * decoding. Published dictionaries sometimes retain legacy-encoded metadata despite + * a {@code SET UTF-8} declaration. Line endings and byte positions remain unchanged, + * while malformed bytes in parsed directives are still reported. + * + * @param bytes The buffered affix file, modified in place. + */ + private static void maskIgnoredAffixLines(byte[] bytes) { + int lineStart = 0; + for (int i = 0; i <= bytes.length; i++) { + if (i == bytes.length || bytes[i] == '\n' || bytes[i] == '\r') { + int fieldStart = lineStart; + while (fieldStart < i && isAsciiFieldSpace(bytes[fieldStart])) { + fieldStart++; + } + int fieldEnd = fieldStart; + while (fieldEnd < i && !isAsciiFieldSpace(bytes[fieldEnd])) { + fieldEnd++; + } + if (fieldStart < fieldEnd && bytes[fieldStart] != '#') { + final String directive = new String(bytes, fieldStart, + fieldEnd - fieldStart, StandardCharsets.US_ASCII); + if (isParsedAffixDirective(directive)) { + maskInlineComment(bytes, fieldEnd, i); + lineStart = i + 1; + continue; + } + } + Arrays.fill(bytes, lineStart, i, (byte) ' '); + lineStart = i + 1; + } + } + } + + /** {@return whether a byte separates fields in an affix line} */ + private static boolean isAsciiFieldSpace(byte value) { + return value == ' ' || value == '\t' || value == '\f'; + } + + /** Replaces an inline comment that starts after an affix field separator. */ + private static void maskInlineComment(byte[] bytes, int from, int to) { + boolean fieldStart = false; + for (int i = from; i < to; i++) { + if (isAsciiFieldSpace(bytes[i])) { + fieldStart = true; + } else if (fieldStart && bytes[i] == '#') { + Arrays.fill(bytes, i, to, (byte) ' '); + return; + } else { + fieldStart = false; + } + } + } + + /** {@return whether this implementation parses a directive's fields} */ + private static boolean isParsedAffixDirective(String directive) { + return switch (directive) { + case SET_TAG, FLAG_TAG, ALIAS_TAG, PREFIX_TAG, SUFFIX_TAG, + "COMPOUNDFLAG", "COMPOUNDBEGIN", "COMPOUNDMIDDLE", "COMPOUNDEND", + "COMPOUNDPERMITFLAG", "COMPOUNDFORBIDFLAG", "NEEDAFFIX", "PSEUDOROOT", + "ONLYINCOMPOUND", "FORBIDDENWORD", "CIRCUMFIX", "COMPOUNDMIN", + "COMPOUNDWORDMAX", "CHECKCOMPOUNDDUP", "CHECKCOMPOUNDCASE", + "CHECKCOMPOUNDTRIPLE", "FULLSTRIP" -> true; + default -> false; + }; + } + + /** + * Converts raw one-byte flags in a UTF-8 affix file to equivalent Unicode code + * points before decoding. Hunspell's default and {@code long} flag modes operate on + * bytes, and published UTF-8 dictionaries can therefore contain non-UTF-8 bytes in + * flag fields. Text, conditions, and affix material remain subject to strict UTF-8 + * decoding. + * + * @param bytes The affix file after unused lines have been masked. + * @param charset The encoding selected by {@code SET}. + * @return The content with raw flag bytes represented as valid UTF-8. + */ + private static byte[] normalizeUtf8ByteFlags(byte[] bytes, Charset charset) { + if (!StandardCharsets.UTF_8.equals(charset) || usesUnicodeOrNumericFlags(bytes)) { + return bytes; + } + final ByteArrayOutputStream normalized = new ByteArrayOutputStream(bytes.length); + int lineStart = 0; + for (int i = 0; i <= bytes.length; i++) { + if (i == bytes.length || bytes[i] == '\n' || bytes[i] == '\r') { + writeNormalizedFlagLine(normalized, bytes, lineStart, i); + if (i < bytes.length) { + normalized.write(bytes[i]); + } + lineStart = i + 1; + } + } + return normalized.toByteArray(); + } + + /** {@return whether {@code FLAG UTF-8} or {@code FLAG num} selects non-byte flags} */ + private static boolean usesUnicodeOrNumericFlags(byte[] bytes) { + final int[] starts = new int[5]; + final int[] fieldEnds = new int[5]; + int lineStart = 0; + for (int i = 0; i <= bytes.length; i++) { + if (i == bytes.length || bytes[i] == '\n' || bytes[i] == '\r') { + final int count = findAsciiFields(bytes, lineStart, i, starts, fieldEnds); + if (count >= 2 && FLAG_TAG.equals( + asciiField(bytes, starts[0], fieldEnds[0]))) { + final String mode = asciiField(bytes, starts[1], fieldEnds[1]); + return "UTF-8".equals(mode) || "num".equals(mode); + } + lineStart = i + 1; + } + } + return false; + } + + /** Writes one affix line, converting high bytes only within raw flag fields. */ + private static void writeNormalizedFlagLine(ByteArrayOutputStream target, + byte[] bytes, int lineStart, int lineEnd) { + final int[] starts = new int[5]; + final int[] fieldEnds = new int[5]; + final int count = findAsciiFields(bytes, lineStart, lineEnd, starts, fieldEnds); + int firstFlagStart = -1; + int firstFlagEnd = -1; + int continuationStart = -1; + int continuationEnd = -1; + if (count >= 2) { + final String directive = asciiField(bytes, starts[0], fieldEnds[0]); + if (ALIAS_TAG.equals(directive) || PREFIX_TAG.equals(directive) + || SUFFIX_TAG.equals(directive) || isSingleFlagDirective(directive)) { + firstFlagStart = starts[1]; + firstFlagEnd = fieldEnds[1]; + } + if (count >= 4 && (PREFIX_TAG.equals(directive) || SUFFIX_TAG.equals(directive))) { + for (int i = starts[3]; i < fieldEnds[3]; i++) { + if (bytes[i] == '/') { + continuationStart = i + 1; + continuationEnd = fieldEnds[3]; + break; + } + } + } + } + for (int i = lineStart; i < lineEnd; i++) { + final boolean flagByte = i >= firstFlagStart && i < firstFlagEnd + || i >= continuationStart && i < continuationEnd; + writeNormalizedByte(target, bytes[i], flagByte); + } + } + + /** + * Converts raw flag bytes after the flag separator of each dictionary entry. + * Word text and morphology fields remain subject to strict UTF-8 decoding. + * + * @param bytes The buffered dictionary file. + * @return The content with raw flag bytes represented as valid UTF-8. + */ + private static byte[] normalizeDictionaryByteFlags(byte[] bytes) { + final ByteArrayOutputStream normalized = new ByteArrayOutputStream(bytes.length); + int lineStart = 0; + for (int i = 0; i <= bytes.length; i++) { + if (i == bytes.length || bytes[i] == '\n' || bytes[i] == '\r') { + int flagStart = -1; + int flagEnd = -1; + for (int cursor = lineStart + 1; cursor < i; cursor++) { + if (bytes[cursor] == '/' && bytes[cursor - 1] != '\\') { + flagStart = cursor + 1; + flagEnd = flagStart; + while (flagEnd < i && bytes[flagEnd] != ' ' && bytes[flagEnd] != '\t') { + flagEnd++; + } + break; + } + } + for (int cursor = lineStart; cursor < i; cursor++) { + writeNormalizedByte(normalized, bytes[cursor], + cursor >= flagStart && cursor < flagEnd); + } + if (i < bytes.length) { + normalized.write(bytes[i]); + } + lineStart = i + 1; + } + } + return normalized.toByteArray(); + } + + /** Writes a raw byte, converting a high flag byte to the matching UTF-8 code point. */ + private static void writeNormalizedByte(ByteArrayOutputStream target, byte source, + boolean flagByte) { + final int value = source & 0xff; + if (flagByte && value >= 0x80) { + target.write(value < 0xc0 ? 0xc2 : 0xc3); + target.write(value < 0xc0 ? value : value - 0x40); + } else { + target.write(value); + } + } + + /** {@return whether the directive value is one Hunspell flag} */ + private static boolean isSingleFlagDirective(String directive) { + return switch (directive) { + case "COMPOUNDFLAG", "COMPOUNDBEGIN", "COMPOUNDMIDDLE", "COMPOUNDEND", + "COMPOUNDPERMITFLAG", "COMPOUNDFORBIDFLAG", "NEEDAFFIX", "PSEUDOROOT", + "ONLYINCOMPOUND", "FORBIDDENWORD", "CIRCUMFIX" -> true; + default -> false; + }; + } + + /** Finds the fields needed to classify one raw line. */ + private static int findAsciiFields(byte[] bytes, int from, int to, + int[] starts, int[] fieldEnds) { + int count = 0; + int cursor = from; + while (cursor < to && count < starts.length) { + while (cursor < to && isAsciiFieldSpace(bytes[cursor])) { + cursor++; + } + if (cursor == to) { + break; + } + starts[count] = cursor; + while (cursor < to && !isAsciiFieldSpace(bytes[cursor])) { + cursor++; + } + fieldEnds[count] = cursor; + count++; + } + return count; + } + + /** Returns one raw ASCII field. */ + private static String asciiField(byte[] bytes, int from, int to) { + return new String(bytes, from, to - from, StandardCharsets.US_ASCII); + } + + /** + * Converts file content without replacing malformed or unmappable input. + * + * @param bytes The encoded file content. + * @param charset The selected character encoding. + * @param label The file label used in the exception message. + * @return The decoded content. + * @throws IOException Thrown if {@code bytes} are invalid in {@code charset}. + */ + private static String decode(byte[] bytes, Charset charset, String label) + throws IOException { + try { + return charset.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(bytes)) + .toString(); + } catch (CharacterCodingException e) { + throw new IOException(label + " is not valid " + charset.name(), e); + } + } + + /** + * Reads an input stream into a byte array, failing when more than {@code maxBytes} + * arrive. + * + * @param in The stream to read. Not closed. + * @param maxBytes The inclusive upper bound on buffered bytes. + * @param label The stream name used in the error message. + * @return The buffered bytes. Never {@code null}. + * @throws IOException Thrown if reading fails or the stream exceeds {@code maxBytes}. + */ + static byte[] readBounded(InputStream in, int maxBytes, String label) + throws IOException { + final byte[] chunk = new byte[8192]; + byte[] buffer = new byte[Math.min(8192, maxBytes)]; + int size = 0; + int n; + while ((n = in.read(chunk)) >= 0) { + if (size + n > maxBytes) { + throw new IOException(label + " size exceeds safe limit of " + maxBytes); + } + if (size + n > buffer.length) { + buffer = Arrays.copyOf(buffer, Math.min(maxBytes, Math.max(buffer.length * 2, size + n))); + } + System.arraycopy(chunk, 0, buffer, size, n); + size += n; + } + return size == buffer.length ? buffer : Arrays.copyOf(buffer, size); + } + + /** + * Looks up a word's flag sets. + * + * @param word The word exactly as listed. + * @return The flag sets of all matching entries, or {@code null} when absent. + */ + List lookup(String word) { + final List found = entries.get(word); + if (found == null) { + return null; + } + final List copy = new ArrayList<>(found.size()); + for (final int[] flags : found) { + copy.add(flags.clone()); + } + return copy; + } + + /** + * The suffix rules whose affix material ends in the given code point, which are the + * only material-bearing rules that can be undone from a word ending in it. + * + * @param last The word's last code point. + * @return The bucket, possibly empty. Never {@code null}. + */ + List suffixesEndingWith(int last) { + return suffixesByLast.bucket(last); + } + + /** {@return the strip-only suffix rules, applicable to any word} Never {@code null}. */ + List suffixesWithoutMaterial() { + return suffixesWithoutMaterial; + } + + /** + * The prefix rules whose affix material starts with the given code point, which are + * the only material-bearing rules that can be undone from a word starting with it. + * + * @param first The word's first code point. + * @return The bucket, possibly empty. Never {@code null}. + */ + List prefixesStartingWith(int first) { + return prefixesByFirst.bucket(first); + } + + /** {@return the strip-only prefix rules, applicable to any word} Never {@code null}. */ + List prefixesWithoutMaterial() { + return prefixesWithoutMaterial; + } + + /** {@return whether the affix file declares any compounding flag at all} */ + boolean compoundsDeclared() { + return compoundFlag != 0 || compoundBegin != 0 || compoundEnd != 0 + || compoundMiddle != 0; + } + + /** {@return the smallest length a compound part may have} At least {@code 1}. */ + int compoundMin() { + return compoundMin; + } + + /** {@return the largest number of parts a compound may have} {@code 0} is unbounded. */ + int compoundWordMax() { + return compoundWordMax; + } + + /** {@return whether {@code CHECKCOMPOUNDDUP} forbids a part repeating its neighbor} */ + boolean checkCompoundDup() { + return checkCompoundDup; + } + + /** {@return whether {@code CHECKCOMPOUNDCASE} forbids uppercase at part boundaries} */ + boolean checkCompoundCase() { + return checkCompoundCase; + } + + /** {@return whether {@code CHECKCOMPOUNDTRIPLE} forbids triple letters at boundaries} */ + boolean checkCompoundTriple() { + return checkCompoundTriple; + } + + /** {@return whether {@code FULLSTRIP} allows an affix rule to strip a whole stem} */ + boolean fullStrip() { + return fullStrip; + } + + /** + * The flag admitting a part at a compound position, next to the general + * compounding flag. + * + * @param position The part's place in the compound. + * @return The dedicated positional flag, or {@code 0} when undeclared. + */ + private int positionalFlag(CompoundPosition position) { + return switch (position) { + case BEGIN -> compoundBegin; + case MIDDLE -> compoundMiddle; + case END -> compoundEnd; + }; + } + + /** + * Checks whether a listed word may stand at a compound position: some homonym's + * flag set contains the general compound flag or the position's dedicated flag + * and is not forbidden. An {@code ONLYINCOMPOUND} entry is valid here, while a + * {@code NEEDAFFIX} entry still requires an affix. + * + * @param flagSets The word's flag sets from {@link #lookup(String)}. + * @param position The part's place in the compound. + * @return {@code true} if the word may stand at the position. + */ + boolean mayStand(List flagSets, CompoundPosition position) { + final int positional = positionalFlag(position); + for (final int[] flags : flagSets) { + if ((contains(flags, compoundFlag) || contains(flags, positional)) + && !contains(flags, forbiddenWord) && !contains(flags, needAffix) + && !forbiddenAtCompoundPosition(flags, position)) { + return true; + } + } + return false; + } + + /** + * Checks whether some homonym supports an affixed compound part. The flag set + * contains the removed affix flag, is not forbidden, and either the affix permits + * the position or the set contains the compound or positional flag. + * + * @param flagSets The part stem's flag sets from {@link #lookup(String)}. + * @param affixFlag The removed affix's flag. + * @param position The part's place in the compound. + * @param affixAdmits Whether the affix's continuation classes admit the position, + * from {@link #affixAdmits(Affix, CompoundPosition)}. + * @return {@code true} if some homonym permits the affixed form at the position. + */ + boolean supportsPart(List flagSets, int affixFlag, CompoundPosition position, + boolean affixAdmits) { + final int positional = positionalFlag(position); + for (final int[] flags : flagSets) { + if (contains(flags, affixFlag) && !contains(flags, forbiddenWord) + && !forbiddenAtCompoundPosition(flags, position) + && (affixAdmits || contains(flags, compoundFlag) + || contains(flags, positional))) { + return true; + } + } + return false; + } + + /** + * Checks whether an affix admits its derived form at a compound position: its + * continuation classes carry the general compounding flag or the position's + * dedicated flag. Published dictionaries position their linking forms this way, + * through zero or dash suffixes whose continuation classes hold the positional + * flags. + * + * @param affix The affix rule applied to the part. + * @param position The part's place in the compound. + * @return {@code true} if the affixed form may stand at the position. + */ + boolean affixAdmits(Affix affix, CompoundPosition position) { + return (compoundFlag != 0 && affix.allowsContinuation(compoundFlag)) + || (positionalFlag(position) != 0 + && affix.allowsContinuation(positionalFlag(position))); + } + + /** + * Checks whether an affix may sit at a compound-internal boundary: it carries the + * {@code COMPOUNDPERMITFLAG} among its continuation classes. Without the flag a + * suffix fits only the last part and a prefix only the first. + * + * @param affix The affix rule applied to the part. + * @return {@code true} if the affix may face another part. + */ + boolean permitsInside(Affix affix) { + return compoundPermit != 0 && affix.allowsContinuation(compoundPermit); + } + + /** + * Checks whether an affix bars its derived form from compounds altogether: it + * carries the {@code COMPOUNDFORBIDFLAG} among its continuation classes. + * + * @param affix The affix rule applied to the part. + * @return {@code true} if the affixed form may not join a compound. + */ + boolean forbidsInCompound(Affix affix) { + return compoundForbid != 0 && affix.allowsContinuation(compoundForbid); + } + + /** + * Checks whether an entry marked with {@code COMPOUNDFORBIDFLAG} is barred from + * this compound position. Hunspell permits such an entry only as the last part, as + * specified by the + * + * format manual and the + * + * regression fixture. + * + * @param flags One entry's flag set. + * @param position The part's place in the compound. + * @return {@code true} if the entry may not stand at the position. + */ + private boolean forbiddenAtCompoundPosition(int[] flags, CompoundPosition position) { + return position != CompoundPosition.END && contains(flags, compoundForbid); + } + + /** + * Checks whether any of a word's flag sets is forbidden, which a dictionary uses + * to block one specific ill-formed compound while its parts stay productive. + * + * @param flagSets The word's flag sets from {@link #lookup(String)}. + * @return {@code true} if some homonym carries the forbidden-word flag. + */ + boolean anyForbidden(List flagSets) { + return hasFlag(flagSets, forbiddenWord); + } + + /** + * Checks whether any of a word's flag sets carries a flag. + * + * @param flagSets The flag sets from {@link #lookup(String)}. + * @param flag The flag to look for. + * @return {@code true} if some flag set contains the flag. + */ + static boolean hasFlag(List flagSets, int flag) { + for (final int[] flags : flagSets) { + if (contains(flags, flag)) { + return true; + } + } + return false; + } + + /** + * Checks one flag set for a flag. An undeclared flag, encoded as {@code 0}, is + * carried by no entry. + * + * @param flags One entry's flag set. + * @param flag The flag to look for. + * @return {@code true} if the set contains the flag. + */ + private static boolean contains(int[] flags, int flag) { + if (flag == 0) { + return false; + } + for (final int candidate : flags) { + if (candidate == flag) { + return true; + } + } + return false; + } + + /** + * Checks whether a listed word is valid on its own: some homonym's flag set carries + * none of the blocking flags. An entry whose every flag set is marked + * {@code NEEDAFFIX} is a virtual stem that exists only to be affixed, one marked + * {@code ONLYINCOMPOUND} appears only inside compounds, and one marked + * {@code FORBIDDENWORD} is listed to be blocked; none of them is a word by itself. + * + * @param flagSets The word's flag sets from {@link #lookup(String)}. + * @return {@code true} if some homonym stands on its own. + */ + boolean validStandalone(List flagSets) { + for (final int[] flags : flagSets) { + if (!contains(flags, needAffix) && !contains(flags, onlyInCompound) + && !contains(flags, forbiddenWord)) { + return true; + } + } + return false; + } + + /** + * Checks whether some homonym supports an affix analysis: its flag set carries the + * affix's flag and is neither compound-only nor forbidden. A {@code NEEDAFFIX} set + * does support the analysis, because the removed affix is exactly what the virtual + * stem needs. + * + * @param flagSets The stem's flag sets from {@link #lookup(String)}. + * @param flag The removed affix's flag. + * @return {@code true} if some homonym carries the flag and may stand affixed. + */ + boolean supports(List flagSets, int flag) { + for (final int[] flags : flagSets) { + if (contains(flags, flag) && !contains(flags, onlyInCompound) + && !contains(flags, forbiddenWord)) { + return true; + } + } + return false; + } + + /** + * Checks whether some homonym supports a cross-product analysis. The root can contain + * both affix flags, or one root flag can select an affix with continuation flags that + * select the other affix. + * + * @param flagSets The stem's flag sets from {@link #lookup(String)}. + * @param prefix The removed prefix. + * @param suffix The removed suffix. + * @return {@code true} if some homonym licenses both affixes. + */ + boolean supportsCrossProduct(List flagSets, Affix prefix, Affix suffix) { + for (final int[] flags : flagSets) { + final boolean rootHasPrefix = contains(flags, prefix.flag()); + final boolean rootHasSuffix = contains(flags, suffix.flag()); + final boolean licensesBoth = (rootHasPrefix + && (rootHasSuffix || prefix.allowsContinuation(suffix.flag()))) + || (rootHasSuffix && suffix.allowsContinuation(prefix.flag())); + if (licensesBoth && !contains(flags, onlyInCompound) + && !contains(flags, forbiddenWord)) { + return true; + } + } + return false; + } + + /** + * Checks whether a form made with this affix alone is still a virtual stem: the + * affix carries the {@code NEEDAFFIX} flag among its continuation classes, so a + * further affix must join before the form is a word. + * + * @param affix The affix rule to inspect. + * @return {@code true} if the affix alone does not finish a word. + */ + boolean needsFurtherAffix(Affix affix) { + return needAffix != 0 && affix.allowsContinuation(needAffix); + } + + /** + * Checks whether an affix applies only inside compounds: it carries the + * {@code ONLYINCOMPOUND} flag among its continuation classes. + * + * @param affix The affix rule to inspect. + * @return {@code true} if the affix never applies to a standalone word. + */ + boolean compoundOnly(Affix affix) { + return onlyInCompound != 0 && affix.allowsContinuation(onlyInCompound); + } + + /** + * Checks whether an affix is one half of a circumfix: it carries the + * {@code CIRCUMFIX} flag among its continuation classes, so it is only valid on a + * word that also carries a circumfix-marked affix of the other kind, the German + * {@code ge...t} participle being the model. + * + * @param affix The affix rule to inspect. + * @return {@code true} if the affix never applies without its other half. + */ + boolean circumfixOnly(Affix affix) { + return circumfix != 0 && affix.allowsContinuation(circumfix); + } + + /** + * Finds the {@code SET} declaration by scanning the raw affix bytes as ASCII, which + * is safe because the declaration itself is ASCII in every supported encoding. Both + * files are then decoded with the declared charset. + * + * @param affixBytes The raw affix file content. + * @return The declared charset, or UTF-8 when no declaration is present. + * @throws IOException Thrown if the declared encoding name is not supported. + */ + private static Charset declaredCharset(byte[] affixBytes) throws IOException { + final String ascii = new String(affixBytes, StandardCharsets.US_ASCII); + for (final String line : splitLines(ascii)) { + final String trimmed = trim(line); + if (trimmed.startsWith(SET_PREFIX) || trimmed.startsWith(SET_TAB_PREFIX)) { + final String name = trim(trimmed.substring(SET_PREFIX.length())); + try { + return Charset.forName(name); + } catch (IllegalCharsetNameException | UnsupportedCharsetException e) { + throw new IOException("unsupported SET encoding: " + name, e); + } + } + } + return StandardCharsets.UTF_8; + } + + /** The flag encodings a dictionary may declare with the {@code FLAG} directive. */ + private enum FlagMode { + /** + * The default: each single character is one flag. Also what {@code FLAG UTF-8} + * declares, which asks for single-character flags in a file the {@code SET} + * declaration already had decoded. + */ + CHAR, + /** Declared as {@code FLAG long}: each pair of characters is one flag. */ + LONG, + /** Declared as {@code FLAG num}: comma-separated decimal numbers are flags. */ + NUM + } + + /** The parsed affix file content. */ + private static final class AffixFile { + private final List prefixes = new ArrayList<>(); + private final List suffixes = new ArrayList<>(); + private final List flagAliases = new ArrayList<>(); + private FlagMode flagMode = FlagMode.CHAR; + private int compoundFlag; + private int compoundBegin; + private int compoundEnd; + private int compoundMin = 3; + private int needAffix; + private int onlyInCompound; + private int forbiddenWord; + private int circumfix; + private int compoundMiddle; + private int compoundPermit; + private int compoundForbid; + private int compoundWordMax; + private boolean checkCompoundDup; + private boolean checkCompoundCase; + private boolean checkCompoundTriple; + private boolean fullStrip; + } + + /** + * Parses the affix file: the {@code FLAG} declaration, the {@code AF} flag alias + * table, the compound and blocking flag declarations, and the {@code PFX} and + * {@code SFX} blocks. Other directives are skipped because this class implements + * affix stemming, not the complete Hunspell spell-checking engine. + * + * @param content The decoded affix file content. + * @return The parsed rules and flag mode. Never {@code null}. + * @throws IOException Thrown if a supported directive is malformed. + */ + private static AffixFile parseAffix(String content) throws IOException { + final AffixFile result = new AffixFile(); + final String[] lines = splitLines(withoutByteOrderMark(content)); + final String[][] fieldsByLine = new String[lines.length][]; + for (int i = 0; i < lines.length; i++) { + fieldsByLine[i] = split(lines[i]); + } + result.flagMode = readFlagMode(fieldsByLine); + result.flagAliases.addAll(readFlagAliases(fieldsByLine, result.flagMode)); + int i = 0; + while (i < lines.length) { + final String[] fields = fieldsByLine[i]; + if (fields.length == 0 || fields[0].startsWith(COMMENT_PREFIX)) { + i++; + continue; + } + switch (fields[0]) { + case FLAG_TAG: + // The file-wide declaration was parsed before any rule fields. + i++; + break; + case "COMPOUNDFLAG": + case "COMPOUNDBEGIN": + case "COMPOUNDMIDDLE": + case "COMPOUNDEND": + case "COMPOUNDPERMITFLAG": + case "COMPOUNDFORBIDFLAG": + case "NEEDAFFIX": + case "PSEUDOROOT": + case "ONLYINCOMPOUND": + case "FORBIDDENWORD": + case "CIRCUMFIX": + if (fields.length < 2) { + throw new IOException(fields[0] + " line without a flag at line " + (i + 1)); + } + final int declared = parseFlag(fields[1], result.flagMode, i + 1); + switch (fields[0]) { + case "COMPOUNDFLAG" -> result.compoundFlag = declared; + case "COMPOUNDBEGIN" -> result.compoundBegin = declared; + case "COMPOUNDMIDDLE" -> result.compoundMiddle = declared; + case "COMPOUNDEND" -> result.compoundEnd = declared; + case "COMPOUNDPERMITFLAG" -> result.compoundPermit = declared; + case "COMPOUNDFORBIDFLAG" -> result.compoundForbid = declared; + // PSEUDOROOT is the directive's name before hunspell renamed it + case "NEEDAFFIX", "PSEUDOROOT" -> result.needAffix = declared; + case "ONLYINCOMPOUND" -> result.onlyInCompound = declared; + case "CIRCUMFIX" -> result.circumfix = declared; + case "FORBIDDENWORD" -> result.forbiddenWord = declared; + default -> throw new IOException( + "unhandled flag directive " + fields[0] + " at line " + (i + 1)); + } + i++; + break; + case "COMPOUNDMIN": + final int compoundMin = parseValue(fields, i + 1); + if (compoundMin < 0) { + throw new IOException("negative COMPOUNDMIN at line " + (i + 1)); + } + if (compoundMin > MAX_COMPOUND_MIN) { + throw new IOException("COMPOUNDMIN exceeds " + MAX_COMPOUND_MIN + + " at line " + (i + 1)); + } + result.compoundMin = Math.max(1, compoundMin); + i++; + break; + case "COMPOUNDWORDMAX": + final int compoundWordMax = parseValue(fields, i + 1); + if (compoundWordMax < 0) { + throw new IOException("negative COMPOUNDWORDMAX at line " + (i + 1)); + } + result.compoundWordMax = compoundWordMax; + i++; + break; + case "CHECKCOMPOUNDDUP": + result.checkCompoundDup = true; + i++; + break; + case "CHECKCOMPOUNDCASE": + result.checkCompoundCase = true; + i++; + break; + case "CHECKCOMPOUNDTRIPLE": + result.checkCompoundTriple = true; + i++; + break; + case "FULLSTRIP": + result.fullStrip = true; + i++; + break; + case ALIAS_TAG: + // The file-wide table was parsed before continuation and entry flags. + i++; + break; + case PREFIX_TAG: + case SUFFIX_TAG: + i = parseAffixBlock(fieldsByLine, i, fields, result); + break; + default: + i++; + break; + } + } + return result; + } + + /** + * Finds the file-wide flag mode before parsing directives that contain flags. + * + * @param fieldsByLine The affix file fields, indexed by source line. + * @return The selected flag mode, or character mode when no declaration is present. + * @throws IOException Thrown if the declaration is missing a mode, unsupported, or + * repeated. + */ + private static FlagMode readFlagMode(String[][] fieldsByLine) throws IOException { + FlagMode mode = FlagMode.CHAR; + boolean foundMode = false; + for (int i = 0; i < fieldsByLine.length; i++) { + final String[] fields = fieldsByLine[i]; + if (fields.length == 0 || fields[0].startsWith(COMMENT_PREFIX) + || !FLAG_TAG.equals(fields[0])) { + continue; + } + if (foundMode) { + throw new IOException("multiple FLAG directives at line " + (i + 1)); + } + if (fields.length < 2) { + throw new IOException("FLAG line without a mode at line " + (i + 1)); + } + mode = switch (fields[1]) { + case "long" -> FlagMode.LONG; + case "num" -> FlagMode.NUM; + case "UTF-8" -> FlagMode.CHAR; + default -> throw new IOException( + "unsupported FLAG mode '" + fields[1] + "' at line " + (i + 1)); + }; + foundMode = true; + } + return mode; + } + + /** + * Parses the file-wide flag alias table before parsing affix continuation flags. + * + * @param fieldsByLine The affix file fields, indexed by source line. + * @param mode The file's flag encoding. + * @return The aliases in their one-based reference order. + * @throws IOException Thrown if the table header, size, or an alias is malformed. + */ + private static List readFlagAliases(String[][] fieldsByLine, FlagMode mode) + throws IOException { + final List aliases = new ArrayList<>(); + int expected = -1; + for (int i = 0; i < fieldsByLine.length; i++) { + final String[] fields = fieldsByLine[i]; + if (fields.length == 0 || fields[0].startsWith(COMMENT_PREFIX) + || !ALIAS_TAG.equals(fields[0])) { + continue; + } + if (fields.length < 2) { + throw new IOException("AF line without a value at line " + (i + 1)); + } + if (expected < 0) { + // The AF header gives the alias count. Later AF lines contain flag runs. + // Numeric dictionary flags reference the one-based position of a run. + expected = parseValue(fields, i + 1); + if (expected < 0) { + throw new IOException("negative AF count at line " + (i + 1)); + } + } else { + aliases.add(parseFlags(fields[1], mode, i + 1)); + } + } + if (expected >= 0 && aliases.size() != expected) { + throw new IOException("AF header specifies " + expected + " aliases but found " + + aliases.size()); + } + return aliases; + } + + /** + * Parses the integer value of a directive that carries exactly one. + * + * @param fields The already-split directive line. + * @param lineNumber The source line, for error messages. + * @return The parsed value. + * @throws IOException Thrown if the value is missing or is not an integer. + */ + private static int parseValue(String[] fields, int lineNumber) throws IOException { + if (fields.length < 2) { + throw new IOException(fields[0] + " line without a value at line " + lineNumber); + } + try { + return Integer.parseInt(fields[1]); + } catch (NumberFormatException e) { + throw new IOException("malformed " + fields[0] + " at line " + lineNumber, e); + } + } + + /** + * Parses one {@code PFX} or {@code SFX} block: the header line naming the flag, the + * cross-product marker, and the rule count, followed by exactly that many rule + * lines. + * + * @param fieldsByLine All affix file fields, indexed by source line. + * @param index The line index of the block header. + * @param header The already-split header fields. + * @param result The parse target the rules are added to. + * @return The index of the first line after the block. + * @throws IOException Thrown if the header or a rule line is malformed. + */ + private static int parseAffixBlock(String[][] fieldsByLine, int index, String[] header, + AffixFile result) throws IOException { + if (header.length < 4) { + throw new IOException("malformed affix header at line " + (index + 1)); + } + final boolean suffix = SUFFIX_TAG.equals(header[0]); + final int flag = parseFlag(header[1], result.flagMode, index + 1); + if (!"Y".equals(header[2]) && !"N".equals(header[2])) { + throw new IOException("invalid cross-product marker at line " + (index + 1)); + } + final boolean crossProduct = "Y".equals(header[2]); + final int count; + try { + count = Integer.parseInt(header[3]); + } catch (NumberFormatException e) { + throw new IOException("malformed affix rule count at line " + (index + 1), e); + } + if (count < 0) { + throw new IOException("negative affix rule count at line " + (index + 1)); + } + int line = index + 1; + for (int rule = 0; rule < count; rule++, line++) { + if (line >= fieldsByLine.length) { + throw new IOException("affix block truncated at line " + (line + 1)); + } + final String[] fields = fieldsByLine[line]; + if (fields.length < 5 || !fields[0].equals(header[0])) { + throw new IOException("malformed affix rule at line " + (line + 1)); + } + if (parseFlag(fields[1], result.flagMode, line + 1) != flag) { + throw new IOException("affix rule flag does not match header at line " + (line + 1)); + } + final String strip = NO_MATERIAL.equals(fields[2]) ? "" : fields[2]; + String affixText = fields[3]; + int[] continuation = new int[0]; + final int slash = affixText.indexOf('/'); + if (slash >= 0) { + continuation = parseAliasedFlags(affixText.substring(slash + 1), + result.flagMode, result.flagAliases, line + 1); + affixText = affixText.substring(0, slash); + } + if (NO_MATERIAL.equals(affixText)) { + affixText = ""; + } + final Affix affix = new Affix(flag, crossProduct, strip, affixText, + AffixCondition.parse(fields[4], suffix, line + 1), continuation); + if (suffix) { + result.suffixes.add(affix); + } else { + result.prefixes.add(affix); + } + } + return line; + } + + /** + * Parses the word list: an optional leading entry count, then one entry per line + * consisting of the word, an optional {@code /flags} run, and optional trailing + * morphological fields, which are ignored. The morphological fields are cut off + * first, because the flag separator is only meaningful in what precedes them; a word + * may itself contain spaces. A slash escaped as {@code \/} belongs to the word itself + * and is unescaped in the stored key. + * + * @param content The decoded word-list content. + * @param flagMode The flag encoding declared by the affix file. + * @param flagAliases The affix file's {@code AF} alias table, possibly empty. When + * it is not empty, a purely numeric flag field is a 1-based + * reference into it rather than a flag run of its own. + * @return The words mapped to the flag sets of their entries. Never {@code null}. + * @throws IOException Thrown if a flag run is malformed or an alias reference is + * out of range. + */ + private static Map> parseWordList(String content, + FlagMode flagMode, List flagAliases) throws IOException { + final String[] lines = splitLines(withoutByteOrderMark(content)); + final Map> entries = new HashMap<>(); + int start = 0; + if (lines.length > 0 && isCount(trim(lines[0]))) { + start = 1; + } + for (int i = start; i < lines.length; i++) { + final String line = trim(lines[i]); + if (line.isEmpty()) { + continue; + } + final int morphology = morphologyIndex(line); + final String entry = morphology < 0 ? line : trim(line.substring(0, morphology)); + String word = entry; + int[] flags = new int[0]; + final int slash = unescapedSlash(entry); + if (slash >= 0) { + word = entry.substring(0, slash); + String flagRun = entry.substring(slash + 1); + // The flag run ends at the first space or tabulator, the separators the + // word-list format defines; whatever follows is a morphological field even + // when it carries no two-letter tag, which hunspell tolerates and so do we. + for (int c = 0; c < flagRun.length(); c++) { + if (isFieldSeparator(flagRun.charAt(c))) { + flagRun = flagRun.substring(0, c); + break; + } + } + flags = parseAliasedFlags(flagRun, flagMode, flagAliases, i + 1); + } + entries.computeIfAbsent(word.replace("\\/", "/"), key -> new ArrayList<>(1)) + .add(flags); + } + return entries; + } + + /** + * Removes a Unicode byte-order mark decoded at the start of a file. + * + * @param content The decoded file content. + * @return The content without an initial byte-order mark. + */ + private static String withoutByteOrderMark(String content) { + return !content.isEmpty() && content.charAt(0) == '\uFEFF' + ? content.substring(1) : content; + } + + /** + * Checks whether a line consists purely of decimal digits, which identifies the + * optional entry-count header of a word list. + * + * @param line The trimmed line to inspect. + * @return {@code true} if the line is a non-empty digit run. + */ + private static boolean isCount(String line) { + if (line.isEmpty()) { + return false; + } + for (int i = 0; i < line.length(); i++) { + if (line.charAt(i) < '0' || line.charAt(i) > '9') { + return false; + } + } + return true; + } + + /** + * Resolves a numeric {@code AF} alias or parses a direct flag run when no alias + * applies. + * + * @param text The flag field without the leading slash. + * @param mode The selected flag encoding. + * @param aliases The affix file's alias table. + * @param lineNumber The source line, for error messages. + * @return The resolved or parsed flags. + * @throws IOException Thrown if the alias is malformed or outside the table, or the + * direct flags do not fit {@code mode}. + */ + private static int[] parseAliasedFlags(String text, FlagMode mode, + List aliases, int lineNumber) throws IOException { + if (!aliases.isEmpty() && isCount(text)) { + final int alias; + try { + alias = Integer.parseInt(text); + } catch (NumberFormatException e) { + throw new IOException("malformed flag alias '" + text + "' at line " + + lineNumber, e); + } + if (alias < 1 || alias > aliases.size()) { + throw new IOException("flag alias " + alias + " at line " + lineNumber + + " is outside the AF table of " + aliases.size() + " aliases"); + } + return aliases.get(alias - 1); + } + return parseFlags(text, mode, lineNumber); + } + + /** + * Finds the first {@code /} that is not escaped as {@code \/}, which separates the + * word from its flag run in a word-list entry. + * + * @param line The word-list line to scan. + * @return The index of the separator, or {@code -1} when the entry has no flags. + */ + private static int unescapedSlash(String line) { + for (int i = 1; i < line.length(); i++) { + if (line.charAt(i) == '/' && line.charAt(i - 1) != '\\') { + return i; + } + } + return -1; + } + + /** + * Finds where the trailing morphological fields of a word-list entry begin, which + * terminates the word and its flag run. A morphological field is either introduced by + * a tabulator, the older separator, or written as a two-letter tag followed by + * {@code :} and preceded by a separator, such as {@code po:verb}. A separator that + * is not followed by such a tag belongs to the word, because a word-list entry may + * name several words. The separators are the space and the tabulator, exactly the + * two characters the reference implementation's {@code hashmgr.cxx} splits on; they + * are format delimiters of the word-list grammar, not a whitespace judgment, so + * wider whitespace such as a no-break space stays part of the word by design. + * + * @param line The trimmed word-list line to scan. + * @return The index at which the morphological fields begin, or {@code -1} if the + * entry carries none. + */ + private static int morphologyIndex(String line) { + int cut = -1; + for (int i = 4; i < line.length(); i++) { + if (line.charAt(i) == ':' && isFieldSeparator(line.charAt(i - 3))) { + int fieldStart = i - 3; + while (fieldStart > 0 && isFieldSeparator(line.charAt(fieldStart - 1))) { + fieldStart--; + } + // a tag with no word in front of it is not a morphological field + cut = fieldStart == 0 ? -1 : fieldStart; + break; + } + } + final int tab = line.indexOf('\t'); + if (tab >= 0 && (cut < 0 || tab < cut)) { + cut = tab; + } + return cut; + } + + /** + * Checks one character against the word-list format's field separators, space and + * tabulator, the exact set the reference implementation splits morphological fields + * on. + * + * @param c The character to test. + * @return {@code true} if {@code c} separates fields in the word-list format. + */ + private static boolean isFieldSeparator(char c) { + return c == ' ' || c == '\t'; + } + + /** + * Removes leading and trailing whitespace, using the whitespace definition the rest + * of the parser scans with. + * + * @param text The text to trim. + * @return The text without leading or trailing whitespace. Never {@code null}. + */ + private static String trim(String text) { + int start = 0; + int end = text.length(); + while (start < end && StringUtil.isWhitespace(text.charAt(start))) { + start++; + } + while (end > start && StringUtil.isWhitespace(text.charAt(end - 1))) { + end--; + } + return text.substring(start, end); + } + + /** + * Parses a flag run according to the declared flag mode: single characters in + * {@code char} mode, character pairs packed into one {@code int} in {@code long} + * mode, and comma-separated decimal numbers in {@code num} mode. + * + * @param text The flag run without its leading {@code /}. An empty run carries no + * flags in every mode. + * @param mode The declared flag encoding. + * @param lineNumber The source line, for error messages. + * @return The parsed flags. Never {@code null}. + * @throws IOException Thrown if the run does not fit the declared encoding. + */ + private static int[] parseFlags(String text, FlagMode mode, int lineNumber) + throws IOException { + if (text.isEmpty()) { + return new int[0]; + } + switch (mode) { + case NUM: { + final String[] parts = splitOn(text, ','); + final int[] flags = new int[parts.length]; + for (int i = 0; i < parts.length; i++) { + final String value = trim(parts[i]); + try { + flags[i] = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new IOException("malformed numeric flag at line " + lineNumber, e); + } + if (flags[i] < 1 || flags[i] > MAX_NUMERIC_FLAG) { + throw new IOException("numeric flag outside 1.." + MAX_NUMERIC_FLAG + " at line " + + lineNumber + ": " + value); + } + } + return flags; + } + case LONG: { + if (text.length() % 2 != 0) { + throw new IOException("odd long-flag run at line " + lineNumber); + } + final int[] flags = new int[text.length() / 2]; + for (int i = 0; i < flags.length; i++) { + flags[i] = (text.charAt(2 * i) << 16) | text.charAt(2 * i + 1); + } + return flags; + } + default: { + // One flag per code point: published dictionaries name affix rules with + // supplementary characters under FLAG UTF-8, and reading per UTF-16 unit + // would split such a flag into a surrogate pair. A variation selector + // (U+FE00..U+FE0F) only selects a flag character's presentation and is + // dropped from flag identity. + final int[] buffer = new int[text.codePointCount(0, text.length())]; + int f = 0; + for (int i = 0; i < text.length(); ) { + final int codePoint = text.codePointAt(i); + i += Character.charCount(codePoint); + if (codePoint >= 0xFE00 && codePoint <= 0xFE0F) { + continue; + } + buffer[f++] = codePoint; + } + return f == buffer.length ? buffer : Arrays.copyOf(buffer, f); + } + } + } + + /** + * Parses a field that must contain exactly one flag, such as the flag name in an + * affix block header. + * + * @param text The flag field. + * @param mode The declared flag encoding. + * @param lineNumber The source line, for error messages. + * @return The single parsed flag. + * @throws IOException Thrown if the field holds no flag or more than one. + */ + private static int parseFlag(String text, FlagMode mode, int lineNumber) + throws IOException { + final int[] flags = parseFlags(text, mode, lineNumber); + if (flags.length != 1) { + throw new IOException("expected exactly one flag at line " + lineNumber); + } + return flags[0]; + } + + /** + * Splits text into lines with a single character scan, tolerating CRLF endings. + * + * @param content The text to split. + * @return The lines without their terminators. Never {@code null}. + */ + private static String[] splitLines(String content) { + final List lines = new ArrayList<>(); + int start = 0; + for (int i = 0; i <= content.length(); i++) { + if (i == content.length() || content.charAt(i) == '\n') { + int end = i; + if (end > start && content.charAt(end - 1) == '\r') { + end--; + } + lines.add(content.substring(start, end)); + start = i + 1; + } + } + return lines.toArray(new String[0]); + } + + /** + * Splits text on a separator character with a single character scan. + * + * @param text The text to split. + * @param separator The separator character. + * @return The parts between the separators, empty ones included. Never {@code null}. + */ + private static String[] splitOn(String text, char separator) { + final List parts = new ArrayList<>(); + int start = 0; + for (int i = 0; i <= text.length(); i++) { + if (i == text.length() || text.charAt(i) == separator) { + parts.add(text.substring(start, i)); + start = i + 1; + } + } + return parts.toArray(new String[0]); + } + + /** + * Splits a line on whitespace with a single character scan. + * + * @param line The line to split. + * @return The whitespace-separated fields, without empty ones. Never {@code null}. + */ + private static String[] split(String line) { + final List parts = new ArrayList<>(); + int start = -1; + for (int i = 0; i <= line.length(); i++) { + if (i == line.length() || StringUtil.isWhitespace(line.charAt(i))) { + if (start >= 0) { + parts.add(line.substring(start, i)); + start = -1; + } + } else if (start < 0) { + start = i; + } + } + return parts.toArray(new String[0]); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownload.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownload.java new file mode 100644 index 0000000000..95f8c62fba --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownload.java @@ -0,0 +1,74 @@ +/* + * 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.stemmer.hunspell; + +import java.io.IOException; +import java.nio.file.Path; + +import opennlp.tools.util.DictionaryCatalog; + +/** + * Opt-in download of Hunspell {@code .aff}/{@code .dic} pairs and their license + * readme from an application-supplied {@link DictionaryCatalog}. Requires + * {@code -Dopennlp.download.remote=true}. OpenNLP bundles neither a catalog nor + * dictionary data. + * + * @since 3.0.0 + */ +public final class HunspellDictionaryDownload { + + /** Prevents construction of this utility class. */ + private HunspellDictionaryDownload() { + } + + /** + * Downloads the cataloged {@code .aff}, {@code .dic}, and readme files for + * {@code dictionaryId} into {@code targetDirectory}. Each file uses its configured + * name or source name, for example {@code en_US.aff}. Existing target files are not + * replaced, so they must be removed before refreshing a dictionary. + * + * @param catalog The application-supplied catalog. Must not be {@code null}. + * @param dictionaryId The catalog dictionary name, for example {@code en_US}. + * Must not be {@code null}. + * @param targetDirectory The directory to write into; created when absent. Must not + * be {@code null}. + * @throws IOException Thrown if remote downloads are disabled, a catalog entry is + * missing, verification fails, or the target already contains one of the + * files. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static void downloadFromCatalog(DictionaryCatalog catalog, String dictionaryId, + Path targetDirectory) throws IOException { + if (catalog == null) { + throw new IllegalArgumentException("catalog must not be null"); + } + if (dictionaryId == null) { + throw new IllegalArgumentException("dictionaryId must not be null"); + } + if (targetDirectory == null) { + throw new IllegalArgumentException("targetDirectory must not be null"); + } + final String prefix = "hunspell." + dictionaryId; + catalog.install(prefix + HunspellDictionary.AFFIX_FILE_SUFFIX, targetDirectory); + catalog.install(prefix + HunspellDictionary.DICTIONARY_FILE_SUFFIX, targetDirectory); + final String readmeId = prefix + ".readme"; + if (catalog.ids().contains(readmeId)) { + catalog.install(readmeId, targetDirectory); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmer.java new file mode 100644 index 0000000000..c5699fa149 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmer.java @@ -0,0 +1,701 @@ +/* + * 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.stemmer.hunspell; + +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.stemmer.Stemmer; +import opennlp.tools.stemmer.hunspell.HunspellDictionary.Affix; +import opennlp.tools.stemmer.hunspell.HunspellDictionary.CompoundPosition; +import opennlp.tools.util.StringUtil; + +/** + * A dictionary-backed {@link Stemmer} over a {@link HunspellDictionary}: a surface form + * is reduced to the dictionary words it can be derived from by removing one suffix, one + * prefix, a cross-product combination of both, or an additional suffix licensed by a + * continuation class. + * + *

    {@link #stem(CharSequence)} returns the first analysis, preferring the word's own + * dictionary entry; {@link #stemAll(CharSequence)} returns every distinct analysis. A + * word with no analysis is returned unchanged, so the stemmer degrades to identity on + * unknown vocabulary. A form containing uppercase characters is also analyzed in its + * lowercase variant, so sentence-initial capitalization does not hide an entry. + * Entries the dictionary marks as virtual stems ({@code NEEDAFFIX}), compound-only + * parts ({@code ONLYINCOMPOUND}), or forbidden words ({@code FORBIDDENWORD}) never + * count as standalone analyses, matching how hunspell reads those flags.

    + * + *

    Compound part search is capped at {@value #PART_CHECK_BUDGET} part-licensing + * attempts per input word; beyond that budget further compound analyses are skipped. + * The {@link Stemmer} interface leaves thread safety to the implementation. This + * implementation reads only the immutable dictionary state, so a single instance is + * safe to share between threads.

    + * + * @since 3.0.0 + */ +@ThreadSafe +public final class HunspellStemmer implements Stemmer { + + /** + * The most part-licensing attempts one decomposition search may spend. Compounding + * searches every split of every tail, which on adversarial input with a + * one-character minimum part length grows without useful bound; the budget stops + * the search there, missing analyses rather than stalling, in line with the + * engine's fail-closed posture. + */ + private static final int PART_CHECK_BUDGET = 2048; + + private final HunspellDictionary dictionary; + + /** + * Initializes the stemmer. + * + * @param dictionary The dictionary to analyze against. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code dictionary} is {@code null}. + */ + public HunspellStemmer(HunspellDictionary dictionary) { + if (dictionary == null) { + throw new IllegalArgumentException("dictionary must not be null"); + } + this.dictionary = dictionary; + } + + /** + * {@inheritDoc} + * + *

    Returns the first analysis, which prefers the word's own dictionary entry.

    + */ + @Override + public CharSequence stem(CharSequence word) { + final List analyses = stemAll(word); + return analyses.get(0); + } + + /** + * {@inheritDoc} + * + *

    Returns every distinct analysis, or a single-element list of the unchanged word + * when it has none.

    + */ + @Override + public List stemAll(CharSequence word) { + if (word == null) { + throw new IllegalArgumentException("word must not be null"); + } + final String surface = word.toString(); + if (surface.isEmpty()) { + // a zero-length word has no morphology; without this guard a strip-only rule + // could restore its strip string onto nothing and answer a non-empty stem + return List.of(surface); + } + final Set analyses = new LinkedHashSet<>(); + for (final String variant : variants(surface)) { + analyze(variant, analyses); + } + if (analyses.isEmpty() && dictionary.compoundsDeclared()) { + for (final String variant : variants(surface)) { + decompose(variant, surface, analyses); + } + } + if (analyses.isEmpty()) { + return List.of(surface); + } + return List.copyOf(analyses); + } + + /** + * Collects the case variants to analyze: the surface form first, then its lowercase + * form when the two differ. Ordering matters because the first analysis found wins + * in {@link #stem(CharSequence)}. + * + * @param surface The surface form. + * @return The variants in analysis order. Never {@code null} or empty. + */ + private List variants(String surface) { + final String lowered = StringUtil.toLowerCase(surface); + return lowered.equals(surface) ? List.of(surface) : List.of(surface, lowered); + } + + /** + * Adds every analysis of one case variant to the result set: the word's own + * dictionary entry, single suffix removal, twofold suffix removal through + * continuation classes, single prefix removal, and cross-product removal of one + * prefix together with one suffix and an optional continuation suffix. Insertion + * order into the set fixes the + * preference order reported by {@link #stemAll(CharSequence)}. + * + * @param word The case variant to analyze. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void analyze(String word, Set analyses) { + final List entries = dictionary.lookup(word); + if (entries != null) { + if (dictionary.anyForbidden(entries)) { + return; + } + if (dictionary.validStandalone(entries)) { + analyses.add(word); + } + } + for (final Affix suffix : dictionary.suffixesEndingWith( + word.codePointBefore(word.length()))) { + undoSuffix(word, suffix, analyses); + } + for (final Affix suffix : dictionary.suffixesWithoutMaterial()) { + undoSuffix(word, suffix, analyses); + } + for (final Affix prefix : dictionary.prefixesStartingWith(word.codePointAt(0))) { + undoPrefix(word, prefix, analyses); + } + for (final Affix prefix : dictionary.prefixesWithoutMaterial()) { + undoPrefix(word, prefix, analyses); + } + } + + /** + * Decomposes a word into listed compound parts when the affix analysis found + * nothing: the first part must be admitted to open a compound, every further part + * to continue or close one, each at least the declared minimum length and counted + * against the declared maximum. A part stands on its own entry or on an entry plus + * one affix, the way published dictionaries position their linking forms through + * zero or dash suffixes. The stems of the parts of every successful splitting are + * reported left to right, so the head-most material comes last. A word the + * dictionary lists as forbidden never decomposes; that is how one specific + * ill-formed compound is blocked while its parts stay productive. + * + * @param word The case variant to decompose. + * @param surface The surface form the variant was derived from; character case at + * junctions is judged against it, so lowercasing a variant cannot + * sidestep a {@code CHECKCOMPOUNDCASE} declaration. + * @param analyses The mutable, insertion-ordered set collecting the part stems. + */ + private void decompose(String word, String surface, Set analyses) { + final List entries = dictionary.lookup(word); + if (entries != null && dictionary.anyForbidden(entries)) { + return; + } + final int codePointCount = word.codePointCount(0, word.length()); + if (codePointCount < 2 * dictionary.compoundMin()) { + return; + } + final int[] codePointOffsets = new int[codePointCount + 1]; + int offset = 0; + for (int i = 0; i < codePointCount; i++) { + codePointOffsets[i] = offset; + offset += Character.charCount(word.codePointAt(offset)); + } + codePointOffsets[codePointCount] = word.length(); + // lowercasing may change the length in exceptional mappings, in which case the + // offsets no longer align and the variant itself is the only usable case source + final String caseSource = surface.length() == word.length() ? surface : word; + search(word, caseSource, codePointOffsets, 0, new ArrayList<>(), + new ArrayList<>(), analyses, new int[] {PART_CHECK_BUDGET}); + } + + /** + * Extends a partial decomposition with the part starting at {@code fromPoint}, trying + * every admissible length and recursing on the remainder. The boundary into this + * part honors the {@code CHECKCOMPOUNDCASE} and {@code CHECKCOMPOUNDTRIPLE} + * declarations, a part repeating its left neighbor honors + * {@code CHECKCOMPOUNDDUP}, and a completed decomposition flushes every part's + * stems into the analyses in part order. + * + * @param word The case variant under decomposition. + * @param caseSource The character-case source for junction checks, the surface + * form when its offsets align with the variant. + * @param codePointOffsets UTF-16 offsets for each code point boundary. + * @param fromPoint The code point index where the next part starts. + * @param surfaces The surface strings of the parts taken so far. + * @param stems The licensed stems of the parts taken so far, one list per part. + * @param analyses The mutable, insertion-ordered set collecting the part stems. + * @param budget The remaining part-licensing attempts, counted down in place. + */ + private void search(String word, String caseSource, int[] codePointOffsets, + int fromPoint, List surfaces, List> stems, + Set analyses, int[] budget) { + final int from = codePointOffsets[fromPoint]; + if (from > 0 && violatesBoundaryChecks(word, caseSource, from)) { + return; + } + final int min = dictionary.compoundMin(); + final int max = dictionary.compoundWordMax(); + final boolean first = from == 0; + final int remaining = codePointOffsets.length - 1 - fromPoint; + // every split leaving room for a further part; a first-position part must also + // leave the closing part, so the whole word is never one part + if (remaining >= 2 * min && (max == 0 || surfaces.size() + 2 <= max)) { + final int lastEndPoint = codePointOffsets.length - 1 - min; + for (int endPoint = fromPoint + min; endPoint <= lastEndPoint; endPoint++) { + if (budget[0] <= 0) { + return; + } + budget[0]--; + final int end = codePointOffsets[endPoint]; + final String part = word.substring(from, end); + if (duplicatesNeighbor(part, surfaces)) { + continue; + } + final List partStems = partStems(part, + first ? CompoundPosition.BEGIN : CompoundPosition.MIDDLE, first, false); + if (partStems.isEmpty()) { + continue; + } + surfaces.add(part); + stems.add(partStems); + search(word, caseSource, codePointOffsets, endPoint, surfaces, stems, + analyses, budget); + surfaces.remove(surfaces.size() - 1); + stems.remove(stems.size() - 1); + } + } + // the closing part takes the whole remainder; a compound has at least two parts + if (first || remaining < min + || (max > 0 && surfaces.size() + 1 > max) || budget[0] <= 0) { + return; + } + budget[0]--; + final String part = word.substring(from); + if (duplicatesNeighbor(part, surfaces)) { + return; + } + final List partStems = partStems(part, CompoundPosition.END, false, true); + if (partStems.isEmpty()) { + return; + } + for (final List earlier : stems) { + analyses.addAll(earlier); + } + analyses.addAll(partStems); + } + + /** + * Applies the {@code CHECKCOMPOUNDDUP} declaration: a part must not repeat the + * part directly before it. + * + * @param part The candidate part. + * @param surfaces The surface strings of the parts taken so far. + * @return {@code true} if the declaration forbids this part here. + */ + private boolean duplicatesNeighbor(String part, List surfaces) { + return dictionary.checkCompoundDup() && !surfaces.isEmpty() + && part.equals(surfaces.get(surfaces.size() - 1)); + } + + /** + * Applies the character-level boundary declarations at the junction before + * {@code from}: {@code CHECKCOMPOUNDCASE} forbids an uppercase character on either + * side of the junction, and {@code CHECKCOMPOUNDTRIPLE} forbids the same character + * three times in a row across it. + * + * @param word The case variant under decomposition. + * @param caseSource The character-case source for the uppercase judgment. + * @param from The index the junction sits before; greater than zero. + * @return {@code true} if a declaration forbids this junction. + */ + private boolean violatesBoundaryChecks(String word, String caseSource, int from) { + final int before = word.codePointBefore(from); + final int after = word.codePointAt(from); + if (dictionary.checkCompoundCase() + && (Character.isUpperCase(caseSource.codePointBefore(from)) + || Character.isUpperCase(caseSource.codePointAt(from)))) { + return true; + } + if (dictionary.checkCompoundTriple() && before == after) { + final int beforeStart = from - Character.charCount(before); + final int afterEnd = from + Character.charCount(after); + if ((beforeStart > 0 && word.codePointBefore(beforeStart) == after) + || (afterEnd < word.length() && word.codePointAt(afterEnd) == after)) { + return true; + } + } + return false; + } + + /** + * Collects the listed stems that admit one part at its compound position: the part + * as its own entry, or an entry plus one suffix or one prefix whose removal leaves + * a listed stem, zero-material rules included, because published dictionaries + * position their linking forms through zero and dash suffixes. An affix at a + * compound-internal boundary must carry the permit flag, a suffix facing the next + * part or a prefix facing the previous one. A part not found as written is also + * tried with its first letter uppercased, the way nouns listed capitalized appear + * lowercase inside a compound. + * + * @param part The part's surface text. + * @param position The part's place in the compound. + * @param first Whether the part opens the word. + * @param last Whether the part closes the word. + * @return The stems admitting the part, in discovery order. Never {@code null}. + */ + private List partStems(String part, CompoundPosition position, + boolean first, boolean last) { + final Set stems = new LinkedHashSet<>(); + collectPartStems(part, position, first, last, stems); + if (stems.isEmpty() && !part.isEmpty()) { + final int initial = part.codePointAt(0); + final int upper = Character.toUpperCase(initial); + if (upper != initial) { + collectPartStems(new StringBuilder().appendCodePoint(upper) + .append(part, Character.charCount(initial), part.length()).toString(), + position, first, last, stems); + } + } + return List.copyOf(stems); + } + + /** + * Collects the stems admitting one spelling of a part, bare and through one affix. + * + * @param part The part spelling to look up. + * @param position The part's place in the compound. + * @param first Whether the part opens the word. + * @param last Whether the part closes the word. + * @param stems The mutable, insertion-ordered set collecting the stems. + */ + private void collectPartStems(String part, CompoundPosition position, + boolean first, boolean last, Set stems) { + final List entries = dictionary.lookup(part); + if (entries != null && dictionary.mayStand(entries, position)) { + stems.add(part); + } + for (final Affix suffix : dictionary.suffixesEndingWith( + part.codePointBefore(part.length()))) { + collectAffixedPartStem(part, suffix, true, position, last, stems); + } + for (final Affix suffix : dictionary.suffixesWithoutMaterial()) { + collectAffixedPartStem(part, suffix, true, position, last, stems); + } + for (final Affix prefix : dictionary.prefixesStartingWith(part.codePointAt(0))) { + collectAffixedPartStem(part, prefix, false, position, first, stems); + } + for (final Affix prefix : dictionary.prefixesWithoutMaterial()) { + collectAffixedPartStem(part, prefix, false, position, first, stems); + } + } + + /** + * Adds the stem of one affixed part reading when the rule and the stem's entry admit + * it at the position. + * + * @param part The part spelling under analysis. + * @param affix The rule to undo. + * @param suffix Whether the rule is a suffix rule. + * @param position The part's place in the compound. + * @param atEdge Whether the part sits at the word end the rule faces, the closing part + * for a suffix rule and the opening part for a prefix rule; an affix + * facing another part instead needs the permit flag. + * @param stems The mutable, insertion-ordered set collecting the stems. + */ + private void collectAffixedPartStem(String part, Affix affix, boolean suffix, + CompoundPosition position, boolean atEdge, Set stems) { + if (dictionary.circumfixOnly(affix) || dictionary.forbidsInCompound(affix) + || (!atEdge && !dictionary.permitsInside(affix))) { + return; + } + final String stem = removeAffixInCompound(part, affix, suffix); + if (stem == null) { + return; + } + final List flagSets = dictionary.lookup(stem); + if (flagSets != null && dictionary.supportsPart(flagSets, affix.flag(), position, + dictionary.affixAdmits(affix, position))) { + stems.add(stem); + } + } + + /** + * Undoes one affix rule on a compound part. Unlike the standalone removals, a rule + * that neither adds nor removes material is undone here, to its own spelling with + * the condition checked, because dictionaries position compound parts through + * exactly such zero rules. + * + * @param part The part spelling under analysis. + * @param affix The rule to undo. + * @param suffix Whether the rule is a suffix rule. + * @return The candidate stem, or {@code null} when the rule does not apply. + */ + private String removeAffixInCompound(String part, Affix affix, boolean suffix) { + return suffix + ? removeSuffixAllowingIdentity(part, affix) + : removePrefixAllowingIdentity(part, affix); + } + + /** + * Undoes one suffix rule and, through continuation classes, one further suffix on + * the intermediate stem, adding dictionary-confirmed analyses. A rule that applies + * only inside compounds or requires the matching circumfix member is not undone + * because no prefix accompanies this path. A rule requiring a further affix produces + * no single-removal analysis. An identity rule also produces no single-removal + * analysis, but it can complete a two-suffix analysis through continuation classes. + * + * @param word The case variant under analysis. + * @param suffix The suffix rule to undo. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void undoSuffix(String word, Affix suffix, Set analyses) { + if (dictionary.compoundOnly(suffix) || dictionary.circumfixOnly(suffix)) { + return; + } + final boolean identity = isIdentityRule(suffix); + final String stem = removeSuffixAllowingIdentity(word, suffix); + if (stem == null) { + return; + } + if (!identity && !dictionary.needsFurtherAffix(suffix)) { + final List flagSets = dictionary.lookup(stem); + if (flagSets != null && dictionary.supports(flagSets, suffix.flag())) { + analyses.add(stem); + } + } + for (final Affix inner : dictionary.suffixesEndingWith( + stem.codePointBefore(stem.length()))) { + undoInnerSuffix(stem, suffix, inner, analyses); + } + for (final Affix inner : dictionary.suffixesWithoutMaterial()) { + undoInnerSuffix(stem, suffix, inner, analyses); + } + } + + /** + * Undoes the inner suffix of a twofold removal when the rule's continuation + * classes allow the outer one. The continuation-linked combination satisfies a + * {@code NEEDAFFIX} marker on either rule. + * + * @param stem The intermediate stem after the outer removal. + * @param outer The already-undone outer suffix rule. + * @param inner The candidate inner suffix rule. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void undoInnerSuffix(String stem, Affix outer, Affix inner, + Set analyses) { + if (!inner.allowsContinuation(outer.flag()) || dictionary.compoundOnly(inner) + || dictionary.circumfixOnly(inner)) { + return; + } + final String doubleStem = removeSuffixAllowingIdentity(stem, inner); + if (doubleStem == null) { + return; + } + final List innerFlags = dictionary.lookup(doubleStem); + if (innerFlags != null && dictionary.supports(innerFlags, inner.flag())) { + analyses.add(doubleStem); + } + } + + /** + * Undoes one prefix rule and, for cross-product rules, one further suffix on the + * intermediate stem, adding dictionary-confirmed analyses. A rule that + * applies only inside compounds is not undone at all. A rule marked as needing a + * further affix or the matching circumfix member produces no single-removal analysis. + * An identity rule also produces no single-removal analysis. A valid cross-product + * suffix can combine with either kind of rule. + * + * @param word The case variant under analysis. + * @param prefix The prefix rule to undo. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void undoPrefix(String word, Affix prefix, Set analyses) { + if (dictionary.compoundOnly(prefix)) { + return; + } + final boolean identity = isIdentityRule(prefix); + final String stem = removePrefixAllowingIdentity(word, prefix); + if (stem == null) { + return; + } + if (!identity && !dictionary.needsFurtherAffix(prefix) + && !dictionary.circumfixOnly(prefix)) { + final List flagSets = dictionary.lookup(stem); + if (flagSets != null && dictionary.supports(flagSets, prefix.flag())) { + analyses.add(stem); + } + } + if (!prefix.crossProduct()) { + return; + } + for (final Affix suffix : dictionary.suffixesEndingWith( + stem.codePointBefore(stem.length()))) { + undoCrossProductSuffix(stem, prefix, suffix, analyses); + } + for (final Affix suffix : dictionary.suffixesWithoutMaterial()) { + undoCrossProductSuffix(stem, prefix, suffix, analyses); + } + } + + /** + * Undoes the suffix half of a cross-product removal when both rules opted in. The + * two rules must agree on circumfixing: a circumfix-marked affix is only valid with + * a marked affix of the other kind, so a pair of which exactly one is marked mixes + * an ordinary affix into a circumfix and is rejected. + * + * @param stem The intermediate stem after the prefix removal. + * @param prefix The already-undone prefix rule. + * @param suffix The candidate suffix rule. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void undoCrossProductSuffix(String stem, Affix prefix, Affix suffix, + Set analyses) { + if (!suffix.crossProduct() || dictionary.compoundOnly(suffix) + || dictionary.circumfixOnly(prefix) != dictionary.circumfixOnly(suffix)) { + return; + } + final String doubleStem = removeSuffix(stem, suffix); + if (doubleStem == null) { + return; + } + // One member can satisfy the other member's needs-further-affix marker. Both rule + // flags must occur in one homonym's flag set. + final List both = dictionary.lookup(doubleStem); + if (both != null && dictionary.supportsCrossProduct(both, prefix, suffix) + && !(dictionary.needsFurtherAffix(prefix) + && dictionary.needsFurtherAffix(suffix))) { + analyses.add(doubleStem); + } + for (final Affix inner : dictionary.suffixesEndingWith( + doubleStem.codePointBefore(doubleStem.length()))) { + undoCrossProductInnerSuffix(doubleStem, prefix, suffix, inner, analyses); + } + for (final Affix inner : dictionary.suffixesWithoutMaterial()) { + undoCrossProductInnerSuffix(doubleStem, prefix, suffix, inner, analyses); + } + } + + /** + * Undoes an inner suffix after a prefix and an outer suffix have been removed. + * The inner suffix must license the outer suffix through the continuation flags. + * The suffix combination satisfies {@code NEEDAFFIX} markers in the derivation. + * + * @param stem The intermediate stem after the prefix and outer suffix removal. + * @param prefix The already-undone prefix rule. + * @param outer The already-undone outer suffix rule. + * @param inner The candidate inner suffix rule. + * @param analyses The mutable, insertion-ordered set collecting the stems found. + */ + private void undoCrossProductInnerSuffix(String stem, Affix prefix, Affix outer, + Affix inner, Set analyses) { + if (!inner.crossProduct() || !inner.allowsContinuation(outer.flag()) + || dictionary.compoundOnly(inner) + || dictionary.circumfixOnly(outer) + || dictionary.circumfixOnly(prefix) != dictionary.circumfixOnly(inner)) { + return; + } + final String root = removeSuffixAllowingIdentity(stem, inner); + if (root == null) { + return; + } + final List flagSets = dictionary.lookup(root); + if (flagSets != null && dictionary.supportsCrossProduct(flagSets, prefix, inner)) { + analyses.add(root); + } + } + + /** + * Undoes one suffix rule: cuts the affix material off the end of the word, restores + * the strip string the rule removed on application, and checks the rule's condition + * against the restored stem. A strip-only rule, whose affix material is empty, is + * undone by restoring its strip string alone. Rules that neither add nor remove + * material and candidates that would leave an empty stem are rejected. A word the + * affix material covers entirely reverses a full-strip application, which hunspell + * only performs when the affix file declares {@code FULLSTRIP}; without that + * declaration the rule does not apply. + * + * @param word The surface form. + * @param suffix The rule to undo. + * @return The candidate stem, or {@code null} when the rule does not apply. + */ + private String removeSuffix(String word, Affix suffix) { + final String affix = suffix.affix(); + final String strip = suffix.strip(); + if (affix.isEmpty() && strip.isEmpty() || !word.endsWith(affix) + || word.length() - affix.length() + strip.length() == 0 + || (word.length() == affix.length() && !dictionary.fullStrip())) { + return null; + } + final String stem = word.substring(0, word.length() - affix.length()) + strip; + return suffix.condition().matches(stem) ? stem : null; + } + + /** + * Undoes a suffix in a continuation sequence, including a rule that changes no + * material. An identity rule still has to satisfy the condition. + * + * @param word The surface form at this point in the sequence. + * @param suffix The rule to undo. + * @return The candidate stem, or {@code null} when the rule does not apply. + */ + private String removeSuffixAllowingIdentity(String word, Affix suffix) { + if (isIdentityRule(suffix)) { + return suffix.condition().matches(word) ? word : null; + } + return removeSuffix(word, suffix); + } + + /** + * Undoes a prefix in a continuation sequence, including a rule that changes no + * material. An identity rule still has to satisfy the condition. + * + * @param word The surface form at this point in the sequence. + * @param prefix The rule to undo. + * @return The candidate stem, or {@code null} when the rule does not apply. + */ + private String removePrefixAllowingIdentity(String word, Affix prefix) { + if (isIdentityRule(prefix)) { + return prefix.condition().matches(word) ? word : null; + } + return removePrefix(word, prefix); + } + + /** + * Checks whether an affix rule adds and strips no material. + * + * @param affix The rule to inspect. + * @return {@code true} if applying the rule does not change the spelling. + */ + private boolean isIdentityRule(Affix affix) { + return affix.affix().isEmpty() && affix.strip().isEmpty(); + } + + /** + * Undoes one prefix rule: cuts the affix material off the start of the word, + * restores the strip string the rule removed on application, and checks the rule's + * condition against the restored stem. A strip-only rule, whose affix material is + * empty, is undone by restoring its strip string alone. Rules that neither add nor + * remove material and candidates that would leave an empty stem are rejected. A + * word the affix material covers entirely reverses a full-strip application, which + * hunspell only performs when the affix file declares {@code FULLSTRIP}; without + * that declaration the rule does not apply. + * + * @param word The surface form. + * @param prefix The rule to undo. + * @return The candidate stem, or {@code null} when the rule does not apply. + */ + private String removePrefix(String word, Affix prefix) { + final String affix = prefix.affix(); + final String strip = prefix.strip(); + if (affix.isEmpty() && strip.isEmpty() || !word.startsWith(affix) + || word.length() - affix.length() + strip.length() == 0 + || (word.length() == affix.length() && !dictionary.fullStrip())) { + return null; + } + final String stem = strip + word.substring(affix.length()); + return prefix.condition().matches(stem) ? stem : null; + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactory.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactory.java new file mode 100644 index 0000000000..b8b7c4e62e --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactory.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.tools.stemmer.hunspell; + +import opennlp.tools.commons.ThreadSafe; +import opennlp.tools.stemmer.Stemmer; +import opennlp.tools.stemmer.StemmerFactory; + +/** + * The shareable handle for Hunspell stemming: holds one immutable + * {@link HunspellDictionary} and hands out {@link HunspellStemmer} instances over it. + * + *

    The factory is immutable and safe to share across threads.

    + * + * @since 3.0.0 + */ +@ThreadSafe +public class HunspellStemmerFactory implements StemmerFactory { + + private final HunspellDictionary dictionary; + + /** + * Initializes the factory. + * + * @param dictionary The dictionary to stem against. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code dictionary} is {@code null}. + */ + public HunspellStemmerFactory(HunspellDictionary dictionary) { + if (dictionary == null) { + throw new IllegalArgumentException("dictionary must not be null"); + } + this.dictionary = dictionary; + } + + /** + * {@inheritDoc} + * + *

    Every call creates a fresh {@link HunspellStemmer} over the same immutable + * dictionary.

    + */ + @Override + public Stemmer newStemmer() { + return new HunspellStemmer(dictionary); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/termvector/TermVectorAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/termvector/TermVectorAnnotator.java new file mode 100644 index 0000000000..4ed12d0d13 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/termvector/TermVectorAnnotator.java @@ -0,0 +1,318 @@ +/* + * 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.termvector; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; +import opennlp.tools.util.normalizer.AlignedText; +import opennlp.tools.util.normalizer.CharSequenceNormalizer; +import opennlp.tools.util.normalizer.OffsetAwareNormalizer; + +/** + * Rolls the token layer up into a term vector layer for index consumers: one + * {@link TermVector} per distinct term, carrying the term string, its occurrence count, + * and (in {@link Mode#FULL full mode}) the occurrence offsets. + * + *

    Term identity comes from the annotator's inputs, not from logic of its own. Without + * a normalizer, the term is the token layer's value as-is, that is, the token's covered + * text in the original document. With a plain {@link CharSequenceNormalizer}, the general + * path, each token's covered text is normalized on its own to produce the term, so any + * normalizer works: case folding, NFC, accent folding, a stemmer-backed normalizer. With + * an {@link OffsetAwareNormalizer}, the whole document text is normalized once with its + * alignment recorded, each token span is mapped forward to the normalized form, and the + * covered normalized text is the term; this path can see across token boundaries but is + * limited to alignment-reporting normalizers. On every path, tokens that differ only by + * a normalization fold (case, an eszett expansion, collapsed whitespace) group together, + * and the occurrence spans emitted in {@link Mode#FULL full mode} are the token layer's + * own spans and therefore always point into the original text. A token whose normalized + * form is empty, for example one the normalizer deleted entirely, is omitted from the + * layer; an empty string is no term, and the token layer still accounts for the + * token.

    + * + *

    The layer is {@link LayerKey.Scope#DOCUMENT document-scoped}: each {@link TermVector} + * is a whole-document statistic, so the annotations carry no span of their own and the + * occurrence offsets live inside the payload. The layer preserves first-occurrence + * order: the first token of a term fixes its position in the layer.

    + * + *

    The annotator holds no per-call state; it is as thread-safe as the normalizer it + * was built with.

    + * + * @since 3.0.0 + */ +public final class TermVectorAnnotator implements DocumentAnnotator { + + /** + * The key of the term vector layer this annotator provides: a document-scoped layer + * of {@link TermVector} values, one per distinct term. + */ + public static final LayerKey TERM_VECTORS = + Layers.documentKey("term-vectors", TermVector.class); + + /** How much each {@link TermVector} records. */ + public enum Mode { + /** + * Counts occurrences and stores every occurrence span in original text coordinates. + */ + FULL, + /** + * Counts occurrences only; the emitted {@link TermVector term vectors} carry no + * spans, so scoring-only consumers do not pay for offset storage. + */ + SCORING_ONLY + } + + private final OffsetAwareNormalizer normalizer; + private final CharSequenceNormalizer tokenNormalizer; + private final Mode mode; + + /** + * Initializes a {@link Mode#FULL full mode} annotator that groups tokens by their + * covered text as-is. + */ + public TermVectorAnnotator() { + this(Mode.FULL); + } + + /** + * Initializes an annotator that groups tokens by their covered text as-is. + * + * @param mode How much each {@link TermVector} records. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code mode} is {@code null}. + */ + public TermVectorAnnotator(Mode mode) { + if (mode == null) { + throw new IllegalArgumentException("mode must not be null"); + } + this.normalizer = null; + this.tokenNormalizer = null; + this.mode = mode; + } + + /** + * Initializes a {@link Mode#FULL full mode} annotator that groups tokens by their + * per-token normalized form. This is the general path: the normalizer is applied to + * each token's covered text on its own, so any {@link CharSequenceNormalizer} works, + * including the folds that cannot report an alignment (case folding, NFC, accent + * folding, stemmer-backed normalizers). The occurrence spans stay the tokens' own + * spans in the original text. + * + * @param normalizer The normalizer that defines term identity, applied to each token's + * covered text. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code normalizer} is {@code null}. + */ + public TermVectorAnnotator(CharSequenceNormalizer normalizer) { + this(normalizer, Mode.FULL); + } + + /** + * Initializes an annotator that groups tokens by their per-token normalized form. This + * is the general path: the normalizer is applied to each token's covered text on its + * own, so any {@link CharSequenceNormalizer} works, including the folds that cannot + * report an alignment (case folding, NFC, accent folding, stemmer-backed normalizers). + * The occurrence spans stay the tokens' own spans in the original text. + * + * @param normalizer The normalizer that defines term identity, applied to each token's + * covered text. Must not be {@code null}. + * @param mode How much each {@link TermVector} records. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code normalizer} or {@code mode} is + * {@code null}. + */ + public TermVectorAnnotator(CharSequenceNormalizer normalizer, Mode mode) { + if (normalizer == null) { + throw new IllegalArgumentException("normalizer must not be null"); + } + if (mode == null) { + throw new IllegalArgumentException("mode must not be null"); + } + this.normalizer = null; + this.tokenNormalizer = normalizer; + this.mode = mode; + } + + /** + * Initializes a {@link Mode#FULL full mode} annotator that groups tokens by their + * normalized form through a whole-document alignment. Prefer the + * {@link #TermVectorAnnotator(CharSequenceNormalizer) plain-normalizer constructor} + * as the general path; this one only accepts alignment-reporting normalizers but can + * see across token boundaries, for example a whitespace collapse spanning two tokens. + * + * @param normalizer The normalizer that defines term identity, applied to the whole + * document text so token spans can be mapped into the normalized + * form through its alignment. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code normalizer} is {@code null}. + */ + public TermVectorAnnotator(OffsetAwareNormalizer normalizer) { + this(normalizer, Mode.FULL); + } + + /** + * Initializes an annotator that groups tokens by their normalized form through a + * whole-document alignment. Prefer the + * {@link #TermVectorAnnotator(CharSequenceNormalizer, Mode) plain-normalizer + * constructor} as the general path; this one only accepts alignment-reporting + * normalizers but can see across token boundaries, for example a whitespace collapse + * spanning two tokens. + * + * @param normalizer The normalizer that defines term identity, applied to the whole + * document text so token spans can be mapped into the normalized + * form through its alignment. Must not be {@code null}. + * @param mode How much each {@link TermVector} records. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code normalizer} or {@code mode} is + * {@code null}. + */ + public TermVectorAnnotator(OffsetAwareNormalizer normalizer, Mode mode) { + if (normalizer == null) { + throw new IllegalArgumentException("normalizer must not be null"); + } + if (mode == null) { + throw new IllegalArgumentException("mode must not be null"); + } + this.normalizer = normalizer; + this.tokenNormalizer = null; + this.mode = mode; + } + + /** + * Aggregates the token layer into the {@link #TERM_VECTORS} layer. A present-but-empty + * token layer yields a present-but-empty term vector layer. + * + * @param document The document to annotate. Must not be {@code null} and must contain + * the {@link Layers#TOKENS} layer. + * @return A new {@link Document} with the {@link #TERM_VECTORS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null}, lacks + * the {@link Layers#TOKENS} layer, or already carries the + * {@link #TERM_VECTORS} layer. + */ + @Override + public Document annotate(Document document) { + if (document == null) { + throw new IllegalArgumentException("document must not be null"); + } + if (!document.layers().contains(Layers.TOKENS)) { + throw new IllegalArgumentException("document lacks the required layer " + + Layers.TOKENS); + } + final List> tokens = document.get(Layers.TOKENS); + final AlignedText aligned = + normalizer != null ? normalizer.normalizeAligned(document.text()) : null; + final String normalized = aligned != null ? aligned.normalizedString() : null; + return document.with(TERM_VECTORS, mode == Mode.FULL + ? fullVectors(tokens, aligned, normalized) + : countVectors(tokens, aligned, normalized)); + } + + /** + * Aggregates the tokens into full term vectors, keeping every occurrence span. + * + * @param tokens The token layer. + * @param aligned The normalized document text with its alignment, or {@code null} when + * no normalizer is present. + * @param normalized The normalized document text, or {@code null} when no normalizer is + * present. + * @return One annotation per distinct term, in first-occurrence order. + */ + private List> fullVectors(List> tokens, + AlignedText aligned, String normalized) { + final Map> spansByTerm = new LinkedHashMap<>(); + for (final Annotation token : tokens) { + final String term = termOf(token, aligned, normalized); + if (!term.isEmpty()) { + spansByTerm.computeIfAbsent(term, key -> new ArrayList<>()).add(token.span()); + } + } + final List> vectors = new ArrayList<>(spansByTerm.size()); + for (final Map.Entry> entry : spansByTerm.entrySet()) { + vectors.add(Annotation.of(TermVector.withSpans(entry.getKey(), entry.getValue()))); + } + return vectors; + } + + /** + * Aggregates the tokens into scoring-only term vectors, so no offset storage is ever + * allocated. + * + * @param tokens The token layer. + * @param aligned The normalized document text with its alignment, or {@code null} when + * no normalizer is present. + * @param normalized The normalized document text, or {@code null} when no normalizer is + * present. + * @return One annotation per distinct term, in first-occurrence order. + */ + private List> countVectors(List> tokens, + AlignedText aligned, String normalized) { + final Map frequencies = new LinkedHashMap<>(); + for (final Annotation token : tokens) { + final String term = termOf(token, aligned, normalized); + if (!term.isEmpty()) { + frequencies.merge(term, 1, Integer::sum); + } + } + final List> vectors = new ArrayList<>(frequencies.size()); + for (final Map.Entry entry : frequencies.entrySet()) { + vectors.add(Annotation.of(TermVector.count(entry.getKey(), entry.getValue()))); + } + return vectors; + } + + /** {@inheritDoc} */ + @Override + public Set> requires() { + return Set.of(Layers.TOKENS); + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(TERM_VECTORS); + } + + /** + * Determines the term one token groups under: its covered text as-is, its covered text + * normalized on its own when a plain per-token normalizer is present, or the covered + * text of its span mapped into the normalized form when an offset-aware normalizer is + * present. + * + * @param token The token annotation. + * @param aligned The normalized document text with its alignment, or {@code null} + * when no offset-aware normalizer is present. + * @param normalized The normalized document text, or {@code null} when no offset-aware + * normalizer is present. + * @return The term string. Never {@code null}, possibly empty. + */ + private String termOf(Annotation token, AlignedText aligned, String normalized) { + if (tokenNormalizer != null) { + return tokenNormalizer.normalize(token.value()).toString(); + } + if (aligned == null) { + return token.value(); + } + final Span span = aligned.toNormalizedSpan(token.span().getStart(), token.span().getEnd()); + return normalized.substring(span.getStart(), span.getEnd()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/TokenizerAnnotator.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/TokenizerAnnotator.java new file mode 100644 index 0000000000..7760b565d3 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/TokenizerAnnotator.java @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnnotator; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +/** + * Adapts a {@link Tokenizer} to the document pipeline: provides {@link Layers#TOKENS}. + * + *

    When {@link Layers#SENTENCES} is present, each sentence is tokenized separately and + * the token spans are shifted back to document coordinates; a present-but-empty sentence + * layer therefore yields a present-but-empty token layer. Only when the sentence layer + * is absent is the whole text tokenized at once. Either way, every token span refers to + * the original document text.

    + * + * @since 3.0.0 + */ +public final class TokenizerAnnotator implements DocumentAnnotator { + + private final Tokenizer tokenizer; + + /** + * Initializes the adapter. + * + * @param tokenizer The tokenizer to delegate to. Must not be {@code null}. + * @throws IllegalArgumentException Thrown if {@code tokenizer} is {@code null}. + */ + public TokenizerAnnotator(Tokenizer tokenizer) { + if (tokenizer == null) { + throw new IllegalArgumentException("tokenizer must not be null"); + } + this.tokenizer = tokenizer; + } + + /** + * Tokenizes the document and adds the {@link Layers#TOKENS} layer, sentence by + * sentence when a sentence layer is present and over the whole text otherwise. + * + * @param document The document to annotate. Must not be {@code null}. + * @return A new {@link Document} with the {@link Layers#TOKENS} layer added. Never + * {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null} or + * already carries the {@link Layers#TOKENS} layer. + */ + @Override + public Document annotate(Document document) { + if (document == null) { + throw new IllegalArgumentException("document must not be null"); + } + final String text = document.text().toString(); + final List> tokens = new ArrayList<>(); + if (!document.layers().contains(Layers.SENTENCES)) { + addTokens(tokens, text, 0); + } else { + for (final Annotation sentence : document.get(Layers.SENTENCES)) { + final Span span = sentence.span(); + addTokens(tokens, text.substring(span.getStart(), span.getEnd()), span.getStart()); + } + } + return document.with(Layers.TOKENS, tokens); + } + + /** + * Tokenizes one stretch of text and appends its tokens, shifted back into document + * coordinates. + * + * @param tokens The layer under construction. + * @param text The stretch to tokenize. + * @param offset The stretch's start offset in the document text. + */ + private void addTokens(List> tokens, String text, int offset) { + for (final Span span : tokenizer.tokenizePos(text)) { + final Span shifted = new Span(span.getStart() + offset, span.getEnd() + offset); + tokens.add(new Annotation<>(shifted, span.getCoveredText(text).toString())); + } + } + + /** {@inheritDoc} */ + @Override + public Set> provides() { + return Set.of(Layers.TOKENS); + } + + /** + * {@return the adapter's simple class name, which names it in pipeline validation + * messages} + */ + @Override + public String toString() { + return getClass().getSimpleName(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstaller.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstaller.java new file mode 100644 index 0000000000..bfd94e34fb --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstaller.java @@ -0,0 +1,271 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.IOException; +import java.net.URI; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.LinkOption; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.Stream; + +import opennlp.tools.util.DictionaryCatalog; +import opennlp.tools.util.ResourceInstaller; + +/** + * Fetches and unpacks a MeCab-format dictionary archive into a local directory, so the + * dictionary is acquired by the user at install time and never ships with this library. + * No dictionary data is bundled. Fetching, verification, and unpacking are done by + * {@link ResourceInstaller} under {@link ResourceInstaller.Limits#DEFAULT}, including + * startup property overrides. An {@code http} or {@code https} archive requires an + * expected checksum, and the ustar, pax, and GNU tar formats are all read. Catalog + * installs are opt-in via + * {@link #installFromCatalog(DictionaryCatalog, String, Path)}. + * + *

    Only the dictionary payload is installed: the {@code *.csv} lexicon files and + * {@code *.def} definition files that a {@link MecabDictionary} reads, plus the + * {@code dicrc} configuration file distributions ship alongside them, taken from the + * archive root only (at most one leading directory deep). Deeper entries are skipped: + * mecab-ko-dic, for example, nests {@code user-dic} templates whose numeric fields are + * empty because they are input for {@code mecab-dict-index}, not loadable lexicon + * data. Installed files are flattened to their base names. A file whose base name + * already exists in the target is not replaced, so it must be removed before refreshing + * a dictionary. The archive unpacks into a hidden scratch directory beneath the target, + * on the target's filesystem, which is removed when the installation ends.

    + * + * @since 3.0.0 + */ +public final class MecabDictionaryInstaller { + + /** The deepest entry path, relative to the archive root, that holds payload. */ + private static final int MAX_PAYLOAD_DEPTH = 2; + + /** The hidden scratch directory beneath the target that the archive unpacks into. */ + private static final String SCRATCH_PREFIX = ".mecab-dict-"; + + /** Prevents construction of this utility class. */ + private MecabDictionaryInstaller() { + } + + /** + * Unpacks a local {@code file:} archive URI. Any other scheme requires + * {@link #install(URI, Path, String)} with an expected checksum. + * + * @param archive The archive location, a gzip-compressed tar. Must not be + * {@code null}. + * @param targetDirectory The directory to unpack into; created when absent. Must not + * be {@code null}. + * @return The number of dictionary files installed. + * @throws IOException Thrown if reading or writing fails, the archive contains no + * dictionary file, an installation limit is exceeded, or the target already + * contains one of the files. + * @throws IllegalArgumentException Thrown if a parameter is {@code null} or + * {@code archive} does not use the {@code file} scheme. + */ + public static int install(URI archive, Path targetDirectory) throws IOException { + return install(archive, targetDirectory, null); + } + + /** + * Downloads a dictionary archive when needed, verifies its checksum, and unpacks it + * through {@link ResourceInstaller#install(URI, Path, String)}. A {@code file:} URI + * may omit the checksum. + * + * @param archive The archive location, a gzip-compressed tar. Must not be + * {@code null}. + * @param targetDirectory The directory to unpack into; created when absent. Must not + * be {@code null}. + * @param expectedChecksum The expected digest of the archive bytes as a hex string, + * 64 characters for SHA-256 or 128 for SHA-512. Required for + * an http or https source; pass {@code null} to skip + * verification for a file source. + * @return The number of dictionary files installed. + * @throws IOException Thrown if fetching, verification, reading, or writing fails, + * the archive contains no dictionary file, an installation limit is + * exceeded, or the target already contains one of the files. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}, the URI is + * not supported by {@link ResourceInstaller}, or an http or https source + * has no checksum. + */ + public static int install(URI archive, Path targetDirectory, String expectedChecksum) + throws IOException { + if (archive == null) { + throw new IllegalArgumentException("archive must not be null"); + } + if (targetDirectory == null) { + throw new IllegalArgumentException("targetDirectory must not be null"); + } + final Path unpacked = createScratch(targetDirectory); + try { + ResourceInstaller.install(archive, unpacked, expectedChecksum); + return promoteDictionaryFiles(unpacked, targetDirectory); + } finally { + deleteRecursively(unpacked); + } + } + + /** + * Downloads a dictionary named in an application-supplied + * {@link DictionaryCatalog} and unpacks it. Requires + * {@code -Dopennlp.download.remote=true}. + * + * @param catalog The application-supplied catalog. Must not be {@code null}. + * @param dictionaryId The catalog id, for example {@code mecab.ipadic} or + * {@code mecab.ko-dic}. Must not be {@code null}. + * @param targetDirectory The directory to unpack into; created when absent. Must not + * be {@code null}. + * @return The number of dictionary files installed. + * @throws IOException Thrown if the catalog entry is missing, remote downloads are + * disabled, or install fails. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public static int installFromCatalog(DictionaryCatalog catalog, String dictionaryId, + Path targetDirectory) throws IOException { + if (catalog == null) { + throw new IllegalArgumentException("catalog must not be null"); + } + if (dictionaryId == null) { + throw new IllegalArgumentException("dictionaryId must not be null"); + } + if (targetDirectory == null) { + throw new IllegalArgumentException("targetDirectory must not be null"); + } + final Path unpacked = createScratch(targetDirectory); + try { + catalog.install(dictionaryId, unpacked); + return promoteDictionaryFiles(unpacked, targetDirectory); + } finally { + deleteRecursively(unpacked); + } + } + + /** + * Creates the scratch directory the archive unpacks into. It lives beneath the target + * so the download, the unpacked tree, and the installed files share one filesystem + * and a large dictionary cannot fill the system temporary directory. Scratch + * directories that an earlier installation left behind, because its process ended + * before cleanup, are removed first. + * + * @param targetDirectory The directory to install into; created when absent. + * @return The new scratch directory. Not {@code null}. + * @throws IOException Thrown if a directory cannot be created or a stale one removed. + */ + private static Path createScratch(Path targetDirectory) throws IOException { + Files.createDirectories(targetDirectory); + final List stale; + try (Stream entries = Files.list(targetDirectory)) { + stale = entries.filter(entry -> entry.getFileName().toString().startsWith(SCRATCH_PREFIX) + && Files.isDirectory(entry, LinkOption.NOFOLLOW_LINKS)).toList(); + } + for (final Path entry : stale) { + deleteRecursively(entry); + } + return Files.createTempDirectory(targetDirectory, SCRATCH_PREFIX); + } + + /** + * Moves the dictionary payload files from an unpacked archive tree into the target + * directory, flattened to their base names. All destinations are checked before the + * first move, so a collision leaves the target unchanged. + * + * @param unpacked The directory the archive was unpacked into. + * @param targetDirectory The directory to install into; created when absent. + * @return The number of dictionary files installed. + * @throws IOException Thrown if the tree holds no dictionary file, two entries + * flatten to the same base name, a target file already exists, or moving + * fails. + */ + private static int promoteDictionaryFiles(Path unpacked, Path targetDirectory) + throws IOException { + final List candidates; + try (Stream files = Files.walk(unpacked, MAX_PAYLOAD_DEPTH)) { + candidates = files.filter(Files::isRegularFile).toList(); + } + final List payload = new ArrayList<>(); + final Set baseNames = new HashSet<>(); + for (final Path file : candidates) { + final String baseName = file.getFileName().toString(); + if (!isDictionaryFile(baseName)) { + continue; + } + if (!baseNames.add(baseName)) { + throw new IOException( + "the archive flattens two entries to the same name: " + baseName); + } + payload.add(file); + } + if (payload.isEmpty()) { + throw new IOException("the archive contains no dictionary file"); + } + for (final Path file : payload) { + final Path destination = targetDirectory.resolve(file.getFileName().toString()); + if (Files.exists(destination, LinkOption.NOFOLLOW_LINKS)) { + throw new IOException("target already contains: " + destination); + } + } + for (final Path file : payload) { + Files.move(file, targetDirectory.resolve(file.getFileName().toString())); + } + return payload.size(); + } + + /** + * Recognizes the file names a {@link MecabDictionary} loads. + * + * @param baseName The file name without any directory prefix. + * @return {@code true} when the name is dictionary payload. + */ + private static boolean isDictionaryFile(String baseName) { + return baseName.endsWith(".csv") || baseName.endsWith(".def") + || "dicrc".equals(baseName); + } + + /** + * Deletes a directory tree, deepest entries first. + * + * @param root The directory to remove. + * @throws IOException Thrown if a deletion fails. + */ + private static void deleteRecursively(Path root) throws IOException { + Files.walkFileTree(root, new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attributes) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path directory, IOException error) + throws IOException { + if (error != null) { + throw error; + } + Files.delete(directory); + return FileVisitResult.CONTINUE; + } + }); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DictionaryCatalog.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DictionaryCatalog.java new file mode 100644 index 0000000000..5e3a2de017 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DictionaryCatalog.java @@ -0,0 +1,197 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.util; + +import java.io.IOException; +import java.io.InputStream; +import java.net.URI; +import java.net.URISyntaxException; +import java.nio.file.Path; +import java.util.Collections; +import java.util.HexFormat; +import java.util.LinkedHashSet; +import java.util.Properties; +import java.util.Set; + +/** + * Opt-in catalog of remote dictionary archives and companion files, loaded from + * application-supplied properties containing URLs and SHA-512 digests. Fetching an + * entry requires {@link #REMOTE_DOWNLOAD_PROPERTY} to be {@code true}. + * + * @since 3.0.0 + */ +public final class DictionaryCatalog { + + private static final int SHA_512_HEX_LENGTH = 128; + private static final String URL_SUFFIX = ".url"; + + /** + * System property that must be {@code true} before a catalog entry may be + * fetched. Direct {@link ResourceInstaller} calls do not require it: there the + * caller already supplied the URI and digest. + */ + public static final String REMOTE_DOWNLOAD_PROPERTY = "opennlp.download.remote"; + private final Properties properties; + + /** + * Initializes a catalog from loaded properties. + * + * @param properties The catalog properties. + */ + private DictionaryCatalog(Properties properties) { + this.properties = properties; + } + + /** + * Loads a catalog from a properties stream. + * + * @param in The properties content. Must not be {@code null}. + * @return The catalog. Never {@code null}. + * @throws IOException Thrown if reading fails. + * @throws IllegalArgumentException Thrown if {@code in} is {@code null}. + */ + public static DictionaryCatalog load(InputStream in) throws IOException { + if (in == null) { + throw new IllegalArgumentException("in must not be null"); + } + final Properties properties = new Properties(); + properties.load(in); + return new DictionaryCatalog(properties); + } + + /** + * {@return an unmodifiable set of catalog entry ids} + */ + public Set ids() { + final Set ids = new LinkedHashSet<>(); + for (final String key : properties.stringPropertyNames()) { + if (key.endsWith(URL_SUFFIX)) { + ids.add(key.substring(0, key.length() - URL_SUFFIX.length())); + } + } + return Collections.unmodifiableSet(ids); + } + + /** + * Looks up one catalog entry. + * + * @param id The entry id, for example {@code mecab.ipadic}. + * @return The entry. Never {@code null}. + * @throws IOException Thrown if the entry is incomplete or invalid. + * @throws IllegalArgumentException Thrown if {@code id} is {@code null}. + */ + public Entry get(String id) throws IOException { + if (id == null) { + throw new IllegalArgumentException("id must not be null"); + } + final String url = properties.getProperty(id + URL_SUFFIX); + final String sha512 = properties.getProperty(id + ".sha512"); + if (url == null || sha512 == null) { + throw new IOException("unknown or incomplete dictionary catalog entry: " + id); + } + final String filename = properties.getProperty(id + ".filename"); + try { + return new Entry(id, new URI(url), sha512.trim(), filename); + } catch (URISyntaxException | IllegalArgumentException e) { + throw new IOException("invalid dictionary catalog entry: " + id, e); + } + } + + /** + * Installs a catalog entry into {@code targetDirectory} after checking that remote + * catalog downloads are enabled. The entry is fetched, digest-verified, and unpacked + * by {@link ResourceInstaller#install(URI, Path, String)}: an archive expands into + * the directory, and a plain file is stored under its source name. + * + * @param id The entry id. Must not be {@code null}. + * @param targetDirectory The directory to install into; created when absent. Must + * not be {@code null}. + * @throws IOException Thrown if the property is not enabled, the entry is missing, + * the download fails verification, or the target already contains an + * installed file. + * @throws IllegalArgumentException Thrown if a parameter is {@code null}. + */ + public void install(String id, Path targetDirectory) throws IOException { + if (id == null) { + throw new IllegalArgumentException("id must not be null"); + } + if (targetDirectory == null) { + throw new IllegalArgumentException("targetDirectory must not be null"); + } + if (!Boolean.getBoolean(REMOTE_DOWNLOAD_PROPERTY)) { + throw new IOException("remote dictionary catalog downloads are disabled; set -D" + + REMOTE_DOWNLOAD_PROPERTY + "=true to enable"); + } + final Entry entry = get(id); + if (entry.filename() == null) { + ResourceInstaller.install(entry.uri(), targetDirectory, entry.sha512()); + } else { + ResourceInstaller.installNamed( + entry.uri(), targetDirectory, entry.sha512(), entry.filename()); + } + } + + /** + * One catalog entry with a URI and the SHA-512 digest of its bytes. + * + * @param id The catalog id. + * @param uri The absolute download URI. + * @param sha512 The expected SHA-512 hex digest. + * @param filename An optional preferred local file name; may be {@code null}. + */ + public record Entry(String id, URI uri, String sha512, String filename) { + /** + * @param id The catalog id. Must not be {@code null}. + * @param uri The absolute download URI. Must not be {@code null}. + * @param sha512 The expected SHA-512 hex digest. Must not be {@code null}. + * @param filename An optional preferred local file name; may be {@code null}. + * @throws IllegalArgumentException Thrown if a required value is {@code null}, + * {@code uri} is relative, {@code sha512} is not 128 hex digits, or + * {@code filename} is not a local file name. + */ + public Entry { + if (id == null) { + throw new IllegalArgumentException("id must not be null"); + } + if (uri == null) { + throw new IllegalArgumentException("uri must not be null"); + } + if (!uri.isAbsolute()) { + throw new IllegalArgumentException("uri must be absolute"); + } + if (sha512 == null) { + throw new IllegalArgumentException("sha512 must not be null"); + } + if (sha512.length() != SHA_512_HEX_LENGTH) { + throw new IllegalArgumentException("sha512 must be 128 hex digits"); + } + try { + HexFormat.of().parseHex(sha512); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("sha512 must be 128 hex digits", e); + } + if (filename != null) { + try { + ResourceInstaller.validateSourceName(filename); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("filename must be a file name", e); + } + } + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DownloadUtil.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DownloadUtil.java index 7554c064b3..0f6f6800d0 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DownloadUtil.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/DownloadUtil.java @@ -62,6 +62,7 @@ public class DownloadUtil { private static final String MODEL_URI_PATH = System.getProperty("OPENNLP_DOWNLOAD_MODEL_PATH", "models/ud-models-1.3/"); private static final String OPENNLP_DOWNLOAD_HOME = "OPENNLP_DOWNLOAD_HOME"; + private static final String CHECKSUM_EXTENSION = ".sha512"; private static Map> availableModels; @@ -89,7 +90,7 @@ static boolean existsModel(String language, ModelType modelType) throws IOExcept boolean exists; if (Files.exists(localFile)) { // if this does not throw the requested model is valid! - validateModel(url + ".sha512", localFile); + validateCachedModel(url + CHECKSUM_EXTENSION, localFile); exists = true; } else { exists = false; @@ -130,7 +131,8 @@ public static T downloadModel(String language, ModelType m * The model is saved to an {@code .opennlp/} directory * located in the user's home directory. This directory will be created * if it does not already exist. If a model to be downloaded already - * exists in that directory, the model will not be re-downloaded. + * exists in that directory, the model will not be re-downloaded, but it is + * verified against its SHA-512 checksum before it is loaded. * * @param url The model's {@link URL}. * @param type The class of the resulting model {@link T}. @@ -159,10 +161,11 @@ public static T downloadModel(URL url, Class type) thro try (final InputStream in = url.openStream()) { Files.copy(in, localFile, StandardCopyOption.REPLACE_EXISTING); } - validateModel(url + ".sha512", localFile); + validateModel(url + CHECKSUM_EXTENSION, localFile); logger.debug("Download complete."); } else { logger.debug("Model file '{}' already exists. Skipping download.", filename); + validateCachedModel(url + CHECKSUM_EXTENSION, localFile); } try { @@ -185,37 +188,108 @@ public static Map> getAvailableModels() { } /** - * Validates a downloaded model via the specified {@link Path downloadedModel path}. + * Validates a freshly downloaded model via the specified {@link Path downloadedModel path} + * and stores the expected checksum next to it, so that subsequent loads of the cached file + * can be verified without contacting the CDN again. * * @param sha512 the url to get the sha512 hash * @param downloadedModel the model file to check - * @throws IOException thrown if the checksum could not be computed + * @throws IOException thrown if the checksum could not be computed or did not match */ private static void validateModel(String sha512, Path downloadedModel) throws IOException { - String expectedChecksum; + final String checksumFile = downloadChecksumFile(sha512, downloadedModel); + verifyChecksum(downloadedModel, parseChecksum(checksumFile)); + storeChecksumFile(downloadedModel, checksumFile); + } + + /** + * Validates a model that is already present in the download home. + *

    + * The expected checksum stored by a previous download is used, so no network access is + * required on this path. When no checksum was stored - the download home was populated by + * an OpenNLP version that predates this check - the published checksum is fetched once and + * then stored. If it cannot be retrieved, the model is loaded and a warning is logged, + * which preserves the behaviour of those earlier versions for offline environments. + * + * @param sha512 the url to get the sha512 hash + * @param cachedModel the cached model file to check + * @throws IOException thrown if the checksum could not be computed or did not match + */ + private static void validateCachedModel(String sha512, Path cachedModel) throws IOException { + final Path checksumFile = checksumPathFor(cachedModel); + + if (Files.exists(checksumFile)) { + verifyChecksum(cachedModel, parseChecksum(Files.readString(checksumFile, StandardCharsets.UTF_8))); + return; + } + + final String publishedChecksumFile; + try { + publishedChecksumFile = downloadChecksumFile(sha512, cachedModel); + } catch (IOException e) { + logger.warn("Could not retrieve the expected checksum for cached model '{}'. " + + "Its integrity has not been verified.", cachedModel.getFileName(), e); + return; + } + + verifyChecksum(cachedModel, parseChecksum(publishedChecksumFile)); + storeChecksumFile(cachedModel, publishedChecksumFile); + } + + /** + * Retrieves the published {@code ".sha512"} file. Its content is returned unmodified so that + * the copy stored next to the model stays interchangeable with the published one. + */ + private static String downloadChecksumFile(String sha512, Path model) throws IOException { try { // Download SHA512 checksum file final URL hashSum = new URI(sha512).toURL(); try (BufferedReader reader = new BufferedReader(new InputStreamReader(hashSum.openStream()))) { - expectedChecksum = reader.readLine(); - - if (expectedChecksum != null) { - expectedChecksum = expectedChecksum.split("\\s")[0].trim(); - } + return reader.readLine(); } } catch (URISyntaxException use) { throw new IOException("Expected SHA512 checksum could not be retrieved for " + - downloadedModel.getFileName(), use); + model.getFileName(), use); + } + } + + /** + * Extracts the hash from the content of a checksum file, which holds the hash followed by the + * name of the file it applies to. + */ + private static String parseChecksum(String checksumFileContent) { + if (checksumFileContent == null) { + return null; } + final String trimmed = checksumFileContent.trim(); + return trimmed.isEmpty() ? null : trimmed.split("\\s")[0]; + } - // Validate SHA512 checksum - final String actualChecksum = calculateSHA512(downloadedModel); + private static void verifyChecksum(Path model, String expectedChecksum) throws IOException { + final String actualChecksum = calculateSHA512(model); if (!actualChecksum.equalsIgnoreCase(expectedChecksum)) { - throw new IOException("SHA512 checksum validation failed for " + downloadedModel.getFileName() + + throw new IOException("SHA512 checksum validation failed for " + model.getFileName() + ". Expected: " + expectedChecksum + ", but got: " + actualChecksum); } } + /** + * Stores the published checksum file alongside the model. A failure to do so is not fatal: + * it only means the next load falls back to retrieving the published checksum again. + */ + private static void storeChecksumFile(Path model, String checksumFileContent) { + final Path checksumFile = checksumPathFor(model); + try { + Files.writeString(checksumFile, checksumFileContent, StandardCharsets.UTF_8); + } catch (IOException e) { + logger.warn("Could not store the expected checksum at {}.", checksumFile, e); + } + } + + private static Path checksumPathFor(Path model) { + return model.resolveSibling(model.getFileName() + CHECKSUM_EXTENSION); + } + private static String calculateSHA512(Path file) throws IOException { try { MessageDigest digest = MessageDigest.getInstance("SHA-512"); diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/ResourceInstaller.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/ResourceInstaller.java new file mode 100644 index 0000000000..54b215b6e3 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/ResourceInstaller.java @@ -0,0 +1,1558 @@ +/* + * 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.util; + +import java.io.BufferedInputStream; +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.HttpURLConnection; +import java.net.URI; +import java.nio.file.FileSystem; +import java.nio.file.FileSystemNotFoundException; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.InvalidPathException; +import java.nio.file.LinkOption; +import java.nio.file.Path; +import java.nio.file.ProviderNotFoundException; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.util.Comparator; +import java.util.Enumeration; +import java.util.HashSet; +import java.util.HexFormat; +import java.util.List; +import java.util.Set; +import java.util.stream.Stream; +import java.util.zip.GZIPInputStream; +import java.util.zip.ZipEntry; +import java.util.zip.ZipException; +import java.util.zip.ZipFile; +import java.util.zip.ZipInputStream; + +import opennlp.tools.util.archive.TarStream; + +/** + * Fetches a third-party resource, such as a training corpus, a dictionary archive, or a + * lexicon, into a local directory. The caller supplies the location and thereby accepts + * that resource's license; no locations are built in and no data is bundled. Only + * {@code http}, {@code https}, and {@code file} locations are accepted. + * + *

    A checksum is required for http and https sources and optional for file sources. + * It is verified against the downloaded bytes before anything is unpacked: a + * 64-character hex digest selects SHA-256, a 128-character one SHA-512. + * The content format is detected from the bytes, not from the name: gzip-compressed + * tar archives and zip archives are unpacked with their relative structure. Invalid, + * escaping, and duplicate file paths are rejected. Plain gzip files are decompressed, + * and other content is stored as a file under the source name. One name rule overrides + * byte detection: a {@code *.bin} source is always stored packed because an OpenNLP + * model file is itself a zip archive that its consumers load packed.

    + * + *

    Each installation is bounded by {@link Limits}: http and https fetches use + * connection and read timeouts, follow at most a fixed number of redirects, reject + * redirects that leave the http and https schemes or downgrade https to http, and + * abort once the download or the expanded content crosses its size limit or the + * archive crosses its entry limit. Compressed content, gzip and zip alike, may + * expand to at most {@link Limits#maxExpansionRatio()} times its compressed size, + * with a floor of {@value #MIN_EXPANSION_BYTES} bytes for small sources. The defaults in + * {@link Limits#DEFAULT} apply when no limits are given, and {@link Limits#builder()} + * starts from them.

    + * + *

    Installation is staged: content is unpacked into a hidden staging directory on + * the same filesystem and moved into the target only after the download was verified + * and every entry unpacked cleanly. A fetch, verification, or unpacking failure + * promotes no files into the target directory. + * Promotion does not replace a file that already exists in the target and detects + * the collision before moving anything, so refreshing a resource means removing its + * old files first. Work files left in the target by an installation that was killed + * are removed at the start of the next installation into that target, so concurrent + * installations into one target directory are not supported.

    + * + * @see DownloadUtil + * @since 3.0.0 + */ +public final class ResourceInstaller { + + private static final String SHA_256 = "SHA-256"; + private static final String SHA_512 = "SHA-512"; + private static final int SHA_256_HEX_LENGTH = 64; + private static final int SHA_512_HEX_LENGTH = 128; + private static final String GZIP_SUFFIX = ".gz"; + + /** OpenNLP model files are packed zip archives; install them packed. */ + private static final String MODEL_SUFFIX = ".bin"; + private static final String DEFAULT_RESOURCE_NAME = "resource"; + private static final String STAGING_PREFIX = ".opennlp-staging"; + private static final String DOWNLOAD_PREFIX = ".opennlp-download"; + private static final String DOWNLOAD_SUFFIX = ".part"; + private static final int BUFFER_SIZE = 8192; + private static final int MAGIC_LENGTH = 4; + private static final int GZIP_MAGIC_FIRST = 0x1F; + private static final int GZIP_MAGIC_SECOND = 0x8B; + private static final int ZIP_MAGIC_FIRST = 'P'; + private static final int ZIP_MAGIC_SECOND = 'K'; + private static final int ZIP_LOCAL_HEADER_THIRD = 3; + private static final int ZIP_LOCAL_HEADER_FOURTH = 4; + private static final int ZIP_END_HEADER_THIRD = 5; + private static final int ZIP_END_HEADER_FOURTH = 6; + private static final int ZIP_END_HEADER_LENGTH = 22; + private static final int ZIP_DISK_OFFSET = 4; + private static final int ZIP_CENTRAL_DISK_OFFSET = 6; + private static final int ZIP_DISK_ENTRIES_OFFSET = 8; + private static final int ZIP_TOTAL_ENTRIES_OFFSET = 10; + private static final int ZIP_CENTRAL_SIZE_OFFSET = 12; + private static final int ZIP_CENTRAL_OFFSET_OFFSET = 16; + private static final int ZIP_COMMENT_LENGTH_OFFSET = 20; + private static final int TAR_END_BLOCKS_LENGTH = 1024; + + /** The expanded size every compressed source may reach regardless of the ratio. */ + private static final long MIN_EXPANSION_BYTES = 1L << 20; + private static final int HTTP_TEMPORARY_REDIRECT = 307; + private static final int HTTP_PERMANENT_REDIRECT = 308; + private static final String SCHEME_HTTP = "http"; + private static final String SCHEME_HTTPS = "https"; + private static final String SCHEME_FILE = "file"; + private static final String MALFORMED_ZIP_ERROR = "malformed zip archive"; + private static final String ZIP_MISMATCH_ERROR = + "zip local headers and central directory list different files"; + + /** + * Safety limits and network behavior for one installation. + * + * @param connectTimeout How long to wait for a connection to be established. Must + * be positive. + * @param readTimeout How long to wait for data on an established connection. Must + * be positive. + * @param maxRedirects How many http redirects to follow before failing. Must not be + * negative; zero rejects all redirects. + * @param maxDownloadBytes The largest download accepted, in bytes. Must be positive. + * @param maxExpandedBytes The largest expanded byte count accepted. For gzip content, + * this counts the entire decompressed stream; otherwise, it + * counts installed file content. Must be positive. + * @param maxEntries The largest number of archive entries accepted, counting every + * entry including directories, so an archive of many tiny files + * cannot exhaust directory entries while staying under the byte + * limits. Must be positive. + * @param maxExpansionRatio The largest expanded size accepted per compressed byte of + * a source, so a small archive cannot expand to the whole + * expansion limit. Applies to gzip and zip content; a source + * may always expand to {@value #MIN_EXPANSION_BYTES} bytes + * regardless of it. Must be positive. + */ + public record Limits(Duration connectTimeout, Duration readTimeout, int maxRedirects, + long maxDownloadBytes, long maxExpandedBytes, long maxEntries, + long maxExpansionRatio) { + + /** The system property overriding the default download limit in bytes. */ + public static final String MAX_DOWNLOAD_BYTES_PROPERTY = "opennlp.download.max.bytes"; + + /** The system property overriding the default expansion limit in bytes. */ + public static final String MAX_EXPANDED_BYTES_PROPERTY = + "opennlp.install.max.total.bytes"; + + /** The system property overriding the default archive entry limit. */ + public static final String MAX_ENTRIES_PROPERTY = "opennlp.install.max.entries"; + + /** The system property overriding the default expansion ratio. */ + public static final String MAX_EXPANSION_RATIO_PROPERTY = + "opennlp.install.max.expansion.ratio"; + + /** + * The limits applied when none are given: 20 second connect timeout, 60 second + * read timeout, at most 5 redirects, a 1 GiB download limit, a 4 GiB expansion + * limit, 100000 archive entries, and an expansion ratio of 100. Each limit can be + * raised or lowered at startup through its system property + * ({@link #MAX_DOWNLOAD_BYTES_PROPERTY}, {@link #MAX_EXPANDED_BYTES_PROPERTY}, + * {@link #MAX_ENTRIES_PROPERTY}, {@link #MAX_EXPANSION_RATIO_PROPERTY}), read once + * at class load; a value that is absent, not a number, or not positive falls back + * to the built-in default. + */ + public static final Limits DEFAULT = new Limits(Duration.ofSeconds(20), + Duration.ofSeconds(60), 5, + longProperty(MAX_DOWNLOAD_BYTES_PROPERTY, 1L << 30), + longProperty(MAX_EXPANDED_BYTES_PROPERTY, 4L << 30), + longProperty(MAX_ENTRIES_PROPERTY, 100_000L), + longProperty(MAX_EXPANSION_RATIO_PROPERTY, 100L)); + + /** + * Reads a limit override from a system property, trimmed before parsing. + * + * @param name The property name. + * @param fallback The built-in default. + * @return The property's value, or {@code fallback} when the property is absent, + * not a number, or not positive. + */ + static long longProperty(String name, long fallback) { + final String value = System.getProperty(name); + if (value == null) { + return fallback; + } + final long parsed; + try { + parsed = Long.parseLong(value.trim()); + } catch (NumberFormatException e) { + return fallback; + } + return parsed > 0 ? parsed : fallback; + } + + /** + * Validates the limit values before constructing an instance. + * + * @param connectTimeout How long to wait for a connection to be established. + * @param readTimeout How long to wait for data on an established connection. + * @param maxRedirects How many http redirects to follow before failing. + * @param maxDownloadBytes The largest download accepted, in bytes. + * @param maxExpandedBytes The largest expanded byte count accepted. + * @param maxEntries The largest number of archive entries accepted. + * @param maxExpansionRatio The largest expanded size accepted per compressed byte. + * @throws IllegalArgumentException Thrown if either timeout is {@code null}, zero, + * or negative, a limit is not positive, or the redirect limit is + * negative. + */ + public Limits(Duration connectTimeout, Duration readTimeout, int maxRedirects, + long maxDownloadBytes, long maxExpandedBytes, long maxEntries, + long maxExpansionRatio) { + if (connectTimeout == null) { + throw new IllegalArgumentException("connectTimeout must not be null"); + } + if (connectTimeout.isZero() || connectTimeout.isNegative()) { + throw new IllegalArgumentException("connectTimeout must be positive"); + } + if (readTimeout == null) { + throw new IllegalArgumentException("readTimeout must not be null"); + } + if (readTimeout.isZero() || readTimeout.isNegative()) { + throw new IllegalArgumentException("readTimeout must be positive"); + } + if (maxRedirects < 0) { + throw new IllegalArgumentException("maxRedirects must not be negative"); + } + if (maxDownloadBytes <= 0) { + throw new IllegalArgumentException("maxDownloadBytes must be positive"); + } + if (maxExpandedBytes <= 0) { + throw new IllegalArgumentException("maxExpandedBytes must be positive"); + } + if (maxEntries <= 0) { + throw new IllegalArgumentException("maxEntries must be positive"); + } + if (maxExpansionRatio <= 0) { + throw new IllegalArgumentException("maxExpansionRatio must be positive"); + } + this.connectTimeout = connectTimeout; + this.readTimeout = readTimeout; + this.maxRedirects = maxRedirects; + this.maxDownloadBytes = maxDownloadBytes; + this.maxExpandedBytes = maxExpandedBytes; + this.maxEntries = maxEntries; + this.maxExpansionRatio = maxExpansionRatio; + } + + /** + * Starts from {@link #DEFAULT} so a caller can state only the limits that differ + * from it, instead of repeating all seven in the canonical constructor. + * + * @return A builder holding the default limits. Not {@code null}. + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Collects limit values and validates them on {@link #build()}. Each setter returns + * this builder. Not thread safe; the {@link Limits} it builds is immutable. + */ + public static final class Builder { + + private Duration connectTimeout = DEFAULT.connectTimeout(); + private Duration readTimeout = DEFAULT.readTimeout(); + private int maxRedirects = DEFAULT.maxRedirects(); + private long maxDownloadBytes = DEFAULT.maxDownloadBytes(); + private long maxExpandedBytes = DEFAULT.maxExpandedBytes(); + private long maxEntries = DEFAULT.maxEntries(); + private long maxExpansionRatio = DEFAULT.maxExpansionRatio(); + + /** Initializes a builder with {@link Limits#DEFAULT}. */ + private Builder() { + } + + /** + * Sets how long to wait for a connection to be established. + * + * @param connectTimeout How long to wait for a connection to be established. + * Must be positive. + * @return This builder. Not {@code null}. + */ + public Builder connectTimeout(Duration connectTimeout) { + this.connectTimeout = connectTimeout; + return this; + } + + /** + * Sets how long to wait for data on an established connection. + * + * @param readTimeout How long to wait for data on an established connection. + * Must be positive. + * @return This builder. Not {@code null}. + */ + public Builder readTimeout(Duration readTimeout) { + this.readTimeout = readTimeout; + return this; + } + + /** + * Sets how many http redirects to follow before failing. + * + * @param maxRedirects How many http redirects to follow before failing. Must not + * be negative; zero rejects all redirects. + * @return This builder. Not {@code null}. + */ + public Builder maxRedirects(int maxRedirects) { + this.maxRedirects = maxRedirects; + return this; + } + + /** + * Sets the largest download accepted. + * + * @param maxDownloadBytes The largest download accepted, in bytes. Must be + * positive. + * @return This builder. Not {@code null}. + */ + public Builder maxDownloadBytes(long maxDownloadBytes) { + this.maxDownloadBytes = maxDownloadBytes; + return this; + } + + /** + * Sets the largest expanded byte count accepted. + * + * @param maxExpandedBytes The largest expanded byte count accepted. For gzip + * content, this counts the entire decompressed stream; + * otherwise, it counts installed file content. Must be + * positive. + * @return This builder. Not {@code null}. + */ + public Builder maxExpandedBytes(long maxExpandedBytes) { + this.maxExpandedBytes = maxExpandedBytes; + return this; + } + + /** + * Sets the largest number of archive entries accepted. + * + * @param maxEntries The largest number of archive entries accepted, counting + * every entry including directories. Must be positive. + * @return This builder. Not {@code null}. + */ + public Builder maxEntries(long maxEntries) { + this.maxEntries = maxEntries; + return this; + } + + /** + * Sets the largest expanded size accepted per compressed byte of a source. + * + * @param maxExpansionRatio The largest expanded size accepted per compressed + * byte of a source. Must be positive. + * @return This builder. Not {@code null}. + */ + public Builder maxExpansionRatio(long maxExpansionRatio) { + this.maxExpansionRatio = maxExpansionRatio; + return this; + } + + /** + * Builds the limits. + * + * @return The limits collected so far. Not {@code null}. + * @throws IllegalArgumentException Thrown if any value is outside its documented + * range. + */ + public Limits build() { + return new Limits(connectTimeout, readTimeout, maxRedirects, maxDownloadBytes, + maxExpandedBytes, maxEntries, maxExpansionRatio); + } + } + } + + /** Prevents construction of this utility class. */ + private ResourceInstaller() { + } + + /** + * Unpacks a resource without checksum verification, under {@link Limits#DEFAULT}. + * This overload treats the source as trusted caller input and performs no + * cryptographic integrity verification, so it accepts only {@code file} sources; an + * http or https source must go through an overload that takes its checksum. + * + * @param source The resource location, a {@code file} URI. Not {@code null}. + * @param targetDirectory The directory to install into; created when absent. Must + * not be {@code null}. + * @return The target directory. Not {@code null}. + * @throws IOException Thrown if fetching or unpacking fails. + * @throws IllegalArgumentException Thrown if {@code source} or + * {@code targetDirectory} is {@code null}, {@code source} contains a scheme + * other than {@code file}, or its last path segment is not a valid local file + * name. + */ + public static Path install(URI source, Path targetDirectory) throws IOException { + return install(source, targetDirectory, null); + } + + /** + * Fetches, verifies, and unpacks a resource under {@link Limits#DEFAULT}. + * + * @param source The resource location, an {@code http}, {@code https}, or + * {@code file} URI. Not {@code null}. + * @param targetDirectory The directory to install into; created when absent. Must + * not be {@code null}. + * @param checksum The expected digest of the downloaded bytes as a hex string, + * compared case-insensitively and ignoring leading and trailing + * whitespace: 64 characters select SHA-256, 128 characters SHA-512. + * Required for an http or https source; pass {@code null} to skip + * verification for a {@code file} source. + * @return The target directory. Not {@code null}. + * @throws IOException Thrown if fetching fails, the checksum does not match, or + * unpacking fails. + * @throws IllegalArgumentException Thrown if {@code source} or + * {@code targetDirectory} is {@code null}, {@code source} contains a scheme + * other than {@code http}, {@code https}, or {@code file}, {@code checksum} + * is not a 64-character or 128-character hex string, an http or https source + * contains no checksum, or the source does not provide a valid local file + * name. + */ + public static Path install(URI source, Path targetDirectory, String checksum) + throws IOException { + return install(source, targetDirectory, checksum, Limits.DEFAULT); + } + + /** + * Fetches, verifies, and unpacks a resource under the given {@link Limits}. + * + * @param source The resource location, an {@code http}, {@code https}, or + * {@code file} URI. Not {@code null}. + * @param targetDirectory The directory to install into; created when absent. Must + * not be {@code null}. + * @param checksum The expected digest of the downloaded bytes as a hex string, + * compared case-insensitively and ignoring leading and trailing + * whitespace: 64 characters select SHA-256, 128 characters SHA-512. + * Required for an http or https source; pass {@code null} to skip + * verification for a {@code file} source. + * @param limits The timeouts, redirect allowance, and size and entry limits to + * enforce. Not {@code null}. + * @return The target directory. Not {@code null}. + * @throws IOException Thrown if fetching fails, a limit is exceeded, the checksum + * does not match, or unpacking fails. + * @throws IllegalArgumentException Thrown if {@code source}, {@code targetDirectory}, + * or {@code limits} is {@code null}, {@code source} contains a scheme other + * than {@code http}, {@code https}, or {@code file}, {@code checksum} is + * not a 64-character or 128-character hex string, an http or https source + * contains no checksum, or the source does not provide a valid local file + * name. + */ + public static Path install(URI source, Path targetDirectory, String checksum, + Limits limits) throws IOException { + return install(source, targetDirectory, checksum, null, limits); + } + + /** + * Installs a catalog entry under its preferred file name when it is not an archive. + * + * @param source The resource location. + * @param targetDirectory The directory to install into. + * @param checksum The expected digest, or {@code null} for a file source. + * @param name The preferred name for non-archive content. + * @return The target directory. + * @throws IOException Thrown if fetching, verification, or unpacking fails. + * @throws IllegalArgumentException Thrown if an argument is invalid. + */ + static Path installNamed(URI source, Path targetDirectory, String checksum, + String name) throws IOException { + if (name == null) { + throw new IllegalArgumentException("name must not be null"); + } + return install(source, targetDirectory, checksum, name, Limits.DEFAULT); + } + + /** + * Validates the request, downloads and verifies the resource, and installs it from a + * staging directory. + * + * @param source The resource location. + * @param targetDirectory The directory to install into. + * @param checksum The expected digest, or {@code null} for a file source. + * @param name The preferred name for non-archive content, or {@code null} to use the + * source name. + * @param limits The limits to enforce. + * @return The target directory. Not {@code null}. + * @throws IOException Thrown if fetching, verification, or installation fails. + * @throws IllegalArgumentException Thrown if an argument is invalid. + */ + private static Path install(URI source, Path targetDirectory, String checksum, + String name, Limits limits) throws IOException { + if (source == null) { + throw new IllegalArgumentException("source must not be null"); + } + if (targetDirectory == null) { + throw new IllegalArgumentException("targetDirectory must not be null"); + } + if (limits == null) { + throw new IllegalArgumentException("limits must not be null"); + } + validateSource(source); + final String expected = validateChecksum(checksum); + if (expected == null && isHttp(source.getScheme())) { + throw new IllegalArgumentException( + "checksum must be given for an http or https source: " + source); + } + final String resourceName = validateSourceName( + name == null ? sourceName(source) : name); + final boolean createdTarget = Files.notExists(targetDirectory); + Files.createDirectories(targetDirectory); + removeStaleWorkFiles(targetDirectory); + final Path downloaded = createDownloadFile(targetDirectory); + try { + download(source, downloaded, limits); + if (expected != null) { + verify(downloaded, expected); + } + installStaged(downloaded, resourceName, targetDirectory, limits); + return targetDirectory; + } catch (IOException e) { + Files.deleteIfExists(downloaded); + if (createdTarget) { + removeIfEmpty(targetDirectory, e); + } + throw e; + } finally { + Files.deleteIfExists(downloaded); + } + } + + /** + * Removes download files and staging directories that an earlier installation left in + * the target because its process ended before cleanup. Only entries with the hidden + * work-file prefixes are touched. + * + * @param targetDirectory The directory to install into. + * @throws IOException Thrown if listing or deleting fails. + */ + private static void removeStaleWorkFiles(Path targetDirectory) throws IOException { + final List stale; + try (Stream entries = Files.list(targetDirectory)) { + stale = entries.filter(ResourceInstaller::isWorkFile).toList(); + } + for (final Path entry : stale) { + if (Files.isDirectory(entry, LinkOption.NOFOLLOW_LINKS)) { + deleteRecursively(entry); + } else { + Files.deleteIfExists(entry); + } + } + } + + /** + * Classifies a target directory entry as a work file of this class. + * + * @param entry The entry to inspect. + * @return {@code true} if the entry name carries a work-file prefix. + */ + private static boolean isWorkFile(Path entry) { + final String fileName = entry.getFileName().toString(); + return fileName.startsWith(STAGING_PREFIX) || fileName.startsWith(DOWNLOAD_PREFIX); + } + + /** + * Removes a target directory this installation created when the failed installation + * left nothing in it, so a failed first attempt leaves the filesystem as it was. + * + * @param targetDirectory The directory this installation created. + * @param failure The failure being reported; a cleanup error is added to it. + */ + private static void removeIfEmpty(Path targetDirectory, IOException failure) { + try (Stream entries = Files.list(targetDirectory)) { + if (entries.findAny().isEmpty()) { + Files.deleteIfExists(targetDirectory); + } + } catch (IOException cleanup) { + failure.addSuppressed(cleanup); + } + } + + /** + * Creates the file the download is written to. It is placed on the target's + * filesystem, not in the system temporary directory, so a large download cannot + * exhaust the system temporary directory while the target has room. Its hidden prefix + * distinguishes it from installed content. + * + * @param targetDirectory The directory to install into. Must already exist. + * @return The newly created, empty download file. Not {@code null}. + * @throws IOException Thrown if the file cannot be created. + */ + static Path createDownloadFile(Path targetDirectory) throws IOException { + return Files.createTempFile(targetDirectory, DOWNLOAD_PREFIX, DOWNLOAD_SUFFIX); + } + + /** + * Validates the checksum argument and normalizes it for comparison. + * + * @param checksum The digest as given by the caller, or {@code null} to skip. + * @return The stripped digest, or {@code null} when verification is skipped. + * @throws IllegalArgumentException Thrown if the digest is not a 64-character or + * 128-character hex string. + */ + private static String validateChecksum(String checksum) { + if (checksum == null) { + return null; + } + final String trimmed = checksum.strip(); + if ((trimmed.length() == SHA_256_HEX_LENGTH || trimmed.length() == SHA_512_HEX_LENGTH) + && isHex(trimmed)) { + return trimmed; + } + throw new IllegalArgumentException( + "checksum must be 64 (SHA-256) or 128 (SHA-512) hex characters; pass null to skip"); + } + + /** + * Checks whether a string is made up entirely of hexadecimal digits. + * + * @param value The string to inspect. + * @return {@code true} if every character is a hexadecimal digit. + */ + private static boolean isHex(String value) { + for (int i = 0; i < value.length(); i++) { + final char c = value.charAt(i); + final boolean digit = c >= '0' && c <= '9'; + final boolean lower = c >= 'a' && c <= 'f'; + final boolean upper = c >= 'A' && c <= 'F'; + if (!digit && !lower && !upper) { + return false; + } + } + return true; + } + + /** + * Fetches the source into the given file, bounded by the download limit. Http and + * https locations are fetched with timeouts and the redirect policy; a {@code file} + * location is read directly. The scheme was accepted by {@link #validateSource(URI)} + * at the public boundary. + * + * @param source The resource location. + * @param file The file receiving the downloaded bytes. + * @param limits The limits to enforce. + * @throws IOException Thrown if fetching fails or a limit is exceeded. + */ + private static void download(URI source, Path file, Limits limits) throws IOException { + final Budget budget = new Budget(limits.maxDownloadBytes(), + "download exceeds the limit of " + limits.maxDownloadBytes() + " bytes"); + if (isHttp(source.getScheme())) { + downloadHttp(source, file, limits, budget); + } else { + try (InputStream in = Files.newInputStream(localFile(source))) { + copyBounded(in, file, budget); + } + } + } + + /** + * Rejects a source the installer will not fetch. Only {@code http}, {@code https}, and + * {@code file} are accepted: any other scheme would be passed to any URL handler + * the runtime happens to have installed, outside the connection timeout, read timeout, + * and redirect policy this class enforces. + * + * @param source The resource location as given by the caller. + * @throws IllegalArgumentException Thrown if the scheme is absent or unsupported. + */ + private static void validateSource(URI source) { + final String scheme = source.getScheme(); + if (!isHttp(scheme) && !SCHEME_FILE.equalsIgnoreCase(scheme)) { + throw new IllegalArgumentException( + "source scheme must be http, https, or file, but was: " + source); + } + } + + /** + * Classifies a scheme as one the http fetch path handles. + * + * @param scheme The URI scheme, or {@code null} when the location has none. + * @return {@code true} for {@code http} and {@code https}, ignoring case. + */ + private static boolean isHttp(String scheme) { + return SCHEME_HTTP.equalsIgnoreCase(scheme) || SCHEME_HTTPS.equalsIgnoreCase(scheme); + } + + /** + * Resolves a {@code file} location to a path on the default filesystem. + * + * @param source The {@code file} location, already validated as such. + * @return The local path. Not {@code null}. + * @throws IOException Thrown if the location does not name a file this runtime can + * open, such as a {@code file} URI naming a remote host. + */ + private static Path localFile(URI source) throws IOException { + try { + return Path.of(source); + } catch (IllegalArgumentException | FileSystemNotFoundException e) { + throw new IOException("not a readable local file location: " + source, e); + } + } + + /** + * Fetches an http or https source with connection and read timeouts, following at + * most the allowed number of redirects under the redirect policy, checking any + * declared content length against the download limit before reading the body, and + * bounding the transferred bytes against the same limit. + * + * @param source The resource location as requested by the caller. + * @param file The file receiving the downloaded bytes. + * @param limits The limits to enforce. + * @param budget The download budget shared with the caller. + * @throws IOException Thrown if fetching fails, the server answers with a status + * other than 200, the redirect policy is violated, or a limit is exceeded. + */ + private static void downloadHttp(URI source, Path file, Limits limits, Budget budget) + throws IOException { + URI current = source; + int redirects = 0; + while (true) { + final HttpURLConnection connection = + (HttpURLConnection) current.toURL().openConnection(); + connection.setInstanceFollowRedirects(false); + connection.setConnectTimeout(timeoutMillis(limits.connectTimeout())); + connection.setReadTimeout(timeoutMillis(limits.readTimeout())); + try { + final int status = connection.getResponseCode(); + if (isRedirect(status)) { + if (redirects >= limits.maxRedirects()) { + throw new IOException( + "more than " + limits.maxRedirects() + " redirects: " + source); + } + current = resolveRedirect(current, connection.getHeaderField("Location")); + redirects++; + continue; + } + if (status != HttpURLConnection.HTTP_OK) { + throw new IOException( + "download failed with HTTP status " + status + ": " + current); + } + final long declared = connection.getContentLengthLong(); + if (declared > limits.maxDownloadBytes()) { + throw new IOException("declared content length " + declared + + " exceeds the download limit of " + limits.maxDownloadBytes() + + " bytes"); + } + try (InputStream in = connection.getInputStream()) { + copyBounded(in, file, budget); + } + return; + } finally { + connection.disconnect(); + } + } + } + + /** + * Classifies a response status as a redirect the installer follows. + * + * @param status The HTTP response status. + * @return {@code true} if the status is one of the redirect statuses 301, 302, 303, + * 307, or 308. + */ + private static boolean isRedirect(int status) { + return status == HttpURLConnection.HTTP_MOVED_PERM + || status == HttpURLConnection.HTTP_MOVED_TEMP + || status == HttpURLConnection.HTTP_SEE_OTHER + || status == HTTP_TEMPORARY_REDIRECT + || status == HTTP_PERMANENT_REDIRECT; + } + + /** + * Resolves a redirect location against the redirected request and enforces the + * redirect policy: the target must be an http or https location, and an https + * request must not be redirected to plain http. + * + * @param from The location that returned the redirect. + * @param location The Location header value, absolute or relative, or {@code null} + * when the header is absent. + * @return The resolved redirect target. Not {@code null}. + * @throws IOException Thrown if the location is absent or malformed, leaves the + * http and https schemes, or downgrades https to http. + */ + static URI resolveRedirect(URI from, String location) throws IOException { + if (location == null || location.isEmpty()) { + throw new IOException("redirect from " + from + " contains no Location header"); + } + final URI target; + try { + target = from.resolve(location); + } catch (IllegalArgumentException e) { + throw new IOException( + "redirect from " + from + " contains a malformed Location: " + location, e); + } + final String scheme = target.getScheme(); + final boolean https = SCHEME_HTTPS.equalsIgnoreCase(scheme); + if (!https && !SCHEME_HTTP.equalsIgnoreCase(scheme)) { + throw new IOException( + "redirect target is not an http or https location: " + target); + } + if (SCHEME_HTTPS.equalsIgnoreCase(from.getScheme()) && !https) { + throw new IOException("redirect downgrades https to http: " + target); + } + return target; + } + + /** + * Converts a timeout to the millisecond form the connection setters take. A positive + * timeout shorter than a millisecond becomes one millisecond because + * {@link HttpURLConnection#setReadTimeout(int) zero disables the timeout}. A timeout + * too large for the int range is capped. + * + * @param timeout The timeout as a duration. Must be positive. + * @return The timeout in milliseconds, at least {@code 1} and at most + * {@link Integer#MAX_VALUE}. + */ + private static int timeoutMillis(Duration timeout) { + final long millis; + try { + millis = timeout.toMillis(); + } catch (ArithmeticException e) { + return Integer.MAX_VALUE; + } + return Math.clamp(millis, 1, Integer.MAX_VALUE); + } + + /** + * Computes the file's digest and compares it with the expected hex digest, ignoring + * hex letter case. The digest length selects the algorithm: 64 characters SHA-256, + * 128 characters SHA-512. + * + * @param file The file to digest. + * @param expected The expected hex digest, already trimmed. + * @throws IOException Thrown if the file cannot be read or the digests differ. + */ + private static void verify(Path file, String expected) throws IOException { + final String algorithm = + expected.length() == SHA_512_HEX_LENGTH ? SHA_512 : SHA_256; + final MessageDigest digest; + try { + digest = MessageDigest.getInstance(algorithm); + } catch (NoSuchAlgorithmException e) { + throw new IOException(algorithm + " is unavailable in this runtime", e); + } + try (InputStream in = Files.newInputStream(file)) { + final byte[] buffer = new byte[BUFFER_SIZE]; + int read; + while ((read = in.read(buffer)) >= 0) { + digest.update(buffer, 0, read); + } + } + final String actual = HexFormat.of().formatHex(digest.digest()); + if (!actual.equalsIgnoreCase(expected)) { + throw new IOException( + "checksum mismatch: expected " + expected + " but downloaded " + actual); + } + } + + /** + * Unpacks the downloaded content into a hidden staging directory beneath the target + * and promotes it into the target only after every entry unpacked cleanly. The + * staging directory lives on the target's filesystem so promotion is a sequence of + * renames, and it is removed whether the installation succeeds or fails. + * + * @param downloaded The fetched and verified file. + * @param name The file name derived from the source location. + * @param target The directory to install into. + * @param limits The limits to enforce while unpacking. + * @throws IOException Thrown if unpacking fails, a limit is exceeded, or promotion + * or staging cleanup fails. + */ + private static void installStaged(Path downloaded, String name, Path target, + Limits limits) throws IOException { + final Path staging = Files.createTempDirectory(target, STAGING_PREFIX); + try { + unpack(downloaded, name, staging, limits); + promote(staging, target); + } catch (IOException e) { + try { + deleteRecursively(staging); + } catch (IOException cleanup) { + e.addSuppressed(cleanup); + } + throw e; + } + deleteRecursively(staging); + } + + /** + * Moves all staged regular files to their relative locations beneath the target + * without replacing anything that already exists there. All destinations are + * checked before the first move, so a collision leaves the target without a mix of + * old and new files, and the move itself refuses an existing destination as well. + * + * @param staging The staging directory holding the fully unpacked content. + * @param target The directory to install into. + * @throws IOException Thrown if a destination already exists, a move fails, or a + * directory on the way to a destination is an existing symbolic link. + */ + private static void promote(Path staging, Path target) throws IOException { + final List files; + try (Stream walk = Files.walk(staging)) { + files = walk.filter(Files::isRegularFile).toList(); + } + for (final Path file : files) { + ensureVacant(target, staging.relativize(file)); + } + for (final Path file : files) { + moveIntoPlace(file, destination(target, staging.relativize(file))); + } + } + + /** + * Moves one staged file to its destination without replacing an existing file. The + * move is not requested atomically: on POSIX filesystems an atomic move renames over + * an existing destination, which would void the vacancy check. + * + * @param file The staged file. + * @param destination The destination beneath the target. + * @throws IOException Thrown if the destination exists or the move fails. + */ + static void moveIntoPlace(Path file, Path destination) throws IOException { + Files.move(file, destination); + } + + /** + * Checks that one staged file's destination is free to receive it, without creating + * anything. A missing directory on the way proves the destination vacant. + * + * @param target The directory to install into. + * @param relative The staged file's path relative to the staging directory. + * @throws IOException Thrown if the destination already exists, or a directory on + * the way is a symbolic link or exists as something other than a directory. + */ + private static void ensureVacant(Path target, Path relative) throws IOException { + final Path destination = destination(target, relative, false); + if (destination != null && Files.exists(destination, LinkOption.NOFOLLOW_LINKS)) { + throw new IOException("target already contains: " + destination); + } + } + + /** + * Resolves one staged file's destination beneath the target, creating the directories + * leading to it one at a time. + * + * @param target The directory to install into. + * @param relative The staged file's path relative to the staging directory. + * @return The destination path beneath the target. Not {@code null}. + * @throws IOException Thrown if a directory on the way is a symbolic link or exists as + * something other than a directory, or if a directory cannot be created. + */ + private static Path destination(Path target, Path relative) throws IOException { + return destination(target, relative, true); + } + + /** + * Walks the directories leading to one staged file's destination without descending + * through a symbolic link that is already there. An entry name that stays inside the + * staging directory can still land outside the target if a directory below the target + * is a link to somewhere else. + * + *

    This covers links present when the installation runs. It is not a defense against + * a link created concurrently, between the check here and the move that follows.

    + * + * @param target The directory to install into. + * @param relative The staged file's path relative to the staging directory. + * @param create Whether to create a missing directory on the way; when {@code false}, + * a missing directory ends the walk. + * @return The destination path beneath the target, or {@code null} when a directory + * on the way is missing and {@code create} is {@code false}. + * @throws IOException Thrown if a directory on the way is a symbolic link or exists as + * something other than a directory, or if a directory cannot be created. + */ + private static Path destination(Path target, Path relative, boolean create) + throws IOException { + Path directory = target; + for (int i = 0; i < relative.getNameCount() - 1; i++) { + directory = directory.resolve(relative.getName(i)); + if (Files.isSymbolicLink(directory)) { + throw new IOException( + "installation path crosses a symbolic link: " + directory); + } + if (!Files.exists(directory)) { + if (!create) { + return null; + } + Files.createDirectory(directory); + } else if (!Files.isDirectory(directory)) { + throw new IOException( + "installation path crosses an existing file: " + directory); + } + } + return directory.resolve(relative.getFileName()); + } + + /** + * Removes the given directory tree, deepest entries first. + * + * @param root The directory to remove. + * @throws IOException Thrown if a deletion fails. + */ + private static void deleteRecursively(Path root) throws IOException { + final List paths; + try (Stream walk = Files.walk(root)) { + paths = walk.sorted(Comparator.reverseOrder()).toList(); + } + for (final Path path : paths) { + Files.deleteIfExists(path); + } + } + + /** + * Detects the content format from its leading bytes and unpacks accordingly, + * bounding the total expanded bytes against the expansion limit. One exception: a + * source named {@code *.bin}, in any letter case, is stored verbatim even when its + * bytes are a zip archive. OpenNLP model consumers load the packed zip artifact. Unpacking it would + * place the internal entries ({@code manifest.properties}, {@code *.model}) in the + * target instead of the model. + * + * @param downloaded The fetched file. + * @param name The file name derived from the source location. + * @param staging The staging directory to unpack into. + * @param limits The limits to enforce. + * @throws IOException Thrown if reading or unpacking fails or the expansion limit + * is exceeded. + */ + private static void unpack(Path downloaded, String name, Path staging, Limits limits) + throws IOException { + final Budget budget = new Budget(limits.maxExpandedBytes(), + "expanded content exceeds the limit of " + limits.maxExpandedBytes() + + " bytes"); + final Budget entryBudget = new Budget(limits.maxEntries(), + "archive entry count exceeds the limit of " + limits.maxEntries() + + " entries"); + try (InputStream raw = new BufferedInputStream(Files.newInputStream(downloaded))) { + raw.mark(MAGIC_LENGTH); + final byte[] magic = raw.readNBytes(MAGIC_LENGTH); + raw.reset(); + if (endsWithIgnoreCase(name, MODEL_SUFFIX)) { + copyBounded(raw, safeChild(staging, name), budget); + } else if (hasMagic(magic, GZIP_MAGIC_FIRST, GZIP_MAGIC_SECOND)) { + unpackGzip(raw, name, staging, + expansionBudget(Files.size(downloaded), limits, budget), entryBudget); + } else if (hasMagic(magic, ZIP_MAGIC_FIRST, ZIP_MAGIC_SECOND, + ZIP_LOCAL_HEADER_THIRD, ZIP_LOCAL_HEADER_FOURTH)) { + final Set unpacked = unpackZip(raw, staging, + expansionBudget(Files.size(downloaded), limits, budget), entryBudget); + validateZip(downloaded, unpacked); + } else if (hasMagic(magic, ZIP_MAGIC_FIRST, ZIP_MAGIC_SECOND, + ZIP_END_HEADER_THIRD, ZIP_END_HEADER_FOURTH)) { + validateEmptyZip(raw); + } else { + copyBounded(raw, safeChild(staging, name), budget); + } + } + } + + /** + * Bounds expansion by the ratio as well as the absolute limit, so a small source + * cannot expand to the whole absolute limit. Deflate reaches roughly 1000 to 1, so + * the absolute limit alone lets a few megabytes fill the target filesystem. + * + * @param compressedSize The size of the compressed source in bytes. + * @param limits The limits holding the accepted expansion ratio. + * @param budget The expansion budget under the absolute limit. + * @return The tighter of the two budgets. Not {@code null}. + */ + private static Budget expansionBudget(long compressedSize, Limits limits, + Budget budget) { + final long ratio = limits.maxExpansionRatio(); + final long ratioCeiling = compressedSize > Long.MAX_VALUE / ratio + ? Long.MAX_VALUE + : Math.max(MIN_EXPANSION_BYTES, compressedSize * ratio); + if (ratioCeiling >= budget.limit()) { + return budget; + } + return new Budget(ratioCeiling, "content expands beyond " + + ratio + " times its compressed size"); + } + + /** + * Compares a name's ending with a suffix, ignoring letter case. + * + * @param name The file name. + * @param suffix The suffix to look for. + * @return {@code true} if the name ends with the suffix in any letter case. + */ + private static boolean endsWithIgnoreCase(String name, String suffix) { + return name.length() >= suffix.length() && name.regionMatches(true, + name.length() - suffix.length(), suffix, 0, suffix.length()); + } + + /** + * Checks whether the bytes at the start of a resource match the given signature. + * + * @param actual The bytes read from the resource. + * @param expected The unsigned byte values in the signature. + * @return {@code true} when the resource begins with the expected values. + */ + private static boolean hasMagic(byte[] actual, int... expected) { + if (actual.length < expected.length) { + return false; + } + for (int i = 0; i < expected.length; i++) { + if ((actual[i] & 0xFF) != expected[i]) { + return false; + } + } + return true; + } + + /** + * Checks that a zip archive contains a valid central directory listing the same files + * the local headers delivered, before staged content is promoted. The two listings + * can disagree in a crafted archive, and the local headers are what was unpacked. + * + * @param archive The downloaded archive. + * @param unpacked The names of the file entries read from the local headers. + * @throws IOException Thrown if the archive is malformed, the listings differ, or the + * archive cannot be read. + */ + private static void validateZip(Path archive, Set unpacked) throws IOException { + final Set listed = new HashSet<>(); + try (ZipFile zip = new ZipFile(archive.toFile())) { + final Enumeration entries = zip.entries(); + while (entries.hasMoreElements()) { + final ZipEntry entry = entries.nextElement(); + if (!entry.isDirectory()) { + listed.add(entry.getName()); + } + } + } catch (UnsupportedOperationException e) { + listed.addAll(listZipOnNonDefaultFileSystem(archive)); + } catch (ZipException e) { + throw new IOException(MALFORMED_ZIP_ERROR, e); + } + if (!listed.equals(unpacked)) { + throw new IOException(ZIP_MISMATCH_ERROR); + } + } + + /** + * Lists the file entries of an archive stored by a file-system provider that cannot + * supply a {@link java.io.File} to {@link ZipFile}. + * + * @param archive The downloaded archive. + * @return The file entry names from the central directory. Not {@code null}. + * @throws IOException Thrown if the archive is malformed or cannot be read. + */ + private static Set listZipOnNonDefaultFileSystem(Path archive) + throws IOException { + final Set listed = new HashSet<>(); + try (FileSystem zip = FileSystems.newFileSystem(archive)) { + for (final Path root : zip.getRootDirectories()) { + try (Stream walk = Files.walk(root)) { + walk.filter(Files::isRegularFile) + .map(path -> root.relativize(path).toString()) + .forEach(listed::add); + } + } + } catch (ZipException | ProviderNotFoundException e) { + throw new IOException(MALFORMED_ZIP_ERROR, e); + } + return listed; + } + + /** + * Validates an empty zip archive from its end-of-central-directory record. An empty + * archive has no local entry headers for {@link ZipInputStream} to validate. + * + * @param raw The zip content, positioned at its first byte. + * @throws IOException Thrown if the record is truncated, declares entries or central + * directory data, or has bytes beyond its declared comment. + */ + private static void validateEmptyZip(InputStream raw) throws IOException { + final byte[] header = raw.readNBytes(ZIP_END_HEADER_LENGTH); + if (header.length != ZIP_END_HEADER_LENGTH + || littleEndianShort(header, ZIP_DISK_OFFSET) != 0 + || littleEndianShort(header, ZIP_CENTRAL_DISK_OFFSET) != 0 + || littleEndianShort(header, ZIP_DISK_ENTRIES_OFFSET) != 0 + || littleEndianShort(header, ZIP_TOTAL_ENTRIES_OFFSET) != 0 + || littleEndianInt(header, ZIP_CENTRAL_SIZE_OFFSET) != 0 + || littleEndianInt(header, ZIP_CENTRAL_OFFSET_OFFSET) != 0) { + throw new IOException(MALFORMED_ZIP_ERROR); + } + final int commentLength = littleEndianShort(header, ZIP_COMMENT_LENGTH_OFFSET); + if (raw.readNBytes(commentLength).length != commentLength || raw.read() >= 0) { + throw new IOException(MALFORMED_ZIP_ERROR); + } + } + + /** + * Reads an unsigned 16-bit little-endian value. + * + * @param bytes The source bytes. + * @param offset The first byte to read. + * @return The decoded value. + */ + private static int littleEndianShort(byte[] bytes, int offset) { + return bytes[offset] & 0xFF | (bytes[offset + 1] & 0xFF) << 8; + } + + /** + * Reads an unsigned 32-bit little-endian value. + * + * @param bytes The source bytes. + * @param offset The first byte to read. + * @return The decoded value. + */ + private static long littleEndianInt(byte[] bytes, int offset) { + return littleEndianShort(bytes, offset) + | (long) littleEndianShort(bytes, offset + 2) << 16; + } + + /** + * Unpacks gzip content: a tar archive inside when present, a plain file otherwise. A + * plain file omits the {@code .gz} suffix of its source name, in any letter case. If the source name is + * only that suffix, the installed file is named {@value #DEFAULT_RESOURCE_NAME}. + * + * @param raw The gzip-compressed content. + * @param name The file name derived from the source location. + * @param staging The staging directory to unpack into. + * @param budget The expansion budget. + * @param entryBudget The entry-count budget. + * @throws IOException Thrown if decompressing or unpacking fails or a limit is + * exceeded. + */ + private static void unpackGzip(InputStream raw, String name, Path staging, + Budget budget, Budget entryBudget) throws IOException { + final InputStream decompressed = new BufferedInputStream( + new BudgetInputStream(new GZIPInputStream(raw), budget), BUFFER_SIZE); + if (TarStream.startsWithHeader(decompressed) || startsWithEmptyTar(decompressed)) { + unpackTar(decompressed, staging, entryBudget); + decompressed.transferTo(OutputStream.nullOutputStream()); + } else { + final String strippedName = endsWithIgnoreCase(name, GZIP_SUFFIX) + ? name.substring(0, name.length() - GZIP_SUFFIX.length()) : name; + final String plainName = strippedName.isEmpty() + ? DEFAULT_RESOURCE_NAME : strippedName; + copy(decompressed, safeChild(staging, plainName)); + } + } + + /** + * Checks for the two zero blocks that make up an empty tar archive. There is no entry + * header for {@link TarStream#startsWithHeader(InputStream)} to recognize in this case. + * + * @param in The decompressed content. Must support mark and reset. + * @return {@code true} when the content starts with two zero tar blocks. + * @throws IOException Thrown if reading or resetting the stream fails. + */ + private static boolean startsWithEmptyTar(InputStream in) throws IOException { + in.mark(TAR_END_BLOCKS_LENGTH); + try { + final byte[] blocks = in.readNBytes(TAR_END_BLOCKS_LENGTH); + if (blocks.length != TAR_END_BLOCKS_LENGTH) { + return false; + } + for (final byte b : blocks) { + if (b != 0) { + return false; + } + } + return true; + } finally { + in.reset(); + } + } + + /** + * Unpacks every regular tar entry to its relative location beneath the staging + * directory. + * + * @param decompressed The uncompressed tar content. + * @param staging The staging directory to unpack into. + * @param entryBudget The archive-header limit. + * @throws IOException Thrown if the archive is malformed, an entry escapes the + * staging directory, or the entry limit is exceeded. + */ + private static void unpackTar(InputStream decompressed, Path staging, + Budget entryBudget) throws IOException { + final TarStream entries = new TarStream(decompressed, entryBudget.limit()); + while (entries.next()) { + if (!entries.isFile()) { + safeChild(staging, entries.name()); + continue; + } + final Path file = newArchiveFile(staging, entries.name()); + copy(entries.entryStream(), file); + } + } + + /** + * Unpacks every regular zip entry to its relative location beneath the staging + * directory. + * + * @param raw The zip content. + * @param staging The staging directory to unpack into. + * @param budget The expansion budget. + * @param entryBudget The entry-count budget, charged for every entry including + * directories. + * @return The names of the file entries unpacked. Not {@code null}. + * @throws IOException Thrown if the archive is malformed, an entry escapes the + * staging directory, or a limit is exceeded. + */ + private static Set unpackZip(InputStream raw, Path staging, Budget budget, + Budget entryBudget) throws IOException { + final ZipInputStream zip = new ZipInputStream(raw); + final Set unpacked = new HashSet<>(); + boolean foundEntry = false; + ZipEntry entry; + while ((entry = zip.getNextEntry()) != null) { + foundEntry = true; + entryBudget.spend(1); + if (entry.isDirectory()) { + safeChild(staging, entry.getName()); + consumeBounded(zip, budget); + continue; + } + final Path file = newArchiveFile(staging, entry.getName()); + copyBounded(zip, file, budget); + unpacked.add(entry.getName()); + } + if (!foundEntry) { + throw new IOException(MALFORMED_ZIP_ERROR); + } + return unpacked; + } + + /** + * Resolves a file entry beneath the staging directory and creates its parent + * directories. A second entry that normalizes to the same path is rejected. + * + * @param staging The staging directory. + * @param entryName The path stored in the archive. + * @return The new file path. Not {@code null}. + * @throws IOException Thrown if the path escapes the staging directory, duplicates + * another file entry, or its parent directories cannot be created. + */ + private static Path newArchiveFile(Path staging, String entryName) throws IOException { + final Path file = safeChild(staging, entryName); + if (Files.exists(file, LinkOption.NOFOLLOW_LINKS)) { + throw new IOException("archive contains duplicate file entry: " + entryName); + } + Files.createDirectories(file.getParent()); + return file; + } + + /** + * Copies the stream into the file, charging every byte against the budget before it + * is written, so an oversized transfer aborts within one buffer of its limit. + * + * @param in The content to copy. + * @param file The file to write. + * @param budget The byte budget to charge. + * @throws IOException Thrown if reading or writing fails or the budget is exceeded. + */ + private static void copyBounded(InputStream in, Path file, Budget budget) + throws IOException { + try (OutputStream out = Files.newOutputStream(file)) { + new BudgetInputStream(in, budget).transferTo(out); + } + } + + /** + * Reads and discards an entry while charging each byte against the expansion limit. + * + * @param in The entry content. + * @param budget The byte budget to charge. + * @throws IOException Thrown if reading fails or the budget is exceeded. + */ + private static void consumeBounded(InputStream in, Budget budget) throws IOException { + new BudgetInputStream(in, budget).transferTo(OutputStream.nullOutputStream()); + } + + /** + * Copies the stream into the file. The input stream must already enforce any byte + * limit. + * + * @param in The content to copy. + * @param file The file to write. + * @throws IOException Thrown if reading or writing fails. + */ + private static void copy(InputStream in, Path file) throws IOException { + try (OutputStream out = Files.newOutputStream(file)) { + in.transferTo(out); + } + } + + /** + * Resolves an archive entry inside the staging directory, rejecting escaping paths. + * + * @param staging The staging directory to unpack into. + * @param entryName The entry name as stored in the archive. + * @return The resolved path beneath the staging directory. Not {@code null}. + * @throws IOException Thrown if the entry resolves outside the staging directory. + */ + private static Path safeChild(Path staging, String entryName) throws IOException { + final Path resolved; + try { + resolved = staging.resolve(entryName).normalize(); + } catch (InvalidPathException e) { + throw new IOException("archive entry has an invalid path: " + entryName, e); + } + if (!resolved.startsWith(staging.normalize())) { + throw new IOException("archive entry escapes the target directory: " + entryName); + } + return resolved; + } + + /** + * Derives a file name from the source URI for non-archive content. + * + * @param source The resource location. + * @return The last path segment, or {@code resource} if the location has none. + */ + private static String sourceName(URI source) { + final String path = source.getPath(); + if (path == null || path.isEmpty()) { + return DEFAULT_RESOURCE_NAME; + } + final int slash = path.lastIndexOf('/'); + final String name = slash < 0 ? path : path.substring(slash + 1); + return name.isEmpty() ? DEFAULT_RESOURCE_NAME : name; + } + + /** + * Rejects names that are empty, path-like, or contain a NUL character. + * + * @param name The candidate local file name. + * @return The validated name. + * @throws IllegalArgumentException Thrown if {@code name} is not a file name. + */ + static String validateSourceName(String name) { + if (name.isEmpty() || ".".equals(name) || "..".equals(name)) { + throw new IllegalArgumentException("name must be a file name"); + } + for (int i = 0; i < name.length(); i++) { + final char c = name.charAt(i); + if (c == '/' || c == '\\' || c == 0) { + throw new IllegalArgumentException("name must be a file name"); + } + } + return name; + } + + /** + * A unit budget, counting bytes or archive entries: {@link #spend(long)} accumulates + * spent units and fails once the limit is crossed. + */ + private static final class Budget { + + private final long limit; + private final String message; + private long used; + + /** + * Creates a budget that has spent zero units. + * + * @param limit The largest total number of units accepted. + * @param message The failure message raised when the limit is crossed. + */ + Budget(long limit, String message) { + this.limit = limit; + this.message = message; + } + + /** {@return the maximum number of units accepted} */ + long limit() { + return limit; + } + + /** + * Charges the given number of units against the budget. + * + * @param units The number of units to charge. + * @throws IOException Thrown if the total charged units exceed the limit. + */ + void spend(long units) throws IOException { + used += units; + if (used > limit) { + throw new IOException(message); + } + } + } + + /** + * Charges every byte read or skipped from an expanded stream against a shared budget. + */ + private static final class BudgetInputStream extends FilterInputStream { + + private final Budget budget; + + /** + * Initializes a budgeted stream. + * + * @param in The expanded stream to read. + * @param budget The budget to charge. + */ + BudgetInputStream(InputStream in, Budget budget) { + super(in); + this.budget = budget; + } + + /** + * {@inheritDoc} + */ + @Override + public int read() throws IOException { + final int value = super.read(); + if (value >= 0) { + budget.spend(1); + } + return value; + } + + /** + * {@inheritDoc} + */ + @Override + public int read(byte[] buffer, int offset, int length) throws IOException { + final int read = super.read(buffer, offset, length); + if (read > 0) { + budget.spend(read); + } + return read; + } + + /** + * {@inheritDoc} + */ + @Override + public long skip(long bytes) throws IOException { + final long skipped = super.skip(bytes); + budget.spend(skipped); + return skipped; + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/archive/TarStream.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/archive/TarStream.java new file mode 100644 index 0000000000..c016a566fa --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/archive/TarStream.java @@ -0,0 +1,659 @@ +/* + * 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.util.archive; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.StandardCharsets; +import java.util.Objects; + +import opennlp.tools.commons.Internal; + +/** + * A forward-only reader for classic v7, POSIX ustar, GNU, and pax tar streams. + * {@link #next()} advances to the following entry and {@link #entryStream()} exposes + * only the current entry's bytes. + * + *

    The reader validates header checksums, supports ustar name prefixes, GNU long + * names, pax {@code path} and {@code size} records, and GNU base-256 sizes. Sparse + * entries and global pax records that change paths or sizes are rejected because this + * reader cannot reproduce their content or global semantics.

    + * + * @since 3.0.0 + */ +@Internal +public final class TarStream { + + private static final int BLOCK = 512; + private static final int NAME_LENGTH = 100; + private static final int CHECKSUM_OFFSET = 148; + private static final int CHECKSUM_LENGTH = 8; + private static final int SIZE_OFFSET = 124; + private static final int SIZE_LENGTH = 12; + private static final int TYPE_OFFSET = 156; + private static final int MAGIC_OFFSET = 257; + private static final int PREFIX_OFFSET = 345; + private static final int PREFIX_LENGTH = 155; + private static final String USTAR_MAGIC = "ustar"; + private static final char TYPE_REGULAR_FILE = '0'; + private static final char TYPE_REGULAR_FILE_CLASSIC = '\0'; + private static final char TYPE_GNU_LONG_NAME = 'L'; + private static final char TYPE_GNU_LONG_LINK = 'K'; + private static final char TYPE_GNU_SPARSE = 'S'; + private static final char TYPE_PAX_EXTENDED = 'x'; + private static final char TYPE_PAX_GLOBAL = 'g'; + private static final int MAX_EXTENSION_BYTES = 1 << 20; + private static final String MALFORMED_RECORD = "malformed pax extended header record"; + private static final String KEYWORD_PATH = "path"; + private static final String KEYWORD_SIZE = "size"; + private static final String SPARSE_PREFIX = "GNU.sparse."; + private static final int BASE_256_MARKER = 0x80; + private static final int BASE_256_NEGATIVE = 0x40; + private static final int BASE_256_FIRST_BYTE_BITS = 0x7F; + + private final InputStream in; + private final long maxEntries; + private final byte[] header = new byte[BLOCK]; + + private String name; + private long size; + private char type; + private long remaining; + private long entries; + private boolean ended; + + /** The name an extension header supplied for the entry that follows it, else null. */ + private String pendingPath; + + /** The size a pax header supplied for the entry that follows it, else {@code -1}. */ + private long pendingSize = -1; + + /** + * Initializes the reader. + * + * @param in The tar content. Not {@code null}. Not closed by this class. + * @throws IllegalArgumentException Thrown if {@code in} is {@code null}. + */ + public TarStream(InputStream in) { + this(in, Long.MAX_VALUE); + } + + /** + * Initializes a reader with an archive-entry limit. Extension headers count toward the + * limit. + * + * @param in The tar content. Not {@code null}. Not closed by this class. + * @param maxEntries The maximum number of archive headers to read. Must be positive. + * @throws IllegalArgumentException Thrown if {@code in} is {@code null} or + * {@code maxEntries} is not positive. + */ + public TarStream(InputStream in, long maxEntries) { + if (in == null) { + throw new IllegalArgumentException("in must not be null"); + } + if (maxEntries <= 0) { + throw new IllegalArgumentException("maxEntries must be positive"); + } + this.in = in; + this.maxEntries = maxEntries; + } + + /** + * Checks whether the given stream is positioned at a tar entry header, leaving its + * position unchanged. + * + * @param in The stream to inspect. Not {@code null} and must support + * {@link InputStream#mark(int) mark} and {@link InputStream#reset() reset}. + * @return {@code true} if the next 512 bytes read as a tar header, {@code false} if + * they do not or if fewer than 512 bytes are available. + * @throws IOException Thrown if reading from or repositioning the stream fails. + * @throws IllegalArgumentException Thrown if {@code in} is {@code null} or does not + * support mark and reset. + */ + public static boolean startsWithHeader(InputStream in) throws IOException { + if (in == null) { + throw new IllegalArgumentException("in must not be null"); + } + if (!in.markSupported()) { + throw new IllegalArgumentException("in must support mark and reset"); + } + in.mark(BLOCK); + try { + final byte[] block = new byte[BLOCK]; + return in.readNBytes(block, 0, BLOCK) == BLOCK && isHeader(block); + } finally { + in.reset(); + } + } + + /** + * Advances to the next entry. + * + * @return {@code true} if an entry is available, {@code false} at the end of the + * archive. + * @throws IOException Thrown if the archive is truncated or a header is malformed. + */ + public boolean next() throws IOException { + if (ended) { + return false; + } + skip(remaining); + skip(padding(size)); + while (true) { + if (!readBlock() || isEndBlock()) { + ended = true; + return false; + } + entries++; + if (entries > maxEntries) { + throw new IOException("archive entry count exceeds the limit of " + + maxEntries + " entries"); + } + if (!hasValidChecksum(header)) { + throw new IOException("malformed tar header checksum"); + } + type = (char) header[TYPE_OFFSET]; + size = parseSize(); + remaining = size; + if (type == TYPE_GNU_SPARSE) { + throw new IOException("sparse tar entries are not supported: " + + "the archived bytes describe file holes, not contiguous content"); + } + if (type == TYPE_PAX_EXTENDED || type == TYPE_PAX_GLOBAL) { + readRecords(readExtensionPayload(), type == TYPE_PAX_GLOBAL); + continue; + } + if (type == TYPE_GNU_LONG_NAME) { + final byte[] payload = readExtensionPayload(); + pendingPath = trimNul(decodeUtf8(payload, 0, payload.length, "GNU long name")); + continue; + } + if (type == TYPE_GNU_LONG_LINK) { + // The link target of an entry this reader does not expose. + readExtensionPayload(); + continue; + } + name = pendingPath != null ? pendingPath : readName(); + pendingPath = null; + if (pendingSize >= 0) { + size = pendingSize; + remaining = size; + pendingSize = -1; + } + if (name.isEmpty()) { + throw new IOException("tar entry header contains an empty name"); + } + return true; + } + } + + /** + * Reads the payload of the extension header just read, leaving the stream positioned on + * the header that follows it. + * + * @return The payload bytes. Not {@code null}. + * @throws IOException Thrown if the payload is larger than + * {@link #MAX_EXTENSION_BYTES} or the archive ends inside it. + */ + private byte[] readExtensionPayload() throws IOException { + if (size > MAX_EXTENSION_BYTES) { + throw new IOException("tar extension header of " + size + " bytes exceeds the " + + MAX_EXTENSION_BYTES + " byte limit"); + } + final byte[] payload = new byte[(int) size]; + if (in.readNBytes(payload, 0, payload.length) < payload.length) { + throw new IOException("truncated tar archive"); + } + skip(padding(size)); + size = 0; + remaining = 0; + return payload; + } + + /** + * Reads a pax extended header payload, which is a sequence of + * {@code " =\n"} records. Each length counts the complete + * record, including the length digits, blank, and newline. Records are parsed from raw + * bytes because the length is a byte count and a multibyte value would shift later records. + * + * @param payload The raw extended header payload. + * @param global Whether this is a global header, which applies to all following entries. + * @throws IOException Thrown if a record is malformed, if a global header contains a + * keyword that would change the entries after it, or if the entry is sparse. + */ + private void readRecords(byte[] payload, boolean global) throws IOException { + int offset = 0; + while (offset < payload.length) { + int blank = offset; + while (blank < payload.length && payload[blank] != ' ') { + blank++; + } + if (blank == payload.length || blank == offset) { + throw new IOException(MALFORMED_RECORD); + } + int length = 0; + for (int i = offset; i < blank; i++) { + final byte b = payload[i]; + if (b < '0' || b > '9') { + throw new IOException(MALFORMED_RECORD); + } + length = length * 10 + (b - '0'); + if (length > payload.length) { + throw new IOException(MALFORMED_RECORD); + } + } + final int end = offset + length; + if (length <= blank - offset || end > payload.length || payload[end - 1] != '\n') { + throw new IOException(MALFORMED_RECORD); + } + int equals = blank + 1; + while (equals < end - 1 && payload[equals] != '=') { + equals++; + } + if (equals >= end - 1 || equals == blank + 1) { + throw new IOException(MALFORMED_RECORD); + } + apply(decodeUtf8(payload, blank + 1, equals - blank - 1, "pax record"), + decodeUtf8(payload, equals + 1, end - equals - 2, "pax record"), global); + offset = end; + } + } + + /** + * Decodes archive text without replacing malformed input, which would silently + * change an archive path. + * + * @param bytes The bytes containing the text. + * @param offset The first byte to decode. + * @param length The number of bytes to decode. + * @param subject The field name to use in an error message. + * @return The decoded text. Not {@code null}. + * @throws IOException Thrown if the bytes are not valid UTF-8. + */ + private static String decodeUtf8(byte[] bytes, int offset, int length, String subject) + throws IOException { + try { + return StandardCharsets.UTF_8.newDecoder() + .decode(ByteBuffer.wrap(bytes, offset, length)).toString(); + } catch (CharacterCodingException e) { + throw new IOException(subject + " is not valid UTF-8", e); + } + } + + /** + * Applies one pax record. Only {@code path} and {@code size} change what this reader + * reports, so other keywords are ignored, except sparse entries, which cannot be + * unpacked. + * + * @param keyword The record's keyword. + * @param value The record's value. + * @param global Whether the record came from a global header. + * @throws IOException Thrown if the entry is sparse, if a global header contains a + * keyword that would change the entries after it, or if {@code size} is not a + * number. + */ + private void apply(String keyword, String value, boolean global) throws IOException { + if (keyword.startsWith(SPARSE_PREFIX)) { + throw new IOException("sparse tar entries are not supported: " + + "the archived bytes describe file holes, not contiguous content"); + } + if (!KEYWORD_PATH.equals(keyword) && !KEYWORD_SIZE.equals(keyword)) { + return; + } + if (global) { + throw new IOException("pax global header contains " + keyword + + ", which would change every entry after it"); + } + if (KEYWORD_PATH.equals(keyword)) { + pendingPath = value; + return; + } + try { + pendingSize = Long.parseLong(value); + } catch (NumberFormatException e) { + throw new IOException("pax size record is not a number: " + value, e); + } + if (pendingSize < 0) { + throw new IOException("pax size record is negative: " + value); + } + } + + /** + * Drops everything from the first NUL onward, which is how a GNU long-name header + * terminates the name it contains. + * + * @param value The decoded payload. + * @return The name without its terminator. Not {@code null}. + */ + private static String trimNul(String value) { + final int nul = value.indexOf('\0'); + return nul < 0 ? value : value.substring(0, nul); + } + + /** + * Reads the current header's entry name. On a POSIX ustar header a non-empty name + * prefix is joined to the name field with {@code /}, which is how a name longer than + * the 100-byte name field is stored when no extension header contains it. + * + * @return The entry name. Not {@code null}. + * @throws IOException Thrown if the stored name or prefix is not valid UTF-8. + */ + private String readName() throws IOException { + final String stored = field(0, NAME_LENGTH, "tar entry name"); + if (!hasPosixUstarMagic(header)) { + return stored; + } + final String prefix = field(PREFIX_OFFSET, PREFIX_LENGTH, "tar entry name prefix"); + return prefix.isEmpty() ? stored : prefix + "/" + stored; + } + + /** + * Reads a NUL-terminated text field of the current header. + * + * @param offset The field's offset in the header block. + * @param length The field's length in bytes. + * @param subject The field name to use in an error message. + * @return The field content up to its first NUL, decoded as UTF-8. Not {@code null}. + * @throws IOException Thrown if the field is not valid UTF-8. + */ + private String field(int offset, int length, String subject) throws IOException { + int end = 0; + while (end < length && header[offset + end] != 0) { + end++; + } + return decodeUtf8(header, offset, end, subject); + } + + /** + * @return The current entry's name as stored in the archive. Not {@code null} + * after a successful {@link #next()}. + */ + public String name() { + return name; + } + + /** + * @return The current entry's size in bytes. + */ + public long size() { + return size; + } + + /** + * @return {@code true} if the current entry is a regular file. + */ + public boolean isFile() { + return type == TYPE_REGULAR_FILE || type == TYPE_REGULAR_FILE_CLASSIC; + } + + /** + * Opens the current entry's content. + * + * @return A stream over exactly this entry's bytes; reading past the end returns end + * of stream, and a zero-length read returns {@code 0} as + * {@link InputStream#read(byte[], int, int)} requires. Not {@code null}. + * Closing it is not required. + */ + public InputStream entryStream() { + return new InputStream() { + @Override + public int read() throws IOException { + if (remaining <= 0) { + return -1; + } + final int b = in.read(); + if (b < 0) { + throw new IOException("truncated tar entry: " + name); + } + remaining--; + return b; + } + + /** + * {@inheritDoc} + * + *

    The range is checked before the entry state, so invalid arguments are reported + * when the entry is exhausted or the requested length is zero. This override uses + * the exceptions specified by {@link InputStream}.

    + * + * @throws NullPointerException Thrown if {@code buffer} is {@code null}. + * @throws IndexOutOfBoundsException Thrown if {@code offset} or {@code length} is + * negative, or {@code length} is greater than + * {@code buffer.length - offset}. + */ + @Override + public int read(byte[] buffer, int offset, int length) throws IOException { + Objects.checkFromIndexSize(offset, length, buffer.length); + if (length == 0) { + return 0; + } + if (remaining <= 0) { + return -1; + } + final int read = in.read(buffer, offset, (int) Math.min(length, remaining)); + if (read < 0) { + throw new IOException("truncated tar entry: " + name); + } + remaining -= read; + return read; + } + }; + } + + /** + * Checks whether a full 512-byte block reads as a tar entry header, which it does when + * it starts with a name and its stored checksum matches the block. Both classic and + * ustar headers contain that checksum, so it identifies both without relying on the + * ustar magic, which arbitrary content can also contain. + * + * @param block The block to inspect. Must be 512 bytes long. + * @return {@code true} if the block reads as a tar header. + */ + private static boolean isHeader(byte[] block) { + return block[0] != 0 && hasValidChecksum(block); + } + + /** + * Verifies a header block against the checksum stored in it. The checksum is the sum + * of every header byte with the checksum field itself read as eight blanks. Historical + * writers used signed-byte sums, so both totals are accepted. + * + * @param block The header block to verify. Must be 512 bytes long. + * @return {@code true} if the stored checksum is well formed and matches the block. + */ + private static boolean hasValidChecksum(byte[] block) { + long stored = 0; + boolean digits = false; + boolean trailingPadding = false; + for (int i = CHECKSUM_OFFSET; i < CHECKSUM_OFFSET + CHECKSUM_LENGTH; i++) { + final byte b = block[i]; + if (b == 0 || b == ' ') { + if (digits) { + trailingPadding = true; + } + continue; + } + if (!isOctalDigit(b) || trailingPadding) { + return false; + } + stored = stored * 8 + (b - '0'); + digits = true; + } + if (!digits) { + return false; + } + int unsigned = 0; + int signed = 0; + for (int i = 0; i < BLOCK; i++) { + final byte b = i >= CHECKSUM_OFFSET && i < CHECKSUM_OFFSET + CHECKSUM_LENGTH + ? (byte) ' ' : block[i]; + unsigned += b & 0xFF; + signed += b; + } + return stored == unsigned || stored == signed; + } + + /** + * Checks for the POSIX ustar magic specifically, which is {@code "ustar"} followed by a + * NUL. GNU writes {@code "ustar"} followed by two blanks and a NUL in the same field, + * and its headers must not be read as ustar: GNU stores {@code atime} at the offset + * ustar gives to the name prefix, so a GNU incremental archive would otherwise deliver + * every entry under a directory named after an octal timestamp. + * + * @param block The block to inspect. Must be 512 bytes long. + * @return {@code true} if the block contains the POSIX ustar magic. + */ + private static boolean hasPosixUstarMagic(byte[] block) { + for (int i = 0; i < USTAR_MAGIC.length(); i++) { + if (block[MAGIC_OFFSET + i] != USTAR_MAGIC.charAt(i)) { + return false; + } + } + return block[MAGIC_OFFSET + USTAR_MAGIC.length()] == 0; + } + + /** + * @param b The byte to classify. + * @return {@code true} if the byte is one of the digits {@code 0} to {@code 7}. + */ + private static boolean isOctalDigit(byte b) { + return b >= '0' && b <= '7'; + } + + /** + * Fills the header buffer with the next 512-byte block. + * + * @return {@code true} when a full block was read, {@code false} at a clean end of + * the stream before any byte of the block. + * @throws IOException Thrown if the stream ends inside the block. + */ + private boolean readBlock() throws IOException { + final int filled = in.readNBytes(header, 0, header.length); + if (filled == 0) { + return false; + } + if (filled < header.length) { + throw new IOException("truncated tar header"); + } + return true; + } + + /** + * @return {@code true} if the current header buffer is one of the all-zero blocks + * that terminate a tar archive. + */ + private boolean isEndBlock() { + for (final byte b : header) { + if (b != 0) { + return false; + } + } + return true; + } + + /** + * Parses the size field of the current header, tolerating NUL and blank padding around + * the octal digits. A field with its leading bit set is in the base-256 encoding + * instead, and is read by {@link #parseBase256Size()}. + * + * @return The entry size in bytes. + * @throws IOException Thrown if the field contains a character that is not an octal + * digit, a blank, or NUL padding, or if the base-256 form is negative or too + * large for a {@code long}. + */ + private long parseSize() throws IOException { + if ((header[SIZE_OFFSET] & BASE_256_MARKER) != 0) { + return parseBase256Size(); + } + long value = 0; + boolean digitSeen = false; + boolean trailingPadding = false; + for (int i = SIZE_OFFSET; i < SIZE_OFFSET + SIZE_LENGTH; i++) { + final byte b = header[i]; + if (b == 0 || b == ' ') { + if (digitSeen) { + trailingPadding = true; + } + continue; + } + if (!isOctalDigit(b) || trailingPadding) { + throw new IOException("malformed tar size field in entry header"); + } + digitSeen = true; + value = value * 8 + (b - '0'); + } + return value; + } + + /** + * Reads a size field in the base-256 encoding, which GNU writes when a value does not + * fit the 11 octal digits the field otherwise contains, so entries of 8 GiB and above + * can state their length. + * + *

    The leading bit marks the encoding, the next bit is the sign, and the remaining + * bits of that byte followed by all later bytes form a big-endian two's complement + * number. Negative sizes are rejected.

    + * + * @return The entry size in bytes. + * @throws IOException Thrown if the encoded value is negative or does not fit a + * {@code long}. + */ + private long parseBase256Size() throws IOException { + if ((header[SIZE_OFFSET] & BASE_256_NEGATIVE) != 0) { + throw new IOException("tar size field is negative"); + } + long value = 0; + for (int i = SIZE_OFFSET; i < SIZE_OFFSET + SIZE_LENGTH; i++) { + final int b = i == SIZE_OFFSET + ? header[i] & BASE_256_FIRST_BYTE_BITS : header[i] & 0xFF; + if (value >>> (Long.SIZE - Byte.SIZE - 1) != 0) { + throw new IOException( + "tar size field exceeds the largest length this reader can represent"); + } + value = value << Byte.SIZE | b; + } + return value; + } + + /** + * @param entrySize The size of an entry's content in bytes. + * @return The number of padding bytes that align the entry to the next 512-byte + * block boundary. + */ + private long padding(long entrySize) { + final long remainder = entrySize % BLOCK; + return remainder == 0 ? 0 : BLOCK - remainder; + } + + /** + * Consumes and discards the given number of bytes from the underlying stream. + * + * @param bytes The number of bytes to discard. + * @throws IOException Thrown if the stream ends before all bytes were consumed. + */ + private void skip(long bytes) throws IOException { + try { + in.skipNBytes(bytes); + } catch (EOFException e) { + throw new IOException("truncated tar archive", e); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizer.java index 56a46b556d..a703e95575 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizer.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizer.java @@ -27,19 +27,36 @@ public class AggregateCharSequenceNormalizer implements CharSequenceNormalizer { private static final long serialVersionUID = 5514902020184083235L; private final CharSequenceNormalizer[] normalizers; - public AggregateCharSequenceNormalizer(CharSequenceNormalizer ... normalizers) { - this.normalizers = normalizers; + /** + * Creates an aggregate that applies the given normalizers in order. + * + * @param normalizers The normalizers to apply, first to last. Must not be + * {@code null} and must not contain {@code null}. The array is + * copied, so later changes to it do not reach this instance. + * @throws IllegalArgumentException Thrown if {@code normalizers} is {@code null} or + * contains {@code null}. + */ + public AggregateCharSequenceNormalizer(CharSequenceNormalizer... normalizers) { + if (normalizers == null) { + throw new IllegalArgumentException("The normalizers must not be null."); + } + for (CharSequenceNormalizer normalizer : normalizers) { + if (normalizer == null) { + throw new IllegalArgumentException("The normalizers must not contain null."); + } + } + this.normalizers = normalizers.clone(); } /** {@inheritDoc} */ @Override - public CharSequence normalize (CharSequence text) { + public CharSequence normalize(CharSequence text) { if (text == null) { throw new IllegalArgumentException("The text must not be null."); } - for (CharSequenceNormalizer normalizers : normalizers) { - text = normalizers.normalize(text); + for (CharSequenceNormalizer normalizer : normalizers) { + text = normalizer.normalize(text); } return text; diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AlignedAggregateCharSequenceNormalizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AlignedAggregateCharSequenceNormalizer.java index efa318f476..1f7a8c14f5 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AlignedAggregateCharSequenceNormalizer.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/AlignedAggregateCharSequenceNormalizer.java @@ -17,11 +17,11 @@ package opennlp.tools.util.normalizer; /** - * An {@link OffsetAwareNormalizer} that applies a chain of offset-aware rungs in order and composes + * An {@link OffsetAwareNormalizer} that applies a chain of offset-aware normalizers in order and composes * their per-stage {@link Alignment}s with {@link Alignment#andThen(Alignment)}, so the result maps a * span found in the fully normalized text back to the original input through every stage. * - *

    Produced by {@code TextNormalizer.Builder.buildAligned()}, which validates that every rung is + *

    Produced by {@code TextNormalizer.Builder.buildAligned()}, which validates that every normalizer is * offset-aware before constructing this.

    */ final class AlignedAggregateCharSequenceNormalizer implements OffsetAwareNormalizer { diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/LineBreakPreservingWhitespaceCharSequenceNormalizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/LineBreakPreservingWhitespaceCharSequenceNormalizer.java index b0fea90320..d79317ff00 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/LineBreakPreservingWhitespaceCharSequenceNormalizer.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/LineBreakPreservingWhitespaceCharSequenceNormalizer.java @@ -36,18 +36,7 @@ public class LineBreakPreservingWhitespaceCharSequenceNormalizer implements Offs private static final CharClass WHITESPACE = CharClass.whitespace(); - // The Unicode mandatory break code points (UAX #14 classes BK/CR/LF/NL): line feed, vertical tab, - // form feed, carriage return, next line, line separator, and paragraph separator. A whitespace run - // that contains any of these collapses to a single newline rather than a space, so line and - // paragraph structure survives while horizontal runs are squished. - private static final CodePointSet LINE_BREAKS = CodePointSet.of( - 0x000A, // line feed - 0x000B, // vertical tab - 0x000C, // form feed - 0x000D, // carriage return - 0x0085, // next line - 0x2028, // line separator - 0x2029); // paragraph separator + private static final CodePointSet LINE_BREAKS = UnicodeWhitespace.lineBreakCodePointSet(); private static final LineBreakPreservingWhitespaceCharSequenceNormalizer INSTANCE = new LineBreakPreservingWhitespaceCharSequenceNormalizer(); diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizer.java new file mode 100644 index 0000000000..44ae476dea --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizer.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.tools.util.normalizer; + +/** + * An {@link OffsetAwareNormalizer} that unwraps hard-wrapped prose while keeping real paragraph + * breaks. Within each whitespace run, at most one logical line break (including {@code CRLF} as a + * single break) collapses to a single space; two or more collapse to one newline + * ({@code U+000A}). Leading and trailing whitespace is trimmed. + * + *

    This is the form wanted before sentence detection on Gutenberg-style or other fixed-width + * plain text: intra-paragraph line wraps become spaces, blank-line paragraph boundaries survive as + * newlines. It reuses the cursor based {@link CharClass#collapseParagraphPreserving(CharSequence, + * CodePointSet, int)} engine, so it recognizes the full Unicode {@code White_Space} set with no + * regular expression.

    + * + *

    For display-oriented text where every line break should survive, use + * {@link LineBreakPreservingWhitespaceCharSequenceNormalizer} instead. For raw markdown with lists, + * code blocks, or other single-newline structure, apply this normalizer only to extracted prose.

    + */ +public class ParagraphPreservingWhitespaceCharSequenceNormalizer implements OffsetAwareNormalizer { + + private static final long serialVersionUID = 7597314924708767414L; + + private static final int NEWLINE = 0x000A; + + private static final CharClass WHITESPACE = CharClass.whitespace(); + + private static final CodePointSet LINE_BREAKS = UnicodeWhitespace.lineBreakCodePointSet(); + + private static final ParagraphPreservingWhitespaceCharSequenceNormalizer INSTANCE = + new ParagraphPreservingWhitespaceCharSequenceNormalizer(); + + /** {@return the shared, stateless instance} */ + public static ParagraphPreservingWhitespaceCharSequenceNormalizer getInstance() { + return INSTANCE; + } + + private Object readResolve() { + return INSTANCE; + } + + /** + * {@inheritDoc} + */ + @Override + public CharSequence normalize(CharSequence text) { + return WHITESPACE.trim(WHITESPACE.collapseParagraphPreserving(text, LINE_BREAKS, NEWLINE)); + } + + /** + * {@inheritDoc} + */ + @Override + public AlignedText normalizeAligned(CharSequence text) { + final AlignedText collapsed = + WHITESPACE.collapseParagraphPreservingAligned(text, LINE_BREAKS, NEWLINE); + final AlignedText trimmed = WHITESPACE.trimAligned(collapsed.normalized()); + return new AlignedText(text, trimmed.normalized(), + collapsed.alignment().andThen(trimmed.alignment())); + } +} diff --git a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/TextNormalizer.java b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/TextNormalizer.java index 65de950f23..8ab0dcd5bd 100644 --- a/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/TextNormalizer.java +++ b/opennlp-core/opennlp-runtime/src/main/java/opennlp/tools/util/normalizer/TextNormalizer.java @@ -105,6 +105,16 @@ public Builder whitespacePreservingLineBreaks() { return add(LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance()); } + /** + * {@return this builder with hard-wrap unwrapping that preserves paragraph breaks appended} + * Within each whitespace run, a single line break (including {@code CRLF}) collapses to a space; + * two or more collapse to a single newline, so Gutenberg-style wraps unwrap without losing + * blank-line paragraph boundaries. + */ + public Builder whitespacePreservingParagraphs() { + return add(ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance()); + } + /** {@return this builder with quotation-mark folding appended} */ public Builder quotes() { return add(QuoteCharSequenceNormalizer.getInstance()); diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/chunker/ChunkerAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/chunker/ChunkerAnnotatorTest.java new file mode 100644 index 0000000000..3214c13897 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/chunker/ChunkerAnnotatorTest.java @@ -0,0 +1,211 @@ +/* + * 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.chunker; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Sequence; +import opennlp.tools.util.Span; + +public class ChunkerAnnotatorTest { + + /** + * A chunker that records every token and tag sequence it receives and answers with + * one {@code NP} chunk per run of {@code N}-initial tags, so slicing and span mapping + * are observable. Tests override {@link #chunkAsSpans(String[], String[])} where a + * deviant answer is the fixture. + */ + private static class RecordingChunker implements Chunker { + + private final List> tokenCalls = new ArrayList<>(); + private final List> tagCalls = new ArrayList<>(); + + @Override + public String[] chunk(String[] toks, String[] tags) { + throw new UnsupportedOperationException("the adapter only calls chunkAsSpans"); + } + + @Override + public Span[] chunkAsSpans(String[] toks, String[] tags) { + tokenCalls.add(List.of(toks)); + tagCalls.add(List.of(tags)); + final List spans = new ArrayList<>(); + int start = -1; + for (int i = 0; i <= tags.length; i++) { + final boolean noun = i < tags.length && tags[i].startsWith("N"); + if (noun && start < 0) { + start = i; + } else if (!noun && start >= 0) { + spans.add(new Span(start, i, "NP")); + start = -1; + } + } + return spans.toArray(new Span[0]); + } + + @Override + public Sequence[] topKSequences(String[] sentence, String[] tags) { + throw new UnsupportedOperationException("the adapter only calls chunkAsSpans"); + } + + @Override + public Sequence[] topKSequences(String[] sentence, String[] tags, + double minSequenceScore) { + throw new UnsupportedOperationException("the adapter only calls chunkAsSpans"); + } + } + + private static List> tokens(String text, String... forms) { + final List> annotations = new ArrayList<>(forms.length); + int cursor = 0; + for (final String form : forms) { + final int start = text.indexOf(form, cursor); + annotations.add(new Annotation<>(new Span(start, start + form.length()), form)); + cursor = start + form.length(); + } + return annotations; + } + + private static List> values(List> tokens, + String... tags) { + final List> annotations = new ArrayList<>(tags.length); + for (int i = 0; i < tags.length; i++) { + annotations.add(new Annotation<>(tokens.get(i).span(), tags[i])); + } + return annotations; + } + + /** Two sentences whose noun runs straddle neither sentence boundary. */ + private static Document twoSentences() { + final String text = "Mary Jones leads Acme. She joined Acme Corp."; + final List> toks = tokens(text, + "Mary", "Jones", "leads", "Acme", ".", "She", "joined", "Acme", "Corp", "."); + return Document.of(text) + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 22), "s"), + new Annotation<>(new Span(23, 44), "s"))) + .with(Layers.TOKENS, toks) + .with(Layers.POS_TAGS, values(toks, + "NNP", "NNP", "VBZ", "NNP", ".", "PRP", "VBD", "NNP", "NNP", ".")); + } + + @Test + void testChunksEachSentenceOntoTokenSpans() { + final RecordingChunker chunker = new RecordingChunker(); + final Document document = new ChunkerAnnotator(chunker).annotate(twoSentences()); + + Assertions.assertEquals(List.of( + List.of("Mary", "Jones", "leads", "Acme", "."), + List.of("She", "joined", "Acme", "Corp", ".")), chunker.tokenCalls); + Assertions.assertEquals(List.of( + List.of("NNP", "NNP", "VBZ", "NNP", "."), + List.of("PRP", "VBD", "NNP", "NNP", ".")), chunker.tagCalls); + final List> chunks = document.get(ChunkerAnnotator.CHUNKS); + Assertions.assertEquals(List.of( + new Annotation<>(new Span(0, 10), "NP"), + new Annotation<>(new Span(17, 21), "NP"), + new Annotation<>(new Span(34, 43), "NP")), chunks); + Assertions.assertEquals("Acme Corp", document.text().subSequence(34, 43).toString()); + } + + @Test + void testEmptyLayersYieldEmptyChunkLayer() { + final Document document = new ChunkerAnnotator(new RecordingChunker()).annotate( + Document.of("") + .with(Layers.SENTENCES, List.of()) + .with(Layers.TOKENS, List.of()) + .with(Layers.POS_TAGS, List.of())); + Assertions.assertTrue(document.layers().contains(ChunkerAnnotator.CHUNKS)); + Assertions.assertTrue(document.get(ChunkerAnnotator.CHUNKS).isEmpty()); + } + + @Test + void testLayerContract() { + final ChunkerAnnotator annotator = new ChunkerAnnotator(new RecordingChunker()); + Assertions.assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS), + annotator.requires()); + Assertions.assertEquals(Set.of(ChunkerAnnotator.CHUNKS), annotator.provides()); + Assertions.assertEquals("opennlp:chunks", ChunkerAnnotator.CHUNKS.id()); + Assertions.assertEquals("ChunkerAnnotator", annotator.toString()); + } + + @Test + void testRejectsNullChunkerAndMissingLayers() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> new ChunkerAnnotator(null)); + final ChunkerAnnotator annotator = new ChunkerAnnotator(new RecordingChunker()); + Assertions.assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(null)); + final Document untagged = Document.of("Mary.") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 5), "s"))) + .with(Layers.TOKENS, tokens("Mary.", "Mary", ".")); + Assertions.assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(untagged)); + } + + @Test + void testRejectsMisalignedTagLayer() { + final String text = "Mary."; + final List> toks = tokens(text, "Mary", "."); + final Document document = Document.of(text) + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 5), "s"))) + .with(Layers.TOKENS, toks) + .with(Layers.POS_TAGS, values(toks, "NNP")); + final ChunkerAnnotator annotator = new ChunkerAnnotator(new RecordingChunker()); + final IllegalArgumentException rejection = Assertions.assertThrows( + IllegalArgumentException.class, () -> annotator.annotate(document)); + Assertions.assertTrue(rejection.getMessage().contains("aligned"), + rejection.getMessage()); + } + + @Test + void testRejectsChunksOutsideSentenceOrWithoutType() { + final ChunkerAnnotator outside = new ChunkerAnnotator(new RecordingChunker() { + @Override + public Span[] chunkAsSpans(String[] toks, String[] tags) { + return new Span[] {new Span(0, toks.length + 1, "NP")}; + } + }); + Assertions.assertThrows(IllegalArgumentException.class, + () -> outside.annotate(twoSentences())); + final ChunkerAnnotator empty = new ChunkerAnnotator(new RecordingChunker() { + @Override + public Span[] chunkAsSpans(String[] toks, String[] tags) { + return new Span[] {new Span(1, 1, "NP")}; + } + }); + Assertions.assertThrows(IllegalArgumentException.class, + () -> empty.annotate(twoSentences())); + final ChunkerAnnotator untyped = new ChunkerAnnotator(new RecordingChunker() { + @Override + public Span[] chunkAsSpans(String[] toks, String[] tags) { + return new Span[] {new Span(0, 1)}; + } + }); + Assertions.assertThrows(IllegalArgumentException.class, + () -> untyped.annotate(twoSentences())); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentAnalyzerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentAnalyzerTest.java new file mode 100644 index 0000000000..7af7a75162 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentAnalyzerTest.java @@ -0,0 +1,206 @@ +/* + * 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.document; + +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.namefind.NameFinderAnnotator; +import opennlp.tools.namefind.TokenNameFinder; +import opennlp.tools.postag.POSTagger; +import opennlp.tools.postag.POSTaggerAnnotator; +import opennlp.tools.sentdetect.SentenceDetectorAnnotator; +import opennlp.tools.tokenize.TokenizerAnnotator; +import opennlp.tools.util.Sequence; +import opennlp.tools.util.Span; + +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 {@link DocumentAnalyzer} pipeline over the adapter annotators, using the + * deterministic components from {@link TestComponents} and a fixed-vocabulary tagger. + * The point under test is the pipeline mechanics and span arithmetic, not model quality. + */ +public class DocumentAnalyzerTest { + + /** Tags the known verbs of the test texts {@code VBZ} and everything else {@code X}. */ + private static final POSTagger TAGGER = new POSTagger() { + + private final Set verbs = Set.of("barks.", "eats."); + + @Override + public String[] tag(String[] sentence) { + final String[] tags = new String[sentence.length]; + for (int i = 0; i < sentence.length; i++) { + tags[i] = verbs.contains(sentence[i]) ? "VBZ" : "X"; + } + return tags; + } + + @Override + public String[] tag(String[] sentence, Object[] additionalContext) { + return tag(sentence); + } + + @Override + public Sequence[] topKSequences(String[] sentence) { + throw new UnsupportedOperationException(); + } + + @Override + public Sequence[] topKSequences(String[] sentence, Object[] additionalContext) { + throw new UnsupportedOperationException(); + } + }; + + /** A finder that finds no names, for tests that only exercise the pipeline plumbing. */ + private static final TokenNameFinder NO_NAMES = new TokenNameFinder() { + + @Override + public Span[] find(String[] tokens) { + return new Span[0]; + } + + @Override + public void clearAdaptiveData() { + } + }; + + /** + * Verifies that the adapters identify themselves by their simple class name, which is + * how a pipeline validation message names the offending annotator. + */ + @Test + void testAdaptersNameThemselvesByClassName() { + assertEquals("SentenceDetectorAnnotator", + new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER).toString()); + assertEquals("TokenizerAnnotator", + new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER).toString()); + assertEquals("POSTaggerAnnotator", new POSTaggerAnnotator(TAGGER).toString()); + assertEquals("NameFinderAnnotator", new NameFinderAnnotator(NO_NAMES).toString()); + } + + @Test + void testPipelineProducesAlignedLayersInOriginalCoordinates() { + final Document document = DocumentAnalyzer.builder() + .add(new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER)) + .add(new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER)) + .add(new POSTaggerAnnotator(TAGGER)) + .build() + .analyze("the dog barks. she eats."); + + final List> sentences = document.get(Layers.SENTENCES); + assertEquals(2, sentences.size()); + assertEquals("she eats.", sentences.get(1).value()); + + final List> tokens = document.get(Layers.TOKENS); + assertEquals(5, tokens.size()); + // token of the second sentence, span in document coordinates + assertEquals("she", tokens.get(3).value()); + assertEquals(new Span(15, 18), tokens.get(3).span()); + + final List> tags = document.get(Layers.POS_TAGS); + assertEquals(5, tags.size()); + assertEquals("VBZ", tags.get(2).value()); + assertEquals(tokens.get(2).span(), tags.get(2).span()); + } + + /** + * Verifies that a full pipeline over empty and whitespace-only input produces a + * document on which every provided layer is present and empty, rather than failing: + * zero sentences legitimately yield zero tokens, zero tags, and zero entities. + */ + @ParameterizedTest + @ValueSource(strings = {"", " "}) + void testEmptyAndBlankInputProduceEmptyLayers(String text) { + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER)) + .add(new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER)) + .add(new POSTaggerAnnotator(TAGGER)) + .add(new NameFinderAnnotator(NO_NAMES)) + .build(); + + final Document document = analyzer.analyze(text); + assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS, Layers.ENTITIES), + document.layers()); + for (final LayerKey layer : document.layers()) { + assertTrue(document.get(layer).isEmpty()); + } + } + + /** + * Verifies that a present-but-empty sentence layer is honored as "no sentences": the + * tokenizer adds a present-but-empty token layer instead of tokenizing the whole text. + */ + @Test + void testTokenizerHonorsPresentButEmptySentenceLayer() { + final Document document = new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER) + .annotate(Document.of("the dog").with(Layers.SENTENCES, List.of())); + assertTrue(document.layers().contains(Layers.TOKENS)); + assertTrue(document.get(Layers.TOKENS).isEmpty()); + } + + @Test + void testTokenizerWorksWithoutSentences() { + final Document document = DocumentAnalyzer.builder() + .add(new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER)) + .build() + .analyze("the dog"); + assertEquals(2, document.get(Layers.TOKENS).size()); + } + + @Test + void testMisorderedPipelineFailsAtBuildTime() { + final DocumentAnalyzer.Builder builder = DocumentAnalyzer.builder() + .add(new POSTaggerAnnotator(TAGGER)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + void testAnnotatorAdaptersRejectNullDelegates() { + assertThrows(IllegalArgumentException.class, () -> new SentenceDetectorAnnotator(null)); + assertThrows(IllegalArgumentException.class, () -> new TokenizerAnnotator(null)); + assertThrows(IllegalArgumentException.class, () -> new POSTaggerAnnotator(null)); + assertThrows(IllegalArgumentException.class, () -> new NameFinderAnnotator(null)); + } + + /** + * Verifies that every adapter rejects a {@code null} document with the shared + * message, whether it checks itself or through + * {@link DocumentAnnotators#requireLayers(Document, LayerKey[])}. + */ + @Test + void testAnnotatorAdaptersRejectNullDocuments() { + final List adapters = List.of( + new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER), + new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER), + new POSTaggerAnnotator(TAGGER), + new NameFinderAnnotator(NO_NAMES)); + for (final DocumentAnnotator adapter : adapters) { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> adapter.annotate(null)); + assertEquals("document must not be null", e.getMessage()); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentPipelineExampleTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentPipelineExampleTest.java new file mode 100644 index 0000000000..f38737d668 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/DocumentPipelineExampleTest.java @@ -0,0 +1,246 @@ +/* + * 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.document; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.postag.POSTagger; +import opennlp.tools.postag.POSTaggerAnnotator; +import opennlp.tools.sentdetect.SentenceDetectorAnnotator; +import opennlp.tools.tokenize.TokenizerAnnotator; +import opennlp.tools.util.Sequence; +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +/** + * Walks through the document pipeline the way a first-time user would: wrap existing + * analysis components in their adapter annotators, add one custom annotator, build a + * {@link DocumentAnalyzer}, analyze a two-sentence text, and read every layer back with + * its spans in original text coordinates. + * + *

    The wrapped components are the tiny deterministic stand-ins from + * {@link TestComponents}, so every expected span and value below follows directly from + * the input text. The point under demonstration is how the layers connect, not the + * quality of any single step.

    + */ +public class DocumentPipelineExampleTest { + + /** + * The key of the custom layer produced by {@link TokenLengthAnnotator}. Any producer + * may introduce such a key in its own code; the container needs no change for it. + */ + private static final LayerKey TOKEN_LENGTHS = + LayerKey.of("token-lengths", Integer.class); + + /** + * A deterministic tagger backed by a fixed dictionary covering exactly the tokens of + * the example text. An unknown token fails the test immediately rather than receiving + * a silent fallback tag. + */ + private static final POSTagger DICTIONARY_TAGGER = new POSTagger() { + + private final Map tagsByToken = Map.of( + "The", "DT", "dog", "NN", "barks.", "VBZ", "It", "PRP", "naps.", "VBZ"); + + @Override + public String[] tag(String[] sentence) { + final String[] tags = new String[sentence.length]; + for (int i = 0; i < sentence.length; i++) { + final String tag = tagsByToken.get(sentence[i]); + if (tag == null) { + throw new IllegalArgumentException("no tag defined for token: " + sentence[i]); + } + tags[i] = tag; + } + return tags; + } + + @Override + public String[] tag(String[] sentence, Object[] additionalContext) { + return tag(sentence); + } + + @Override + public Sequence[] topKSequences(String[] sentence) { + throw new UnsupportedOperationException("the adapter only calls tag"); + } + + @Override + public Sequence[] topKSequences(String[] sentence, Object[] additionalContext) { + throw new UnsupportedOperationException("the adapter only calls tag"); + } + }; + + /** + * A custom pipeline step written directly against {@link DocumentAnnotator}: it reads + * the token layer and provides {@link #TOKEN_LENGTHS}, one annotation per token on the + * token's span, whose value is the character length of the token text. + */ + private static final class TokenLengthAnnotator implements DocumentAnnotator { + + /** + * Adds the {@link #TOKEN_LENGTHS} layer computed from {@link Layers#TOKENS}. + * + * @param document The document to annotate. Must not be {@code null} and must + * contain the token layer, which may be empty. + * @return A new {@link Document} carrying the token length layer. Never {@code null}. + * @throws IllegalArgumentException Thrown if {@code document} is {@code null} or the + * token layer is absent. + */ + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.TOKENS); + final List> tokens = document.get(Layers.TOKENS); + final List> lengths = new ArrayList<>(tokens.size()); + for (final Annotation token : tokens) { + lengths.add(new Annotation<>(token.span(), token.value().length())); + } + return document.with(TOKEN_LENGTHS, lengths); + } + + @Override + public Set> requires() { + return Set.of(Layers.TOKENS); + } + + @Override + public Set> provides() { + return Set.of(TOKEN_LENGTHS); + } + } + + /** + * Mirrors the manual's document-scoped layer example: a language id rides a + * document-scoped key as a span-less value with exactly the id, value, and null span + * the chapter shows. + */ + @Test + void testDocumentScopedLayerExample() { + final LayerKey language = LayerKey.document("app:language", String.class); + final Document document = Document.of("The dog barks. It naps."); + + final Document tagged = document.with(language, List.of(Annotation.of("eng"))); + assertEquals("eng", tagged.get(language).get(0).value()); + assertNull(tagged.get(language).get(0).span()); + } + + /** + * Runs the full pipeline story: sentences, tokens, part-of-speech tags, and one custom + * layer over a two-sentence text, then verifies every annotation of every layer, span + * by span, in original text coordinates. + */ + @Test + void testFullPipelineStory() { + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER)) + .add(new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER)) + .add(new POSTaggerAnnotator(DICTIONARY_TAGGER)) + .add(new TokenLengthAnnotator()) + .build(); + + final Document document = analyzer.analyze("The dog barks. It naps."); + + // The document carries exactly the four layers the pipeline provides. + assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS, TOKEN_LENGTHS), + document.layers()); + assertEquals("The dog barks. It naps.", document.text()); + + // Sentence layer: one annotation per sentence, covering it in document coordinates. + final List> sentences = document.get(Layers.SENTENCES); + assertEquals(2, sentences.size()); + assertEquals(new Span(0, 14), sentences.get(0).span()); + assertEquals("The dog barks.", sentences.get(0).value()); + assertEquals(new Span(15, 23), sentences.get(1).span()); + assertEquals("It naps.", sentences.get(1).value()); + + // Token layer: five tokens; the second sentence's spans are shifted back to + // document coordinates, so every span can index into the original text. + final List> tokens = document.get(Layers.TOKENS); + assertEquals(5, tokens.size()); + assertEquals(new Span(0, 3), tokens.get(0).span()); + assertEquals("The", tokens.get(0).value()); + assertEquals(new Span(4, 7), tokens.get(1).span()); + assertEquals("dog", tokens.get(1).value()); + assertEquals(new Span(8, 14), tokens.get(2).span()); + assertEquals("barks.", tokens.get(2).value()); + assertEquals(new Span(15, 17), tokens.get(3).span()); + assertEquals("It", tokens.get(3).value()); + assertEquals(new Span(18, 23), tokens.get(4).span()); + assertEquals("naps.", tokens.get(4).value()); + + // Tag layer: aligned with the token layer by position, each tag on its token's span. + final List> tags = document.get(Layers.POS_TAGS); + assertEquals(5, tags.size()); + assertEquals("DT", tags.get(0).value()); + assertEquals("NN", tags.get(1).value()); + assertEquals("VBZ", tags.get(2).value()); + assertEquals("PRP", tags.get(3).value()); + assertEquals("VBZ", tags.get(4).value()); + for (int i = 0; i < tags.size(); i++) { + assertEquals(tokens.get(i).span(), tags.get(i).span()); + } + + // Custom layer: the container returns it as List>, so the + // values are used as numbers without a cast. + final List> lengths = document.get(TOKEN_LENGTHS); + assertEquals(5, lengths.size()); + assertEquals(3, lengths.get(0).value()); + assertEquals(3, lengths.get(1).value()); + assertEquals(6, lengths.get(2).value()); + assertEquals(2, lengths.get(3).value()); + assertEquals(5, lengths.get(4).value()); + for (int i = 0; i < lengths.size(); i++) { + assertEquals(tokens.get(i).span(), lengths.get(i).span()); + } + + // Every span refers to the original text, so covered text always round-trips. + for (final Annotation token : tokens) { + assertEquals(token.value(), + token.span().getCoveredText(document.text()).toString()); + } + } + + /** + * Verifies that the analyzer leaves the input untouched between calls: analyzing two + * texts with the same analyzer yields two independent documents. + */ + @Test + void testAnalyzerIsReusableAcrossTexts() { + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new SentenceDetectorAnnotator(TestComponents.PERIOD_SPLITTER)) + .add(new TokenizerAnnotator(TestComponents.SPACE_TOKENIZER)) + .build(); + + final Document first = analyzer.analyze("The dog barks."); + final Document second = analyzer.analyze("It naps."); + + assertEquals(3, first.get(Layers.TOKENS).size()); + assertEquals(2, second.get(Layers.TOKENS).size()); + assertEquals("The dog barks.", first.text()); + assertEquals("It naps.", second.text()); + assertEquals(1, first.get(Layers.SENTENCES).size()); + assertEquals(1, second.get(Layers.SENTENCES).size()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/TestComponents.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/TestComponents.java new file mode 100644 index 0000000000..6145e5b790 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/document/TestComponents.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.tools.document; + +import java.util.ArrayList; +import java.util.List; + +import opennlp.tools.sentdetect.SentenceDetector; +import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.util.Span; + +/** + * Deterministic stand-in components shared by the document pipeline tests, so every + * expected span in those tests follows directly from the definitions here. + */ +final class TestComponents { + + /** + * A deterministic sentence detector that ends a sentence after every period and + * expects a single space between sentences. Only the span-producing method is + * implemented because the adapter calls no other method. + */ + static final SentenceDetector PERIOD_SPLITTER = new SentenceDetector() { + + @Override + public String[] sentDetect(CharSequence s) { + throw new UnsupportedOperationException("the adapter only calls sentPosDetect"); + } + + @Override + public Span[] sentPosDetect(CharSequence s) { + final String text = s.toString(); + final List spans = new ArrayList<>(); + int start = 0; + for (int i = 0; i < text.length(); i++) { + if (text.charAt(i) == '.') { + spans.add(new Span(start, i + 1)); + start = i + 2; + } + } + return spans.toArray(new Span[0]); + } + }; + + /** + * A deterministic tokenizer that splits on single space characters and keeps all + * other characters, including sentence-final periods, attached to their token. Only + * the span-producing method is implemented because the adapter calls no other method. + */ + static final Tokenizer SPACE_TOKENIZER = new Tokenizer() { + + @Override + public String[] tokenize(String s) { + throw new UnsupportedOperationException("the adapter only calls tokenizePos"); + } + + @Override + public Span[] tokenizePos(String s) { + final List spans = new ArrayList<>(); + int start = -1; + for (int i = 0; i <= s.length(); i++) { + final boolean boundary = i == s.length() || s.charAt(i) == ' '; + if (boundary && start >= 0) { + spans.add(new Span(start, i)); + start = -1; + } else if (!boundary && start < 0) { + start = i; + } + } + return spans.toArray(new Span[0]); + } + }; + + private TestComponents() { + // Not instantiated; this class provides shared test fixtures only. + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/DictionaryLemmatizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/DictionaryLemmatizerTest.java index add85b2d5e..e4f53c4b24 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/DictionaryLemmatizerTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/DictionaryLemmatizerTest.java @@ -17,14 +17,26 @@ package opennlp.tools.lemmatizer; +import java.io.ByteArrayInputStream; +import java.nio.charset.StandardCharsets; +import java.util.Locale; + +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; public class DictionaryLemmatizerTest { + /** + * Turkish folds {@code 'I'} to the dotless {@code 'ı'} (U+0131) instead of {@code 'i'}. + */ + private static final Locale TURKISH = Locale.of("tr", "TR"); + private static DictionaryLemmatizer dictionaryLemmatizer; + private final Locale defaultLocale = Locale.getDefault(); + @BeforeAll static void loadDictionary() throws Exception { dictionaryLemmatizer = new DictionaryLemmatizer( @@ -49,4 +61,25 @@ void testForNullPointerException() { } } + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + + /** + * The dictionary is read verbatim but looked up folded, so the fold has to match the + * casing of the dictionary file rather than whatever locale the JVM happens to run in. + */ + @Test + void testLookupIsIndependentOfDefaultLocale() throws Exception { + final String entries = "illinois\tNNP\tIllinois\n" + "indices\tNNS\tindex\n"; + final DictionaryLemmatizer lemmatizer = new DictionaryLemmatizer( + new ByteArrayInputStream(entries.getBytes(StandardCharsets.UTF_8))); + + Locale.setDefault(TURKISH); + + Assertions.assertArrayEquals(new String[] {"Illinois", "index"}, + lemmatizer.lemmatize(new String[] {"Illinois", "INDICES"}, new String[] {"NNP", "NNS"})); + } + } diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerAnnotatorTest.java new file mode 100644 index 0000000000..63654bbada --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerAnnotatorTest.java @@ -0,0 +1,247 @@ +/* + * 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.lemmatizer; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class LemmatizerAnnotatorTest { + + /** Lowercases verbs and keeps everything else, enough to observe the adapter. */ + private static final Lemmatizer FIXTURE = new Lemmatizer() { + @Override + public String[] lemmatize(String[] toks, String[] tags) { + final String[] lemmas = new String[toks.length]; + for (int i = 0; i < toks.length; i++) { + lemmas[i] = "VERB".equals(tags[i]) ? "run" : toks[i]; + } + return lemmas; + } + + @Override + public List> lemmatize(List toks, List tags) { + throw new UnsupportedOperationException(); + } + }; + + @Test + void testLemmasAlignWithTokens() { + final Document document = Document.of("She ran home") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 12), "She ran home"))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "She"), + new Annotation<>(new Span(4, 7), "ran"), + new Annotation<>(new Span(8, 12), "home"))) + .with(Layers.POS_TAGS, List.of( + new Annotation<>(new Span(0, 3), "PRON"), + new Annotation<>(new Span(4, 7), "VERB"), + new Annotation<>(new Span(8, 12), "NOUN"))); + + final Document lemmatized = new LemmatizerAnnotator(FIXTURE).annotate(document); + + final List> lemmas = lemmatized.get(LemmatizerAnnotator.LEMMAS); + assertEquals(3, lemmas.size()); + assertEquals("run", lemmas.get(1).value()); + assertEquals(new Span(4, 7), lemmas.get(1).span()); + assertEquals("home", lemmas.get(2).value()); + } + + /** + * Verifies that the lemmatizer is invoked once per sentence with exactly that + * sentence's tokens and tags, so lemmatization decisions never see material from a + * neighboring sentence, and that the lemma layer still aligns with the token layer. + */ + @Test + void testLemmatizesPerSentence() { + final List> calls = new ArrayList<>(); + final Lemmatizer recording = new Lemmatizer() { + @Override + public String[] lemmatize(String[] toks, String[] tags) { + calls.add(List.of(toks)); + return toks.clone(); + } + + @Override + public List> lemmatize(List toks, List tags) { + throw new UnsupportedOperationException(); + } + }; + final Document document = Document.of("Ana runs. Bob sits.") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 9), "Ana runs."), + new Annotation<>(new Span(10, 19), "Bob sits."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(10, 13), "Bob"), + new Annotation<>(new Span(14, 19), "sits."))) + .with(Layers.POS_TAGS, List.of( + new Annotation<>(new Span(0, 3), "PROPN"), + new Annotation<>(new Span(4, 9), "VERB"), + new Annotation<>(new Span(10, 13), "PROPN"), + new Annotation<>(new Span(14, 19), "VERB"))); + + final Document lemmatized = new LemmatizerAnnotator(recording).annotate(document); + + assertEquals(List.of( + List.of("Ana", "runs."), + List.of("Bob", "sits.")), calls); + assertEquals(4, lemmatized.get(LemmatizerAnnotator.LEMMAS).size()); + assertEquals(new Span(10, 13), + lemmatized.get(LemmatizerAnnotator.LEMMAS).get(2).span()); + } + + /** + * Verifies that the adapter identifies itself by its simple class name, which is how a + * pipeline validation message names the offending annotator. + */ + @Test + void testAdapterNamesItselfByClassName() { + assertEquals("LemmatizerAnnotator", new LemmatizerAnnotator(FIXTURE).toString()); + } + + @Test + void testInvalidArguments() { + assertThrows(IllegalArgumentException.class, + () -> new LemmatizerAnnotator(null)); + final LemmatizerAnnotator annotator = new LemmatizerAnnotator(FIXTURE); + assertThrows(IllegalArgumentException.class, () -> annotator.annotate(null)); + final Document misaligned = Document.of("a b") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 3), "a b"))) + .with(Layers.TOKENS, List.of(new Annotation<>(new Span(0, 1), "a"))) + .with(Layers.POS_TAGS, List.of()); + assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(misaligned)); + } + + /** + * Verifies that a document lacking a required layer is rejected with a message naming + * the missing layer, for each of the three required layers in declaration order. + */ + @Test + void testAbsentRequiredLayerThrowsWithExactMessage() { + final LemmatizerAnnotator annotator = new LemmatizerAnnotator(FIXTURE); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(Document.of("no layers"))); + assertEquals("document lacks the required layer opennlp:sentences", + e.getMessage()); + + final Document sentencesOnly = Document.of("a") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 1), "a"))); + final IllegalArgumentException tokenless = assertThrows( + IllegalArgumentException.class, () -> annotator.annotate(sentencesOnly)); + assertEquals("document lacks the required layer opennlp:tokens", + tokenless.getMessage()); + + final Document untagged = sentencesOnly + .with(Layers.TOKENS, List.of(new Annotation<>(new Span(0, 1), "a"))); + final IllegalArgumentException tagless = assertThrows( + IllegalArgumentException.class, () -> annotator.annotate(untagged)); + assertEquals("document lacks the required layer opennlp:pos", + tagless.getMessage()); + } + + /** + * Verifies that present-but-empty layers yield a present-but-empty lemma layer rather + * than an exception. + */ + @Test + void testEmptyPresentLayersYieldEmptyLemmaLayer() { + final Document document = Document.of("") + .with(Layers.SENTENCES, List.of()) + .with(Layers.TOKENS, List.of()) + .with(Layers.POS_TAGS, List.of()); + final Document lemmatized = new LemmatizerAnnotator(FIXTURE).annotate(document); + assertTrue(lemmatized.layers().contains(LemmatizerAnnotator.LEMMAS)); + assertTrue(lemmatized.get(LemmatizerAnnotator.LEMMAS).isEmpty()); + } + + /** + * Verifies that a token lying outside every sentence is rejected loudly, matching the + * walk contract of the other per-sentence adapters. + */ + @Test + void testTokenOutsideEverySentenceThrowsWithExactMessage() { + final Document document = Document.of("Ana runs. Bob") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 9), "Ana runs."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(10, 13), "Bob"))) + .with(Layers.POS_TAGS, List.of( + new Annotation<>(new Span(0, 3), "PROPN"), + new Annotation<>(new Span(4, 9), "VERB"), + new Annotation<>(new Span(10, 13), "PROPN"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new LemmatizerAnnotator(FIXTURE).annotate(document)); + assertEquals("token at [10..13) lies outside every sentence", e.getMessage()); + } + + /** + * Verifies that a lemmatizer returning a wrong number of lemmas for a sentence is + * rejected loudly instead of silently misaligning the lemma layer. + */ + @Test + void testWrongLemmaCountFailsLoud() { + final Lemmatizer shortLemmatizer = new Lemmatizer() { + @Override + public String[] lemmatize(String[] toks, String[] tags) { + return new String[] {"a"}; + } + + @Override + public List> lemmatize(List toks, List tags) { + throw new UnsupportedOperationException(); + } + }; + final Document document = Document.of("a b") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 3), "a b"))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 1), "a"), + new Annotation<>(new Span(2, 3), "b"))) + .with(Layers.POS_TAGS, List.of( + new Annotation<>(new Span(0, 1), "X"), + new Annotation<>(new Span(2, 3), "X"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new LemmatizerAnnotator(shortLemmatizer).annotate(document)); + assertEquals("lemmatizer returned 1 lemmas for 2 tokens", e.getMessage()); + } + + /** + * Verifies that the adapter declares all three consumed layers as required, so a + * pipeline without a sentence step fails at build time. + */ + @Test + void testRequiresSentencesTokensAndTags() { + assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS, Layers.POS_TAGS), + new LemmatizerAnnotator(FIXTURE).requires()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerMELocaleTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerMELocaleTest.java new file mode 100644 index 0000000000..70f20ce07c --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/lemmatizer/LemmatizerMELocaleTest.java @@ -0,0 +1,72 @@ +/* + * 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.lemmatizer; + +import java.util.Locale; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Tests that the lemma class encoding and decoding of {@link LemmatizerME} do not depend + * on the JVM's default {@link Locale}. + *

    + * Encoding produces the outcome labels that are persisted in a trained model, decoding + * consumes them at inference time. Both sides therefore have to agree across JVMs. + */ +public class LemmatizerMELocaleTest { + + /** + * Turkish folds {@code 'I'} to the dotless {@code 'ı'} (U+0131) instead of {@code 'i'}. + */ + private static final Locale TURKISH = Locale.of("tr", "TR"); + + private static final String[] TOKENS = {"MICE", "INDICES"}; + private static final String[] LEMMAS = {"mouse", "index"}; + + /** The lemma classes an existing, English trained model contains for {@link #TOKENS}. */ + private static final String[] LEMMA_CLASSES = {"R2ioR1cuI1s", "D3iD2cR0sx"}; + + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + + /** + * Training on a Turkish JVM must not write different outcome labels into the model. + */ + @Test + void testEncodeLemmasIsIndependentOfDefaultLocale() { + Locale.setDefault(TURKISH); + + Assertions.assertArrayEquals(LEMMA_CLASSES, LemmatizerME.encodeLemmas(TOKENS, LEMMAS)); + } + + /** + * Serving an existing model on a Turkish JVM must not corrupt the reconstructed lemmas. + */ + @Test + void testDecodeLemmasIsIndependentOfDefaultLocale() { + Locale.setDefault(TURKISH); + + Assertions.assertArrayEquals(LEMMAS, LemmatizerME.decodeLemmas(TOKENS, LEMMA_CLASSES)); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchEquivalenceTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchEquivalenceTest.java new file mode 100644 index 0000000000..45efbc741b --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchEquivalenceTest.java @@ -0,0 +1,612 @@ +/* + * 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.ml; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.PriorityQueue; +import java.util.Queue; +import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import opennlp.tools.ml.model.MaxentModel; +import opennlp.tools.util.BeamSearchContextGenerator; +import opennlp.tools.util.Cache; +import opennlp.tools.util.Sequence; +import opennlp.tools.util.SequenceValidator; + +/** + * Equivalence tests for the {@code BeamSearch.bestSequences} chain-node implementation + * introduced with OPENNLP-1903. + *

    + * Every test runs the current {@link BeamSearch} side by side with + * {@link #referenceBestSequences}, a faithful port of the per-candidate + * {@link Sequence}-copying implementation that OPENNLP-1903 replaced, and demands + * identical output: same number of sequences, identical outcome lists + * (order-sensitive), and bit-identical scores and per-position probabilities. + */ +public class BeamSearchEquivalenceTest { + + /** Mirror of the private {@code BeamSearch.ZERO_LOG} default threshold. */ + private static final double ZERO_LOG = -100000; + + private static final String[] OUTCOMES = {"o0", "o1", "o2", "o3"}; + private static final long MODEL_SEED = 0x5eedL; + + private static final int[] BEAM_SIZES = {1, 2, 3, 5, 10}; // 10 > OUTCOMES.length on purpose + private static final int[] INPUT_LENGTHS = {0, 1, 2, 7, 33, 128}; + private static final int[] CACHE_SIZES = {0, 64}; + + private static final SequenceValidator ACCEPT_ALL = + (i, input, outcomes, outcome) -> true; + + /** A {@link SequenceValidator} paired with a stable name for test display. */ + record NamedValidator(String name, SequenceValidator validator) { + @Override + public String toString() { + return name; + } + } + + private static final List VALIDATORS = List.of( + new NamedValidator("rejectOneOutcome", + (i, input, outcomes, outcome) -> !"o2".equals(outcome)), + // Rejects every outcome at position 2: at that position the threshold loop adds + // nothing, so the next.isEmpty() fallback runs (and also rejects everything, + // killing the search for inputs longer than 2). + new NamedValidator("rejectAllAtPosition2", + (i, input, outcomes, outcome) -> i != 2), + // Rejects everything except "o3" at position 2: the fallback actually populates + // next with the sub-threshold "o3" candidate whenever "o3" fell below the min. + new NamedValidator("onlyO3AtPosition2", + (i, input, outcomes, outcome) -> i != 2 || "o3".equals(outcome)), + new NamedValidator("rejectEverything", + (i, input, outcomes, outcome) -> false)); + + // --------------------------------------------------------------------------- + // Context generator + // --------------------------------------------------------------------------- + + /** + * Builds contexts from the index, the token, the previous outcome and the first + * additional-context element. Identical context content is interned to the same + * {@code String[]} instance so the identity-keyed contexts {@link Cache} in + * {@link BeamSearch} actually produces hits. Thread-safe. + */ + static final class SeededContextGenerator implements BeamSearchContextGenerator { + + private final ConcurrentHashMap intern = new ConcurrentHashMap<>(); + private final AtomicInteger callCount = new AtomicInteger(); + + @Override + public String[] getContext(int index, String[] sequence, + String[] priorDecisions, Object[] additionalContext) { + callCount.incrementAndGet(); + String prev = index > 0 ? priorDecisions[index - 1] : ""; + String ac = additionalContext != null && additionalContext.length > 0 + ? String.valueOf(additionalContext[0]) : "-"; + String[] ctx = {"ix=" + index, "tok=" + sequence[index], "prev=" + prev, "ac=" + ac}; + String key = String.join("", ctx); + String[] existing = intern.putIfAbsent(key, ctx); + return existing != null ? existing : ctx; + } + + /** + * @return The number of {@link #getContext} invocations so far. + */ + int callCount() { + return callCount.get(); + } + } + + // --------------------------------------------------------------------------- + // Reference implementation: faithful port of the bestSequences implementation + // prior to OPENNLP-1903 + // --------------------------------------------------------------------------- + + /** + * Port of the {@code BeamSearch.bestSequences} control flow prior to OPENNLP-1903: + * PriorityQueue over {@link Sequence}, per-candidate + * {@code new Sequence(top, out, scores[p])} copies, tempScores sort/min, the + * {@code next.isEmpty()} advance-all-valid fallback, the queue swap, and the winner + * removal order. The cache path (a {@code Cache} exactly like the + * per-thread one in {@link BeamSearch}) is used when {@code cacheSize > 0}; otherwise + * the uncached eval path is taken. + */ + static Sequence[] referenceBestSequences( + final int numSequences, final T[] sequence, final Object[] additionalContext, + final double minSequenceScore, final BeamSearchContextGenerator cg, + final SequenceValidator validator, final MaxentModel model, + final int beamSize, final int cacheSize) { + + // Local equivalents of the per-thread CacheState in BeamSearch. + final double[] probs = new double[model.getNumOutcomes()]; + final double[] tempScores = new double[model.getNumOutcomes()]; + final Cache cache = cacheSize > 0 ? new Cache<>(cacheSize) : null; + + Queue prev = new PriorityQueue<>(beamSize); + Queue next = new PriorityQueue<>(beamSize); + Queue tmp; + prev.add(new Sequence()); + + Object[] context = additionalContext; + if (context == null) { + context = new Object[0]; + } + + for (int i = 0; i < sequence.length; i++) { + final int sz = StrictMath.min(beamSize, prev.size()); + + for (int sc = 0; prev.size() > 0 && sc < sz; sc++) { + final Sequence top = prev.remove(); + final List tmpOutcomes = top.getOutcomes(); + final String[] outcomes = tmpOutcomes.toArray(new String[0]); + final String[] contexts = cg.getContext(i, sequence, outcomes, context); + final double[] scores; + if (cache != null) { + scores = cache.computeIfAbsent(contexts, c -> { + double[] res = model.eval(c, probs); + double[] copy = new double[res.length]; + System.arraycopy(res, 0, copy, 0, res.length); + return copy; + }); + } else { + scores = model.eval(contexts, probs); + } + + System.arraycopy(scores, 0, tempScores, 0, scores.length); + Arrays.sort(tempScores); + + final double min = tempScores[StrictMath.max(0, scores.length - beamSize)]; + + for (int p = 0; p < scores.length; p++) { + if (scores[p] >= min) { + final String out = model.getOutcome(p); + if (validator.validSequence(i, sequence, outcomes, out)) { + final Sequence ns = new Sequence(top, out, scores[p]); + if (ns.getScore() > minSequenceScore) { + next.add(ns); + } + } + } + } + + if (next.isEmpty()) { // if no advanced sequences, advance all valid + for (int p = 0; p < scores.length; p++) { + final String out = model.getOutcome(p); + if (validator.validSequence(i, sequence, outcomes, out)) { + final Sequence ns = new Sequence(top, out, scores[p]); + if (ns.getScore() > minSequenceScore) { + next.add(ns); + } + } + } + } + } + + // make prev = next; and re-init next (reuse existing prev set once cleared) + prev.clear(); + tmp = prev; + prev = next; + next = tmp; + } + + final int numSeq = StrictMath.min(numSequences, prev.size()); + final Sequence[] topSequences = new Sequence[numSeq]; + + for (int seqIndex = 0; seqIndex < numSeq; seqIndex++) { + topSequences[seqIndex] = prev.remove(); + } + + return topSequences; + } + + /** Reference twin of the two-arg overload (default {@code ZERO_LOG} threshold). */ + static Sequence[] referenceBestSequences( + int numSequences, T[] sequence, Object[] additionalContext, + BeamSearchContextGenerator cg, SequenceValidator validator, + MaxentModel model, int beamSize, int cacheSize) { + return referenceBestSequences(numSequences, sequence, additionalContext, ZERO_LOG, + cg, validator, model, beamSize, cacheSize); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + /** + * @return A new model over {@link #OUTCOMES} seeded with {@link #MODEL_SEED}. + */ + private static MaxentModel model() { + return new SeededMaxentModel(OUTCOMES, MODEL_SEED); + } + + /** + * @return A seeded pseudo-random token sequence of {@code length} tokens. + */ + private static String[] randomInput(int length, long seed) { + Random rnd = new Random(seed); + String[] input = new String[length]; + for (int i = 0; i < length; i++) { + input[i] = "t" + rnd.nextInt(17); + } + return input; + } + + private static void assertBitIdentical(double expected, double actual, String what) { + Assertions.assertEquals(Double.doubleToRawLongBits(expected), + Double.doubleToRawLongBits(actual), + what + " (expected=" + expected + ", actual=" + actual + ")"); + } + + private static void assertSequencesEqual(Sequence[] expected, Sequence[] actual, + String caseDesc) { + Assertions.assertNotNull(actual, caseDesc + ": result array must not be null"); + Assertions.assertEquals(expected.length, actual.length, + caseDesc + ": number of returned sequences"); + for (int s = 0; s < expected.length; s++) { + String seqDesc = caseDesc + ", sequence[" + s + "]"; + Assertions.assertEquals(expected[s].getOutcomes(), actual[s].getOutcomes(), + seqDesc + ": outcomes"); + assertBitIdentical(expected[s].getScore(), actual[s].getScore(), + seqDesc + ": score"); + Assertions.assertEquals(expected[s].getSize(), actual[s].getSize(), + seqDesc + ": size"); + double[] expectedProbs = expected[s].getProbs(); + double[] actualProbs = actual[s].getProbs(); + Assertions.assertEquals(expectedProbs.length, actualProbs.length, + seqDesc + ": probs length"); + for (int p = 0; p < expectedProbs.length; p++) { + assertBitIdentical(expectedProbs[p], actualProbs[p], + seqDesc + ": prob[" + p + "]"); + assertBitIdentical(expected[s].getProb(p), actual[s].getProb(p), + seqDesc + ": getProb(" + p + ")"); + } + } + } + + private static String caseDesc(String test, int beam, int length, int cache) { + return test + "[beam=" + beam + ", len=" + length + ", cache=" + cache + "]"; + } + + // --------------------------------------------------------------------------- + // Parameter sources + // --------------------------------------------------------------------------- + + /** + * @return The cartesian product of {@link #BEAM_SIZES}, {@link #INPUT_LENGTHS} + * and {@link #CACHE_SIZES}. + */ + static Stream beamLengthCacheMatrix() { + Stream.Builder cases = Stream.builder(); + for (int beam : BEAM_SIZES) { + for (int length : INPUT_LENGTHS) { + for (int cache : CACHE_SIZES) { + cases.add(Arguments.of(beam, length, cache)); + } + } + } + return cases.build(); + } + + /** + * @return The cartesian product of {@link #VALIDATORS}, {@link #BEAM_SIZES}, + * {@link #INPUT_LENGTHS} and {@link #CACHE_SIZES}. + */ + static Stream validatorBeamLengthCacheMatrix() { + Stream.Builder cases = Stream.builder(); + for (NamedValidator nv : VALIDATORS) { + for (int beam : BEAM_SIZES) { + for (int length : INPUT_LENGTHS) { + for (int cache : CACHE_SIZES) { + cases.add(Arguments.of(nv, beam, length, cache)); + } + } + } + } + return cases.build(); + } + + /** + * @return The cartesian product of {@link #INPUT_LENGTHS} and {@link #CACHE_SIZES}. + */ + static Stream lengthCacheMatrix() { + Stream.Builder cases = Stream.builder(); + for (int length : INPUT_LENGTHS) { + for (int cache : CACHE_SIZES) { + cases.add(Arguments.of(length, cache)); + } + } + return cases.build(); + } + + /** + * @return All values of {@link #CACHE_SIZES}. + */ + static IntStream cacheSizes() { + return IntStream.of(CACHE_SIZES); + } + + // --------------------------------------------------------------------------- + // 1. Equivalence matrix: beam sizes x input lengths x cache sizes, + // default (ZERO_LOG) threshold via the two-arg overload, accept-all validator + // --------------------------------------------------------------------------- + + @ParameterizedTest(name = "beam={0}, len={1}, cache={2}") + @MethodSource("beamLengthCacheMatrix") + void equivalenceAcrossBeamSizesLengthsAndCaches(int beam, int length, int cache) { + MaxentModel model = model(); + String[] input = randomInput(length, 1000L + length); + String desc = caseDesc("matrix", beam, length, cache); + SeededContextGenerator cg = new SeededContextGenerator(); + + Sequence[] expected = referenceBestSequences(1, input, null, cg, ACCEPT_ALL, + model, beam, cache); + Sequence[] actual = new BeamSearch(beam, model, cache) + .bestSequences(1, input, null, cg, ACCEPT_ALL); + + assertSequencesEqual(expected, actual, desc); + if (length == 0) { + Assertions.assertEquals(0, cg.callCount(), + desc + ": context generator must not be called for empty input"); + } + } + + // --------------------------------------------------------------------------- + // 2. Equivalence with a tight minSequenceScore that actually filters candidates + // --------------------------------------------------------------------------- + + @ParameterizedTest(name = "beam={0}, len={1}, cache={2}") + @MethodSource("beamLengthCacheMatrix") + void equivalenceWithTightMinSequenceScore(int beam, int length, int cache) { + MaxentModel model = model(); + String[] input = randomInput(length, 2000L + length); + String desc = caseDesc("threshold", beam, length, cache); + + // Derive a threshold that bites: run uncapped, then cut between the best + // and worst candidate scores (or just above the best when only one exists). + Sequence[] uncapped = referenceBestSequences(beam, input, null, + new SeededContextGenerator(), ACCEPT_ALL, model, beam, cache); + final double threshold; + if (uncapped.length == 0) { + threshold = 0; + } else { + double best = uncapped[0].getScore(); + double worst = uncapped[uncapped.length - 1].getScore(); + threshold = (uncapped.length > 1 && worst < best) + ? (best + worst) / 2.0 : best + 0.5; + } + + Sequence[] expected = referenceBestSequences(1, input, null, threshold, + new SeededContextGenerator(), ACCEPT_ALL, model, beam, cache); + Sequence[] actual = new BeamSearch(beam, model, cache) + .bestSequences(1, input, null, threshold, new SeededContextGenerator(), + ACCEPT_ALL); + + assertSequencesEqual(expected, actual, desc + ", threshold=" + threshold); + } + + // --------------------------------------------------------------------------- + // 3. Equivalence under restrictive validators, including the next.isEmpty() + // advance-all-valid fallback and the reject-everything empty-result path + // --------------------------------------------------------------------------- + + @ParameterizedTest(name = "{0}, beam={1}, len={2}, cache={3}") + @MethodSource("validatorBeamLengthCacheMatrix") + void equivalenceWithRestrictiveValidators(NamedValidator nv, int beam, int length, + int cache) { + MaxentModel model = model(); + String[] input = randomInput(length, 3000L + length); + String desc = caseDesc("validator-" + nv.name(), beam, length, cache); + + Sequence[] expected = referenceBestSequences(1, input, null, + new SeededContextGenerator(), nv.validator(), model, beam, cache); + Sequence[] actual = new BeamSearch(beam, model, cache) + .bestSequences(1, input, null, new SeededContextGenerator(), + nv.validator()); + + assertSequencesEqual(expected, actual, desc); + if ("rejectEverything".equals(nv.name()) && length > 0) { + Assertions.assertEquals(0, actual.length, + desc + ": reject-everything must yield an empty (non-null) array"); + Assertions.assertEquals(0, expected.length, + desc + ": reference reject-everything must also be empty"); + } + } + + // --------------------------------------------------------------------------- + // 4. numSequences > 1: winner order and scores match the reference exactly + // --------------------------------------------------------------------------- + + @ParameterizedTest(name = "len={0}, cache={1}") + @MethodSource("lengthCacheMatrix") + void multiWinnerOrderingMatchesReference(int length, int cache) { + MaxentModel model = model(); + int beam = 5; + int numSequences = 3; + Object[] additionalContext = {"ac-ctx"}; + String[] input = randomInput(length, 4000L + length); + String desc = caseDesc("multiWinner[k=3]", beam, length, cache); + + Sequence[] expected = referenceBestSequences(numSequences, input, + additionalContext, new SeededContextGenerator(), ACCEPT_ALL, + model, beam, cache); + Sequence[] actual = new BeamSearch(beam, model, cache) + .bestSequences(numSequences, input, additionalContext, + new SeededContextGenerator(), ACCEPT_ALL); + + assertSequencesEqual(expected, actual, desc); + if (length > 0) { + Assertions.assertEquals(numSequences, actual.length, + desc + ": expected a full k-best list"); + // Winners must come out in non-increasing score order. + for (int s = 1; s < actual.length; s++) { + Assertions.assertTrue(actual[s - 1].getScore() >= actual[s].getScore(), + desc + ": winner order not non-increasing at index " + s); + } + } + } + + // --------------------------------------------------------------------------- + // 5. Concurrency determinism: one shared BeamSearch, 8 worker threads + // --------------------------------------------------------------------------- + + @Test + void concurrentResultsMatchSerialReference() throws Exception { + final int numInputs = 64; + final int rounds = 4; // 64 x 4 = 256 decode evaluations across the pool + final int numThreads = 8; + final int beam = 3; + final int cache = 64; + + MaxentModel model = model(); + BeamSearch shared = new BeamSearch(beam, model, cache); + + String[][] inputs = new String[numInputs][]; + for (int n = 0; n < numInputs; n++) { + inputs[n] = randomInput(n, 5000L + n); // lengths 0..63 + } + SeededContextGenerator cg = new SeededContextGenerator(); + + // Serial reference results, one per input. + Sequence[][] reference = new Sequence[numInputs][]; + for (int n = 0; n < numInputs; n++) { + reference[n] = referenceBestSequences(1, inputs[n], null, cg, ACCEPT_ALL, + model, beam, cache); + } + + // Run 1: 8 workers, each with a disjoint subset of inputs, `rounds` passes each. + Sequence[][][] runResults = new Sequence[rounds][numInputs][]; + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + try { + List> futures = new ArrayList<>(); + for (int w = 0; w < numThreads; w++) { + final int worker = w; + futures.add(pool.submit(() -> { + for (int round = 0; round < rounds; round++) { + for (int n = worker; n < numInputs; n += numThreads) { + runResults[round][n] = shared.bestSequences(1, inputs[n], null, cg, + ACCEPT_ALL); + } + } + })); + } + for (Future f : futures) { + f.get(); + } + } finally { + pool.shutdown(); + } + + for (int round = 0; round < rounds; round++) { + for (int n = 0; n < numInputs; n++) { + assertSequencesEqual(reference[n], runResults[round][n], + "concurrent[round=" + round + ", input=" + n + ", len=" + inputs[n].length + + "]"); + } + } + + // Run 2: the SAME input set twice concurrently; the two passes must agree with + // each other (and with the serial reference). + Sequence[][] passA = new Sequence[numInputs][]; + Sequence[][] passB = new Sequence[numInputs][]; + ExecutorService pool2 = Executors.newFixedThreadPool(numThreads); + try { + Future fa = pool2.submit(() -> { + for (int n = 0; n < numInputs; n++) { + passA[n] = shared.bestSequences(1, inputs[n], null, cg, ACCEPT_ALL); + } + }); + Future fb = pool2.submit(() -> { + for (int n = numInputs - 1; n >= 0; n--) { // reverse order, still racing passA + passB[n] = shared.bestSequences(1, inputs[n], null, cg, ACCEPT_ALL); + } + }); + fa.get(); + fb.get(); + } finally { + pool2.shutdown(); + } + + for (int n = 0; n < numInputs; n++) { + assertSequencesEqual(passA[n], passB[n], + "concurrent-agreement[input=" + n + "]"); + assertSequencesEqual(reference[n], passA[n], + "concurrent-vs-reference[input=" + n + "]"); + } + } + + // --------------------------------------------------------------------------- + // 6. Winner materialization: outcomes, per-position probs and score of the + // winning Sequence are consistent with the model's eval outputs + // --------------------------------------------------------------------------- + + @ParameterizedTest(name = "cache={0}") + @MethodSource("cacheSizes") + void winnerMaterializationMatchesModelOutputs(int cache) { + MaxentModel model = model(); + String[] input = randomInput(7, 6000L); + int beam = 3; + + String desc = "materialization[cache=" + cache + "]"; + BeamSearch bs = new BeamSearch(beam, model, cache); + Sequence winner = bs.bestSequence(input, null, new SeededContextGenerator(), + ACCEPT_ALL); + Assertions.assertNotNull(winner, desc); + Assertions.assertEquals(input.length, winner.getSize(), desc + ": size"); + + // The winner must equal the reference winner. + Sequence refWinner = referenceBestSequences(1, input, null, + new SeededContextGenerator(), ACCEPT_ALL, model, beam, cache)[0]; + Assertions.assertEquals(refWinner.getOutcomes(), winner.getOutcomes(), + desc + ": outcomes vs reference"); + + // Walk the winning path and recompute the expected probs/score independently. + List outcomes = winner.getOutcomes(); + double[] probs = winner.getProbs(); + double expectedScore = 0d; + SeededContextGenerator cg = new SeededContextGenerator(); + for (int i = 0; i < outcomes.size(); i++) { + String[] prefix = outcomes.subList(0, i).toArray(new String[0]); + String[] contexts = cg.getContext(i, input, prefix, new Object[0]); + double[] eval = model.eval(contexts); + int outcomeIndex = model.getIndex(outcomes.get(i)); + Assertions.assertTrue(outcomeIndex >= 0, desc + ": outcome known to model"); + double expectedProb = eval[outcomeIndex]; + + assertBitIdentical(expectedProb, probs[i], desc + ": getProbs()[" + i + "]"); + assertBitIdentical(expectedProb, winner.getProb(i), + desc + ": getProb(" + i + ")"); + expectedScore += StrictMath.log(expectedProb); + } + assertBitIdentical(expectedScore, winner.getScore(), desc + ": score"); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchTest.java index 2c9ee2ea65..948c50cc17 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/BeamSearchTest.java @@ -17,7 +17,9 @@ package opennlp.tools.ml; +import java.util.Arrays; import java.util.HashMap; +import java.util.List; import java.util.Map; import org.junit.jupiter.api.Assertions; @@ -133,11 +135,83 @@ void testBestSequenceZeroLengthInput() { Assertions.assertEquals(sequence.length, seq.getOutcomes().size()); } + /** + * Uniform model: every outcome is equally probable in every context, so all + * candidates at a given depth tie on score exactly. + */ + static class UniformModel implements MaxentModel { + + private final String[] outcomes; + + UniformModel(String[] outcomes) { + this.outcomes = outcomes; + } + + public double[] eval(String[] context) { + double[] probs = new double[outcomes.length]; + Arrays.fill(probs, 1.0d / outcomes.length); + return probs; + } + + public double[] eval(String[] context, double[] probs) { + Arrays.fill(probs, 1.0d / outcomes.length); + return probs; + } + + public double[] eval(String[] context, float[] values) { + return eval(context); + } + + public String getAllOutcomes(double[] outcomes) { + return null; + } + + public String getBestOutcome(double[] outcomes) { + return null; + } + + public int getIndex(String outcome) { + return 0; + } + + public int getNumOutcomes() { + return outcomes.length; + } + + public String getOutcome(int i) { + return outcomes[i]; + } + } + + /** + * Tests that exact score ties resolve in a canonical outcome order rather than + * in whatever order the priority queue's heap layout produces. + */ + @Test + void testBestSequencesBreakScoreTiesDeterministically() { + String[] sequence = {"t1", "t2"}; + BeamSearchContextGenerator cg = new IdentityFeatureGenerator(sequence); + + // The model's outcome iteration order ("b" before "a") must not leak into the tie order. + BeamSearch bs = new BeamSearch(4, new UniformModel(new String[] {"b", "a"})); + + Sequence[] best = bs.bestSequences(4, sequence, null, cg, + (int i, String[] inputSequence, String[] outcomesSequence, + String outcome) -> true); + + Assertions.assertEquals(4, best.length); + Assertions.assertEquals(List.of("a", "a"), best[0].getOutcomes()); + Assertions.assertEquals(List.of("a", "b"), best[1].getOutcomes()); + Assertions.assertEquals(List.of("b", "a"), best[2].getOutcomes()); + Assertions.assertEquals(List.of("b", "b"), best[3].getOutcomes()); + } + /** * Tests finding a sequence of length one. */ @Test void testBestSequenceOneElementInput() { + String[] sequence = {"1"}; BeamSearchContextGenerator cg = new IdentityFeatureGenerator(sequence); diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/SeededMaxentModel.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/SeededMaxentModel.java new file mode 100644 index 0000000000..1d2ddb6b2a --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/ml/SeededMaxentModel.java @@ -0,0 +1,120 @@ +/* + * 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.ml; + +import opennlp.tools.ml.model.MaxentModel; + +/** + * A deterministic pseudo-random {@link MaxentModel} test fixture. The probability of an + * outcome is derived from a hash of the joined context strings, the outcome index and a + * fixed seed with splitmix64-style mixing, so repeated evals of the same context return + * identical values in (0.01, 0.99]. Values are intentionally not normalized. The + * {@code eval(context, probs)} buffer contract is honored: values are written into the + * passed array and that same array is returned. Stateless and therefore thread-safe. + */ +class SeededMaxentModel implements MaxentModel { + + private final String[] outcomes; + private final long seed; + + /** + * Initializes a {@link SeededMaxentModel} instance. + * + * @param outcomes The outcome labels; the outcome index is the array index. + * @param seed The seed all probabilities are derived from. + */ + SeededMaxentModel(String[] outcomes, long seed) { + this.outcomes = outcomes; + this.seed = seed; + } + + /** + * @return A pseudo-random probability in (0.01, 0.99], a pure function + * of {@code context}, {@code outcomeIndex} and the seed. + */ + private double prob(String[] context, int outcomeIndex) { + long h = seed; + for (String c : context) { + h = mix(h, c.hashCode()); + } + h = mix(h, outcomeIndex); + // splitmix64 finalizer for avalanche + h ^= h >>> 30; + h *= 0xBF58476D1CE4E5B9L; + h ^= h >>> 27; + h *= 0x94D049BB133111EBL; + h ^= h >>> 31; + double u = (h >>> 11) * (1.0 / (1L << 53)); // [0, 1) + return 0.01 + 0.98 * u; // (0.01, 0.99] + } + + /** + * @return {@code h} combined with {@code v} FNV-style. + */ + private static long mix(long h, long v) { + return (h ^ (v + 0x9E3779B97F4A7C15L)) * 0x100000001B3L; + } + + @Override + public double[] eval(String[] context) { + return eval(context, new double[outcomes.length]); + } + + @Override + public double[] eval(String[] context, double[] probs) { + for (int i = 0; i < outcomes.length; i++) { + probs[i] = prob(context, i); + } + return probs; // buffer contract: write into the passed array AND return it + } + + @Override + public double[] eval(String[] context, float[] values) { + return eval(context); + } + + @Override + public String getOutcome(int i) { + return outcomes[i]; + } + + @Override + public int getNumOutcomes() { + return outcomes.length; + } + + @Override + public String getAllOutcomes(double[] outcomes) { + return null; + } + + @Override + public String getBestOutcome(double[] outcomes) { + return null; + } + + @Override + public int getIndex(String outcome) { + for (int i = 0; i < outcomes.length; i++) { + if (outcomes[i].equals(outcome)) { + return i; + } + } + return -1; + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/NameFinderAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/NameFinderAnnotatorTest.java new file mode 100644 index 0000000000..14f294c8ce --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/NameFinderAnnotatorTest.java @@ -0,0 +1,298 @@ +/* + * 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.namefind; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +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 that {@link NameFinderAnnotator} maps token-index mentions to character spans on + * the original text, carries the entity type as the annotation value on an untyped span, + * and clears the finder's adaptive data per document. + */ +public class NameFinderAnnotatorTest { + + /** + * Builds a finder over a fixed find function, with adaptive-data clearing counted + * in the given counter when one is supplied. + * + * @param find Maps a sentence's tokens to its mentions. + * @param cleared Counts {@code clearAdaptiveData} calls, or {@code null} to ignore. + * @return The finder fixture. Never {@code null}. + */ + private static TokenNameFinder finder(Function find, + AtomicInteger cleared) { + return new TokenNameFinder() { + + @Override + public Span[] find(String[] tokens) { + return find.apply(tokens); + } + + @Override + public void clearAdaptiveData() { + if (cleared != null) { + cleared.incrementAndGet(); + } + } + }; + } + + /** + * @return A two-sentence document with sentence and token layers over + * {@code "Ana runs. Bob sits."}. Never {@code null}. + */ + private static Document twoSentenceDocument() { + return Document.of("Ana runs. Bob sits.") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 9), "Ana runs."), + new Annotation<>(new Span(10, 19), "Bob sits."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(10, 13), "Bob"), + new Annotation<>(new Span(14, 19), "sits."))); + } + + @Test + void testTokenIndexSpansBecomeCharacterSpans() { + final AtomicInteger cleared = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> { + // "New York" as a two-token location mention + return new Span[] {new Span(1, 3, "location")}; + }, cleared); + + final Document document = Document.of("in New York today") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 17), "in New York today"))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 2), "in"), + new Annotation<>(new Span(3, 6), "New"), + new Annotation<>(new Span(7, 11), "York"), + new Annotation<>(new Span(12, 17), "today"))); + + final Document annotated = new NameFinderAnnotator(finder).annotate(document); + final List> entities = annotated.get(Layers.ENTITIES); + assertEquals(1, entities.size()); + assertEquals(new Span(3, 11), entities.get(0).span()); + // the annotation value is the single source of the entity type; the span is untyped + assertEquals("location", entities.get(0).value()); + assertNull(entities.get(0).span().getType()); + assertEquals("New York", + entities.get(0).span().getCoveredText(annotated.text()).toString()); + assertEquals(1, cleared.get()); + } + + /** + * Verifies that a document carrying sentences but no token layer is rejected with a + * message naming the token layer, so the token check is exercised on its own rather + * than being shadowed by the sentence check. + */ + @Test + void testMissingTokenLayerThrows() { + final TokenNameFinder finder = finder(tokens -> new Span[0], null); + final Document document = Document.of("no tokens") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 9), "no tokens"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new NameFinderAnnotator(finder).annotate(document)); + assertEquals("document lacks the required layer opennlp:tokens", e.getMessage()); + } + + /** + * Verifies that a mention the finder returns without a type is recorded with the + * {@link NameFinderAnnotator#UNTYPED} label, and that the label is the name-sample + * default type, so downstream consumers can rely on the two being interchangeable. + */ + @Test + void testUntypedMentionRecordedAsUntyped() { + final TokenNameFinder finder = finder(tokens -> new Span[] {new Span(0, 1)}, null); + final Document document = Document.of("Ana runs.") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 9), "Ana runs."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."))); + + final Document annotated = new NameFinderAnnotator(finder).annotate(document); + final List> entities = annotated.get(Layers.ENTITIES); + assertEquals(1, entities.size()); + assertEquals(NameSample.DEFAULT_TYPE, NameFinderAnnotator.UNTYPED); + assertEquals(NameFinderAnnotator.UNTYPED, entities.get(0).value()); + assertEquals(new Span(0, 3), entities.get(0).span()); + assertNull(entities.get(0).span().getType()); + } + + /** + * Verifies that a mention whose token indices reach beyond its sentence's tokens is + * rejected loudly instead of silently taking its character span from the following + * sentence's tokens, and that the adaptive data is still cleared on that failure. + */ + @Test + void testMentionOutsideSentenceTokensFailsLoud() { + final AtomicInteger cleared = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> { + // two tokens in the sentence, but the mention claims three + return new Span[] {new Span(0, 3, "person")}; + }, cleared); + final Document document = twoSentenceDocument(); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new NameFinderAnnotator(finder).annotate(document)); + assertEquals("finder returned mention [0..3) person outside the sentence's 2 tokens", + e.getMessage()); + assertEquals(1, cleared.get()); + } + + /** + * Verifies that a zero-length mention is rejected loudly. {@link Span} permits + * {@code start == end}, but such a mention covers no token, so mapping its end + * through {@code end - 1} would read the previous sentence's last token instead of + * failing. The finder here returns the empty mention for the second sentence, the + * case that would otherwise be mapped silently wrong, and the adaptive data is still + * cleared on the failure. + */ + @Test + void testZeroLengthMentionFailsLoud() { + final AtomicInteger cleared = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> + "Bob".equals(tokens[0]) ? new Span[] {new Span(0, 0)} : new Span[0], cleared); + final Document document = twoSentenceDocument(); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new NameFinderAnnotator(finder).annotate(document)); + assertEquals("finder returned mention [0..0) outside the sentence's 2 tokens", + e.getMessage()); + assertEquals(1, cleared.get()); + } + + /** + * Verifies that a token lying outside every sentence is rejected loudly and that the + * adaptive data is still cleared on that failure, so a rejected document cannot leak + * finder state into the next one. + */ + @Test + void testTokenOutsideEverySentenceThrowsAndStillClears() { + final AtomicInteger cleared = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> new Span[0], cleared); + final Document document = Document.of("Ana runs. Bob") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 9), "Ana runs."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "Ana"), + new Annotation<>(new Span(4, 9), "runs."), + new Annotation<>(new Span(10, 13), "Bob"))); + + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new NameFinderAnnotator(finder).annotate(document)); + assertEquals("token at [10..13) lies outside every sentence", e.getMessage()); + assertEquals(1, cleared.get()); + } + + /** + * Verifies that the finder is invoked once per sentence with exactly that sentence's + * tokens, that sentence-local mention indices are mapped through the sentence's first + * token position into document character spans, and that the adaptive data is cleared + * exactly once after the whole document. + */ + @Test + void testFindsPerSentenceAndMapsSentenceLocalIndices() { + final List> calls = new ArrayList<>(); + final AtomicInteger cleared = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> { + calls.add(List.of(tokens)); + // the first token of every sentence is a person mention, in sentence-local indices + return new Span[] {new Span(0, 1, "person")}; + }, cleared); + + final Document annotated = new NameFinderAnnotator(finder).annotate(twoSentenceDocument()); + + assertEquals(List.of( + List.of("Ana", "runs."), + List.of("Bob", "sits.")), calls); + + final List> entities = annotated.get(Layers.ENTITIES); + assertEquals(2, entities.size()); + assertEquals(new Span(0, 3), entities.get(0).span()); + assertEquals("person", entities.get(0).value()); + assertEquals(new Span(10, 13), entities.get(1).span()); + assertEquals("person", entities.get(1).value()); + assertEquals("Bob", + entities.get(1).span().getCoveredText(annotated.text()).toString()); + assertEquals(1, cleared.get()); + } + + /** + * Verifies that the annotator declares both the sentence layer and the token layer as + * required, so a pipeline without a sentence step fails at build time. + */ + @Test + void testRequiresSentencesAndTokens() { + final TokenNameFinder finder = finder(tokens -> new Span[0], null); + assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS), + new NameFinderAnnotator(finder).requires()); + } + + /** + * Verifies that present-but-empty sentence and token layers yield a present-but-empty + * entity layer without invoking the finder, rather than an exception. + */ + @Test + void testEmptyPresentLayersYieldEmptyEntityLayer() { + final AtomicInteger found = new AtomicInteger(); + final TokenNameFinder finder = finder(tokens -> { + found.incrementAndGet(); + return new Span[0]; + }, null); + final Document document = Document.of("") + .with(Layers.SENTENCES, List.of()) + .with(Layers.TOKENS, List.of()); + + final Document annotated = new NameFinderAnnotator(finder).annotate(document); + + assertTrue(annotated.layers().contains(Layers.ENTITIES)); + assertTrue(annotated.get(Layers.ENTITIES).isEmpty()); + assertEquals(0, found.get()); + } + + /** + * Verifies that a document without a sentence layer is rejected with a message naming + * the missing layer. + */ + @Test + void testAbsentSentenceLayerThrowsWithExactMessage() { + final TokenNameFinder finder = finder(tokens -> new Span[0], null); + final Document document = Document.of("Ana") + .with(Layers.TOKENS, List.of(new Annotation<>(new Span(0, 3), "Ana"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new NameFinderAnnotator(finder).annotate(document)); + assertEquals("document lacks the required layer opennlp:sentences", e.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/RegexNameFinderFactoryTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/RegexNameFinderFactoryTest.java index 809429925b..9bc72dbb2b 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/RegexNameFinderFactoryTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/namefind/RegexNameFinderFactoryTest.java @@ -17,12 +17,18 @@ package opennlp.tools.namefind; +import java.time.Duration; import java.util.Arrays; import java.util.List; +import java.util.stream.Stream; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Named; 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.WhitespaceTokenizer; import opennlp.tools.util.Span; @@ -90,6 +96,107 @@ void testLatLong() { Assertions.assertEquals("45", tokens[latLongSpan2.getStart()]); } + /** + * Crafted inputs that used to drive the built-in EMAIL and URL patterns into + * catastrophic backtracking / deep recursion (ReDoS, CWE-1333 / CWE-400 / CWE-674). + * The hardened patterns must finish quickly regardless of input length. + *

    + * Each attack runs as its own test so a regression names the exact input that + * got slow. Each case does one warmup run outside the timed section, so the + * budget measures the match itself rather than JVM cold start on throttled + * machines. + */ + private static Stream reDoSAttacks() { + return Stream.of( + Arguments.of(Named.of("emailLocalBlowup", "a".repeat(100_000) + "@ ")), + Arguments.of(Named.of("emailDomainBlowup", "x@a" + "-a".repeat(60_000) + " ")), + Arguments.of(Named.of("urlPathRecursion", "http://a.com/" + "a".repeat(100_000) + " ")), + Arguments.of(Named.of("urlNestedBlowup", "http://a.com" + "/a".repeat(50_000) + "%z")), + // The reported StackOverflowError repro: a long ?a&a&a&... query string, which + // drove the nested (&(...)+ ... )* group in the old URL pattern into deep recursion. + Arguments.of(Named.of("urlQueryRecursion", "http://a.com/p?a" + "&a".repeat(50_000) + "= "))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("reDoSAttacks") + void testBuiltinPatternsAreNotVulnerableToReDoS(String attack) { + String[] tokens = WhitespaceTokenizer.INSTANCE.tokenize(attack); + // Warmup: JIT-compile the pattern engine before the timed run. + regexNameFinder.find(tokens); + + Assertions.assertTimeoutPreemptively(Duration.ofSeconds(2), + () -> regexNameFinder.find(tokens)); + } + + /** + * Regression tests for the hardened built-in patterns: a trailing slash or a + * sentence-final period must not cause the match to be abandoned entirely. + */ + @Test + void testUrlWithTrailingSlashOrSentencePunctuation() { + RegexNameFinder urlFinder = RegexNameFinderFactory.getDefaultRegexNameFinders( + RegexNameFinderFactory.DEFAULT_REGEX_NAME_FINDER.URL); + final String input = "see http://a.com/ or http://example.com/path/ or www.google.com."; + Span[] spans = urlFinder.find(input); + Assertions.assertEquals(3, spans.length); + Assertions.assertEquals("http://a.com", + input.substring(spans[0].getStart(), spans[0].getEnd())); + Assertions.assertEquals("http://example.com/path", + input.substring(spans[1].getStart(), spans[1].getEnd())); + Assertions.assertEquals("www.google.com", + input.substring(spans[2].getStart(), spans[2].getEnd())); + } + + @Test + void testUrlKeepsPortWhenFollowedBySlash() { + RegexNameFinder urlFinder = RegexNameFinderFactory.getDefaultRegexNameFinders( + RegexNameFinderFactory.DEFAULT_REGEX_NAME_FINDER.URL); + final String input = "http://example.com:8080/"; + Span[] spans = urlFinder.find(input); + Assertions.assertEquals(1, spans.length); + Assertions.assertEquals("http://example.com:8080", + input.substring(spans[0].getStart(), spans[0].getEnd())); + } + + @Test + void testEmailAtEndOfSentence() { + RegexNameFinder emailFinder = RegexNameFinderFactory.getDefaultRegexNameFinders( + RegexNameFinderFactory.DEFAULT_REGEX_NAME_FINDER.EMAIL); + final String input = "mail me at a@b.com. and a@sub.b.co.uk."; + Span[] spans = emailFinder.find(input); + Assertions.assertEquals(2, spans.length); + Assertions.assertEquals("a@b.com", input.substring(spans[0].getStart(), spans[0].getEnd())); + Assertions.assertEquals("a@sub.b.co.uk", + input.substring(spans[1].getStart(), spans[1].getEnd())); + } + + /** + * The hardened URL pattern must still match a URL that is followed by a trailing + * delimiter (a path slash, a sentence-final period, {@code #} or {@code &}), stopping + * the span before the delimiter rather than abandoning the match. + */ + @Test + void testUrlTrailingDelimiters() { + assertFirstUrl("http://a.com", "http://a.com/"); + assertFirstUrl("http://example.com/path", "http://example.com/path/"); + assertFirstUrl("ftp://files.example.com/pub", "ftp://files.example.com/pub/"); + assertFirstUrl("www.google.com", "check www.google.com/ now"); + assertFirstUrl("http://example.com", "I saw http://example.com. Then"); + assertFirstUrl("www.google.com", "www.google.com."); + assertFirstUrl("http://a.com/path", "http://a.com/path."); + assertFirstUrl("http://a.com", "http://a.com/#"); + assertFirstUrl("http://example.com:8080", "http://example.com:8080/"); + assertFirstUrl("http://a.com/p?q=1", "http://a.com/p?q=1&"); + } + + private static void assertFirstUrl(String expected, String input) { + RegexNameFinder urlFinder = RegexNameFinderFactory.getDefaultRegexNameFinders( + RegexNameFinderFactory.DEFAULT_REGEX_NAME_FINDER.URL); + Span[] spans = urlFinder.find(input); + Assertions.assertTrue(spans.length > 0, "no URL match for: " + input); + Assertions.assertEquals(expected, input.substring(spans[0].getStart(), spans[0].getEnd())); + } + @Test void testMgrs() { String[] tokens = WhitespaceTokenizer.INSTANCE.tokenize(text); diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/parser/ParserAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/parser/ParserAnnotatorTest.java new file mode 100644 index 0000000000..dee78a2ef7 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/parser/ParserAnnotatorTest.java @@ -0,0 +1,158 @@ +/* + * 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.parser; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.Layers; +import opennlp.tools.parser.ParserAnnotator.Phrase; +import opennlp.tools.util.Span; + +public class ParserAnnotatorTest { + + /** + * A parser that builds one fixed bracketing over any sentence of six tokens, + * {@code (S (NP (NP 0 1) (PP 2 (NP 3))) (VP 4) 5)}, with explicit heads, so the + * span and head mapping is observable without a model. + */ + private static class FixedParser implements Parser { + + @Override + public Parse[] parse(Parse tokens, int numParses) { + return new Parse[] {parse(tokens)}; + } + + @Override + public Parse parse(Parse tokens) { + final Parse[] toks = tokens.getChildren(); + if (toks.length != 6) { + return tokens; + } + final String[] tags = {"DT", "NN", "IN", "NNP", "VBD", "."}; + final Parse[] pos = new Parse[toks.length]; + for (int i = 0; i < toks.length; i++) { + pos[i] = node(tokens, tags[i], toks[i], toks[i], toks[i]); + tokens.insert(pos[i]); + } + final Parse innerNp = node(tokens, "NP", pos[0], pos[1], toks[1]); + tokens.insert(innerNp); + final Parse maryNp = node(tokens, "NP", pos[3], pos[3], toks[3]); + tokens.insert(maryNp); + final Parse pp = node(tokens, "PP", pos[2], maryNp, toks[2]); + tokens.insert(pp); + final Parse outerNp = node(tokens, "NP", innerNp, pp, toks[1]); + tokens.insert(outerNp); + final Parse vp = node(tokens, "VP", pos[4], pos[4], toks[4]); + tokens.insert(vp); + final Parse s = node(tokens, "S", outerNp, pos[5], toks[4]); + tokens.insert(s); + return tokens; + } + + private static Parse node(Parse root, String type, Parse from, Parse to, Parse head) { + return new Parse(root.getText(), + new Span(from.getSpan().getStart(), to.getSpan().getEnd()), type, 1.0, head); + } + } + + private static List> tokens(String text, String... forms) { + final List> annotations = new ArrayList<>(forms.length); + int cursor = 0; + for (final String form : forms) { + final int start = text.indexOf(form, cursor); + annotations.add(new Annotation<>(new Span(start, start + form.length()), form)); + cursor = start + form.length(); + } + return annotations; + } + + /** One six-token sentence whose text carries a double space the parse text lacks. */ + private static Document sentence() { + final String text = "The dog of Mary ran."; + return Document.of(text) + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 21), "s"))) + .with(Layers.TOKENS, tokens(text, "The", "dog", "of", "Mary", "ran", ".")); + } + + @Test + void testEmitsPhrasesInPreOrderOnTokenSpansWithHeads() { + final Document document = new ParserAnnotator(new FixedParser()).annotate(sentence()); + + final List> phrases = document.get(ParserAnnotator.PHRASES); + Assertions.assertEquals(List.of("S", "NP", "NP", "PP", "NP", "VP"), + phrases.stream().map(a -> a.value().label()).toList()); + Assertions.assertEquals(List.of( + new Span(0, 21), new Span(0, 16), new Span(0, 7), new Span(9, 16), + new Span(12, 16), new Span(17, 20)), + phrases.stream().map(Annotation::span).toList()); + Assertions.assertEquals("The dog of Mary", document.text().subSequence(0, 16).toString()); + final Span dog = new Span(4, 7); + final Span ran = new Span(17, 20); + Assertions.assertEquals(List.of(ran, dog, dog, new Span(9, 11), new Span(12, 16), ran), + phrases.stream().map(a -> a.value().head()).toList()); + } + + @Test + void testEmptyLayersYieldEmptyPhraseLayer() { + final Document document = new ParserAnnotator(new FixedParser()).annotate( + Document.of("").with(Layers.SENTENCES, List.of()).with(Layers.TOKENS, List.of())); + Assertions.assertTrue(document.layers().contains(ParserAnnotator.PHRASES)); + Assertions.assertTrue(document.get(ParserAnnotator.PHRASES).isEmpty()); + } + + @Test + void testLayerContract() { + final ParserAnnotator annotator = new ParserAnnotator(new FixedParser()); + Assertions.assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS), annotator.requires()); + Assertions.assertEquals(Set.of(ParserAnnotator.PHRASES), annotator.provides()); + Assertions.assertEquals("opennlp:phrases", ParserAnnotator.PHRASES.id()); + Assertions.assertEquals("ParserAnnotator", annotator.toString()); + } + + @Test + void testRejectsNullParserMissingLayersAndNullParse() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> new ParserAnnotator(null)); + final ParserAnnotator annotator = new ParserAnnotator(new FixedParser()); + Assertions.assertThrows(IllegalArgumentException.class, () -> annotator.annotate(null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(Document.of("x"))); + final ParserAnnotator silent = new ParserAnnotator(new FixedParser() { + @Override + public Parse parse(Parse tokens) { + return null; + } + }); + Assertions.assertThrows(IllegalArgumentException.class, + () -> silent.annotate(sentence())); + } + + @Test + void testPhraseRejectsBlankLabelAndNullHead() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> new Phrase(" ", new Span(0, 1))); + Assertions.assertThrows(IllegalArgumentException.class, () -> new Phrase("NP", null)); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/postag/POSTaggerAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/postag/POSTaggerAnnotatorTest.java new file mode 100644 index 0000000000..72fea7f524 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/postag/POSTaggerAnnotatorTest.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.tools.postag; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Sequence; +import opennlp.tools.util.Span; + +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 that {@link POSTaggerAnnotator} tags one sentence per {@link POSTagger#tag(String[])} + * call, keeps the tag layer aligned with the token layer, distinguishes a present-but-empty + * required layer from an absent one, and rejects tokens outside every sentence. + */ +public class POSTaggerAnnotatorTest { + + /** + * A tagger that records the exact token sequence of every call and answers with one + * {@code "X"} tag per token, so the per-call slicing is observable. Tests override + * {@link #tag(String[])} where a deviant answer is the fixture. + */ + private static class RecordingTagger implements POSTagger { + + private final List> calls = new ArrayList<>(); + + @Override + public String[] tag(String[] sentence) { + calls.add(List.of(sentence)); + final String[] tags = new String[sentence.length]; + Arrays.fill(tags, "X"); + return tags; + } + + @Override + public String[] tag(String[] sentence, Object[] additionalContext) { + return tag(sentence); + } + + @Override + public Sequence[] topKSequences(String[] sentence) { + throw new UnsupportedOperationException("the adapter only calls tag"); + } + + @Override + public Sequence[] topKSequences(String[] sentence, Object[] additionalContext) { + throw new UnsupportedOperationException("the adapter only calls tag"); + } + } + + /** + * @return A two-sentence document with sentence and token layers over + * {@code "The dog barks. It naps."}. Never {@code null}. + */ + private static Document twoSentenceDocument() { + return Document.of("The dog barks. It naps.") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 14), "The dog barks."), + new Annotation<>(new Span(15, 23), "It naps."))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "The"), + new Annotation<>(new Span(4, 7), "dog"), + new Annotation<>(new Span(8, 14), "barks."), + new Annotation<>(new Span(15, 17), "It"), + new Annotation<>(new Span(18, 23), "naps."))); + } + + /** + * Verifies that the tagger is invoked once per sentence with exactly that sentence's + * tokens, and that the resulting tag layer stays aligned with the token layer by + * position, each tag on its token's span. + */ + @Test + void testTagsEachSentenceSeparately() { + final RecordingTagger tagger = new RecordingTagger(); + final Document annotated = new POSTaggerAnnotator(tagger).annotate(twoSentenceDocument()); + + assertEquals(List.of( + List.of("The", "dog", "barks."), + List.of("It", "naps.")), tagger.calls); + + final List> tokens = annotated.get(Layers.TOKENS); + final List> tags = annotated.get(Layers.POS_TAGS); + assertEquals(tokens.size(), tags.size()); + for (int i = 0; i < tags.size(); i++) { + assertEquals(tokens.get(i).span(), tags.get(i).span()); + assertEquals("X", tags.get(i).value()); + } + } + + /** + * Verifies that the annotator declares both the sentence layer and the token layer as + * required, so a pipeline without a sentence step fails at build time. + */ + @Test + void testRequiresSentencesAndTokens() { + assertEquals(Set.of(Layers.SENTENCES, Layers.TOKENS), + new POSTaggerAnnotator(new RecordingTagger()).requires()); + } + + /** + * Verifies that present-but-empty sentence and token layers yield a present-but-empty + * tag layer without invoking the tagger, rather than an exception. + */ + @Test + void testEmptyPresentLayersYieldEmptyTagLayer() { + final RecordingTagger tagger = new RecordingTagger(); + final Document document = Document.of("") + .with(Layers.SENTENCES, List.of()) + .with(Layers.TOKENS, List.of()); + + final Document annotated = new POSTaggerAnnotator(tagger).annotate(document); + + assertTrue(annotated.layers().contains(Layers.POS_TAGS)); + assertTrue(annotated.get(Layers.POS_TAGS).isEmpty()); + assertTrue(tagger.calls.isEmpty()); + } + + /** + * Verifies that a document without a sentence layer is rejected with a message naming + * the missing layer. + */ + @Test + void testAbsentSentenceLayerThrowsWithExactMessage() { + final Document document = Document.of("The dog") + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "The"), + new Annotation<>(new Span(4, 7), "dog"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new POSTaggerAnnotator(new RecordingTagger()).annotate(document)); + assertEquals("document lacks the required layer opennlp:sentences", e.getMessage()); + } + + /** + * Verifies that a document without a token layer is rejected with a message naming the + * missing layer. + */ + @Test + void testAbsentTokenLayerThrowsWithExactMessage() { + final Document document = Document.of("The dog") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 7), "The dog"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new POSTaggerAnnotator(new RecordingTagger()).annotate(document)); + assertEquals("document lacks the required layer opennlp:tokens", e.getMessage()); + } + + /** + * Verifies that a token whose span no sentence encloses is rejected with a message + * naming the token's span. + */ + @Test + void testTokenOutsideEverySentenceThrowsWithExactMessage() { + final Document document = Document.of("The dog") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 3), "The"))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "The"), + new Annotation<>(new Span(4, 7), "dog"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new POSTaggerAnnotator(new RecordingTagger()).annotate(document)); + assertEquals("token at [4..7) lies outside every sentence", e.getMessage()); + } + + /** + * Verifies that a sentence containing no tokens contributes nothing: the tagger is + * never called with an empty sequence and the tag layer still matches the token layer. + */ + @Test + void testSentenceWithoutTokensContributesNothing() { + final RecordingTagger tagger = new RecordingTagger(); + final Document document = Document.of("The ???") + .with(Layers.SENTENCES, List.of( + new Annotation<>(new Span(0, 3), "The"), + new Annotation<>(new Span(4, 7), "???"))) + .with(Layers.TOKENS, List.of(new Annotation<>(new Span(0, 3), "The"))); + + final Document annotated = new POSTaggerAnnotator(tagger).annotate(document); + + assertEquals(List.of(List.of("The")), tagger.calls); + assertEquals(1, annotated.get(Layers.POS_TAGS).size()); + } + + /** + * Verifies that a tagger returning a wrong number of tags for a sentence is rejected + * loudly instead of silently misaligning the tag layer with the token layer. + */ + @Test + void testWrongTagCountFailsLoud() { + // one tag regardless of sentence length, so a two-token sentence trips the check + final POSTagger shortTagger = new RecordingTagger() { + + @Override + public String[] tag(String[] sentence) { + return new String[] {"X"}; + } + }; + final Document document = Document.of("The dog") + .with(Layers.SENTENCES, List.of(new Annotation<>(new Span(0, 7), "The dog"))) + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 3), "The"), + new Annotation<>(new Span(4, 7), "dog"))); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> new POSTaggerAnnotator(shortTagger).annotate(document)); + assertEquals("tagger returned 1 tags for 2 tokens", e.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/LegacyAbbreviationSentenceDetectorME.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/LegacyAbbreviationSentenceDetectorME.java new file mode 100644 index 0000000000..a04a208f78 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/LegacyAbbreviationSentenceDetectorME.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.tools.sentdetect; + +import opennlp.tools.dictionary.Dictionary; +import opennlp.tools.util.StringList; +import opennlp.tools.util.StringUtil; + +/** + * A {@link SentenceDetectorME} whose abbreviation veto is the full text scan that + * {@link SentenceDetectorME#isAcceptableBreak(CharSequence, int, int)} used before it was given a + * bounded window. It is the oracle of {@link SentenceDetectorMEAbbreviationEquivalenceTest} and + * the baseline of {@code SentenceDetectorMEAbbreviationBenchmark}, which is why it lives in its + * own class rather than in either of them. + */ +class LegacyAbbreviationSentenceDetectorME extends SentenceDetectorME { + + private final Dictionary abbDict; + private int vetoes; + + /** + * @param model The {@link SentenceModel} to be used. + * @param abbDict The {@link Dictionary} to veto with, may be {@code null}. + */ + LegacyAbbreviationSentenceDetectorME(SentenceModel model, Dictionary abbDict) { + super(model, abbDict); + this.abbDict = abbDict; + } + + @Override + protected boolean isAcceptableBreak(CharSequence s, int fromIndex, int candidateIndex) { + final boolean acceptable = decide(abbDict, s, fromIndex, candidateIndex); + if (!acceptable) { + vetoes++; + } + return acceptable; + } + + /** + * @return How often this instance vetoed a break, so a comparison run can show that it + * exercised the veto at all. + */ + int vetoes() { + return vetoes; + } + + /** + * The pre-rewrite implementation, kept verbatim. + * + * @param abbDict The abbreviation {@link Dictionary}, may be {@code null}. + * @param s The {@link CharSequence} in which the break occurred. + * @param fromIndex The start of the segment currently being evaluated. + * @param candidateIndex The index of the candidate sentence ending. + * @return {@code true} if the break is acceptable, {@code false} otherwise. + */ + static boolean decide(Dictionary abbDict, CharSequence s, int fromIndex, int candidateIndex) { + if (abbDict == null) + return true; + + final String text = s.toString(); + final boolean caseSensitive = abbDict.isCaseSensitive(); + final String searchText = caseSensitive ? text : StringUtil.toLowerCase(text); + for (StringList abb : abbDict) { + final String abbToken = caseSensitive ? abb.getToken(0) + : StringUtil.toLowerCase(abb.getToken(0)); + final int tokenLength = abbToken.length(); + int tokenStartPos = searchText.indexOf(abbToken, fromIndex); + while (tokenStartPos != -1) { + if (tokenStartPos > candidateIndex) { + break; // past candidate position, no point searching further + } + if (tokenStartPos == fromIndex + && searchText.substring(tokenStartPos, candidateIndex + 1).equals(abbToken)) { + return false; // full abbreviation match at segment start -> no acceptable break + } + final char prevChar = + s.charAt(tokenStartPos == fromIndex ? tokenStartPos : tokenStartPos - 1); + if (tokenStartPos + tokenLength >= candidateIndex + && (Character.isWhitespace(prevChar) || isApostrophe(prevChar) || prevChar == '(')) { + return false; // in case of a valid abbreviation: the (sentence) break is not accepted + } + // Try next occurrence of this abbreviation in the text + tokenStartPos = searchText.indexOf(abbToken, tokenStartPos + 1); + } + } + return true; // no abbreviation(s) at given positions: valid sentence boundary + } + + /** + * @param c The character to check. + * @return {@code true} if the character represents an apostrophe, {@code false} otherwise. + */ + private static boolean isApostrophe(char c) { + return c == '\'' || c == '`' || c == '´'; + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationEquivalenceTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationEquivalenceTest.java new file mode 100644 index 0000000000..ffd1251c5f --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/sentdetect/SentenceDetectorMEAbbreviationEquivalenceTest.java @@ -0,0 +1,481 @@ +/* + * 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.sentdetect; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Random; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import opennlp.tools.dictionary.Dictionary; +import opennlp.tools.util.Span; +import opennlp.tools.util.StringList; + +/** + * Proves that the bounded-window abbreviation veto of + * {@link SentenceDetectorME#isAcceptableBreak(CharSequence, int, int)} decides exactly what the + * previous full-text scan decided. + *

    + * {@link LegacyAbbreviationSentenceDetectorME} holds that previous implementation verbatim and + * is the oracle here. Agreement is asserted twice over: for every single {@code (fromIndex, + * candidateIndex)} pair of a set of hand-written and seeded random texts, and for the sentence + * spans a full {@link SentenceDetectorME#sentPosDetect(CharSequence)} run produces on a corpus. + */ +public class SentenceDetectorMEAbbreviationEquivalenceTest extends AbstractSentenceDetectorTest { + + private static SentenceModel sentdetectModel; + + /** + * A model without an abbreviation dictionary, the shape every sentence model published by the + * project has. Detectors built from it never enter the veto at all. + */ + private static SentenceModel modelWithoutAbbreviations; + private static String corpus; + + @BeforeAll + public static void prepareResources() throws IOException { + final SentenceDetectorFactory factory = new SentenceDetectorFactory( + "eng", true, loadAbbDictionary(Locale.ENGLISH), null); + sentdetectModel = train(factory, Locale.ENGLISH); + Assertions.assertNotNull(sentdetectModel); + modelWithoutAbbreviations = train( + new SentenceDetectorFactory("eng", true, null, null), Locale.ENGLISH); + Assertions.assertNull(modelWithoutAbbreviations.getAbbreviations()); + corpus = readCorpus(); + Assertions.assertTrue(corpus.length() > 10_000, "corpus too small to be meaningful"); + } + + /** + * The pre-rewrite implementation, the oracle every comparison below is made against. + * + * @see LegacyAbbreviationSentenceDetectorME#decide + */ + private static boolean legacyIsAcceptableBreak(Dictionary abbDict, CharSequence s, + int fromIndex, int candidateIndex) { + return LegacyAbbreviationSentenceDetectorME.decide(abbDict, s, fromIndex, candidateIndex); + } + + /* + * ------------------------------------------------------------------------------------------ + * Differential harness. + * ------------------------------------------------------------------------------------------ + */ + + private static Dictionary dictionaryOf(boolean caseSensitive, String... entries) { + final Dictionary dict = new Dictionary(caseSensitive); + for (String entry : entries) { + dict.put(new StringList(entry)); + } + return dict; + } + + /** + * Compares old and new for every {@code 0 <= fromIndex <= candidateIndex < text.length()}, + * that is over the whole contract domain of the method. + * + * @return The number of decisions compared. + */ + private static long assertAgreesEverywhere(Dictionary dict, String text) { + final SentenceDetectorME detector = new SentenceDetectorME(sentdetectModel, dict); + long compared = 0; + for (int fromIndex = 0; fromIndex < text.length(); fromIndex++) { + for (int candidateIndex = fromIndex; candidateIndex < text.length(); candidateIndex++) { + final boolean expected = legacyIsAcceptableBreak(dict, text, fromIndex, candidateIndex); + final boolean actual = detector.isAcceptableBreak(text, fromIndex, candidateIndex); + if (expected != actual) { + Assertions.fail(String.format(Locale.ROOT, + "disagreement at fromIndex=%d candidateIndex=%d: legacy=%b, new=%b, text=<%s>", + fromIndex, candidateIndex, expected, actual, text)); + } + compared++; + } + } + return compared; + } + + /** + * As {@link #assertAgreesEverywhere}, and additionally requires that the oracle vetoes at + * least once, so a case cannot pass by never reaching the interesting branch. + */ + private static long assertAgreesEverywhereAndVetoes(Dictionary dict, String text) { + final long compared = assertAgreesEverywhere(dict, text); + int vetoes = 0; + for (int fromIndex = 0; fromIndex < text.length(); fromIndex++) { + for (int candidateIndex = fromIndex; candidateIndex < text.length(); candidateIndex++) { + if (!legacyIsAcceptableBreak(dict, text, fromIndex, candidateIndex)) { + vetoes++; + } + } + } + Assertions.assertTrue(vetoes > 0, + "the oracle never vetoes here, so the case proves nothing: " + text); + return compared; + } + + /* + * ------------------------------------------------------------------------------------------ + * Edge cases. + * ------------------------------------------------------------------------------------------ + */ + + @Test + void testAbbreviationAtDocumentStart() { + final Dictionary dict = dictionaryOf(true, "Mr.", "Dr."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Mr. Smith left.") > 0); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Dr.") > 0); + } + + @Test + void testAbbreviationAtDocumentEnd() { + final Dictionary dict = dictionaryOf(true, "etc.", "Mr."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Bring apples, pears, etc.") > 0); + // The entry reaches exactly the last index, and one that reaches past the end. + Assertions.assertTrue(assertAgreesEverywhere(dict, "Bring apples, pears, etc") > 0); + } + + @Test + void testOverlappingCandidates() { + // "U.S." and "S." overlap, and "..." makes several candidate ends adjacent. + final Dictionary dict = dictionaryOf(true, "U.S.", "S.", "e.g.", "g."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "The U.S. and e.g. Spain...") > 0); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "U.S.S.R. is gone.") > 0); + } + + @Test + void testEmptyDictionary() { + final Dictionary dict = dictionaryOf(true); + Assertions.assertTrue(assertAgreesEverywhere(dict, "A sentence. Another one.") > 0); + final Dictionary insensitive = dictionaryOf(false); + Assertions.assertTrue(assertAgreesEverywhere(insensitive, "A sentence. Another one.") > 0); + } + + @Test + void testNullDictionaryAlwaysAccepts() { + final SentenceDetectorME detector = new SentenceDetectorME(sentdetectModel, (Dictionary) null); + final String text = "Mr. Smith left. Dr. Jones stayed."; + for (int fromIndex = 0; fromIndex < text.length(); fromIndex++) { + for (int candidateIndex = fromIndex; candidateIndex < text.length(); candidateIndex++) { + Assertions.assertTrue(detector.isAcceptableBreak(text, fromIndex, candidateIndex)); + } + } + } + + @Test + void testDictionaryEntryNotPresentInText() { + final Dictionary dict = dictionaryOf(true, "Zzz.", "Qqq.", "Mr."); + Assertions.assertTrue(assertAgreesEverywhere(dict, "Nothing here matches at all.") > 0); + } + + @Test + void testEntryLongerThanText() { + final Dictionary dict = dictionaryOf(true, "averyveryverylongabbreviation."); + Assertions.assertTrue(assertAgreesEverywhere(dict, "short.") > 0); + } + + @Test + void testMultiCharacterAndUnicodeAbbreviations() { + final Dictionary dict = dictionaryOf(true, "z.B.", "Abb.", "Straße.", "№.", "ç."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, + "Siehe z.B. Abb. 3 in der Straße. Und №. 7 sowie ç. hier.") > 0); + } + + @Test + void testSupplementaryCodePointsInTextAndDictionary() { + // U+1D400 MATHEMATICAL BOLD CAPITAL A and U+10400 DESERET CAPITAL LONG I, the latter has a + // lower case mapping, so it exercises the case-folded window on a surrogate pair. + final String bold = new String(Character.toChars(0x1D400)); + final String deseret = new String(Character.toChars(0x10400)); + final String deseretLower = new String(Character.toChars(0x10428)); + final Dictionary sensitive = dictionaryOf(true, bold + ".", deseret + ".", "Mr."); + final String text = "A " + bold + ". B " + deseret + ". C " + deseretLower + ". Mr. D."; + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(sensitive, text) > 0); + final Dictionary insensitive = dictionaryOf(false, bold + ".", deseret + ".", "Mr."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(insensitive, text) > 0); + // A lone high surrogate, so the window can start or end on an unpaired half. + final Dictionary loneHalf = dictionaryOf(false, "\uD801.", "Mr."); + Assertions.assertTrue(assertAgreesEverywhere(loneHalf, "x \uD801. y " + deseret + ". Mr. z") > 0); + } + + @Test + void testRepeatedAbbreviations() { + final Dictionary dict = dictionaryOf(true, "Mr."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, + "Mr. A, Mr. B, Mr. C, Mr. D, and Mr. E.") > 0); + } + + @Test + void testCaseInsensitiveDictionary() { + // "i̇." is the decomposed lower case of "İ" (LATIN CAPITAL LETTER I WITH DOT + // ABOVE), whose single code point lower case mapping is plain "i", so the two do not match. + // The point is only that both implementations say so. + final Dictionary dict = dictionaryOf(false, "Mr.", "TEL.", "i̇."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, + "mr. Smith, MR. Jones and Tel. 555 plus İ. and I. here.") > 0); + } + + @Test + void testEmptyStringEntry() { + // An empty entry matches at every position; the oracle can still veto through the + // preceding-character branch, so the window has to probe length zero as well. + final Dictionary onlyEmpty = dictionaryOf(true, ""); + Assertions.assertFalse(legacyIsAcceptableBreak(onlyEmpty, "a b", 0, 2), + "the empty entry has to be able to veto on its own, otherwise this proves nothing"); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(onlyEmpty, "a b c d.") > 0); + final Dictionary dict = dictionaryOf(true, "", "Mr."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Mr. A said no.") > 0); + } + + @Test + void testMultiTokenEntryUsesFirstTokenOnly() { + final Dictionary dict = new Dictionary(true); + dict.put(new StringList("Mr.", "Smith")); + dict.put(new StringList("St.", "Louis")); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Mr. Smith of St. Louis left.") > 0); + } + + @Test + void testApostropheAndOpeningBracketPrefixes() { + final Dictionary dict = dictionaryOf(true, "Mr.", "cf."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, + "He said 'Mr. X' and (cf. Y) and `Mr. Z` and ´cf. W´ and xMr. V.") > 0); + } + + @Test + void testSegmentStartInsideAnAbbreviation() { + // fromIndex walking through the middle of an abbreviation is what the sentence loop does + // after accepting a break, and it is where the "match at segment start" branch fires. + final Dictionary dict = dictionaryOf(true, "Mr.", "r.", "."); + Assertions.assertTrue(assertAgreesEverywhereAndVetoes(dict, "Mr. Mr. Mr.") > 0); + } + + /* + * ------------------------------------------------------------------------------------------ + * Seeded random sweep. + * ------------------------------------------------------------------------------------------ + */ + + @Test + void testSeededRandomTexts() { + final String[] pieces = { + "Mr.", "Dr.", "e.g.", "i.e.", "U.S.A.", "etc.", "St.", "Nr.", "z.B.", "vs.", + " ", " ", "\n", "\t", " ", "(", ")", "'", "`", "´", "\"", + "word", "Word", "WORD", "x", ".", "?", "!", "3.5", "ç", "ä", "Straße", + new String(Character.toChars(0x10400)), new String(Character.toChars(0x1D400)), + }; + final Dictionary[] dicts = { + dictionaryOf(true, "Mr.", "Dr.", "e.g.", "etc.", "St.", "U.S.A.", ".", "Nr."), + dictionaryOf(false, "mr.", "DR.", "E.g.", "etc.", "st.", "u.s.a.", "z.b.", "vs.", + "straße", "ç.", new String(Character.toChars(0x10400)) + "."), + dictionaryOf(false, "notpresent.", "alsonot."), + }; + + final Random random = new Random(20260803L); + long compared = 0; + int vetoing = 0; + for (int run = 0; run < 250; run++) { + final StringBuilder text = new StringBuilder(); + while (text.length() < 70) { + text.append(pieces[random.nextInt(pieces.length)]); + } + final String sample = text.toString(); + for (Dictionary dict : dicts) { + compared += assertAgreesEverywhere(dict, sample); + if (!legacyIsAcceptableBreak(dict, sample, 0, sample.length() - 1)) { + vetoing++; + } + } + } + Assertions.assertTrue(compared > 1_500_000L, + "the sweep degenerated, only " + compared + " decisions compared"); + Assertions.assertTrue(vetoing > 0, "no sampled text ever vetoed"); + } + + /* + * ------------------------------------------------------------------------------------------ + * Corpus level: identical spans out of sentPosDetect. + * ------------------------------------------------------------------------------------------ + */ + + /** + * Differential check of {@link SentenceDetectorME#sentPosDetect(CharSequence)} against the + * legacy veto. The English sentence model is paired with each shipped abbreviation dictionary + * (including the German one); this is not German-language coverage. + */ + @Test + void testSentPosDetectAgreesOnCorpus() throws IOException { + for (Locale locale : new Locale[] {Locale.ENGLISH, Locale.GERMAN}) { + assertCorpusRunsAgree(loadAbbDictionary(locale), + "the shipped abbreviation dictionary of " + locale + " on the English model"); + } + } + + @Test + void testSentPosDetectAgreesOnCaseSensitiveDictionary() { + assertCorpusRunsAgree(dictionaryOf(true, "Mr.", "Mrs.", "Ms.", "Dr.", "St.", "etc.", + "e.g.", "i.e.", "vs.", "No.", "Jr.", "Sr.", "Prof.", "Inc.", "Ltd."), + "a case-sensitive dictionary"); + } + + /** + * The path every user of a published sentence model is on: no abbreviation dictionary at all, + * so the veto is never entered. This is the case an optimisation of the veto is most likely to + * regress by accident, and the least likely to be noticed. + */ + @Test + void testNoDictionaryConfiguredIsUnchangedOnCorpus() { + // A model published without an abbreviation dictionary, used through the one argument + // constructor, and a model that has one but is asked to ignore it. + assertNoDictionaryRunAgrees(modelWithoutAbbreviations, + new SentenceDetectorME(modelWithoutAbbreviations)); + assertNoDictionaryRunAgrees(sentdetectModel, + new SentenceDetectorME(sentdetectModel, (Dictionary) null)); + } + + private static void assertNoDictionaryRunAgrees(SentenceModel model, SentenceDetectorME fixed) { + final LegacyAbbreviationSentenceDetectorME legacy = + new LegacyAbbreviationSentenceDetectorME(model, null); + + final Span[] expected = legacy.sentPosDetect(corpus); + final double[] expectedProbs = legacy.probs(); + final Span[] actual = fixed.sentPosDetect(corpus); + final double[] actualProbs = fixed.probs(); + + Assertions.assertTrue(expected.length > 100, + "corpus produced too few sentences to be meaningful: " + expected.length); + Assertions.assertEquals(0, legacy.vetoes(), + "the veto must not be reachable without a dictionary"); + Assertions.assertArrayEquals(expected, actual); + Assertions.assertArrayEquals(expectedProbs, actualProbs, 0.0d); + } + + private static void assertCorpusRunsAgree(Dictionary dict, String description) { + final SentenceDetectorME fixed = new SentenceDetectorME(sentdetectModel, dict); + final LegacyAbbreviationSentenceDetectorME legacy = + new LegacyAbbreviationSentenceDetectorME(sentdetectModel, dict); + + final Span[] expected = legacy.sentPosDetect(corpus); + final double[] expectedProbs = legacy.probs(); + final Span[] actual = fixed.sentPosDetect(corpus); + final double[] actualProbs = fixed.probs(); + + Assertions.assertTrue(expected.length > 100, + "corpus produced too few sentences to be meaningful: " + expected.length); + Assertions.assertTrue(legacy.vetoes() > 0, + "the veto never fired for " + description + ", so this run proves nothing"); + Assertions.assertArrayEquals(expected, actual, "span mismatch for " + description); + Assertions.assertArrayEquals(expectedProbs, actualProbs, 0.0d, + "probability mismatch for " + description); + } + + /* + * ------------------------------------------------------------------------------------------ + * Documented differences, outside the contract of the method. + * ------------------------------------------------------------------------------------------ + */ + + @Test + void testCandidateIndexOutsideTextNoLongerThrows() { + final Dictionary dict = dictionaryOf(false, "Mr.", "etc."); + final SentenceDetectorME detector = new SentenceDetectorME(sentdetectModel, dict); + + // Match at fromIndex with candidateIndex == length: legacy clause A substring-throws; + // the bounded window accepts instead. + final String matchAtFromIndex = "Mr. x"; + Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> legacyIsAcceptableBreak(dict, matchAtFromIndex, 0, matchAtFromIndex.length())); + Assertions.assertTrue(detector.isAcceptableBreak(matchAtFromIndex, 0, matchAtFromIndex.length())); + + // Match elsewhere with candidateIndex == length: legacy short-circuits clause A and + // clause B rejects the break. The bounded window must agree (not silently accept). + final String matchElsewhere = "Bring apples, pears, etc."; + Assertions.assertFalse(legacyIsAcceptableBreak(dict, matchElsewhere, 0, matchElsewhere.length())); + Assertions.assertFalse(detector.isAcceptableBreak(matchElsewhere, 0, matchElsewhere.length())); + + // Past the end of the text cannot carry an abbreviation; both accept, and the new path + // must not throw. + Assertions.assertTrue(detector.isAcceptableBreak(matchAtFromIndex, 0, + matchAtFromIndex.length() + 100)); + } + + @Test + void testNegativeFromIndexIsStillUndefined() { + final Dictionary dict = dictionaryOf(true, "Mr."); + final SentenceDetectorME detector = new SentenceDetectorME(sentdetectModel, dict); + // A negative segment start makes the character before the first match unreadable. Both + // implementations fail on it the same way when the match sits at the start of the text. + final String atStart = "Mr. x"; + Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> legacyIsAcceptableBreak(dict, atStart, -1, 3)); + Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> detector.isAcceptableBreak(atStart, -1, 3)); + // They part company when the match is too far left to be relevant: the previous + // implementation still read index -1 for it, the bounded window never looks there. + final String far = "Mr. and something else."; + Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> legacyIsAcceptableBreak(dict, far, -1, far.length() - 1)); + Assertions.assertTrue(detector.isAcceptableBreak(far, -1, far.length() - 1)); + } + + @Test + void testCandidateBeforeSegmentStartAcceptsAsBefore() { + final Dictionary dict = dictionaryOf(true, "Mr."); + final SentenceDetectorME detector = new SentenceDetectorME(sentdetectModel, dict); + final String text = "Mr. x"; + Assertions.assertTrue(legacyIsAcceptableBreak(dict, text, 4, 2)); + Assertions.assertTrue(detector.isAcceptableBreak(text, 4, 2)); + } + + /* + * ------------------------------------------------------------------------------------------ + * Helpers. + * ------------------------------------------------------------------------------------------ + */ + + private static String readCorpus() throws IOException { + final List lines = new ArrayList<>(); + for (String resource : new String[] {"/opennlp/tools/sentdetect/Sentences.txt", + "/opennlp/tools/sentdetect/Sentences_DE.txt"}) { + try (InputStream in = SentenceDetectorMEAbbreviationEquivalenceTest.class + .getResourceAsStream(resource)) { + Assertions.assertNotNull(in, resource + " is not on the test classpath"); + final String all = new String(in.readAllBytes(), StandardCharsets.UTF_8); + for (String line : all.split("\n")) { + if (!line.isBlank()) { + lines.add(line.strip()); + } + } + } + } + // The shipped training corpora happen to contain almost no abbreviations, which would make + // the comparison runs above vacuous. This tail supplies them, in the shape they occur in. + lines.add("Mr. Smith called tel. 555 1234 this morning. Mrs. Smith did not."); + lines.add("Ms. Adams, Mr. Brown and Mrs. Clark met in the hall."); + lines.add("Ask Mr. Brown or, failing that, Ms. Adams (cf. tel. 555 1234)."); + lines.add("Er wohnt in der S. Bahnstrasse, vgl. S. 12, ca. 30 Minuten entfernt."); + lines.add("Das gilt z.B. fuer Bek. 4 ff. und lt. V. 7 ugs. auch sonst."); + lines.add("Siehe ca. 20 Stueck, z. B. Bek. 9, S. 3 f. und ff."); + return String.join(" ", lines); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/StemmerAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/StemmerAnnotatorTest.java new file mode 100644 index 0000000000..668acd901c --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/StemmerAnnotatorTest.java @@ -0,0 +1,97 @@ +/* + * 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.stemmer; + +import java.util.List; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnalyzer; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class StemmerAnnotatorTest { + + @Test + void testStemsAlignWithTokens() { + final Document document = Document.of("running dogs") + .with(Layers.TOKENS, List.of( + new Annotation<>(new Span(0, 7), "running"), + new Annotation<>(new Span(8, 12), "dogs"))); + + final Document stemmed = new StemmerAnnotator( + new PorterStemmer()).annotate(document); + + final List> stems = stemmed.get(StemmerAnnotator.STEMS); + assertEquals(2, stems.size()); + assertEquals("run", stems.get(0).value()); + assertEquals(new Span(0, 7), stems.get(0).span()); + assertEquals("dog", stems.get(1).value()); + } + + @Test + void testInvalidArguments() { + assertThrows(IllegalArgumentException.class, + () -> new StemmerAnnotator(null)); + final StemmerAnnotator annotator = new StemmerAnnotator(new PorterStemmer()); + assertThrows(IllegalArgumentException.class, () -> annotator.annotate(null)); + } + + /** + * Verifies that a misordered pipeline names the adapter by its simple class name, not + * by its default identity string. + */ + @Test + void testPipelineValidationNamesTheAdapter() { + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> DocumentAnalyzer.builder() + .add(new StemmerAnnotator(new PorterStemmer())).build()); + assertEquals("annotator StemmerAnnotator requires layer opennlp:tokens," + + " which no earlier annotator provides", e.getMessage()); + } + + /** + * Verifies that a document without a token layer is rejected with a message naming the + * missing layer, instead of silently producing an empty stem layer. + */ + @Test + void testAbsentTokenLayerThrowsWithExactMessage() { + final StemmerAnnotator annotator = new StemmerAnnotator(new PorterStemmer()); + final IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> annotator.annotate(Document.of("no tokens"))); + assertEquals("document lacks the required layer opennlp:tokens", e.getMessage()); + } + + /** + * Verifies that a present-but-empty token layer yields a present-but-empty stem layer + * rather than an exception. + */ + @Test + void testEmptyPresentTokenLayerYieldsEmptyStemLayer() { + final Document document = Document.of("").with(Layers.TOKENS, List.of()); + final Document stemmed = new StemmerAnnotator(new PorterStemmer()).annotate(document); + assertTrue(stemmed.layers().contains(StemmerAnnotator.STEMS)); + assertTrue(stemmed.get(StemmerAnnotator.STEMS).isEmpty()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownloadTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownloadTest.java new file mode 100644 index 0000000000..7c219fd4f3 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellDictionaryDownloadTest.java @@ -0,0 +1,163 @@ +/* + * 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.stemmer.hunspell; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.DictionaryCatalog; +import opennlp.tools.util.DigestTestUtil; + +/** + * Pins the Hunspell catalog download gate; network fetches are not exercised here. + */ +public class HunspellDictionaryDownloadTest { + + /** + * Verifies that a catalog download without the remote-download property fails with + * the property name in the message, leaving the previous property value restored. + * + * @param target A scratch directory managed by the test framework. + * @throws IOException Thrown if the local catalog cannot be prepared. + */ + @Test + void testDownloadRequiresRemoteProperty(@TempDir Path target) throws IOException { + final DictionaryCatalog catalog = localCatalog(target); + final String previous = + System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + try { + final IOException e = Assertions.assertThrows(IOException.class, + () -> HunspellDictionaryDownload.downloadFromCatalog(catalog, "demo", target)); + Assertions.assertTrue( + e.getMessage().contains(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY)); + } finally { + restore(previous); + } + } + + /** + * Verifies that an application-supplied catalog downloads a Hunspell pair and its + * license readme under the names declared by that catalog. + * + * @param target A scratch directory managed by the test framework. + * @throws IOException Thrown if the local catalog cannot be prepared or downloaded. + */ + @Test + void testDownloadsFromApplicationCatalog(@TempDir Path target) throws IOException { + final DictionaryCatalog catalog = localCatalog(target); + final Path output = target.resolve("output"); + final String previous = + System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.setProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY, "true"); + try { + HunspellDictionaryDownload.downloadFromCatalog(catalog, "demo", output); + Assertions.assertEquals("SET UTF-8\n", + Files.readString(output.resolve("demo" + HunspellDictionary.AFFIX_FILE_SUFFIX))); + Assertions.assertEquals("1\nword\n", + Files.readString(output.resolve("demo" + HunspellDictionary.DICTIONARY_FILE_SUFFIX))); + Assertions.assertEquals("license\n", Files.readString(output.resolve("README.txt"))); + } finally { + restore(previous); + } + } + + /** + * Verifies that each required parameter is checked before a download begins. + * + * @param argument The invalid method parameter. + * @param target A scratch directory managed by the test framework. + * @throws IOException Thrown if the local catalog cannot be prepared. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"catalog", "dictionaryId", "targetDirectory"}) + void testRejectsNullParameters(String argument, @TempDir Path target) + throws IOException { + final DictionaryCatalog catalog = localCatalog(target); + final Executable download = switch (argument) { + case "catalog" -> () -> + HunspellDictionaryDownload.downloadFromCatalog(null, "demo", target); + case "dictionaryId" -> () -> + HunspellDictionaryDownload.downloadFromCatalog(catalog, null, target); + case "targetDirectory" -> () -> + HunspellDictionaryDownload.downloadFromCatalog(catalog, "demo", null); + default -> throw new IllegalArgumentException("unknown argument: " + argument); + }; + + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, download); + Assertions.assertEquals(argument + " must not be null", thrown.getMessage()); + } + + /** + * Creates an application-supplied catalog backed by local files. + * + * @param directory The directory to hold the source files. + * @return The loaded catalog. Never {@code null}. + * @throws IOException Thrown if the source files cannot be written. + */ + private static DictionaryCatalog localCatalog(Path directory) throws IOException { + final byte[] affix = "SET UTF-8\n".getBytes(StandardCharsets.UTF_8); + final byte[] dictionary = "1\nword\n".getBytes(StandardCharsets.UTF_8); + final byte[] readme = "license\n".getBytes(StandardCharsets.UTF_8); + final Path affixSource = directory.resolve( + "source" + HunspellDictionary.AFFIX_FILE_SUFFIX); + final Path dictionarySource = directory.resolve( + "source" + HunspellDictionary.DICTIONARY_FILE_SUFFIX); + final Path readmeSource = directory.resolve("source-readme.txt"); + Files.write(affixSource, affix); + Files.write(dictionarySource, dictionary); + Files.write(readmeSource, readme); + + final String prefix = "hunspell.demo"; + final String catalog = entry(prefix + HunspellDictionary.AFFIX_FILE_SUFFIX, + affixSource, affix, "demo" + HunspellDictionary.AFFIX_FILE_SUFFIX) + + entry(prefix + HunspellDictionary.DICTIONARY_FILE_SUFFIX, + dictionarySource, dictionary, + "demo" + HunspellDictionary.DICTIONARY_FILE_SUFFIX) + + entry(prefix + ".readme", readmeSource, readme, "README.txt"); + return DictionaryCatalog.load( + new ByteArrayInputStream(catalog.getBytes(StandardCharsets.UTF_8))); + } + + /** Builds one properties entry for a local source file. */ + private static String entry(String id, Path source, byte[] content, String filename) { + return id + ".url=" + source.toUri() + "\n" + + id + ".sha512=" + DigestTestUtil.sha512(content) + "\n" + + id + ".filename=" + filename + "\n"; + } + + /** Restores the remote-download property to its previous value. */ + private static void restore(String previous) { + if (previous == null) { + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + } else { + System.setProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY, previous); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellManualExampleTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellManualExampleTest.java new file mode 100644 index 0000000000..21daec7e0b --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellManualExampleTest.java @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.stemmer.hunspell; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import opennlp.tools.stemmer.Stemmer; + +/** + * Runs the manual's Hunspell examples (docbkx {@code stemmer.xml}) verbatim: every + * value the chapter states is asserted here, so a change breaking this test breaks the + * manual. The fixture dictionary is authored inside this class; no external dictionary + * data is involved. + */ +public class HunspellManualExampleTest { + + /** + * Affix fixture matching the chapter: agentive {@code -er} with continuation class + * {@code S}, and the plural {@code -s}. + */ + private static final String AFFIX = String.join("\n", + "SET UTF-8", + "SFX E Y 1", + "SFX E 0 er/S .", + "SFX S Y 1", + "SFX S 0 s [^sxy]", + ""); + + /** Word-list fixture: {@code work} accepts both suffixes. */ + private static final String WORDS = "1\nwork/ES\n"; + + /** + * Loads the chapter's miniature dictionary, stems through a factory-minted stemmer, + * and asserts the exact stems the manual prints. + * + * @throws IOException Thrown if the in-memory fixture fails to load. + */ + @Test + void testLoadAndStemWorkers() throws IOException { + final HunspellDictionary dictionary = HunspellDictionary.load( + new ByteArrayInputStream(AFFIX.getBytes(StandardCharsets.UTF_8)), + new ByteArrayInputStream(WORDS.getBytes(StandardCharsets.UTF_8))); + final Stemmer stemmer = new HunspellStemmerFactory(dictionary).newStemmer(); + + Assertions.assertEquals("work", stemmer.stem("workers").toString()); + Assertions.assertEquals("work", stemmer.stem("worker").toString()); + Assertions.assertEquals(List.of("work"), + stemmer.stemAll("workers").stream().map(CharSequence::toString).toList()); + // unknown vocabulary passes through unchanged + Assertions.assertEquals("table", stemmer.stem("table").toString()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellRealDictionaryTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellRealDictionaryTest.java new file mode 100644 index 0000000000..9221658ddd --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellRealDictionaryTest.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.tools.stemmer.hunspell; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; + +/** + * Gated checks against published dictionaries, which are never bundled: the tests run + * only when {@code -Dopennlp.hunspell.dict.dir} names a directory holding + * {@code .aff}/{@code .dic} pairs, and each test additionally skips when + * its dictionary pair is absent. The download helper in {@code dev/} fetches the pairs + * together with their license files; see {@code dev/README-hunspell-dictionaries.md}. + * + *

    The assertions are limited to morphology stable across dictionary revisions: + * everyday inflections, and for German the decomposability of ordinary compounds.

    + */ +public class HunspellRealDictionaryTest { + + private static final String DICT_DIR_PROPERTY = "opennlp.hunspell.dict.dir"; + + /** + * Loads one dictionary pair from the gated directory, skipping the test when the + * gate or the pair is absent. + * + * @param name The dictionary base name, such as {@code en_US}. + * @return A stemmer over the loaded pair. Never {@code null}. + * @throws IOException Thrown if a present pair fails to load, which is a failure, + * not a skip. + */ + private static HunspellStemmer loadOrSkip(String name) throws IOException { + final String dir = System.getProperty(DICT_DIR_PROPERTY); + Assumptions.assumeTrue(dir != null && !dir.isBlank(), + "no " + DICT_DIR_PROPERTY + " given"); + final Path affix = Path.of(dir, name + HunspellDictionary.AFFIX_FILE_SUFFIX); + final Path words = Path.of(dir, name + HunspellDictionary.DICTIONARY_FILE_SUFFIX); + Assumptions.assumeTrue(Files.isReadable(affix) && Files.isReadable(words), + name + " pair not present under " + dir); + return new HunspellStemmer(HunspellDictionary.load(affix, words)); + } + + /** + * Checks everyday English inflections against {@code en_US}, plus the identity + * fallback on vocabulary no dictionary lists. + * + * @throws IOException Thrown if a present dictionary pair fails to load. + */ + @Test + void testEnglishInflections() throws IOException { + final HunspellStemmer stemmer = loadOrSkip("en_US"); + Assertions.assertEquals("worker", stemmer.stem("workers").toString()); + Assertions.assertEquals("cat", stemmer.stem("cats").toString()); + Assertions.assertEquals("unhappy", stemmer.stem("unhappiest").toString()); + Assertions.assertEquals("quick", stemmer.stem("quickly").toString()); + Assertions.assertEquals("look", stemmer.stem("looked").toString()); + // unknown vocabulary degrades to identity + Assertions.assertEquals("zyzzyvax", stemmer.stem("zyzzyvax").toString()); + } + + /** + * Checks everyday German inflections against {@code de_DE_frami}: a plural, an + * umlauted plural, and a superlative. + * + * @throws IOException Thrown if a present dictionary pair fails to load. + */ + @Test + void testGermanInflections() throws IOException { + final HunspellStemmer stemmer = loadOrSkip("de_DE_frami"); + Assertions.assertEquals("Kind", stemmer.stem("Kinder").toString()); + // Haeuser, written with a-umlaut, stems to Haus + Assertions.assertEquals("Haus", stemmer.stem("H\u00E4user").toString()); + Assertions.assertEquals("schnell", stemmer.stem("schnellsten").toString()); + } + + /** + * Checks that ordinary German compounds decompose against {@code de_DE_frami}. Only + * the part count is asserted: the exact part spellings follow the dictionary's own + * entries and may shift between its revisions. + * + * @throws IOException Thrown if a present dictionary pair fails to load. + */ + @Test + void testGermanCompoundsDecompose() throws IOException { + final HunspellStemmer stemmer = loadOrSkip("de_DE_frami"); + // Haustuer, written with u-umlaut, is Haus + Tuer + Assertions.assertTrue(stemmer.stemAll("Haust\u00FCr").size() >= 2); + Assertions.assertTrue(stemmer.stemAll("Kinderzimmer").size() >= 2); + Assertions.assertTrue(stemmer.stemAll("Abbildungsverzeichnis").size() >= 2); + } + + /** + * Checks everyday Hungarian inflections against {@code hu_HU}: a plural and two + * case-suffixed forms. + * + * @throws IOException Thrown if a present dictionary pair fails to load. + */ + @Test + void testHungarianInflections() throws IOException { + final HunspellStemmer stemmer = loadOrSkip("hu_HU"); + // kutyak, written with a-acute, is the plural of kutya + Assertions.assertEquals("kutya", stemmer.stem("kuty\u00E1k").toString()); + Assertions.assertEquals("asztal", stemmer.stem("asztalon").toString()); + // konyveket, written with o-umlaut, is an inflected form of konyv + Assertions.assertEquals("k\u00F6nyv", stemmer.stem("k\u00F6nyveket").toString()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactoryTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactoryTest.java new file mode 100644 index 0000000000..6907fd1430 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerFactoryTest.java @@ -0,0 +1,183 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.stemmer.hunspell; + +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.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.stemmer.Stemmer; + +/** + * Demonstrates the intended end-to-end usage of the Hunspell stemming classes: a user + * writes (or ships) a {@code .aff}/{@code .dic} file pair, loads it once into a + * {@link HunspellDictionary}, wraps the dictionary in a {@link HunspellStemmerFactory}, + * and obtains {@link Stemmer} instances from the factory wherever stemming is needed. + * The fixture dictionary is authored inside this test class, so no external dictionary + * data is involved. + */ +public class HunspellStemmerFactoryTest { + + /** + * The affix fixture: the prefix {@code re-}, the suffix {@code -er} whose continuation + * class {@code S} lets the plural {@code -s} stack on top of it, and the plural + * {@code -s} itself, restricted to stems not ending in {@code s}, {@code x}, or + * {@code y}. All three rules opt into cross-product combination. + */ + private static final String AFFIX = String.join("\n", + "# project-authored test fixture", + "SET UTF-8", + "", + "PFX R Y 1", + "PFX R 0 re .", + "", + "SFX E Y 1", + "SFX E 0 er/S .", + "", + "SFX S Y 1", + "SFX S 0 s [^sxy]", + ""); + + /** + * The word-list fixture: {@code work} accepts the prefix and both suffixes, + * {@code paint} accepts only the agentive {@code -er}. + */ + private static final String WORDS = String.join("\n", + "2", + "work/RES", + "paint/E", + ""); + + /** + * Writes the fixture dictionary pair into a directory and loads it through the + * file-based {@link HunspellDictionary#load(Path, Path)} entry point. + * + * @param directory The directory to write into. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if writing or loading fails. + */ + private static HunspellDictionary writeAndLoadFixture(Path directory) throws IOException { + final Path affixFile = directory.resolve( + "fixture" + HunspellDictionary.AFFIX_FILE_SUFFIX); + final Path dictionaryFile = directory.resolve( + "fixture" + HunspellDictionary.DICTIONARY_FILE_SUFFIX); + Files.write(affixFile, AFFIX.getBytes(StandardCharsets.UTF_8)); + Files.write(dictionaryFile, WORDS.getBytes(StandardCharsets.UTF_8)); + return HunspellDictionary.load(affixFile, dictionaryFile); + } + + /** + * Walks the whole intended flow on a single thread: files on disk, one dictionary, + * one factory, one stemmer, and exact stems for a prefixed form, a suffixed form, a + * twofold suffix chain, a cross-product form, an in-dictionary word, and an unknown + * word. + * + * @param tempDir A scratch directory managed by the test framework. + * @throws IOException Thrown if the fixture cannot be written or loaded. + */ + @Test + void testEndToEndUsageFromFiles(@TempDir Path tempDir) throws IOException { + final HunspellDictionary dictionary = writeAndLoadFixture(tempDir); + final HunspellStemmerFactory factory = new HunspellStemmerFactory(dictionary); + final Stemmer stemmer = factory.newStemmer(); + + // one suffix removed + Assertions.assertEquals("work", stemmer.stem("worker").toString()); + Assertions.assertEquals("paint", stemmer.stem("painter").toString()); + // twofold suffixes: -s stacks on -er through the continuation class S + Assertions.assertEquals("work", stemmer.stem("workers").toString()); + // one prefix removed + Assertions.assertEquals("work", stemmer.stem("rework").toString()); + // cross product: the prefix re- and the suffix -s on the same stem + Assertions.assertEquals("work", stemmer.stem("reworks").toString()); + // a word that is itself listed stems to itself + Assertions.assertEquals("work", stemmer.stem("work").toString()); + // unknown vocabulary passes through unchanged + Assertions.assertEquals("table", stemmer.stem("table").toString()); + } + + /** + * Shares one factory between two threads: each thread obtains its own stemmer + * instance from the factory and stems the same inputs. The test asserts that the two + * instances are distinct objects and that their results are identical to each other + * and to the expected stems. + * + * @param tempDir A scratch directory managed by the test framework. + * @throws Exception Thrown if the fixture cannot be loaded or a worker fails. + */ + @Test + void testFactorySharedAcrossThreads(@TempDir Path tempDir) throws Exception { + final HunspellStemmerFactory factory = + new HunspellStemmerFactory(writeAndLoadFixture(tempDir)); + final List inputs = List.of("workers", "reworks", "painter", "table"); + final List expected = List.of("work", "work", "paint", "table"); + + final Stemmer[] created = new Stemmer[2]; + final ExecutorService pool = Executors.newFixedThreadPool(2); + try { + final List>> futures = new ArrayList<>(2); + for (int worker = 0; worker < 2; worker++) { + final int slot = worker; + futures.add(pool.submit(() -> { + final Stemmer stemmer = factory.newStemmer(); + created[slot] = stemmer; + final List stems = new ArrayList<>(inputs.size()); + for (final String input : inputs) { + stems.add(stemmer.stem(input).toString()); + } + return stems; + })); + } + final List first = futures.get(0).get(); + final List second = futures.get(1).get(); + Assertions.assertEquals(expected, first); + Assertions.assertEquals(expected, second); + } finally { + pool.shutdownNow(); + } + Assertions.assertNotSame(created[0], created[1]); + } + + /** + * Verifies that the file-based entry point rejects each {@code null} path with the + * documented exception naming the offending argument. + * + * @param tempDir A scratch directory managed by the test framework. + */ + @Test + void testNullPathsAreRejected(@TempDir Path tempDir) { + final Path present = tempDir.resolve("present" + HunspellDictionary.AFFIX_FILE_SUFFIX); + IllegalArgumentException e = Assertions.assertThrows(IllegalArgumentException.class, + () -> HunspellDictionary.load(null, present)); + Assertions.assertEquals("affixFile must not be null", e.getMessage()); + + e = Assertions.assertThrows(IllegalArgumentException.class, + () -> HunspellDictionary.load(present, null)); + Assertions.assertEquals("dictionaryFile must not be null", e.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerTest.java new file mode 100644 index 0000000000..312d785483 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/stemmer/hunspell/HunspellStemmerTest.java @@ -0,0 +1,1916 @@ +/* + * 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.stemmer.hunspell; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +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.stemmer.Stemmer; + +/** + * Tests the affix engine against a project-authored miniature dictionary; no external + * dictionary data is involved. + */ +public class HunspellStemmerTest { + + private static final byte TRUNCATED_UTF8_LEAD_BYTE = (byte) 0xC3; + + private static final String AFFIX = String.join("\n", + "# project-authored test fixture", + "SET UTF-8", + "", + "PFX U Y 1", + "PFX U 0 un .", + "", + "SFX S Y 3", + "SFX S 0 s [^sxy]", + "SFX S y ies y", + "SFX S 0 es [sx]", + "", + "SFX G Y 2", + "SFX G 0 ing [^e]", + "SFX G e ing e", + ""); + + private static final String WORDS = String.join("\n", + "6", + "lock/USG", + "pony/S", + "make/G", + "cat/S", + "box/S", + "fish", + ""); + + private static HunspellStemmer stemmer; + + /** + * Loads the shared fixture dictionary once for the tests that stem against it. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @BeforeAll + static void loadDictionary() throws IOException { + stemmer = new HunspellStemmer(load(AFFIX, WORDS)); + } + + /** + * Loads a dictionary from in-memory affix and word-list content, both encoded as + * UTF-8, through the stream-based entry point. + * + * @param affix The {@code .aff} content. + * @param words The {@code .dic} content. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if the content is malformed. + */ + private static HunspellDictionary load(String affix, String words) throws IOException { + return load(affix, words, StandardCharsets.UTF_8); + } + + /** + * Loads a dictionary from in-memory affix and word-list content encoded in the given + * charset, through the stream-based entry point. + * + * @param affix The {@code .aff} content. + * @param words The {@code .dic} content. + * @param charset The charset both contents are encoded with. + * @return The loaded dictionary. Never {@code null}. + * @throws IOException Thrown if the content is malformed. + */ + private static HunspellDictionary load(String affix, String words, Charset charset) + throws IOException { + return HunspellDictionary.load(new ByteArrayInputStream(affix.getBytes(charset)), + new ByteArrayInputStream(words.getBytes(charset))); + } + + /** + * Verifies the fixture's suffix rules: the plural {@code -s}, the {@code y} to + * {@code ies} replacement, the {@code -es} plural after a sibilant, and the + * progressive {@code -ing} with and without the silent {@code e}. + * + * @param word The surface form to stem. + * @param expected The stem the fixture licenses. + */ + @ParameterizedTest + @CsvSource({"cats,cat", "ponies,pony", "boxes,box", "making,make", "locking,lock"}) + void testSuffixRules(String word, String expected) { + Assertions.assertEquals(expected, stemmer.stem(word).toString()); + } + + /** + * Verifies prefix removal alone and combined with a suffix through the cross-product + * marker both rules declare. + * + * @param word The surface form to stem. + * @param expected The stem the fixture licenses. + */ + @ParameterizedTest + @CsvSource({"unlock,lock", "unlocks,lock", "unlocking,lock"}) + void testPrefixAndCrossProduct(String word, String expected) { + Assertions.assertEquals(expected, stemmer.stem(word).toString()); + } + + /** + * Verifies that an analysis a rule condition or a missing flag rejects is not + * reported: {@code boxs} fails the {@code [^sxy]} condition of the {@code -s} rule, + * {@code cat} carries no {@code G} flag, and {@code fish} carries no flag at all, so + * each surface form falls through unchanged. + * + * @param word The surface form to stem. + */ + @ParameterizedTest + @CsvSource({"boxs", "cating", "fishs"}) + void testConditionsBlockWrongAnalyses(String word) { + Assertions.assertEquals(word, stemmer.stem(word).toString()); + } + + /** + * Verifies that a listed word stems to itself and that a capitalized surface form is + * analyzed through its lowercase variant. + * + * @param word The surface form to stem. + * @param expected The stem the fixture licenses. + */ + @ParameterizedTest + @CsvSource({"fish,fish", "Cats,cat", "Unlocks,lock"}) + void testDirectLookupAndCase(String word, String expected) { + Assertions.assertEquals(expected, stemmer.stem(word).toString()); + } + + /** Verifies that a word with no analysis is returned unchanged as its only analysis. */ + @Test + void testUnknownWordsPassThroughUnchanged() { + Assertions.assertEquals("zebras", stemmer.stem("zebras").toString()); + Assertions.assertEquals(1, stemmer.stemAll("zebras").size()); + } + + /** + * Verifies that {@link HunspellStemmer#stemAll(CharSequence)} reports the analyses + * and that {@link HunspellStemmer#stem(CharSequence)} answers the first of them. + */ + @Test + void testStemAllReportsEveryAnalysis() { + Assertions.assertEquals(1, stemmer.stemAll("unlocks").size()); + Assertions.assertEquals("lock", stemmer.stemAll("unlocks").get(0).toString()); + // the surface form itself is an entry AND an analysis target + Assertions.assertEquals("lock", stemmer.stemAll("lock").get(0).toString()); + } + + /** + * Verifies twofold suffix removal: the plural {@code -s} stacks on the comparative + * {@code -er} through the continuation class the outer rule declares, while the inner + * flag alone licenses nothing because no entry carries it. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testTwofoldSuffixesThroughContinuationClasses() throws IOException { + final HunspellStemmer twofold = new HunspellStemmer(load(String.join("\n", + "SET UTF-8", + "SFX A Y 1", + "SFX A 0 er/B .", + "SFX B Y 1", + "SFX B 0 s .", + ""), String.join("\n", "1", "kind/A", ""))); + + Assertions.assertEquals("kind", twofold.stem("kinder").toString()); + Assertions.assertEquals("kind", twofold.stem("kinders").toString()); + // B alone never applies: no entry carries it directly + Assertions.assertEquals("kinds", twofold.stem("kinds").toString()); + } + + /** + * Verifies {@code FLAG num} mode: a comma-separated run of decimal numbers is the + * entry's flag set, and an affix block named by one of them applies. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testNumericFlagMode() throws IOException { + final HunspellDictionary dictionary = load(String.join("\n", + "SET UTF-8", + "FLAG num", + "SFX 100 Y 1", + "SFX 100 0 s .", + ""), String.join("\n", "1", "walk/100,7", "")); + Assertions.assertEquals("walk", + new HunspellStemmer(dictionary).stem("walks").toString()); + } + + /** + * Verifies {@code FLAG long} mode: each pair of characters in the run is one flag, + * and an affix block named by such a pair applies. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testLongFlagMode() throws IOException { + final HunspellDictionary dictionary = load(String.join("\n", + "SET UTF-8", + "FLAG long", + "SFX Aa Y 1", + "SFX Aa 0 s .", + ""), String.join("\n", "1", "walk/AaBb", "")); + Assertions.assertEquals("walk", + new HunspellStemmer(dictionary).stem("walks").toString()); + } + + /** + * Verifies that a stemmer minted by the factory analyzes against the same dictionary. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testFactoryHandsOutWorkingStemmers() throws IOException { + final Stemmer fresh = new HunspellStemmerFactory(load(AFFIX, WORDS)).newStemmer(); + Assertions.assertEquals("pony", fresh.stem("ponies").toString()); + } + + /** + * Verifies that cross-product combination of a prefix with a suffix only happens + * when both rules declare the cross-product marker {@code Y}. Removing just the one + * affix whose rule exists keeps working; the combined form must not be analyzed. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCrossProductRequiresBothRulesOptIn() throws IOException { + // the prefix rule declares N, so it never combines with the suffix + final HunspellStemmer prefixOptedOut = new HunspellStemmer(load(String.join("\n", + "PFX U N 1", + "PFX U 0 un .", + "SFX S Y 1", + "SFX S 0 s .", + ""), "1\nlock/US\n")); + Assertions.assertEquals("lock", prefixOptedOut.stem("unlock").toString()); + Assertions.assertEquals("lock", prefixOptedOut.stem("locks").toString()); + Assertions.assertEquals("unlocks", prefixOptedOut.stem("unlocks").toString()); + + // the suffix rule declares N, so the combined form is likewise not analyzed + final HunspellStemmer suffixOptedOut = new HunspellStemmer(load(String.join("\n", + "PFX U Y 1", + "PFX U 0 un .", + "SFX S N 1", + "SFX S 0 s .", + ""), "1\nlock/US\n")); + Assertions.assertEquals("lock", suffixOptedOut.stem("unlock").toString()); + Assertions.assertEquals("lock", suffixOptedOut.stem("locks").toString()); + Assertions.assertEquals("unlocks", suffixOptedOut.stem("unlocks").toString()); + } + + /** + * Verifies that a non-negated character class rejects a candidate stem: the + * {@code es} rule requires a stem ending in {@code s} or {@code x}, so removing + * {@code es} from {@code cates} produces {@code cat}, which the class rejects, and + * the surface form falls through unchanged. + */ + @Test + void testPositiveCharacterClassRejectsCandidate() { + Assertions.assertEquals("cates", stemmer.stem("cates").toString()); + Assertions.assertEquals(1, stemmer.stemAll("cates").size()); + } + + /** + * Verifies that the {@code SET} declaration selects the charset both files are + * decoded with: a word list holding the byte {@code 0xE9} only maps to the word + * caf\u00E9 (e with acute accent) when decoded as ISO-8859-1, as the affix file declares. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testSetDeclarationSelectsEncoding() throws IOException { + final Charset latin1 = StandardCharsets.ISO_8859_1; + final HunspellDictionary dictionary = load(String.join("\n", + "SET ISO8859-1", + "SFX S Y 1", + "SFX S 0 s .", + ""), "1\ncaf\u00E9/S\n", latin1); + final HunspellStemmer latin1Stemmer = new HunspellStemmer(dictionary); + Assertions.assertEquals("caf\u00E9", latin1Stemmer.stem("caf\u00E9s").toString()); + Assertions.assertEquals("caf\u00E9", latin1Stemmer.stem("caf\u00E9").toString()); + } + + /** + * Verifies that continuation classes also work in {@code FLAG long} mode, where a + * flag is a two-character run: the plural {@code Bb} stacks on the agentive + * {@code Aa} to analyze a twofold suffix chain. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testLongFlagContinuation() throws IOException { + final HunspellStemmer longFlags = new HunspellStemmer(load(String.join("\n", + "FLAG long", + "SFX Aa Y 1", + "SFX Aa 0 er/Bb .", + "SFX Bb Y 1", + "SFX Bb 0 s .", + ""), "1\nkind/Aa\n")); + Assertions.assertEquals("kind", longFlags.stem("kinder").toString()); + Assertions.assertEquals("kind", longFlags.stem("kinders").toString()); + } + + /** + * Verifies that cross-product prefix and suffix combination also works in + * {@code FLAG num} mode, where flags are comma-separated decimal numbers. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testNumericFlagCrossProduct() throws IOException { + final HunspellStemmer numericFlags = new HunspellStemmer(load(String.join("\n", + "FLAG num", + "PFX 1 Y 1", + "PFX 1 0 un .", + "SFX 2 Y 1", + "SFX 2 0 s .", + ""), "1\nlock/1,2\n")); + Assertions.assertEquals("lock", numericFlags.stem("unlock").toString()); + Assertions.assertEquals("lock", numericFlags.stem("locks").toString()); + Assertions.assertEquals("lock", numericFlags.stem("unlocks").toString()); + } + + /** + * Verifies the exact exception and message for each malformed {@code FLAG} + * declaration the parser detects: a missing mode and an unrecognized mode name. + */ + @Test + void testMalformedFlagDeclarationMessages() { + IOException e = Assertions.assertThrows(IOException.class, + () -> load("FLAG\n", "0\n")); + Assertions.assertEquals("FLAG line without a mode at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, () -> load("FLAG short\n", "0\n")); + Assertions.assertEquals("unsupported FLAG mode 'short' at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("FLAG num\nFLAG UTF-8\n", "0\n")); + Assertions.assertEquals("multiple FLAG directives at line 2", e.getMessage()); + } + + /** + * Verifies the exact exception and message for each malformed affix block the + * parser detects: a header with too few fields, a non-numeric rule count, a block + * with fewer rule lines than its count announces, a rule line whose type tag does + * not match its header, and an unterminated character class in a condition. + */ + @Test + void testMalformedAffixBlockMessages() { + IOException e = Assertions.assertThrows(IOException.class, + () -> load("PFX U Y\n", "0\n")); + Assertions.assertEquals("malformed affix header at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y many\nSFX S 0 s .\n", "0\n")); + Assertions.assertEquals("malformed affix rule count at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y 2\nSFX S 0 s .", "0\n")); + Assertions.assertEquals("affix block truncated at line 3", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S X 1\nSFX S 0 s .\n", "0\n")); + Assertions.assertEquals("invalid cross-product marker at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y 1\nPFX S 0 s .\n", "0\n")); + Assertions.assertEquals("malformed affix rule at line 2", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y 1\nSFX T 0 s .\n", "0\n")); + Assertions.assertEquals("affix rule flag does not match header at line 2", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y -1\n", "0\n")); + Assertions.assertEquals("negative affix rule count at line 1", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SFX S Y 1\nSFX S 0 s [ab\n", "0\n")); + Assertions.assertEquals("unterminated character class at line 2", e.getMessage()); + } + + /** + * Verifies the exact exception and message for each malformed flag value the parser + * detects: an odd-length flag run in {@code FLAG long} mode, a non-numeric flag in + * {@code FLAG num} mode, an affix header naming more than one flag, and a + * {@code SET} declaration naming an unknown encoding. + */ + @Test + void testMalformedFlagValueMessages() { + IOException e = Assertions.assertThrows(IOException.class, + () -> load("FLAG long\n", "1\nwalk/AaB\n")); + Assertions.assertEquals("odd long-flag run at line 2", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("FLAG num\n", "1\nwalk/12,x\n")); + Assertions.assertEquals("malformed numeric flag at line 2", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("FLAG long\nSFX AaBb Y 1\nSFX AaBb 0 s .\n", "0\n")); + Assertions.assertEquals("expected exactly one flag at line 2", e.getMessage()); + + e = Assertions.assertThrows(IOException.class, + () -> load("SET NO-SUCH-ENCODING\n", "0\n")); + Assertions.assertEquals("unsupported SET encoding: NO-SUCH-ENCODING", e.getMessage()); + } + + /** + * Verifies that a morphological field is cut off the entry before the flag separator + * is looked for, so a slash inside a morphological field is not mistaken for the + * separator: the entry {@code walk po:verb/noun} registers the word {@code walk} + * with no flags in every flag mode, and its morphology is ignored. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testMorphologicalFieldsAreCutBeforeTheFlagSeparator() throws IOException { + final HunspellDictionary chars = load("SFX G Y 1\nSFX G 0 ing .\n", + "1\nwalk po:verb/noun\n"); + Assertions.assertNotNull(chars.lookup("walk")); + Assertions.assertEquals(0, chars.lookup("walk").get(0).length); + Assertions.assertNull(chars.lookup("walk po:verb")); + + final HunspellDictionary numbers = load("FLAG num\nSFX 1 Y 1\nSFX 1 0 ing .\n", + "1\nwalk po:verb/noun\n"); + Assertions.assertNotNull(numbers.lookup("walk")); + Assertions.assertEquals(0, numbers.lookup("walk").get(0).length); + + // the tabulator is the older morphological field separator + final HunspellDictionary tabbed = load("SFX G Y 1\nSFX G 0 ing .\n", + "1\nwalk\tpo:verb/noun\n"); + Assertions.assertNotNull(tabbed.lookup("walk")); + Assertions.assertEquals(0, tabbed.lookup("walk").get(0).length); + } + + /** + * Verifies that an entry keeps its flags when it carries both a flag run and a + * morphological field holding a slash, in every flag mode. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testFlaggedEntriesKeepTheirFlagsBesideMorphology() throws IOException { + final HunspellDictionary chars = load("SFX A Y 1\nSFX A 0 ing .\n", + "1\nwalk/AB po:verb/noun\n"); + Assertions.assertArrayEquals(new int[] {'A', 'B'}, chars.lookup("walk").get(0)); + Assertions.assertEquals("walk", + new HunspellStemmer(chars).stem("walking").toString()); + + final HunspellDictionary numbers = load("FLAG num\nSFX 1 Y 1\nSFX 1 0 ing .\n", + "1\nwalk/1,2 po:verb/noun\n"); + Assertions.assertArrayEquals(new int[] {1, 2}, numbers.lookup("walk").get(0)); + Assertions.assertEquals("walk", + new HunspellStemmer(numbers).stem("walking").toString()); + } + + /** + * Verifies that a multi-word entry keeps both its spaces and its flags: the word of + * a word-list entry runs up to its morphological fields, not up to its first space. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testMultiWordEntriesKeepTheirSpacesAndFlags() throws IOException { + final HunspellDictionary dictionary = load("FLAG num\nSFX 39 Y 1\nSFX 39 0 s .\n", + "1\nall right/39\n"); + Assertions.assertArrayEquals(new int[] {39}, dictionary.lookup("all right").get(0)); + Assertions.assertNull(dictionary.lookup("all")); + } + + /** + * Verifies that the parser trims word-list entries with the same whitespace + * definition it uses to find their fields: an entry edged by Unicode whitespace, + * leading or trailing, is registered under its real word, both with and without a + * flag run. + * + * @param space The whitespace character at the line edges: the no-break space + * U+00A0 and the ideographic space U+3000, both whitespace to + * {@code StringUtil.isWhitespace} but not to {@code String.trim()}. + * @throws IOException Thrown if a fixture fails to load. + */ + @ParameterizedTest + @ValueSource(strings = {"\u00A0", "\u3000"}) + void testEntriesEdgedByUnicodeWhitespaceAreTrimmed(String space) throws IOException { + final HunspellDictionary dictionary = load("SFX S Y 1\nSFX S 0 s .\n", + "2\n" + space + "fish" + space + "\n" + space + "cat/S" + space + "\n"); + Assertions.assertNotNull(dictionary.lookup("fish")); + Assertions.assertNotNull(dictionary.lookup("cat")); + Assertions.assertNull(dictionary.lookup("")); + Assertions.assertEquals("cat", new HunspellStemmer(dictionary).stem("cats").toString()); + } + + /** + * Verifies that {@code FLAG UTF-8}, which declares single-character flags, is + * accepted and read exactly like the default single-character mode, including a flag + * outside ASCII. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testUtf8FlagModeDeclaresSingleCharacterFlags() throws IOException { + final HunspellStemmer plain = new HunspellStemmer(load(String.join("\n", + "FLAG UTF-8", + "SFX S Y 1", + "SFX S 0 s .", + ""), "1\nwalk/S\n")); + Assertions.assertEquals("walk", plain.stem("walks").toString()); + + // \u00E9 is e with an acute accent, a single-character flag outside ASCII + final HunspellStemmer accented = new HunspellStemmer(load(String.join("\n", + "FLAG UTF-8", + "SFX \u00E9 Y 1", + "SFX \u00E9 0 s .", + ""), "1\nwalk/\u00E9\n")); + Assertions.assertEquals("walk", accented.stem("walks").toString()); + } + + /** + * Verifies that a strip-only rule, whose affix material is empty and which therefore + * only removes stem material, is undone: the suffix rule turns the entry + * {@code bake} into the surface form {@code bak}, and the prefix rule turns + * {@code apple} into {@code pple}. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testStripOnlyAffixRulesAreUndone() throws IOException { + final HunspellStemmer suffixStripping = new HunspellStemmer(load(String.join("\n", + "SFX A Y 1", + "SFX A e 0 e", + ""), "1\nbake/A\n")); + Assertions.assertEquals("bake", suffixStripping.stem("bak").toString()); + + final HunspellStemmer prefixStripping = new HunspellStemmer(load(String.join("\n", + "PFX B Y 1", + "PFX B a 0 a", + ""), "1\napple/B\n")); + Assertions.assertEquals("apple", prefixStripping.stem("pple").toString()); + } + + /** + * Verifies that an entry written with an empty flag run loads and carries no flags in + * every flag mode, rather than failing the load in {@code FLAG num} mode alone. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testEmptyFlagRunYieldsNoFlagsInEveryMode() throws IOException { + Assertions.assertEquals(0, load("", "1\nword/\n").lookup("word").get(0).length); + Assertions.assertEquals(0, + load("FLAG long\n", "1\nword/\n").lookup("word").get(0).length); + Assertions.assertEquals(0, + load("FLAG num\n", "1\nword/\n").lookup("word").get(0).length); + } + + /** Verifies that a malformed affix file aborts the load instead of loading partially. */ + @Test + void testMalformedInputFailsLoud() { + Assertions.assertThrows(IOException.class, + () -> load("SFX S Y 2\nSFX S 0 s .\n", "1\ncat/S\n")); + Assertions.assertThrows(IOException.class, + () -> load("SET NO-SUCH-ENCODING\n", "0\n")); + Assertions.assertThrows(IOException.class, () -> load("SFX S 0 s [a\n", "0\n")); + } + + /** + * Verifies that every entry point rejects a {@code null} argument with the documented + * exception, and that the stream-based loader names the offending argument the way + * its file-based sibling does. + */ + @Test + void testNullArgumentsAreRejected() { + final InputStream present = new ByteArrayInputStream(new byte[0]); + IllegalArgumentException e = Assertions.assertThrows(IllegalArgumentException.class, + () -> HunspellDictionary.load(null, present)); + Assertions.assertEquals("affixStream must not be null", e.getMessage()); + + e = Assertions.assertThrows(IllegalArgumentException.class, + () -> HunspellDictionary.load(present, null)); + Assertions.assertEquals("dictionaryStream must not be null", e.getMessage()); + + e = Assertions.assertThrows(IllegalArgumentException.class, + () -> new HunspellStemmer(null)); + Assertions.assertEquals("dictionary must not be null", e.getMessage()); + + e = Assertions.assertThrows(IllegalArgumentException.class, + () -> new HunspellStemmerFactory(null)); + Assertions.assertEquals("dictionary must not be null", e.getMessage()); + + e = Assertions.assertThrows(IllegalArgumentException.class, + () -> stemmer.stemAll(null)); + Assertions.assertEquals("word must not be null", e.getMessage()); + } + + /** + * Verifies hunspell's tolerance for trailing text after a numeric or long flag run: + * the flag run ends at the first space, and whatever follows is a morphological + * field even without a two-letter tag, so such an entry loads instead of aborting + * the whole dictionary. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testTrailingTextAfterNumericFlagRunIsMorphologyNotAnError() throws IOException { + final HunspellDictionary numbers = load("FLAG num\n", + "2\nwalk/39 blah\nrun/7,9 xyz abc\n"); + Assertions.assertNotNull(numbers.lookup("walk")); + Assertions.assertTrue(HunspellDictionary.hasFlag(numbers.lookup("walk"), 39)); + Assertions.assertTrue(HunspellDictionary.hasFlag(numbers.lookup("run"), 7)); + Assertions.assertTrue(HunspellDictionary.hasFlag(numbers.lookup("run"), 9)); + + final HunspellDictionary longs = load("FLAG long\n", "1\nwalk/AB cd\n"); + Assertions.assertTrue(HunspellDictionary.hasFlag(longs.lookup("walk"), + ('A' << 16) | 'B')); + } + + /** + * Verifies that stemming the empty word answers the empty word: a zero-length + * surface has no morphology, and a strip-only rule must not restore its strip + * string onto nothing and answer a non-empty stem. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testEmptyWordStemsToItself() throws IOException { + final HunspellStemmer stripOnly = new HunspellStemmer(load( + "PFX P Y 1\nPFX P xy 0 .\n", + "1\nxy/P\n")); + Assertions.assertEquals("", stripOnly.stem("").toString()); + Assertions.assertEquals(List.of(""), stripOnly.stemAll("")); + } + + /** + * Verifies the escaped-slash feature: {@code \/} belongs to the word, so an entry + * naming a slashed term keeps its slash while the first unescaped slash still + * separates the flag run. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testEscapedSlashBelongsToTheWord() throws IOException { + final HunspellDictionary slashed = load("FLAG num\n", + "2\nTCP\\/IP/39\nAC\\/DC\n"); + Assertions.assertNotNull(slashed.lookup("TCP/IP")); + Assertions.assertTrue(HunspellDictionary.hasFlag(slashed.lookup("TCP/IP"), 39)); + Assertions.assertNotNull(slashed.lookup("AC/DC")); + Assertions.assertNull(slashed.lookup("TCP")); + } + + /** + * Verifies the sharpest combination of the morphology cut: an entry that is both a + * multi-word term and carries trailing tag morphology keeps the whole multi-word + * surface and its flags, and the tags stay out of the word. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testMultiWordEntryWithTrailingTagMorphology() throws IOException { + final HunspellDictionary phrases = load("FLAG num\n", + "1\nall right/39 po:phrase st:allright\n"); + Assertions.assertNotNull(phrases.lookup("all right")); + Assertions.assertTrue(HunspellDictionary.hasFlag(phrases.lookup("all right"), 39)); + Assertions.assertNull(phrases.lookup("all right po:phrase st:allright")); + } + + /** + * Pins FLAG UTF-8 for a supplementary flag character: a flag is one code point, so + * a character above U+FFFF is one flag carrying its code point value, never two + * surrogate-unit flags. The Spanish dictionary of the LibreOffice collection names + * affix rules with such characters, so an affix keyed by a supplementary flag must + * connect to the entries that carry it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testSupplementaryFlagCharacterIsOneCodePointFlag() throws IOException { + // U+1F600 as a flag, written as its surrogate pair + final HunspellDictionary emoji = load("FLAG UTF-8\n", + "1\nwalk/\uD83D\uDE00\n"); + Assertions.assertTrue(HunspellDictionary.hasFlag(emoji.lookup("walk"), 0x1F600)); + Assertions.assertFalse(HunspellDictionary.hasFlag(emoji.lookup("walk"), 0xD83D)); + + final HunspellStemmer supplementaryFlags = new HunspellStemmer(load( + "FLAG UTF-8\nSFX \uD83D\uDE00 Y 1\nSFX \uD83D\uDE00 0 s .\n", + "1\nwalk/\uD83D\uDE00\n")); + Assertions.assertEquals("walk", supplementaryFlags.stem("walks").toString()); + } + + /** + * Pins the variation-selector rule the Spanish dictionary of the LibreOffice + * collection relies on: a variation selector after a flag character selects its + * presentation and is no flag of its own, so an affix rule named with the emoji + * form of a character connects to entries flagged with either spelling. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testVariationSelectorIsDroppedFromFlagIdentity() throws IOException { + // U+260E BLACK TELEPHONE followed by U+FE0F VARIATION SELECTOR-16, the exact + // shape of a prefix flag in the published es_ES affix file + final HunspellStemmer stemmer = new HunspellStemmer(load( + "FLAG UTF-8\nPFX \u260E\uFE0F Y 1\nPFX \u260E\uFE0F 0 tele .\n", + "1\nfono/\u260E\n")); + Assertions.assertEquals("fono", stemmer.stem("telefono").toString()); + } + + /** + * Pins the documented rejection of rules that neither add nor remove material: a + * suffix rule with strip {@code 0} and affix {@code 0} loads without error and + * never fires, so stemming a flagged dictionary word answers that word exactly + * once. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testRuleThatNeitherAddsNorRemovesLoadsAndNeverFires() throws IOException { + final HunspellStemmer identity = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 0 .\n", + "1\nwalk/X\n")); + Assertions.assertEquals(List.of("walk"), identity.stemAll("walk")); + } + + /** + * Verifies the AF flag alias table: the first AF line declares the count, every + * further AF line is one flag run, and a purely numeric flag field in the word + * list is a 1-based reference into that table, the layout the published Hungarian + * dictionary uses for all of its ninety-seven thousand entries. Alias lines may + * carry trailing comments, which the field split already discards. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testNumericDictionaryFlagsResolveThroughTheAliasTable() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "AF 2", + "AF S # 1", + "AF SP # 2", + "SFX S Y 1", + "SFX S 0 s .", + "PFX P Y 1", + "PFX P 0 re .", + ""), + "2\nwalk/1\nplay/2\n")); + Assertions.assertEquals("walk", stemmer.stem("walks").toString()); + Assertions.assertEquals("play", stemmer.stem("plays").toString()); + Assertions.assertEquals("play", stemmer.stem("replay").toString()); + // walk carries alias 1, the suffix-only run, so the prefix must not apply + Assertions.assertEquals("rewalk", stemmer.stem("rewalk").toString()); + } + + /** + * Verifies that an alias reference outside the AF table fails loud with the line + * and the table size, instead of silently flagging the entry with nothing, and that + * a digit run too large for an alias number fails loud as well. + */ + @Test + void testAliasReferenceOutsideTheTableFailsLoud() { + IOException e = Assertions.assertThrows(IOException.class, () -> load( + "AF 1\nAF S # 1\nSFX S Y 1\nSFX S 0 s .\n", + "1\nwalk/2\n")); + Assertions.assertEquals( + "flag alias 2 at line 2 is outside the AF table of 1 aliases", + e.getMessage()); + + e = Assertions.assertThrows(IOException.class, () -> load( + "AF 1\nAF S # 1\nSFX S Y 1\nSFX S 0 s .\n", + "1\nwalk/99999999999999999999\n")); + Assertions.assertEquals( + "malformed flag alias '99999999999999999999' at line 2", + e.getMessage()); + } + + /** + * Verifies that numeric flag fields stay ordinary flags when no AF table exists: + * under FLAG num a digit run is a flag value, not an alias reference. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testNumericFlagsWithoutAliasTableStayFlags() throws IOException { + final HunspellDictionary numbers = load("FLAG num\n", "1\nwalk/39\n"); + Assertions.assertTrue(HunspellDictionary.hasFlag(numbers.lookup("walk"), 39)); + } + + /** + * Verifies two-part compound decomposition under the general compounding flag: a + * word the affix analysis cannot explain splits into two listed parts that both + * carry the flag, reported left to right, while a part without the flag blocks the + * split and the word stays unanalyzed. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundFlagDecomposesUnanalyzedWords() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\n", + "3\ndog/Z\nhouse/Z\ncat\n")); + Assertions.assertEquals(List.of("dog", "house"), stemmer.stemAll("doghouse")); + // cat is listed without the compounding flag, so no split may use it + Assertions.assertEquals(List.of("cathouse"), stemmer.stemAll("cathouse")); + // a listed word never decomposes; it is its own analysis + Assertions.assertEquals(List.of("dog"), stemmer.stemAll("dog")); + } + + /** + * Verifies the positional compound flags: the begin flag only opens and the end + * flag only closes, so the parts compose in one order and refuse the other. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundBeginAndEndFlagsArePositional() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDBEGIN B\nCOMPOUNDEND E\nCOMPOUNDMIN 3\n", + "2\ndog/B\nhouse/E\n")); + Assertions.assertEquals(List.of("dog", "house"), stemmer.stemAll("doghouse")); + Assertions.assertEquals(List.of("housedog"), stemmer.stemAll("housedog")); + } + + /** + * Verifies the minimum part length: a split leaving a side shorter than + * COMPOUNDMIN is never taken, although both sides are listed and flagged. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundMinBoundsThePartLength() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 4\n", + "2\ndog/Z\nhouse/Z\n")); + // the left side would be three characters, below the declared minimum of four + Assertions.assertEquals(List.of("doghouse"), stemmer.stemAll("doghouse")); + } + + /** + * Verifies the NEEDAFFIX flag on entries: a virtual stem exists only to be affixed, + * the linking forms of the published German dictionary being the model, so its bare + * form is no analysis of itself while its affixed forms still reduce to it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testNeedAffixEntryIsNoStandaloneAnalysis() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "NEEDAFFIX h\nSFX S Y 1\nSFX S 0 s .\nSFX K Y 1\nSFX K 0 k .\n", + "2\nlink/hS\nlin/K\n")); + // the virtual entry no longer explains the bare form; only the k analysis remains + Assertions.assertEquals(List.of("lin"), stemmer.stemAll("link")); + // affixed, the virtual stem is exactly what the s removal lands on + Assertions.assertEquals(List.of("link"), stemmer.stemAll("links")); + } + + /** + * Verifies NEEDAFFIX against homonyms: the flag blocks one entry's flag set, not + * the word, so a second listing without the flag keeps the bare form valid. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testNeedAffixHomonymKeepsTheBareWord() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "NEEDAFFIX h\nSFX S Y 1\nSFX S 0 s .\n", + "2\nlink/hS\nlink\n")); + Assertions.assertEquals(List.of("link"), stemmer.stemAll("link")); + } + + /** + * Verifies the historical PSEUDOROOT alias, the directive's name before hunspell + * renamed it to NEEDAFFIX; older dictionaries still declare it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testPseudoRootIsNeedAffixByItsOldName() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "PSEUDOROOT h\nSFX S Y 1\nSFX S 0 s .\nSFX K Y 1\nSFX K 0 k .\n", + "2\nlink/hS\nlin/K\n")); + Assertions.assertEquals(List.of("lin"), stemmer.stemAll("link")); + } + + /** + * Verifies the NEEDAFFIX flag on affix rules: a rule carrying the flag among its + * continuation classes makes a form that still needs another affix, so its + * single-removal analysis is suppressed while a twofold removal, whose inner affix + * is the further one required, still reports the stem. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testNeedAffixOnAnAffixRequiresAnotherAffix() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "NEEDAFFIX h", + "SFX A Y 1", + "SFX A 0 er/hB .", + "SFX B Y 1", + "SFX B 0 s .", + ""), + "1\nwork/A\n")); + // work + er alone is virtual, so worker has no analysis and passes through + Assertions.assertEquals(List.of("worker"), stemmer.stemAll("worker")); + // work + er + s is complete; the twofold removal reaches the listed stem + Assertions.assertEquals(List.of("work"), stemmer.stemAll("workers")); + } + + /** + * Verifies that a cross-product analysis satisfies an affix's NEEDAFFIX marker: + * the prefix is the further affix the marked suffix requires, mirroring how + * hunspell accepts a prefix plus a needs-affix suffix together. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCrossProductSatisfiesNeedAffixOnTheSuffix() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "NEEDAFFIX h", + "PFX P Y 1", + "PFX P 0 un .", + "SFX A Y 1", + "SFX A 0 er/h .", + ""), + "1\nwork/AP\n")); + Assertions.assertEquals(List.of("worker"), stemmer.stemAll("worker")); + Assertions.assertEquals(List.of("work"), stemmer.stemAll("unworker")); + } + + /** + * Verifies the ONLYINCOMPOUND flag: an entry carrying it appears only inside + * compounds, the ordinal parts of the published US English dictionary being the + * model, so neither its bare form nor its affixed forms are standalone analyses, + * while compound decomposition may still use it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testOnlyInCompoundEntrySupportsNoStandaloneAnalyses() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "ONLYINCOMPOUND c\nCOMPOUNDFLAG Z\nCOMPOUNDMIN 3\nSFX S Y 1\nSFX S 0 s .\n", + "3\npart/cSZ\nhouse/Z\nwalk/S\n")); + // the affix analysis is suppressed because part's only flag set is compound-only + Assertions.assertEquals(List.of("parts"), stemmer.stemAll("parts")); + Assertions.assertEquals(List.of("part"), stemmer.stemAll("part")); + // inside a compound the entry serves exactly its declared purpose + Assertions.assertEquals(List.of("part", "house"), stemmer.stemAll("parthouse")); + Assertions.assertEquals(List.of("walk"), stemmer.stemAll("walks")); + } + + /** + * Verifies the FORBIDDENWORD flag: an entry carrying it is listed to be blocked, + * so it supports no analysis and no compound part. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testForbiddenWordSupportsNothing() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "FORBIDDENWORD w\nCOMPOUNDFLAG Z\nCOMPOUNDMIN 3\nSFX S Y 1\nSFX S 0 s .\n", + "3\nfoo/wSZ\nhouse/Z\nbar/S\n")); + Assertions.assertEquals(List.of("foo"), stemmer.stemAll("foo")); + Assertions.assertEquals(List.of("foos"), stemmer.stemAll("foos")); + Assertions.assertEquals(List.of("foohouse"), stemmer.stemAll("foohouse")); + Assertions.assertEquals(List.of("bar"), stemmer.stemAll("bars")); + } + + /** The circumfix fixture: the German {@code ge...t} participle in miniature. */ + private static final String CIRCUMFIX_AFFIX = String.join("\n", + "CIRCUMFIX f", + "PFX G Y 1", + "PFX G 0 ge/f .", + "SFX T Y 1", + "SFX T en et/f en", + "PFX U Y 1", + "PFX U 0 un .", + "SFX S Y 1", + "SFX S 0 s .", + ""); + + /** + * Verifies the CIRCUMFIX flag: two marked halves analyze together and neither + * analyzes alone, so the participle reduces to its verb while the half-applied + * forms stay unexplained. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCircumfixHalvesOnlyAnalyzeTogether() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + CIRCUMFIX_AFFIX, "1\narbeiten/GT\n")); + Assertions.assertEquals(List.of("arbeiten"), stemmer.stemAll("gearbeitet")); + // the suffix half alone is no word, although the stem carries its flag + Assertions.assertEquals(List.of("arbeitet"), stemmer.stemAll("arbeitet")); + // the prefix half alone is no word either + Assertions.assertEquals(List.of("gearbeiten"), stemmer.stemAll("gearbeiten")); + } + + /** + * Verifies that circumfixing rejects mixed pairs: a marked half never combines + * with an unmarked affix of the other kind, in either direction, while a fully + * unmarked cross-product in the same dictionary still analyzes. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCircumfixRejectsMixedPairs() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + CIRCUMFIX_AFFIX, "2\narbeiten/GTUS\nlauf/US\n")); + // unmarked prefix with the marked suffix half + Assertions.assertEquals(List.of("unarbeitet"), stemmer.stemAll("unarbeitet")); + // the marked prefix half with an unmarked suffix + Assertions.assertEquals(List.of("gearbeitens"), stemmer.stemAll("gearbeitens")); + // both halves marked still analyze beside the rejected mixtures + Assertions.assertEquals(List.of("arbeiten"), stemmer.stemAll("gearbeitet")); + // a fully unmarked cross-product is untouched by the circumfix declaration + Assertions.assertEquals(List.of("lauf"), stemmer.stemAll("unlaufs")); + } + + /** + * Verifies decomposition beyond two parts: the positional flags admit a begin, a + * middle, and an end part, a part fit only for the middle neither opens nor closes, + * and repeated middles fold into the reported set. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundMiddleAdmitsInnerParts() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDBEGIN B\nCOMPOUNDMIDDLE M\nCOMPOUNDEND E\nCOMPOUNDMIN 3\n", + "3\ndog/B\ncat/M\nhouse/E\n")); + Assertions.assertEquals(List.of("dog", "cat", "house"), + stemmer.stemAll("dogcathouse")); + Assertions.assertEquals(List.of("dog", "cat", "house"), + stemmer.stemAll("dogcatcathouse")); + Assertions.assertEquals(List.of("dog", "house"), stemmer.stemAll("doghouse")); + // cat holds only the middle flag, so it neither opens nor closes + Assertions.assertEquals(List.of("cathouse"), stemmer.stemAll("cathouse")); + Assertions.assertEquals(List.of("dogcat"), stemmer.stemAll("dogcat")); + } + + /** + * Verifies COMPOUNDWORDMAX: a decomposition needing more parts than declared is + * rejected while one within the bound still analyzes. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundWordMaxBoundsThePartCount() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDBEGIN B\nCOMPOUNDMIDDLE M\nCOMPOUNDEND E\nCOMPOUNDMIN 3\n" + + "COMPOUNDWORDMAX 2\n", + "3\ndog/B\ncat/M\nhouse/E\n")); + Assertions.assertEquals(List.of("dog", "house"), stemmer.stemAll("doghouse")); + Assertions.assertEquals(List.of("dogcathouse"), stemmer.stemAll("dogcathouse")); + } + + /** + * Verifies affixed compound parts, the German linking form being the model: a part + * is its entry plus one suffix whose continuation classes position the derived form + * and permit it at the internal boundary, and the reported analysis is the entry, + * not the linking form. The lowercase interior spelling of a capitalized entry is + * found through the part's uppercased retry. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testLinkingSuffixJoinsCompoundParts() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "COMPOUNDBEGIN x", + "COMPOUNDEND z", + "COMPOUNDPERMITFLAG c", + "COMPOUNDMIN 2", + "SFX j Y 1", + "SFX j 0 s/xc .", + ""), + "2\nAbbildung/j\nVerzeichnis/z\n")); + Assertions.assertEquals(List.of("Abbildung", "Verzeichnis"), + stemmer.stemAll("Abbildungsverzeichnis")); + // without the linking s the first part has no admitting reading + Assertions.assertEquals(List.of("Abbildungverzeichnis"), + stemmer.stemAll("Abbildungverzeichnis")); + } + + /** + * Verifies zero-suffix part positioning, the pattern the published German + * dictionary uses: a virtual stem enters compounds through a rule that adds no + * material but whose continuation classes carry the positional and permit flags. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testZeroSuffixPositionsAVirtualStemInCompounds() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "NEEDAFFIX h", + "COMPOUNDBEGIN x", + "COMPOUNDEND z", + "COMPOUNDPERMITFLAG c", + "COMPOUNDMIN 3", + "SFX j Y 1", + "SFX j 0 0/xc .", + ""), + "2\nfugen/hj\nwerk/z\n")); + Assertions.assertEquals(List.of("fugen", "werk"), stemmer.stemAll("fugenwerk")); + // the virtual stem alone is still no word + Assertions.assertEquals(List.of("fugen"), stemmer.stemAll("fugen")); + } + + /** + * Verifies COMPOUNDFORBIDFLAG: an affixed form whose rule carries the flag stays + * out of compounds although its positioning otherwise admits it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCompoundForbidFlagBarsAnAffixedPart() throws IOException { + final String words = "2\ndog/ZS\nhouse/Z\n"; + final HunspellStemmer barred = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDPERMITFLAG c\nCOMPOUNDFORBIDFLAG F\nCOMPOUNDMIN 3\n" + + "SFX S Y 1\nSFX S 0 s/cF .\n", + words)); + Assertions.assertEquals(List.of("dogshouse"), barred.stemAll("dogshouse")); + final HunspellStemmer allowed = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDPERMITFLAG c\nCOMPOUNDMIN 3\n" + + "SFX S Y 1\nSFX S 0 s/c .\n", + words)); + Assertions.assertEquals(List.of("dog", "house"), allowed.stemAll("dogshouse")); + } + + /** + * Verifies that an affix without the permit flag keeps off internal boundaries: a + * suffixed reading fits the last part but not an earlier one. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testAffixWithoutPermitFlagStaysAtTheEdge() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\nSFX S Y 1\nSFX S 0 s .\n", + "2\ndog/ZS\nhouse/ZS\n")); + // the suffix closes the word, so the last part may carry it + Assertions.assertEquals(List.of("dog", "house"), stemmer.stemAll("doghouses")); + // an internal suffix without the permit flag blocks the split + Assertions.assertEquals(List.of("dogshouse"), stemmer.stemAll("dogshouse")); + } + + /** + * Verifies CHECKCOMPOUNDDUP: a part must not repeat its left neighbor, while the + * same dictionary without the declaration accepts the repetition. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCheckCompoundDupForbidsRepeatedParts() throws IOException { + final String words = "2\ndog/Z\nhouse/Z\n"; + final HunspellStemmer checked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\nCHECKCOMPOUNDDUP\n", words)); + Assertions.assertEquals(List.of("dogdoghouse"), checked.stemAll("dogdoghouse")); + final HunspellStemmer unchecked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\n", words)); + Assertions.assertEquals(List.of("dog", "house"), unchecked.stemAll("dogdoghouse")); + } + + /** + * Verifies CHECKCOMPOUNDCASE: an uppercase character on either side of a junction + * forbids the split, while the same dictionary without the declaration accepts it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCheckCompoundCaseForbidsUppercaseJunctions() throws IOException { + final String words = "2\ndog/Z\nHouse/Z\n"; + final HunspellStemmer checked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\nCHECKCOMPOUNDCASE\n", words)); + Assertions.assertEquals(List.of("dogHouse"), checked.stemAll("dogHouse")); + final HunspellStemmer unchecked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\n", words)); + Assertions.assertEquals(List.of("dog", "House"), unchecked.stemAll("dogHouse")); + } + + /** + * Verifies CHECKCOMPOUNDTRIPLE: the same character three times in a row across a + * junction forbids the split, while the same dictionary without the declaration + * accepts it. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testCheckCompoundTripleForbidsTripleLetters() throws IOException { + final String words = "2\nbell/Z\nlow/Z\n"; + final HunspellStemmer checked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\nCHECKCOMPOUNDTRIPLE\n", words)); + Assertions.assertEquals(List.of("belllow"), checked.stemAll("belllow")); + final HunspellStemmer unchecked = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 3\n", words)); + Assertions.assertEquals(List.of("bell", "low"), unchecked.stemAll("belllow")); + } + + /** + * Verifies that a listed forbidden word never decomposes: the dictionary blocks + * one specific ill-formed compound while its parts stay productive elsewhere. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testForbiddenEntryBlocksItsDecomposition() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "FORBIDDENWORD w\nCOMPOUNDFLAG Z\nCOMPOUNDMIN 3\n", + "4\ndog/Z\nhouse/Z\ncat/Z\ndoghouse/w\n")); + Assertions.assertEquals(List.of("doghouse"), stemmer.stemAll("doghouse")); + Assertions.assertEquals(List.of("cat", "house"), stemmer.stemAll("cathouse")); + } + + /** + * Verifies that directives outside the affix-stemming subset do not prevent use of + * the rules this implementation supports. + * + * @param line The affix file line. + */ + @ParameterizedTest + @ValueSource(strings = { + "ICONV 1", + "OCONV 1", + "COMPLEXPREFIXES", + "COMPOUNDRULE 1", + "COMPOUNDMORESUFFIXES", + "COMPOUNDROOT R", + "CHECKCOMPOUNDREP", + "SIMPLIFIEDTRIPLE", + "CHECKCOMPOUNDPATTERN 1", + "FORCEUCASE U", + "COMPOUNDSYLLABLE 6 aeiou", + "SYLLABLENUM ABC", + "LANG tr", + "CHECKSHARPS", + "BREAK 1", + "FORBIDWARN", + "IGNORE x", + "KEEPCASE k" + }) + void testUnsupportedDirectiveDoesNotBlockSupportedRules(String line) + throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + line + "\nSFX A Y 1\nSFX A 0 s .\n", "1\ndog/A\n")); + + Assertions.assertEquals("dog", stemmer.stem("dogs").toString()); + } + + /** + * Verifies that a full-strip suffix rule is not applied unless the affix file + * declares {@code FULLSTRIP}. Hunspell applies a rule whose strip string consumes + * the whole stem only under that declaration; without it, inventing a stem from + * such a rule contradicts the fail-closed loader policy. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testFullStripRuleRequiresFullStripDirective() throws IOException { + final String rule = "SFX A Y 1\nSFX A work ed .\n"; + final String words = "1\nwork/A\n"; + Assertions.assertEquals(List.of("work"), + new HunspellStemmer(load("FULLSTRIP\n" + rule, words)).stemAll("ed")); + Assertions.assertEquals(List.of("ed"), + new HunspellStemmer(load(rule, words)).stemAll("ed")); + } + + /** + * Verifies that {@link HunspellDictionary#load(InputStream, InputStream)} rejects an + * affix stream larger than {@link HunspellDictionary#MAX_STREAM_BYTES}. + */ + @Test + void testLoadRejectsOversizedAffixStream() { + final IOException e = Assertions.assertThrows(IOException.class, + () -> HunspellDictionary.load( + filledStream(HunspellDictionary.MAX_STREAM_BYTES + 1), + new ByteArrayInputStream("1\nlock\n".getBytes(StandardCharsets.UTF_8)))); + Assertions.assertEquals( + "affix stream size exceeds safe limit of " + HunspellDictionary.MAX_STREAM_BYTES, + e.getMessage()); + } + + /** + * Verifies that {@link HunspellDictionary#load(InputStream, InputStream)} rejects a + * dictionary stream larger than {@link HunspellDictionary#MAX_STREAM_BYTES}. + */ + @Test + void testLoadRejectsOversizedDictionaryStream() { + final byte[] affix = "SET UTF-8\n".getBytes(StandardCharsets.UTF_8); + final IOException e = Assertions.assertThrows(IOException.class, + () -> HunspellDictionary.load( + new ByteArrayInputStream(affix), + filledStream(HunspellDictionary.MAX_STREAM_BYTES + 1))); + Assertions.assertEquals( + "dictionary stream size exceeds safe limit of " + + HunspellDictionary.MAX_STREAM_BYTES, + e.getMessage()); + } + + /** + * Pins the inclusive stream-byte ceiling: a stream of exactly {@code limit} bytes + * succeeds, and {@code limit + 1} fails. Uses a small limit so the test does not + * allocate the production ceiling. + * + * @throws IOException Thrown if reading the in-bound stream fails. + */ + @Test + void testBoundedReadCeilingIsInclusive() throws IOException { + final int limit = 64; + final byte[] bytes = HunspellDictionary.readBounded(filledStream(limit), limit, + "affix stream"); + Assertions.assertEquals(limit, bytes.length); + final IOException e = Assertions.assertThrows(IOException.class, + () -> HunspellDictionary.readBounded(filledStream(limit + 1), limit, + "affix stream")); + Assertions.assertEquals( + "affix stream size exceeds safe limit of " + limit, e.getMessage()); + } + + /** + * Pins affix conditions and boundary bucketing to code points, matching FLAG UTF-8: + * a condition of two dots needs two code points, so a stem that is one supplementary + * character must not match, while a one-dot condition and a supplementary affix + * character still analyze. + * + * @throws IOException Thrown if a fixture fails to load. + */ + @Test + void testAffixConditionAndBoundaryUseCodePoints() throws IOException { + final HunspellStemmer twoDots = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 s ..\n", + "1\n\uD83D\uDE00/X\n")); + Assertions.assertEquals("\uD83D\uDE00s", twoDots.stem("\uD83D\uDE00s").toString()); + + final HunspellStemmer oneDot = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 s .\n", + "1\n\uD83D\uDE00/X\n")); + Assertions.assertEquals("\uD83D\uDE00", oneDot.stem("\uD83D\uDE00s").toString()); + + final HunspellStemmer emojiSuffix = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 \uD83D\uDE00 .\n", + "1\nwalk/X\n")); + Assertions.assertEquals("walk", emojiSuffix.stem("walk\uD83D\uDE00").toString()); + + final HunspellStemmer classCondition = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 s [\uD83D\uDE00]\n", + "1\nwalk\uD83D\uDE00/X\n")); + Assertions.assertEquals("walk\uD83D\uDE00", + classCondition.stem("walk\uD83D\uDE00s").toString()); + } + + /** + * Returns a stream of {@code size} zero bytes. + * + * @param size The number of bytes the stream yields. + * @return The stream. Never {@code null}. + */ + private static InputStream filledStream(int size) { + return new InputStream() { + private int remaining = size; + + @Override + public int read() { + if (remaining <= 0) { + return -1; + } + remaining--; + return 0; + } + + @Override + public int read(byte[] buffer, int offset, int length) { + if (remaining <= 0) { + return -1; + } + final int n = Math.min(length, remaining); + Arrays.fill(buffer, offset, offset + n, (byte) 0); + remaining -= n; + return n; + } + }; + } + + /** + * Verifies that a cosmetic unsupported directive such as {@code REP} is skipped so + * the dictionary still loads. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testCosmeticUnsupportedDirectiveIsSkipped() throws IOException { + final HunspellDictionary dictionary = load("REP 1\nREP alot a lot\n", "1\nlock\n"); + Assertions.assertNotNull(dictionary.lookup("lock")); + } + + /** + * Verifies that an {@code AF} table applies to affix rules above the table. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testAliasTableAppliesToAffixesThatPrecedeIt() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "FLAG num", + "SFX 1 Y 1", + "SFX 1 0 er/1 .", + "SFX 2 Y 1", + "SFX 2 0 s .", + "AF 2", + "AF 2", + "AF 1", + ""), + "1\nkind/2\n")); + + Assertions.assertEquals(List.of("kind"), stemmer.stemAll("kinders")); + } + + /** + * Verifies that {@code COMPOUNDMIN} counts Unicode code points. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testCompoundMinCountsSupplementaryCharactersOnce() throws IOException { + final String first = "\uD840\uDC00"; + final String rightPart = "\uD840\uDC01"; + final String words = "2\n" + first + "/Z\n" + rightPart + "/Z\n"; + + final HunspellStemmer minimumTwo = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 2\n", words)); + Assertions.assertEquals(List.of(first + rightPart), minimumTwo.stemAll(first + rightPart)); + + final HunspellStemmer minimumOne = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 1\n", words)); + Assertions.assertEquals(List.of(first, rightPart), minimumOne.stemAll(first + rightPart)); + } + + /** + * Verifies cross-product analysis with stacked suffixes. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testCrossProductSupportsTwofoldSuffixes() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "PFX U Y 1", + "PFX U 0 un .", + "SFX A Y 1", + "SFX A 0 s/B .", + "SFX B Y 1", + "SFX B 0 bar .", + ""), + "1\nfoo/AU\n")); + + Assertions.assertEquals("foo", stemmer.stem("unfoosbar").toString()); + } + + /** Verifies that an unrecognized directive does not block supported affix rules. */ + @Test + void testUnknownAffixDirectiveIsSkipped() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "UNRECOGNIZED value\nSFX A Y 1\nSFX A 0 s .\n", "1\ndog/A\n")); + + Assertions.assertEquals("dog", stemmer.stem("dogs").toString()); + } + + /** Verifies validation of the {@code AF} count line. */ + @ParameterizedTest + @ValueSource(strings = {"count-mismatch", "malformed", "negative"}) + void testAliasTableCountIsValidated(String fixture) { + final String affix; + final String message; + switch (fixture) { + case "count-mismatch" -> { + affix = "AF 2\nAF A\n"; + message = "AF header specifies 2 aliases but found 1"; + } + case "malformed" -> { + affix = "AF count\n"; + message = "malformed AF at line 1"; + } + case "negative" -> { + affix = "AF -1\n"; + message = "negative AF count at line 1"; + } + default -> throw new AssertionError(fixture); + } + final IOException exception = Assertions.assertThrows(IOException.class, + () -> load(affix, "0\n")); + + Assertions.assertEquals(message, exception.getMessage()); + } + + /** + * Verifies the numeric flag range. + * + * @param flag The invalid numeric flag. + */ + @ParameterizedTest + @ValueSource(strings = {"-1", "0", "65001"}) + void testNumericFlagOutsideRangeIsRejected(String flag) { + final IOException e = Assertions.assertThrows(IOException.class, + () -> load("FLAG num\n", "1\nword/" + flag + "\n")); + + Assertions.assertEquals("numeric flag outside 1..65000 at line 2: " + flag, + e.getMessage()); + } + + /** + * Verifies that an identity suffix can license an outer suffix. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testZeroMaterialInnerSuffixLicensesOuterSuffix() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "SFX A Y 1", + "SFX A 0 0/B .", + "SFX B Y 1", + "SFX B 0 baz .", + ""), + "1\nbar/A\n")); + + Assertions.assertEquals(List.of("bar"), stemmer.stemAll("barbaz")); + } + + /** + * Verifies that both {@code NEEDAFFIX} markers cannot satisfy one another. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testCrossProductNeedAffixMarkersDoNotSatisfyEachOther() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "NEEDAFFIX X", + "PFX P Y 1", + "PFX P 0 pseudo/X .", + "SFX A Y 1", + "SFX A 0 pseudo/X .", + ""), + "1\nfoo/AP\n")); + + Assertions.assertEquals(List.of("pseudofoopseudo"), + stemmer.stemAll("pseudofoopseudo")); + } + + /** + * Verifies that a slash at the start of a dictionary entry is word text. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testLeadingSlashIsPartOfWord() throws IOException { + final HunspellStemmer slashWord = new HunspellStemmer(load( + "SFX X Y 1\nSFX X 0 s .\n", + "2\n/foo\n/foo/X\n")); + + Assertions.assertEquals(List.of("/foo"), slashWord.stemAll("/foos")); + } + + /** + * Verifies that an identity continuation completes a virtual suffix. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testZeroMaterialContinuationCompletesVirtualSuffix() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "PSEUDOROOT X", + "SFX A Y 1", + "SFX A 0 0 .", + "SFX C Y 1", + "SFX C 0 baz/XA .", + ""), + "1\nbar/C\n")); + + Assertions.assertEquals(List.of("bar"), stemmer.stemAll("barbaz")); + } + + /** + * Verifies cross-product licensing from either member's continuation flags. + * + * @param licensingRule The member that identifies the partner. + * @throws IOException Thrown if a fixture fails to load. + */ + @ParameterizedTest + @ValueSource(strings = {"prefix", "suffix"}) + void testContinuationFlagLicensesCrossProductPartner(String licensingRule) + throws IOException { + final boolean prefixLicenses = "prefix".equals(licensingRule); + final String affix = prefixLicenses + ? String.join("\n", + "PFX P Y 1", + "PFX P 0 un/S .", + "SFX S Y 1", + "SFX S 0 s .", + "") + : String.join("\n", + "PFX P Y 1", + "PFX P 0 un .", + "SFX R Y 1", + "SFX R 0 able/P .", + ""); + final String words = prefixLicenses ? "1\nlock/P\n" : "1\ndrink/R\n"; + final String surface = prefixLicenses ? "unlocks" : "undrinkable"; + final String expected = prefixLicenses ? "lock" : "drink"; + + final HunspellStemmer stemmer = new HunspellStemmer(load(affix, words)); + + Assertions.assertEquals(expected, stemmer.stem(surface).toString()); + } + + /** + * Verifies that {@code FLAG} applies to affix rules above the declaration. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testFlagModeAppliesToRulesThatPrecedeTheDeclaration() throws IOException { + final HunspellStemmer lateFlagMode = new HunspellStemmer(load(String.join("\n", + "SFX 1 Y 1", + "SFX 1 0 s .", + "FLAG num", + ""), String.join("\n", "1", "dog/1", ""))); + + Assertions.assertEquals("dog", lateFlagMode.stem("dogs").toString()); + } + + /** + * Verifies that {@code COMPOUNDFORBIDFLAG} rejects nonfinal dictionary entries. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testCompoundForbidFlagBarsDictionaryEntryBeforeEnd() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDFORBIDFLAG F\nCOMPOUNDMIN 3\n", + "2\ndog/ZF\nhouse/Z\n")); + + Assertions.assertEquals(List.of("doghouse"), stemmer.stemAll("doghouse")); + Assertions.assertEquals(List.of("house", "dog"), stemmer.stemAll("housedog")); + } + + /** + * Verifies that a forbidden surface form is not analyzed through an affix rule. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testForbiddenSurfaceOverridesAffixAnalysis() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "FORBIDDENWORD X\nSFX A Y 1\nSFX A 0 s .\n", + "2\nfoo/A\nfoos/X\n")); + + Assertions.assertEquals(List.of("foos"), stemmer.stemAll("foos")); + } + + /** + * Verifies that a forbidden homonym blocks affix analysis even when another entry + * for the same surface is valid as a standalone entry. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testForbiddenHomonymOverridesStandaloneEntry() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + "FORBIDDENWORD X\nSFX A Y 1\nSFX A 0 s .\n", + "3\nfoo/A\nfoos\nfoos/X\n")); + + Assertions.assertEquals(List.of("foos"), stemmer.stemAll("foos")); + } + + /** + * Verifies that malformed UTF-8 is rejected in semantic affix content and in the + * dictionary file. + * + * @param file The malformed input file. + */ + @ParameterizedTest + @ValueSource(strings = {"affix", "dictionary"}) + void testMalformedFileEncodingIsRejected(String file) { + final byte[] affix = "SET UTF-8\n".getBytes(StandardCharsets.UTF_8); + final byte[] words = Arrays.copyOf("1\n".getBytes(StandardCharsets.UTF_8), 3); + words[2] = TRUNCATED_UTF8_LEAD_BYTE; + final byte[] affixPrefix = "SET UTF-8\nSFX A Y 1\nSFX A 0 " + .getBytes(StandardCharsets.UTF_8); + final byte[] malformedAffix = Arrays.copyOf(affixPrefix, affixPrefix.length + 1); + malformedAffix[malformedAffix.length - 1] = TRUNCATED_UTF8_LEAD_BYTE; + final byte[] selectedAffix = "affix".equals(file) ? malformedAffix : affix; + final byte[] selectedWords = "dictionary".equals(file) + ? words : "0\n".getBytes(StandardCharsets.UTF_8); + + final IOException exception = Assertions.assertThrows(IOException.class, + () -> HunspellDictionary.load(new ByteArrayInputStream(selectedAffix), + new ByteArrayInputStream(selectedWords))); + + Assertions.assertEquals(file + " stream is not valid UTF-8", exception.getMessage()); + } + + /** + * Verifies that invalid bytes in a comment do not prevent loading an otherwise valid + * UTF-8 affix file. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testMalformedCommentEncodingIsIgnored() throws IOException { + final byte[] prefix = "SET UTF-8\n# ".getBytes(StandardCharsets.UTF_8); + final byte[] affix = Arrays.copyOf(prefix, prefix.length + 1); + affix[affix.length - 1] = TRUNCATED_UTF8_LEAD_BYTE; + + final HunspellDictionary dictionary = HunspellDictionary.load( + new ByteArrayInputStream(affix), + new ByteArrayInputStream("1\ndog\n".getBytes(StandardCharsets.UTF_8))); + + Assertions.assertNotNull(dictionary.lookup("dog")); + } + + /** + * Verifies raw one-byte flags in a file where word text uses UTF-8. + * + * @param representation Whether the dictionary entry uses an alias or a direct flag. + * @throws IOException Thrown if the fixture fails to load. + */ + @ParameterizedTest + @ValueSource(strings = {"alias", "direct"}) + void testDefaultFlagModePreservesRawBytesInUtf8File(String representation) + throws IOException { + final boolean alias = "alias".equals(representation); + final String aliasTable = alias ? "AF 1\nAF \u00D7\n" : ""; + final byte[] affix = ("SET UTF-8\n" + aliasTable + + "SFX \u00D7 Y 1\nSFX \u00D7 0 s .\n") + .getBytes(StandardCharsets.ISO_8859_1); + final byte[] words = (alias ? "1\ndog/1\n" : "1\ndog/\u00D7\n") + .getBytes(alias ? StandardCharsets.UTF_8 : StandardCharsets.ISO_8859_1); + final HunspellStemmer stemmer = new HunspellStemmer(HunspellDictionary.load( + new ByteArrayInputStream(affix), new ByteArrayInputStream(words))); + + Assertions.assertEquals("dog", stemmer.stem("dogs").toString()); + } + + /** + * Verifies {@code AF} references in affix continuation fields. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testAffixContinuationFlagsResolveThroughTheAliasTable() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "AF 2", + "AF AB", + "AF A", + "SFX A Y 1", + "SFX A 0 x .", + "SFX B Y 1", + "SFX B 0 y/2 .", + ""), + "1\nfoo/1\n")); + + Assertions.assertEquals(List.of("foo"), stemmer.stemAll("fooyx")); + } + + /** + * Verifies that a stacked suffix satisfies {@code NEEDAFFIX} in a cross-product. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testStackedSuffixSatisfiesNeedAffixWithinCrossProduct() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "NEEDAFFIX X", + "PFX P Y 1", + "PFX P 0 pseudo/X .", + "SFX A Y 1", + "SFX A 0 pseudo/XB .", + "SFX B Y 1", + "SFX B 0 bar/X .", + ""), + "1\nfoo/AP\n")); + + Assertions.assertEquals(List.of("foo"), + stemmer.stemAll("pseudofoopseudobar")); + } + + /** + * Verifies cross-product analysis with an identity inner suffix. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testZeroMaterialInnerSuffixSupportsCrossProduct() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "PFX P Y 1", + "PFX P 0 un .", + "SFX A Y 1", + "SFX A 0 0/B .", + "SFX B Y 1", + "SFX B 0 baz .", + ""), + "1\nbar/AP\n")); + + Assertions.assertEquals(List.of("bar"), stemmer.stemAll("unbarbaz")); + } + + /** + * Verifies that compound limits cannot be negative. + * + * @param directive The compound limit directive. + */ + @ParameterizedTest + @ValueSource(strings = {"COMPOUNDMIN", "COMPOUNDWORDMAX"}) + void testNegativeCompoundLimitIsRejected(String directive) { + final IOException e = Assertions.assertThrows(IOException.class, + () -> load(directive + " -1\n", "0\n")); + + Assertions.assertEquals("negative " + directive + " at line 1", e.getMessage()); + } + + /** Verifies that {@code COMPOUNDMIN} cannot overflow the doubled length check. */ + @Test + void testCompoundMinAboveSafeRangeIsRejected() { + final IOException exception = Assertions.assertThrows(IOException.class, + () -> load("COMPOUNDMIN 1073741824\n", "0\n")); + + Assertions.assertEquals("COMPOUNDMIN exceeds 1073741823 at line 1", + exception.getMessage()); + } + + /** + * Verifies cross-product analysis with an identity prefix. + * + * @throws IOException Thrown if the fixture fails to load. + */ + @Test + void testZeroMaterialPrefixParticipatesInCrossProduct() throws IOException { + final HunspellStemmer stemmer = new HunspellStemmer(load( + String.join("\n", + "PFX P Y 1", + "PFX P 0 0/S .", + "SFX S Y 1", + "SFX S 0 s .", + ""), + "1\nroot/P\n")); + + Assertions.assertEquals(List.of("root"), stemmer.stemAll("roots")); + } + + /** + * Verifies supplementary characters in compound boundary checks. + * + * @param check The boundary check to exercise. + * @throws IOException Thrown if a fixture fails to load. + */ + @ParameterizedTest + @ValueSource(strings = {"case", "triple"}) + void testCompoundBoundaryChecksUseCodePoints(String check) throws IOException { + final boolean triple = "triple".equals(check); + final String codePoint = triple ? "\uD840\uDC00" : "\uD801\uDC00"; + final String words = triple + ? "2\na" + codePoint + codePoint + "/Z\n" + codePoint + "b/Z\n" + : "2\na/Z\n" + codePoint + "b/Z\n"; + final String declaration = triple ? "CHECKCOMPOUNDTRIPLE" : "CHECKCOMPOUNDCASE"; + final String surface = triple + ? "a" + codePoint + codePoint + codePoint + "b" + : "a" + codePoint + "b"; + final HunspellStemmer stemmer = new HunspellStemmer(load( + "COMPOUNDFLAG Z\nCOMPOUNDMIN 1\n" + declaration + "\n", words)); + + Assertions.assertEquals(List.of(surface), stemmer.stemAll(surface)); + } + +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/SingleSpaceTokens.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/SingleSpaceTokens.java new file mode 100644 index 0000000000..4f7beb4643 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/SingleSpaceTokens.java @@ -0,0 +1,90 @@ +/* + * 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.termvector; + +import java.util.ArrayList; +import java.util.List; + +import opennlp.tools.document.Annotation; +import opennlp.tools.tokenize.Tokenizer; +import opennlp.tools.util.Span; + +/** + * Deterministic single-space tokenization shared by the term vector tests: splits on + * single space characters and keeps all other characters, including sentence-final + * periods, attached to their token. Runs of spaces yield no empty tokens, so every + * expected span follows directly from the input text. + */ +final class SingleSpaceTokens { + + /** + * A {@link Tokenizer} view of {@link #spans(String)}. Only the span-producing method + * is implemented because the annotator adapter calls no other method. + */ + static final Tokenizer TOKENIZER = new Tokenizer() { + + @Override + public String[] tokenize(String s) { + throw new UnsupportedOperationException("the adapter only calls tokenizePos"); + } + + @Override + public Span[] tokenizePos(String s) { + return spans(s).toArray(new Span[0]); + } + }; + + private SingleSpaceTokens() { + } + + /** + * Computes the token spans of a text split on single space characters. + * + * @param text The text to split. + * @return One span per token, in text order. + */ + static List spans(String text) { + final List spans = new ArrayList<>(); + int start = -1; + for (int i = 0; i <= text.length(); i++) { + final boolean boundary = i == text.length() || text.charAt(i) == ' '; + if (boundary && start >= 0) { + spans.add(new Span(start, i)); + start = -1; + } else if (!boundary && start < 0) { + start = i; + } + } + return spans; + } + + /** + * Builds a token layer from {@link #spans(String)}, each token valued with its covered + * text. + * + * @param text The text to split. + * @return One annotation per token, in text order. + */ + static List> tokens(String text) { + final List> tokens = new ArrayList<>(); + for (final Span span : spans(text)) { + tokens.add(new Annotation<>(span, text.substring(span.getStart(), span.getEnd()))); + } + return tokens; + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorAnnotatorTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorAnnotatorTest.java new file mode 100644 index 0000000000..33559371f9 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorAnnotatorTest.java @@ -0,0 +1,424 @@ +/* + * 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.termvector; + +import java.io.Serial; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.LayerKey; +import opennlp.tools.document.Layers; +import opennlp.tools.util.Span; +import opennlp.tools.util.normalizer.AlignedText; +import opennlp.tools.util.normalizer.Alignment; +import opennlp.tools.util.normalizer.CharSequenceNormalizer; +import opennlp.tools.util.normalizer.OffsetAwareNormalizer; + +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; + +/** + * Verifies the {@link TermVectorAnnotator} roll-up: term identity with and without a + * normalizer, both recording modes, and the graceful-degradation rules of the document + * pipeline. The token layer is built directly here; the wiring through a + * {@code DocumentAnalyzer} is covered by {@code TermVectorPipelineTest}. + */ +public class TermVectorAnnotatorTest { + + /** + * A deterministic {@link OffsetAwareNormalizer} that collapses every run of + * whitespace to one space and applies a full case fold, including the German eszett + * expansion {@code ß -> ss}. Every edit is recorded in an {@link Alignment}, so the + * expected spans below follow directly from the input text. + */ + private static final class WhitespaceCaseFoldNormalizer implements OffsetAwareNormalizer { + + @Serial + private static final long serialVersionUID = 1L; + + @Override + public CharSequence normalize(CharSequence text) { + return normalizeAligned(text).normalized(); + } + + @Override + public AlignedText normalizeAligned(CharSequence text) { + final String source = text.toString(); + final StringBuilder out = new StringBuilder(source.length()); + final Alignment.Builder alignment = new Alignment.Builder(source.length()); + int i = 0; + while (i < source.length()) { + final char c = source.charAt(i); + if (Character.isWhitespace(c)) { + int end = i; + while (end < source.length() && Character.isWhitespace(source.charAt(end))) { + end++; + } + out.append(' '); + alignment.replace(end - i, 1); + i = end; + } else if (c == 'ß') { + out.append("ss"); + alignment.replace(1, 2); + i++; + } else { + final String folded = String.valueOf(c).toLowerCase(Locale.ROOT); + out.append(folded); + if (folded.length() == 1 && folded.charAt(0) == c) { + alignment.equal(1); + } else { + alignment.replace(1, folded.length()); + } + i++; + } + } + return new AlignedText(source, out.toString(), alignment.build(source.length())); + } + } + + /** + * A deterministic {@link OffsetAwareNormalizer} that deletes every digit and copies + * every other character unchanged, so a token made up of digits alone normalizes to + * the empty string. + */ + private static final class DigitDeletingNormalizer implements OffsetAwareNormalizer { + + @Serial + private static final long serialVersionUID = 1L; + + @Override + public CharSequence normalize(CharSequence text) { + return normalizeAligned(text).normalized(); + } + + @Override + public AlignedText normalizeAligned(CharSequence text) { + final String source = text.toString(); + final StringBuilder out = new StringBuilder(source.length()); + final Alignment.Builder alignment = new Alignment.Builder(source.length()); + for (int i = 0; i < source.length(); i++) { + final char c = source.charAt(i); + if (Character.isDigit(c)) { + alignment.replace(1, 0); + } else { + out.append(c); + alignment.equal(1); + } + } + return new AlignedText(source, out.toString(), alignment.build(source.length())); + } + } + + private static final OffsetAwareNormalizer FOLD = new WhitespaceCaseFoldNormalizer(); + + private static final OffsetAwareNormalizer DROP_DIGITS = new DigitDeletingNormalizer(); + + /** + * A plain, alignment-free case folder for the per-token path: it stands in for the + * shipped normalizers (case fold, NFC, accent fold) that cannot report offsets and + * therefore cannot implement {@link OffsetAwareNormalizer}. + */ + private static final CharSequenceNormalizer PLAIN_LOWER = + text -> text.toString().toLowerCase(Locale.ROOT); + + /** + * A plain normalizer that deletes every digit, so a token made up of digits alone + * normalizes to the empty string on the per-token path. + */ + private static final CharSequenceNormalizer PLAIN_DROP_DIGITS = text -> { + final StringBuilder out = new StringBuilder(text.length()); + for (int i = 0; i < text.length(); i++) { + if (!Character.isDigit(text.charAt(i))) { + out.append(text.charAt(i)); + } + } + return out.toString(); + }; + + private static Document documentWithTokens(String text) { + return Document.of(text).with(Layers.TOKENS, SingleSpaceTokens.tokens(text)); + } + + @Test + void testFullModeGroupsByCoveredTextAsIs() { + final Document document = new TermVectorAnnotator() + .annotate(documentWithTokens("The dog barks. The dog naps.")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(4, vectors.size()); + // The layer preserves first-occurrence order. + assertEquals(new TermVector("The", 2, List.of(new Span(0, 3), new Span(15, 18))), + vectors.get(0).value()); + assertEquals(new TermVector("dog", 2, List.of(new Span(4, 7), new Span(19, 22))), + vectors.get(1).value()); + assertEquals(new TermVector("barks.", 1, List.of(new Span(8, 14))), + vectors.get(2).value()); + assertEquals(new TermVector("naps.", 1, List.of(new Span(23, 28))), + vectors.get(3).value()); + } + + @Test + void testLayerIsDocumentScopedAndCarriesNoAnnotationSpans() { + final Document document = new TermVectorAnnotator() + .annotate(documentWithTokens("The dog barks.")); + assertEquals(LayerKey.Scope.DOCUMENT, TermVectorAnnotator.TERM_VECTORS.scope()); + for (final Annotation vector : document.get(TermVectorAnnotator.TERM_VECTORS)) { + assertNull(vector.span()); + } + } + + @Test + void testScoringOnlyModeOmitsSpans() { + final Document document = new TermVectorAnnotator(TermVectorAnnotator.Mode.SCORING_ONLY) + .annotate(documentWithTokens("The dog barks. The dog naps.")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(4, vectors.size()); + assertEquals(TermVector.count("The", 2), vectors.get(0).value()); + assertEquals(TermVector.count("dog", 2), vectors.get(1).value()); + assertEquals(TermVector.count("barks.", 1), vectors.get(2).value()); + assertEquals(TermVector.count("naps.", 1), vectors.get(3).value()); + for (final Annotation vector : vectors) { + assertTrue(vector.value().spans().isEmpty()); + } + } + + /** + * The offset-fidelity case: whitespace collapse shifts offsets and the eszett case + * fold grows the text, yet every emitted occurrence span must land on the original + * text. {@code "Groß groß GROSS"} normalizes to {@code "gross gross gross"}, so all + * three tokens group under one term while their spans keep pointing at the original + * surface forms. + */ + @Test + void testNormalizationGroupsFoldedTokensWithOriginalOffsets() { + final String text = "Groß groß GROSS"; + final Document document = new TermVectorAnnotator(FOLD) + .annotate(documentWithTokens(text)); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + final TermVector vector = vectors.get(0).value(); + assertEquals("gross", vector.term()); + assertEquals(3, vector.frequency()); + assertEquals(List.of(new Span(0, 4), new Span(6, 10), new Span(12, 17)), vector.spans()); + + // Every occurrence span covers the original surface form, not the normalized one. + final List surfaceForms = + vector.spans().stream().map(s -> s.getCoveredText(text).toString()).toList(); + assertEquals(List.of("Groß", "groß", "GROSS"), surfaceForms); + + // The same spans round-trip through the alignment: a token's normalized span maps + // back to exactly its original span, through both edits. + final AlignedText aligned = FOLD.normalizeAligned(text); + assertEquals("gross gross gross", aligned.normalizedString()); + assertEquals(new Span(0, 4), aligned.toOriginalSpan(0, 5)); + assertEquals(new Span(6, 10), aligned.toOriginalSpan(6, 11)); + assertEquals(new Span(12, 17), aligned.toOriginalSpan(12, 17)); + } + + @Test + void testNormalizationKeepsDistinctTermsApart() { + final String text = "Das große Haus ist groß"; + final Document document = new TermVectorAnnotator(FOLD) + .annotate(documentWithTokens(text)); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(5, vectors.size()); + assertEquals(new TermVector("das", 1, List.of(new Span(0, 3))), vectors.get(0).value()); + assertEquals(new TermVector("grosse", 1, List.of(new Span(5, 10))), vectors.get(1).value()); + assertEquals(new TermVector("haus", 1, List.of(new Span(12, 16))), vectors.get(2).value()); + assertEquals(new TermVector("ist", 1, List.of(new Span(18, 21))), vectors.get(3).value()); + assertEquals(new TermVector("gross", 1, List.of(new Span(23, 27))), vectors.get(4).value()); + } + + @Test + void testScoringOnlyModeWithNormalizerCountsWithoutOffsets() { + final Document document = new TermVectorAnnotator(FOLD, + TermVectorAnnotator.Mode.SCORING_ONLY).annotate(documentWithTokens("Groß groß GROSS")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(TermVector.count("gross", 3), vectors.get(0).value()); + } + + /** + * The general per-token path: a plain {@link CharSequenceNormalizer} defines term + * identity by folding each token's covered text, no alignment involved, while every + * occurrence span stays the token's own span in the original text. This admits the + * folds {@code buildAligned()} rejects (case fold, NFC, accent fold). + */ + @Test + void testPlainNormalizerGroupsFoldedTokensWithOriginalSpans() { + final String text = "Word word WORD"; + final Document document = new TermVectorAnnotator(PLAIN_LOWER) + .annotate(documentWithTokens(text)); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(new TermVector("word", 3, + List.of(new Span(0, 4), new Span(5, 9), new Span(10, 14))), vectors.get(0).value()); + + // Every occurrence span covers the original surface form, not the folded one. + final List surfaceForms = vectors.get(0).value().spans().stream() + .map(s -> s.getCoveredText(text).toString()).toList(); + assertEquals(List.of("Word", "word", "WORD"), surfaceForms); + } + + @Test + void testPlainNormalizerScoringOnlyModeCountsWithoutOffsets() { + final Document document = new TermVectorAnnotator(PLAIN_LOWER, + TermVectorAnnotator.Mode.SCORING_ONLY).annotate(documentWithTokens("Word word WORD")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(TermVector.count("word", 3), vectors.get(0).value()); + } + + /** + * Empty-term omission behaves identically on the per-token path: tokens a plain + * normalizer folds to the empty string are left out of the layer, matching + * {@link #testTokensNormalizedAwayAreOmitted()}. + */ + @Test + void testPlainNormalizerOmitsDeletedTokens() { + final Document document = new TermVectorAnnotator(PLAIN_DROP_DIGITS) + .annotate(documentWithTokens("dog 42 dog 7")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(new TermVector("dog", 2, List.of(new Span(0, 3), new Span(7, 10))), + vectors.get(0).value()); + } + + @Test + void testNullPlainNormalizerIsRejected() { + final CharSequenceNormalizer noNormalizer = null; + assertThrows(IllegalArgumentException.class, () -> new TermVectorAnnotator(noNormalizer)); + assertThrows(IllegalArgumentException.class, + () -> new TermVectorAnnotator(noNormalizer, TermVectorAnnotator.Mode.FULL)); + assertThrows(IllegalArgumentException.class, + () -> new TermVectorAnnotator(PLAIN_LOWER, null)); + } + + /** + * A token the normalizer deletes entirely is omitted from the layer: an empty string + * is no term, and the token layer still accounts for the token itself. + */ + @Test + void testTokensNormalizedAwayAreOmitted() { + final Document document = new TermVectorAnnotator(DROP_DIGITS) + .annotate(documentWithTokens("dog 42 dog 7")); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(new TermVector("dog", 2, List.of(new Span(0, 3), new Span(7, 10))), + vectors.get(0).value()); + } + + /** + * Spans are UTF-16 offsets, so a supplementary-plane character occupies two positions: + * {@code "𝕏 x 𝕏"} tokenizes to spans of width two around the surrogate pairs, and both + * occurrences group under one term whose spans still cover the original text exactly. + */ + @Test + void testSupplementaryPlaneTokensKeepUtf16Offsets() { + final String text = "𝕏 x 𝕏"; + final Document document = new TermVectorAnnotator() + .annotate(documentWithTokens(text)); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(2, vectors.size()); + assertEquals(new TermVector("𝕏", 2, List.of(new Span(0, 2), new Span(5, 7))), + vectors.get(0).value()); + assertEquals(new TermVector("x", 1, List.of(new Span(3, 4))), vectors.get(1).value()); + for (final Annotation vector : vectors) { + for (final Span span : vector.value().spans()) { + assertEquals(vector.value().term(), span.getCoveredText(text).toString()); + } + } + } + + /** + * A document whose every token normalizes to the empty string yields the layer + * present but empty, the same graceful degradation as an empty token layer. + */ + @Test + void testAllTokensNormalizedAwayYieldPresentButEmptyLayer() { + final Document document = new TermVectorAnnotator(DROP_DIGITS) + .annotate(documentWithTokens("42 7")); + assertTrue(document.layers().contains(TermVectorAnnotator.TERM_VECTORS)); + assertTrue(document.get(TermVectorAnnotator.TERM_VECTORS).isEmpty()); + } + + @Test + void testEmptyTokenLayerYieldsPresentButEmptyLayer() { + final Document document = new TermVectorAnnotator() + .annotate(Document.of("").with(Layers.TOKENS, List.of())); + assertTrue(document.layers().contains(TermVectorAnnotator.TERM_VECTORS)); + assertTrue(document.get(TermVectorAnnotator.TERM_VECTORS).isEmpty()); + } + + @Test + void testMissingTokenLayerIsRejected() { + final TermVectorAnnotator annotator = new TermVectorAnnotator(); + final Document bare = Document.of("The dog barks."); + assertThrows(IllegalArgumentException.class, () -> annotator.annotate(bare)); + } + + @Test + void testNullDocumentIsRejected() { + final TermVectorAnnotator annotator = new TermVectorAnnotator(); + assertThrows(IllegalArgumentException.class, () -> annotator.annotate(null)); + } + + @Test + void testNullConstructorArgumentsAreRejected() { + final TermVectorAnnotator.Mode noMode = null; + final OffsetAwareNormalizer noNormalizer = null; + assertThrows(IllegalArgumentException.class, () -> new TermVectorAnnotator(noMode)); + assertThrows(IllegalArgumentException.class, () -> new TermVectorAnnotator(noNormalizer)); + assertThrows(IllegalArgumentException.class, () -> new TermVectorAnnotator(FOLD, null)); + assertThrows(IllegalArgumentException.class, + () -> new TermVectorAnnotator(null, TermVectorAnnotator.Mode.FULL)); + } + + @Test + void testRequiresAndProvides() { + final TermVectorAnnotator annotator = new TermVectorAnnotator(); + assertEquals(Set.of(Layers.TOKENS), annotator.requires()); + assertEquals(Set.of(TermVectorAnnotator.TERM_VECTORS), annotator.provides()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorNormalizedExampleTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorNormalizedExampleTest.java new file mode 100644 index 0000000000..2868b021e7 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorNormalizedExampleTest.java @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.termvector; + +import java.util.List; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnalyzer; +import opennlp.tools.tokenize.TokenizerAnnotator; +import opennlp.tools.tokenize.WhitespaceTokenizer; +import opennlp.tools.util.Span; +import opennlp.tools.util.normalizer.CharSequenceNormalizer; +import opennlp.tools.util.normalizer.TextNormalizer; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Mirrors the normalized-term workflow shown in the term vector section of the manual: + * a whitespace tokenizer feeding a {@link TermVectorAnnotator} built with a shipped, + * plain {@link CharSequenceNormalizer} case folder. The folder defines term identity per + * token, so case variants group under one term, while every occurrence span stays the + * token's own span in the original text. + */ +public class TermVectorNormalizedExampleTest { + + /** + * The documented example: {@code "Word word WORD"} yields one term, {@code "word"}, + * with three occurrence spans, each the token's exact original span. + */ + @Test + void testCaseFoldedTermsKeepOriginalSpans() { + final CharSequenceNormalizer folder = TextNormalizer.builder().caseFold().build(); + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new TokenizerAnnotator(WhitespaceTokenizer.INSTANCE)) + .add(new TermVectorAnnotator(folder)) + .build(); + + final Document document = analyzer.analyze("Word word WORD"); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(1, vectors.size()); + assertEquals(new TermVector("word", 3, + List.of(new Span(0, 4), new Span(5, 9), new Span(10, 14))), vectors.get(0).value()); + + // The spans point at the original surface forms, not the folded term. + final List surfaceForms = vectors.get(0).value().spans().stream() + .map(s -> s.getCoveredText(document.text()).toString()).toList(); + assertEquals(List.of("Word", "word", "WORD"), surfaceForms); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorPipelineTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorPipelineTest.java new file mode 100644 index 0000000000..8510cfdac9 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/termvector/TermVectorPipelineTest.java @@ -0,0 +1,109 @@ +/* + * 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.termvector; + +import java.util.List; +import java.util.Set; + +import org.junit.jupiter.api.Test; + +import opennlp.tools.document.Annotation; +import opennlp.tools.document.Document; +import opennlp.tools.document.DocumentAnalyzer; +import opennlp.tools.document.Layers; +import opennlp.tools.tokenize.TokenizerAnnotator; +import opennlp.tools.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Wires the {@link TermVectorAnnotator} into a {@link DocumentAnalyzer} behind a + * {@link TokenizerAnnotator}: the token layer goes in, the term vector layer comes out, + * and nothing else about the document changes. The tokenizer is the shared deterministic + * single-space fixture, so every expected span follows directly from the input text. + */ +public class TermVectorPipelineTest { + + /** + * Runs tokenizer plus term vector roll-up over a text with repeated tokens and reads + * the aggregated layer back, span by span, in original text coordinates. + */ + @Test + void testTokenizerAndTermVectorPipeline() { + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new TokenizerAnnotator(SingleSpaceTokens.TOKENIZER)) + .add(new TermVectorAnnotator()) + .build(); + + final Document document = analyzer.analyze("The dog barks. The dog naps."); + + // The document carries exactly the two layers the pipeline provides. + assertEquals(Set.of(Layers.TOKENS, TermVectorAnnotator.TERM_VECTORS), document.layers()); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(4, vectors.size()); + assertEquals(new TermVector("The", 2, List.of(new Span(0, 3), new Span(15, 18))), + vectors.get(0).value()); + assertEquals(new TermVector("dog", 2, List.of(new Span(4, 7), new Span(19, 22))), + vectors.get(1).value()); + assertEquals(new TermVector("barks.", 1, List.of(new Span(8, 14))), + vectors.get(2).value()); + assertEquals(new TermVector("naps.", 1, List.of(new Span(23, 28))), + vectors.get(3).value()); + + // Every occurrence span indexes into the original text and covers its own term. + for (final Annotation vector : vectors) { + for (final Span span : vector.value().spans()) { + assertEquals(vector.value().term(), span.getCoveredText(document.text()).toString()); + } + } + } + + /** + * The same pipeline in scoring-only mode: counts survive, offsets are never stored. + */ + @Test + void testScoringOnlyPipeline() { + final DocumentAnalyzer analyzer = DocumentAnalyzer.builder() + .add(new TokenizerAnnotator(SingleSpaceTokens.TOKENIZER)) + .add(new TermVectorAnnotator(TermVectorAnnotator.Mode.SCORING_ONLY)) + .build(); + + final Document document = analyzer.analyze("The dog barks. The dog naps."); + + final List> vectors = + document.get(TermVectorAnnotator.TERM_VECTORS); + assertEquals(4, vectors.size()); + assertEquals(TermVector.count("The", 2), vectors.get(0).value()); + assertEquals(TermVector.count("dog", 2), vectors.get(1).value()); + } + + /** + * The analyzer validates the pipeline at build time: the term vector annotator + * requires the token layer, so a pipeline without a tokenizer fails when it is + * assembled. + */ + @Test + void testPipelineWithoutTokenizerIsRejectedAtBuildTime() { + final DocumentAnalyzer.Builder builder = DocumentAnalyzer.builder() + .add(new TermVectorAnnotator()); + assertThrows(IllegalArgumentException.class, builder::build); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java deleted file mode 100644 index 070834316b..0000000000 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/BertTokenizerTest.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package opennlp.tools.tokenize; - -import java.util.HashSet; -import java.util.List; -import java.util.Set; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; - -import static org.junit.jupiter.api.Assertions.assertArrayEquals; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; - -/** - * Pins the deprecated {@link BertTokenizer} shim: constructor validation, delegation to - * {@link WordpieceEncoder#encodeToPieces(CharSequence)}, the default special tokens, and - * the unsupported {@link BertTokenizer#tokenizePos(String)}. - */ -@SuppressWarnings("removal") // Exercises BertTokenizer deliberately until its removal in 3.1. -class BertTokenizerTest { - - private static final List VOCABULARY = List.of( - "[PAD]", "[UNK]", "[CLS]", "[SEP]", "hello", "world", "##s", "ca", "##fe", ",", "!"); - - private static Set vocabularySet() { - return new HashSet<>(VOCABULARY); - } - - @ParameterizedTest - @ValueSource(strings = {"Hello, world!", "Caf\u00E9 hellos", ""}) - void testTokenizeReturnsTheEncoderPieceSequence(String input) { - final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - assertArrayEquals(encoder.encodeToPieces(input), tokenizer.tokenize(input), - "delegation broke on: " + input); - } - - @Test - void testDefaultConstructorsChainToBertSpecialTokensAndLowerCasing() { - final String[] uncasedDefault = new BertTokenizer(vocabularySet()).tokenize("Hello worldS"); - assertArrayEquals(new String[] {"[CLS]", "hello", "world", "##s", "[SEP]"}, uncasedDefault); - assertArrayEquals(uncasedDefault, - new BertTokenizer(vocabularySet(), true).tokenize("Hello worldS")); - assertArrayEquals(uncasedDefault, - new BertTokenizer(vocabularySet(), true, WordpieceTokenizer.BERT_CLS_TOKEN, - WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN) - .tokenize("Hello worldS")); - } - - @Test - void testCasedTokenizerKeepsCase() { - // Without lower casing, the capitalized word misses the lowercase-only vocabulary. - assertArrayEquals(new String[] {"[CLS]", "[UNK]", "[SEP]"}, - new BertTokenizer(vocabularySet(), false).tokenize("Hello")); - } - - @Test - void testConstructorsRejectNullArguments() { - assertThrows(IllegalArgumentException.class, () -> new BertTokenizer(null)); - assertThrows(IllegalArgumentException.class, () -> new BertTokenizer(null, true)); - assertThrows(IllegalArgumentException.class, - () -> new BertTokenizer(null, true, "[CLS]", "[SEP]", "[UNK]")); - assertThrows(IllegalArgumentException.class, - () -> new BertTokenizer(vocabularySet(), true, null, "[SEP]", "[UNK]")); - assertThrows(IllegalArgumentException.class, - () -> new BertTokenizer(vocabularySet(), true, "[CLS]", null, "[UNK]")); - assertThrows(IllegalArgumentException.class, - () -> new BertTokenizer(vocabularySet(), true, "[CLS]", "[SEP]", null)); - } - - @Test - void testTokenizeRejectsNullText() { - final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); - assertThrows(IllegalArgumentException.class, () -> tokenizer.tokenize(null)); - } - - @Test - void testTokenizePosIsUnsupportedWithTheDocumentedMessage() { - final BertTokenizer tokenizer = new BertTokenizer(vocabularySet()); - final UnsupportedOperationException e = assertThrows(UnsupportedOperationException.class, - () -> tokenizer.tokenizePos("hello world")); - assertEquals("Wordpiece tokens cannot be mapped to character spans of the original text", - e.getMessage()); - } -} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java index bbaeadaca3..9ecb227f68 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderReferenceSequencesTest.java @@ -26,15 +26,11 @@ import org.junit.jupiter.params.provider.MethodSource; /** - * Reference token-sequence expectations for {@link WordpieceEncoder}, covering lower casing, - * accent stripping, punctuation and CJK isolation, and text cleaning. - *

    - * All expected token sequences in this test were generated with the HuggingFace - * {@code tokenizers} reference implementation ({@code BertWordPieceTokenizer}) - * using the same vocabulary, so they are verified to be identical to the - * reference BERT tokenization. The encoder requires its special tokens to be - * present in the vocabulary (every piece must have an id), so the vocabularies - * here include them; the token sequences are unchanged. + * Checks token sequences defined by the BERT tokenizer for lower casing, accent stripping, + * punctuation, CJK isolation, and text cleaning. + * + * @see + * BERT reference tokenizer */ class WordpieceEncoderReferenceSequencesTest { @@ -43,16 +39,18 @@ class WordpieceEncoderReferenceSequencesTest { "the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", "em", "##bed", "##ding", "##s", "wurttemberg", "strasse", "grosse", - "don", "t", "wait", "what", ".", ",", "?", "!", "'", + "\u03C3\u03BF\u03C6\u03BF\u03C2", "\u03C3", "\u03BF\u03C3\u03B1", + "\u03BF\u03C3", "\u03BF\u03C2", "\u03B1", "\u03BF\u03C2\u302E\u03C3", + "don", "t", "wait", "what", ".", ",", "?", "!", "'", "\"", "\u6211", "\u7231", // CJK - "natural", "language", "processing"); + "natural", "language", "processing", "foxjumps"); /** * The reference input and expected-sequence pairs, one argument set per pipeline behavior. * * @return The (input, expected pieces) pairs. */ - static Stream referenceSequences() { + private static Stream referenceSequences() { return Stream.of( // Lower cases capitalized words. Arguments.of("The quick brown fox jumps over the lazy dog.", @@ -61,8 +59,31 @@ static Stream referenceSequences() { // Lower cases before wordpiece splitting. Arguments.of("Embeddings", new String[] {"[CLS]", "em", "##bed", "##ding", "##s", "[SEP]"}), + // Lowercases a word-final Greek sigma to U+03C2. + Arguments.of("\u03A3\u039F\u03A6\u039F\u03A3", + new String[] {"[CLS]", "\u03C3\u03BF\u03C6\u03BF\u03C2", "[SEP]"}), + // A sigma without a preceding cased letter is not a final sigma. + Arguments.of("\u03A3", new String[] {"[CLS]", "\u03C3", "[SEP]"}), + // A sigma followed by a cased letter is not a final sigma. + Arguments.of("\u039F\u03A3\u0391", + new String[] {"[CLS]", "\u03BF\u03C3\u03B1", "[SEP]"}), + // Case context skips the apostrophe before punctuation isolation. + Arguments.of("\u039F\u03A3'\u0391", + new String[] {"[CLS]", "\u03BF\u03C3", "'", "\u03B1", "[SEP]"}), + // A full stop is also case-ignorable before punctuation isolation. + Arguments.of("\u039F\u03A3.\u0391", + new String[] {"[CLS]", "\u03BF\u03C3", ".", "\u03B1", "[SEP]"}), + // A quotation mark terminates the case context, so the sigma before it is word-final. + Arguments.of("\u039F\u03A3\"\u0391", + new String[] {"[CLS]", "\u03BF\u03C2", "\"", "\u03B1", "[SEP]"}), + // Case context skips a combining mark, which accent stripping later removes. + Arguments.of("\u039F\u03A3\u0301", + new String[] {"[CLS]", "\u03BF\u03C2", "[SEP]"}), + // A combining spacing mark terminates the case context. + Arguments.of("\u039F\u03A3\u302E\u03A3", + new String[] {"[CLS]", "\u03BF\u03C2\u302E\u03C3", "[SEP]"}), // The u-umlaut decomposes to u plus a combining diaeresis and the mark is stripped; - // the sharp s is not a combining mark and must survive, leaving an OOV token. + // the sharp s is not a combining mark and is retained, leaving an OOV token. Arguments.of("W\u00fcrttemberg Stra\u00dfe", new String[] {"[CLS]", "wurttemberg", "[UNK]", "[SEP]"}), // Splits punctuation runs into single characters. @@ -79,11 +100,11 @@ static Stream referenceSequences() { // joining "brown" and "fox" into one out-of-vocabulary token. Arguments.of("the\tquick\u00a0brown\u0000fox", new String[] {"[CLS]", "the", "quick", "[UNK]", "[SEP]"}), - // The reference implementation treats all C* categories as control - // characters: private use (U+E000, Co) and noncharacters (U+FDD0, Cn) - // are removed, joining the surrounding text into one OOV token. - Arguments.of("fox\ue000jumps and fox\ufdd0jumps", - new String[] {"[CLS]", "[UNK]", "[UNK]", "[UNK]", "[SEP]"})); + // BERT removes all Unicode control categories, including private-use, + // unassigned, and surrogate code units. + Arguments.of("fox\ue000jumps and fox\ufdd0jumps fox\ud800jumps", + new String[] {"[CLS]", "foxjumps", "[UNK]", "foxjumps", "foxjumps", + "[SEP]"})); } @ParameterizedTest @@ -91,18 +112,7 @@ static Stream referenceSequences() { void testEncodesTheReferenceSequence(String input, String[] expected) { final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); Assertions.assertArrayEquals(expected, encoder.encodeToPieces(input), - "sequence broke on: " + input); - } - - @Test - void testRejectsNullSpecialTokens() { - // The encoder's contract throws IllegalArgumentException for null special tokens. - Assertions.assertThrows(IllegalArgumentException.class, - () -> new WordpieceEncoder(VOCABULARY, true, null, "[SEP]", "[UNK]")); - Assertions.assertThrows(IllegalArgumentException.class, - () -> new WordpieceEncoder(VOCABULARY, true, "[CLS]", null, "[UNK]")); - Assertions.assertThrows(IllegalArgumentException.class, - () -> new WordpieceEncoder(VOCABULARY, true, "[CLS]", "[SEP]", null)); + "unexpected sequence for: " + input); } @Test diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java index 997036b092..551fef6419 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceEncoderTest.java @@ -16,31 +16,37 @@ */ package opennlp.tools.tokenize; +import java.text.Normalizer; import java.util.ArrayList; -import java.util.HashSet; +import java.util.HashMap; +import java.util.LinkedHashSet; import java.util.List; +import java.util.Locale; +import java.util.Map; import java.util.Random; +import java.util.Set; import java.util.stream.Stream; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; -import opennlp.tools.util.Span; - import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -/** - * The encoder held against the deprecated {@link BertTokenizer} for piece-sequence parity (the - * shim's contract is "the encoder's piece sequence, nothing more"), plus exact hand-computed span - * assertions through every normalization step that changes, inserts, or removes characters. - */ +/** Checks token ids and original-text offsets through each normalization stage. */ class WordpieceEncoderTest { + private static final long REFERENCE_RANDOM_SEED = 1885L; + private static final int REFERENCE_SAMPLE_COUNT = 500; + private static final long SPAN_RANDOM_SEED = 42L; + private static final int SPAN_SAMPLE_COUNT = 400; + // Ids are indices: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, hello=4, world=5, ##s=6, won=7, // ##der=8, ##ful=9, ca=10, ##fe=11, istanbul=12, U+4E2D=13, U+56FD=14, .=15, ,=16, !=17, // he=18, ##llo=19, Greek "sofos" with a final sigma=20. @@ -59,72 +65,131 @@ private static void assertPiece(SubwordPiece piece, String expectedPiece, int ex assertEquals(expectedId, piece.id()); assertEquals(expectedStart, piece.start(), "start of " + piece); assertEquals(expectedEnd, piece.end(), "end of " + piece); - assertEquals(new Span(expectedStart, expectedEnd), piece.span(), "span of " + piece); } - /** - * The curated parity inputs, each exercising a normalization step of the pipeline. - * - * @return The inputs. - */ - static Stream curatedInputs() { - return Stream.of( - "", - " ", - "Hello, WORLD!", - "Wonderful", - "hellos", - // An accented e, stripped by NFD decomposition. - "Caf\u00E9", - // The Turkish dotted capital I: lower cases to two chars, then the dot strips away. - "\u0130stanbul", - // CJK ideographs are isolated into single-character tokens. - "\u4E2D\u56FD is CJK", - // Greek upper case: the trailing sigma takes the contextual final-sigma mapping. - "\u03A3\u039F\u03A6\u039F\u03A3", - // The NBSP is whitespace in the BERT sense. - "hello\u00A0world", - // NUL and the zero-width space are removed by the cleaning stage. - "a\u0000b\u200Bc", - // An emoji: unknown to the vocabulary, and a surrogate pair. - "\uD83D\uDE00", - "!!!", - "a".repeat(101), - "he said: \u00ABhello\u00BB."); + @Test + void testGeneratedInputsMatchReferenceNormalizationOrder() { + final int[] pool = {'a', 'B', ' ', '\t', '\n', 0x00A0, 0x00C9, 0x0130, 0x0301, + 0x0391, 0x03BF, 0x1715, 0x200B, 0x2028, 0x2029, 0x302E, 0x4E2D, 0x20000, + '.', '\'', 0x10100, 0x1F600, 0xFFFD, 0}; + final Random random = new Random(REFERENCE_RANDOM_SEED); + for (final boolean lowerCase : new boolean[] {true, false}) { + for (int sampleIndex = 0; sampleIndex < REFERENCE_SAMPLE_COUNT; sampleIndex++) { + final StringBuilder input = new StringBuilder(); + final int length = random.nextInt(24); + for (int i = 0; i < length; i++) { + input.appendCodePoint(pool[random.nextInt(pool.length)]); + } + + final List expectedWords = referenceBasicTokens(input, lowerCase); + final Set vocabulary = new LinkedHashSet<>( + List.of("[UNK]", "[CLS]", "[SEP]")); + vocabulary.addAll(expectedWords); + final WordpieceEncoder encoder = new WordpieceEncoder( + new ArrayList<>(vocabulary), lowerCase); + + final List expected = new ArrayList<>(expectedWords.size() + 2); + expected.add("[CLS]"); + expected.addAll(expectedWords); + expected.add("[SEP]"); + assertEquals(expected, List.of(encoder.encodeToPieces(input)), + "unexpected normalization for: " + codePointList(input)); + } + } } - @ParameterizedTest - @MethodSource("curatedInputs") - @SuppressWarnings("removal") // BertTokenizer is pinned to the encoder until its removal in 3.1. - void testPieceSequenceMatchesBertTokenizerOnCuratedInputs(String input) { - final BertTokenizer bertTokenizer = new BertTokenizer(new HashSet<>(VOCAB), true); - final WordpieceEncoder encoder = uncased(); - assertArrayEquals(bertTokenizer.tokenize(input), encoder.encodeToPieces(input), - "parity broke on: " + input); + private static List codePointList(CharSequence text) { + return text.codePoints() + .mapToObj(codePoint -> "U+" + Integer.toHexString(codePoint).toUpperCase(Locale.ROOT)) + .toList(); + } + + private static List referenceBasicTokens(CharSequence input, boolean lowerCase) { + final StringBuilder cleaned = new StringBuilder(input.length() + 16); + input.codePoints().forEach(codePoint -> { + if (codePoint == 0 || codePoint == 0xFFFD || BertNormalization.isControl(codePoint)) { + return; + } + if (BertNormalization.isWhitespace(codePoint)) { + cleaned.append(' '); + } else if (BertNormalization.isCjk(codePoint)) { + cleaned.append(' ').appendCodePoint(codePoint).append(' '); + } else { + cleaned.appendCodePoint(codePoint); + } + }); + + final List words = splitReferenceWhitespace(cleaned); + final List pieces = new ArrayList<>(); + for (String word : words) { + if (lowerCase) { + word = stripReferenceAccents(word.toLowerCase(Locale.ROOT)); + } + final StringBuilder current = new StringBuilder(); + word.codePoints().forEach(codePoint -> { + if (BertNormalization.isPunctuation(codePoint)) { + addIfNotEmpty(pieces, current); + pieces.add(new String(Character.toChars(codePoint))); + } else { + current.appendCodePoint(codePoint); + } + }); + addIfNotEmpty(pieces, current); + } + return pieces; + } + + private static List splitReferenceWhitespace(CharSequence text) { + final List words = new ArrayList<>(); + final StringBuilder current = new StringBuilder(); + text.codePoints().forEach(codePoint -> { + if (codePoint == ' ' || codePoint == 0x2028 || codePoint == 0x2029) { + addIfNotEmpty(words, current); + } else { + current.appendCodePoint(codePoint); + } + }); + addIfNotEmpty(words, current); + return words; + } + + private static String stripReferenceAccents(String text) { + final StringBuilder stripped = new StringBuilder(text.length()); + Normalizer.normalize(text, Normalizer.Form.NFD).codePoints().forEach(codePoint -> { + if (Character.getType(codePoint) != Character.NON_SPACING_MARK) { + stripped.appendCodePoint(codePoint); + } + }); + return stripped.toString(); + } + + private static void addIfNotEmpty(List words, StringBuilder current) { + if (!current.isEmpty()) { + words.add(current.toString()); + current.setLength(0); + } } @Test - @SuppressWarnings("removal") // BertTokenizer is pinned to the encoder until its removal in 3.1. - void testPieceSequenceMatchesBertTokenizerOnRandomInputs() { + void testGeneratedUnicodeProducesOrderedBoundedSpans() { final int[] pool = {'a', 'b', 'A', 'B', 'z', ' ', ' ', '\t', 0x00A0, '.', '!', ',', 0x0301, 0x00E9, 0x0130, 0x03A3, 0x03C3, 0x03BF, 0x4E2D, 0xFFFD, 0x200B, 0x1F600, 0}; - final Random random = new Random(42); + final Random random = new Random(SPAN_RANDOM_SEED); for (final boolean lowerCase : new boolean[] {true, false}) { - final BertTokenizer bertTokenizer = new BertTokenizer(new HashSet<>(VOCAB), lowerCase); final WordpieceEncoder encoder = new WordpieceEncoder(VOCAB, lowerCase); - for (int round = 0; round < 400; round++) { + for (int sampleIndex = 0; sampleIndex < SPAN_SAMPLE_COUNT; sampleIndex++) { final StringBuilder text = new StringBuilder(); final int length = random.nextInt(25); for (int i = 0; i < length; i++) { text.appendCodePoint(pool[random.nextInt(pool.length)]); } final String input = text.toString(); - assertArrayEquals(bertTokenizer.tokenize(input), encoder.encodeToPieces(input), - "parity broke on: " + input); + final List pieces = encoder.encode(input); + assertPiece(pieces.getFirst(), "[CLS]", 2, 0, 0); + assertPiece(pieces.getLast(), "[SEP]", 3, input.length(), input.length()); - // Span invariants: within bounds and never moving backwards. int previousStart = 0; - for (final SubwordPiece piece : encoder.encode(input)) { + for (final SubwordPiece piece : pieces) { assertTrue(piece.start() >= previousStart && piece.end() <= input.length(), "span out of order or bounds in " + input + ": " + piece); previousStart = piece.start(); @@ -134,7 +199,7 @@ void testPieceSequenceMatchesBertTokenizerOnRandomInputs() { } @Test - void testSpansSurvivePunctuationIsolationAndCaseFolding() { + void testSpansPreservePunctuationIsolationAndCaseMapping() { final List pieces = uncased().encode("Hello, WORLD!"); assertEquals(6, pieces.size()); assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); @@ -146,7 +211,7 @@ void testSpansSurvivePunctuationIsolationAndCaseFolding() { } @Test - void testSpansSurviveAccentStripping() { + void testSpansPreserveAccentStripping() { // The accent is stripped by NFD, yet ##fe still covers the accented surface. final List pieces = uncased().encode("Caf\u00E9"); assertEquals(4, pieces.size()); @@ -155,9 +220,29 @@ void testSpansSurviveAccentStripping() { } @Test - void testSpansSurviveLengthChangingLowerCasing() { - // The Turkish dotted capital I lower cases to two chars before the combining dot strips; - // the piece still covers the original eight chars. + void testSpansPreserveDecomposedAccentStripping() { + final List pieces = uncased().encode("Cafe\u0301"); + + assertEquals(4, pieces.size()); + assertPiece(pieces.get(1), "ca", 10, 0, 2); + assertPiece(pieces.get(2), "##fe", 11, 2, 5); + } + + @Test + void testNormalizationReordersCombiningMarksAcrossCodePoints() { + final String normalized = "a\u1715\u302E"; + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", normalized)); + + final List pieces = encoder.encode("a\u302E\u1715"); + + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), normalized, 3, 0, 3); + } + + @Test + void testSpansPreserveLengthChangingLowerCasing() { + // The Turkish dotted capital I normalizes to i; the piece still covers the full input range. final List pieces = uncased().encode("\u0130stanbul"); assertEquals(3, pieces.size()); assertPiece(pieces.get(1), "istanbul", 12, 0, 8); @@ -171,11 +256,23 @@ void testCjkIsolationYieldsOnePieceAndSpanPerIdeograph() { assertPiece(pieces.get(2), "\u56FD", 14, 1, 2); } + @Test + void testSupplementaryCjkOffsetsUseUtf16Indices() { + final String ideograph = new String(Character.toChars(0x20000)); + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", "a", ideograph, "b")); + + final List pieces = encoder.encode("A" + ideograph + "B"); + + assertPiece(pieces.get(1), "a", 3, 0, 1); + assertPiece(pieces.get(2), ideograph, 4, 1, 3); + assertPiece(pieces.get(3), "b", 5, 3, 4); + } + @Test void testLineAndParagraphSeparatorsSplitWords() { // Zl and Zp are not whitespace in the BERT _is_whitespace sense, but the reference - // pipeline's whitespace_tokenize (Python's str.split()) breaks words on them, as did the - // previous OpenNLP pipeline via WhitespaceTokenizer. + // pipeline's whitespace_tokenize (Python's str.split()) breaks words on them. final List pieces = uncased().encode("hello\u2028world\u2029hello"); assertEquals(5, pieces.size()); assertPiece(pieces.get(1), "hello", 4, 0, 5); @@ -184,7 +281,7 @@ void testLineAndParagraphSeparatorsSplitWords() { } @Test - void testUnknownWordCoversItsWholeSurfaceIncludingRemovedChars() { + void testUnknownWordCoversCompleteSourceRangeIncludingRemovedChars() { // NUL and the zero-width space are removed by cleaning, so one word "abc" remains; it is // not representable and becomes the unknown piece spanning the full original surface. final List pieces = uncased().encode("a\u0000b\u200Bc"); @@ -193,10 +290,18 @@ void testUnknownWordCoversItsWholeSurfaceIncludingRemovedChars() { } @Test - void testContextualCaseMappingFallsBackToWordWideSpans() { - // Greek final sigma is a contextual mapping the per-char rerun cannot reproduce, so the - // word's pieces fall back to spanning the whole word; the piece content is asserted - // exactly below. + void testPartialWordpieceMatchBecomesOneUnknownPiece() { + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", "won", "##der")); + + final List pieces = encoder.encode("Wonderful"); + + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), "[UNK]", 0, 0, 9); + } + + @Test + void testFinalSigmaMappingPreservesSourceRange() { final List pieces = uncased().encode("\u03A3\u039F\u03A6\u039F\u03A3"); assertEquals(3, pieces.size()); @@ -205,12 +310,150 @@ void testContextualCaseMappingFallsBackToWordWideSpans() { } @Test - void testEncodeToIdsCarriesVocabularyLineNumbers() { + void testEncodeToIdsReturnsVocabularyLineNumbers() { assertArrayEquals(new int[] {2, 4, 5, 6, 3}, uncased().encodeToIds("Hello worldS")); } @Test - void testCasedEncoderKeepsCase() { + void testExplicitVocabularyIdsArePreserved() { + final WordpieceEncoder encoder = new WordpieceEncoder( + Map.of("", 101, "", 205, "", 999, "alice", 42), true, + "", "", ""); + + assertArrayEquals(new int[] {101, 42, 999, 205}, + encoder.encodeToIds("Alice rabbit")); + } + + @Test + void testCopiesVocabularyInputs() { + final List list = new ArrayList<>(VOCAB); + final WordpieceEncoder fromList = new WordpieceEncoder(list); + list.set(4, "changed"); + + final Map map = new HashMap<>(Map.of( + "[CLS]", 10, "[SEP]", 11, "[UNK]", 12, "hello", 13)); + final WordpieceEncoder fromMap = new WordpieceEncoder( + map, true, "[CLS]", "[SEP]", "[UNK]"); + map.remove("hello"); + + assertArrayEquals(new String[] {"[CLS]", "hello", "[SEP]"}, + fromList.encodeToPieces("Hello")); + assertArrayEquals(new int[] {10, 13, 11}, fromMap.encodeToIds("Hello")); + } + + @Test + void testAliceExamplePreservesOriginalOffsets() { + final List vocabulary = List.of( + "[UNK]", "[CLS]", "[SEP]", "alice", "was", "begin", "##ning", "to", "get", + "very", "tired", "."); + final WordpieceEncoder encoder = new WordpieceEncoder(vocabulary); + final String text = "Alice was beginning to get very tired."; + + final List pieces = encoder.encode(text); + + // The manual's WordPiece example prints one line per piece: id, piece, and source text. + final List lines = new ArrayList<>(); + for (final SubwordPiece piece : pieces) { + final CharSequence source = text.subSequence(piece.start(), piece.end()); + lines.add(piece.id() + "\t" + piece.piece() + "\t" + source); + } + assertEquals(List.of( + "1\t[CLS]\t", "3\talice\tAlice", "4\twas\twas", "5\tbegin\tbegin", "6\t##ning\tning", + "7\tto\tto", "8\tget\tget", "9\tvery\tvery", "10\ttired\ttired", "11\t.\t.", + "2\t[SEP]\t"), lines); + assertArrayEquals(new String[] { + "[CLS]", "alice", "was", "begin", "##ning", "to", "get", "very", "tired", ".", + "[SEP]"}, encoder.encodeToPieces(text)); + assertArrayEquals(new int[] {1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 2}, + encoder.encodeToIds(text)); + assertPiece(pieces.get(1), "alice", 3, 0, 5); + assertPiece(pieces.get(2), "was", 4, 6, 9); + assertPiece(pieces.get(3), "begin", 5, 10, 15); + assertPiece(pieces.get(4), "##ning", 6, 15, 19); + assertPiece(pieces.get(8), "tired", 10, 32, 37); + assertPiece(pieces.get(9), ".", 11, 37, 38); + } + + @Test + void testMaximumWordLengthCountsUnicodeCodePoints() { + final String face = "\uD83D\uDE00"; + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", face, "##" + face), false); + + final String input = face.repeat(100); + final List pieces = encoder.encode(input); + + assertEquals(102, pieces.size()); + assertPiece(pieces.get(1), face, 3, 0, 2); + assertPiece(pieces.get(100), "##" + face, 4, 198, 200); + } + + @Test + void testDoesNotSplitSupplementaryCodePoint() { + final String face = "\uD83D\uDE00"; + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", "\uD83D", "##\uDE00"), false); + + final List pieces = encoder.encode(face); + + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), "[UNK]", 0, 0, 2); + } + + @Test + void testWordBeyondMaximumLengthBecomesUnknown() { + final String face = "\uD83D\uDE00"; + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", face, "##" + face), false); + + final String input = face.repeat(101); + final List pieces = encoder.encode(input); + + assertEquals(3, pieces.size()); + assertPiece(pieces.get(1), "[UNK]", 0, 0, 202); + } + + @Test + void testCustomMaximumWordLength() { + final WordpieceEncoder encoder = new WordpieceEncoder( + List.of("[UNK]", "[CLS]", "[SEP]", "a", "##a"), false, 2); + + assertArrayEquals(new String[] {"[CLS]", "a", "##a", "[SEP]"}, + encoder.encodeToPieces("aa")); + assertArrayEquals(new String[] {"[CLS]", "[UNK]", "[SEP]"}, + encoder.encodeToPieces("aaa")); + } + + @Test + void testRejectsNegativeMaximumWordLength() { + assertEquals("maxWordCodePoints must not be negative", + assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(VOCAB, true, -1)).getMessage()); + } + + @ParameterizedTest + @MethodSource("invalidSubwordPieces") + void testSubwordPieceRejectsInvalidFields(String expectedMessage, Executable constructor) { + assertEquals(expectedMessage, + assertThrows(IllegalArgumentException.class, constructor).getMessage()); + } + + private static Stream invalidSubwordPieces() { + return Stream.of( + Arguments.of("piece must not be null", + (Executable) () -> new SubwordPiece(null, 0, 0, 0)), + Arguments.of("piece must not be empty", + (Executable) () -> new SubwordPiece("", 0, 0, 0)), + Arguments.of("id must not be negative", + (Executable) () -> new SubwordPiece("piece", -1, 0, 0)), + Arguments.of("start must not be negative", + (Executable) () -> new SubwordPiece("piece", 0, -1, 0)), + Arguments.of("end must be at least start", + (Executable) () -> new SubwordPiece("piece", 0, 2, 1))); + } + + @Test + void testCasedEncoderPreservesCase() { final List vocabulary = new ArrayList<>(VOCAB); vocabulary.add("Hello"); final WordpieceEncoder cased = new WordpieceEncoder(vocabulary, false); @@ -220,24 +463,112 @@ void testCasedEncoderKeepsCase() { } @ParameterizedTest - @ValueSource(strings = {"", " "}) + @ValueSource(strings = {"", " ", "\u0000", "\u200B", "\u0301"}) void testEmptyAndBlankTextEncodeToTheFramePiecesOnly(String input) { final List pieces = uncased().encode(input); - assertEquals(2, pieces.size(), "frame pieces broke on <" + input + ">"); + assertEquals(2, pieces.size(), "unexpected frame pieces for <" + input + ">"); assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); assertPiece(pieces.get(1), "[SEP]", 3, input.length(), input.length()); } - @Test - void testValidationFailsLoudly() { - assertThrows(IllegalArgumentException.class, () -> new WordpieceEncoder(null)); - assertThrows(IllegalArgumentException.class, - () -> new WordpieceEncoder(List.of("[CLS]", "[SEP]"))); - assertThrows(IllegalArgumentException.class, - () -> new WordpieceEncoder(List.of("[CLS]", "[SEP]", "[UNK]", "dup", "dup"))); + @ParameterizedTest + @MethodSource("invalidListVocabularies") + void testRejectsInvalidListVocabulary(String expectedMessage, Executable constructor) { + assertEquals(expectedMessage, + assertThrows(IllegalArgumentException.class, constructor).getMessage()); + } + + private static Stream invalidListVocabularies() { final List withNull = new ArrayList<>(VOCAB); withNull.add(null); - assertThrows(IllegalArgumentException.class, () -> new WordpieceEncoder(withNull)); - assertThrows(IllegalArgumentException.class, () -> uncased().encode(null)); + return Stream.of( + Arguments.of("vocabulary must not be null", + (Executable) () -> new WordpieceEncoder(null)), + Arguments.of("vocabulary must not contain duplicate piece 'dup'", + (Executable) () -> new WordpieceEncoder( + List.of("[CLS]", "[SEP]", "[UNK]", "dup", "dup"))), + Arguments.of("vocabulary must not contain an empty piece at index 3", + (Executable) () -> new WordpieceEncoder( + List.of("[CLS]", "[SEP]", "[UNK]", ""))), + Arguments.of("vocabulary must not contain null at index " + (withNull.size() - 1), + (Executable) () -> new WordpieceEncoder(withNull))); + } + + @ParameterizedTest + @MethodSource("invalidMapVocabularies") + void testRejectsInvalidMapVocabulary(String expectedMessage, Executable constructor) { + assertEquals(expectedMessage, + assertThrows(IllegalArgumentException.class, constructor).getMessage()); + } + + private static Stream invalidMapVocabularies() { + final Map nullPiece = new HashMap<>(); + nullPiece.put("[CLS]", 0); + nullPiece.put("[SEP]", 1); + nullPiece.put("[UNK]", 2); + nullPiece.put(null, 0); + final Map nullId = new HashMap<>(); + nullId.put("[SEP]", 1); + nullId.put("[UNK]", 2); + nullId.put("[CLS]", null); + return Stream.of( + Arguments.of("vocabularyIds must not be null", + (Executable) () -> new WordpieceEncoder((Map) null, true, + "[CLS]", "[SEP]", "[UNK]")), + Arguments.of("vocabularyIds must not contain a negative id for piece '[UNK]'", + (Executable) () -> new WordpieceEncoder( + Map.of("[CLS]", 0, "[SEP]", 1, "[UNK]", -1), true, + "[CLS]", "[SEP]", "[UNK]")), + Arguments.of("vocabularyIds must not contain an empty piece", + (Executable) () -> new WordpieceEncoder( + Map.of("[CLS]", 0, "[SEP]", 1, "[UNK]", 2, "", 3), true, + "[CLS]", "[SEP]", "[UNK]")), + Arguments.of("vocabularyIds must not contain null pieces or ids", + (Executable) () -> new WordpieceEncoder( + nullPiece, true, "[CLS]", "[SEP]", "[UNK]")), + Arguments.of("vocabularyIds must not contain null pieces or ids", + (Executable) () -> new WordpieceEncoder( + nullId, true, "[CLS]", "[SEP]", "[UNK]"))); + } + + @ParameterizedTest + @ValueSource(strings = {"[CLS]", "[SEP]", "[UNK]"}) + void testRejectsMissingSpecialToken(String missingToken) { + final List vocabulary = new ArrayList<>(VOCAB); + vocabulary.remove(missingToken); + + assertEquals("vocabulary must contain special token '" + missingToken + "'", + assertThrows(IllegalArgumentException.class, + () -> new WordpieceEncoder(vocabulary)).getMessage()); + } + + @ParameterizedTest + @MethodSource("invalidSpecialTokens") + void testRejectsNullAndEmptySpecialTokens(String expectedMessage, Executable constructor) { + assertEquals(expectedMessage, + assertThrows(IllegalArgumentException.class, constructor).getMessage()); + } + + private static Stream invalidSpecialTokens() { + return Stream.of( + Arguments.of("classificationToken must not be null", + (Executable) () -> new WordpieceEncoder(VOCAB, true, null, "[SEP]", "[UNK]")), + Arguments.of("separatorToken must not be null", + (Executable) () -> new WordpieceEncoder(VOCAB, true, "[CLS]", null, "[UNK]")), + Arguments.of("unknownToken must not be null", + (Executable) () -> new WordpieceEncoder(VOCAB, true, "[CLS]", "[SEP]", null)), + Arguments.of("classificationToken must not be empty", + (Executable) () -> new WordpieceEncoder(VOCAB, true, "", "[SEP]", "[UNK]")), + Arguments.of("separatorToken must not be empty", + (Executable) () -> new WordpieceEncoder(VOCAB, true, "[CLS]", "", "[UNK]")), + Arguments.of("unknownToken must not be empty", + (Executable) () -> new WordpieceEncoder(VOCAB, true, "[CLS]", "[SEP]", ""))); + } + + @Test + void testRejectsNullText() { + assertEquals("text must not be null", + assertThrows(IllegalArgumentException.class, + () -> uncased().encode(null)).getMessage()); } } diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceTokenizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceTokenizerTest.java index 87362bde15..f070e17105 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceTokenizerTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/WordpieceTokenizerTest.java @@ -18,10 +18,16 @@ package opennlp.tools.tokenize; import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; import java.util.Set; +import java.util.stream.Stream; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; public class WordpieceTokenizerTest { @@ -91,6 +97,77 @@ void testRejectsNegativeMaxTokenLength() { } + @Test + void testDefaultMaximumMatchesTheReferencePipeline() { + final WordpieceTokenizer tokenizer = new WordpieceTokenizer(Set.of("a", "##a")); + + final String[] tokens = tokenizer.tokenize("a".repeat(100)); + + Assertions.assertEquals(102, tokens.length); + Assertions.assertEquals("a", tokens[1]); + Assertions.assertEquals("##a", tokens[100]); + } + + @Test + void testMaximumLengthCountsCodePoints() { + final String face = "\uD83D\uDE00"; + final WordpieceTokenizer tokenizer = new WordpieceTokenizer( + Set.of(face, "##" + face), 2); + + final String[] tokens = tokenizer.tokenize(face.repeat(2)); + + Assertions.assertArrayEquals( + new String[] {"[CLS]", face, "##" + face, "[SEP]"}, tokens); + } + + @Test + void testSegmentationDoesNotSplitSupplementaryCodePoints() { + final WordpieceTokenizer tokenizer = new WordpieceTokenizer( + Set.of("\uD83D", "##\uDE00")); + + final String[] tokens = tokenizer.tokenize("\uD83D\uDE00"); + + Assertions.assertArrayEquals(new String[] {"[CLS]", "[UNK]", "[SEP]"}, tokens); + } + + @ParameterizedTest + @MethodSource("invalidVocabularies") + void testRejectsInvalidVocabulary(Set vocabulary) { + Assertions.assertThrows(IllegalArgumentException.class, + () -> new WordpieceTokenizer(vocabulary)); + } + + private static Stream> invalidVocabularies() { + final Set withNull = new LinkedHashSet<>(List.of("word")); + withNull.add(null); + return Stream.of(null, Set.of(""), withNull); + } + + @ParameterizedTest + @MethodSource("invalidSpecialTokens") + void testRejectsInvalidSpecialTokens(String classificationToken, String separatorToken, + String unknownToken) { + Assertions.assertThrows(IllegalArgumentException.class, () -> new WordpieceTokenizer( + getVocabulary(), classificationToken, separatorToken, unknownToken)); + } + + private static Stream invalidSpecialTokens() { + return Stream.of( + Arguments.of(null, "[SEP]", "[UNK]"), + Arguments.of("", "[SEP]", "[UNK]"), + Arguments.of("[CLS]", null, "[UNK]"), + Arguments.of("[CLS]", "", "[UNK]"), + Arguments.of("[CLS]", "[SEP]", null), + Arguments.of("[CLS]", "[SEP]", "")); + } + + @Test + void testTokenizeRejectsNullText() { + final WordpieceTokenizer tokenizer = new WordpieceTokenizer(getVocabulary()); + + Assertions.assertThrows(IllegalArgumentException.class, () -> tokenizer.tokenize(null)); + } + @Test void testTokenizePosIsUnsupported() { diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeTokenizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeTokenizerTest.java new file mode 100644 index 0000000000..c0cf0a6445 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeTokenizerTest.java @@ -0,0 +1,1101 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.IOException; +import java.nio.charset.MalformedInputException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import org.junit.jupiter.api.Assertions; +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.ResourceLimits; +import opennlp.tools.util.Span; + +/** + * Tests the lattice segmenter against a project-authored miniature dictionary; no + * external dictionary data is involved. + * + *

    Source strings are written as Unicode escapes to keep this file ASCII-only; the + * class works over the same miniature Japanese dictionary as the sibling usage + * example, whose javadoc spells out each fixture word.

    + */ +public class LatticeTokenizerTest { + + private static final String LEXICON_CSV = "lexicon.csv"; + private static final String MATRIX_DEF = "matrix.def"; + private static final String CHAR_DEF = "char.def"; + private static final String UNK_DEF = "unk.def"; + + /** UTF-8 lead byte with the required continuation byte omitted. */ + private static final byte TRUNCATED_UTF8_LEAD_BYTE = (byte) 0xC3; + + /** A one by one connection matrix charging cost zero, for single-context fixtures. */ + private static final String UNIT_MATRIX = "1 1\n0 0 0\n"; + + /** + * The {@code char.def} line defining the DEFAULT category: it does not invoke + * unknown-word handling beside a lexicon match, it groups a whole run into one + * candidate, and it offers no fixed-length candidates. + */ + private static final String DEFAULT_CATEGORY_LINE = "DEFAULT 0 1 0"; + + /** The {@code unk.def} template line for the DEFAULT category. */ + private static final String DEFAULT_UNKNOWN_TEMPLATE = "DEFAULT,0,0,10000,symbol,unknown"; + + @TempDir + static Path directory; + + private static LatticeTokenizer tokenizer; + + @BeforeAll + static void loadDictionary() throws IOException { + write(LEXICON_CSV, String.join("\n", + "\u6771\u4EAC,0,0,3000,noun,proper", + "\u4EAC\u90FD,0,0,3000,noun,proper", + "\u6771,0,0,6000,noun,common", + "\u90FD,0,0,4000,noun,suffix", + "\u306B,0,0,1000,particle,case", + "\u884C\u304F,0,0,3000,verb,base", + "")); + write(MATRIX_DEF, UNIT_MATRIX); + write(CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "KANJI 0 0 2", + "HIRAGANA 0 1 0", + "LATIN 1 1 0", + "", + "0x3041..0x3096 HIRAGANA", + "0x4E00..0x9FFF KANJI", + "0x0041..0x005A LATIN", + "0x0061..0x007A LATIN", + "")); + write(UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "LATIN,0,0,4000,noun,foreign", + "KANJI,0,0,8000,noun,unknown", + "HIRAGANA,0,0,9000,particle,unknown", + "")); + tokenizer = new LatticeTokenizer(MecabDictionary.load(directory)); + } + + /** Writes one UTF-8 dictionary file into the shared dictionary directory. */ + private static void write(String name, String content) throws IOException { + write(directory, name, content); + } + + /** Writes one UTF-8 dictionary file into a test-supplied directory. */ + private static void write(Path target, String name, String content) throws IOException { + Files.write(target.resolve(name), content.getBytes(StandardCharsets.UTF_8)); + } + + @Test + void testLatticePrefersTheCheaperSegmentation() { + // Tokyo plus the metropolis suffix must beat the competing reading east plus Kyoto. + final String text = "\u6771\u4EAC\u90FD\u306B\u884C\u304F"; + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u90FD", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 2), new Span(2, 3), new Span(3, 4), new Span(4, 6)}, + tokenizer.tokenizePos(text)); + } + + @Test + void testMorphemesCarryDictionaryFeatures() { + final List morphemes = + tokenizer.analyze("\u6771\u4EAC\u90FD\u306B\u884C\u304F"); + Assertions.assertEquals(4, morphemes.size()); + Assertions.assertEquals(List.of("noun", "proper"), morphemes.get(0).features()); + Assertions.assertEquals(List.of("particle", "case"), morphemes.get(2).features()); + Assertions.assertFalse(morphemes.get(0).unknown()); + } + + @Test + void testUnknownLatinRunGroupsIntoOneMorpheme() { + final List morphemes = tokenizer.analyze("ABC\u306B\u884C\u304F"); + Assertions.assertEquals(3, morphemes.size()); + Assertions.assertEquals("ABC", morphemes.get(0).surface()); + Assertions.assertTrue(morphemes.get(0).unknown()); + Assertions.assertEquals(List.of("noun", "foreign"), morphemes.get(0).features()); + } + + @Test + void testUnknownKanjiPreferOneMorphemeOverTwo() { + final List morphemes = tokenizer.analyze("\u5CE0\u9053\u306B\u884C\u304F"); + Assertions.assertEquals(3, morphemes.size()); + Assertions.assertEquals("\u5CE0\u9053", morphemes.get(0).surface()); + Assertions.assertTrue(morphemes.get(0).unknown()); + } + + /** + * Verifies that an unknown-word candidate never spans a character category boundary. + * An unlisted kanji directly followed by a Latin letter must be analyzed as two + * morphemes of their own categories, never as one KANJI morpheme whose surface glues + * the kanji to the letter. + */ + @Test + void testUnknownCandidatesNeverSpanCategoryBoundaries() { + final String text = "\u5CE0a"; + Assertions.assertArrayEquals(new String[] {"\u5CE0", "a"}, tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] {new Span(0, 1), new Span(1, 2)}, + tokenizer.tokenizePos(text)); + final List morphemes = tokenizer.analyze(text); + Assertions.assertEquals(List.of("noun", "unknown"), morphemes.get(0).features()); + Assertions.assertEquals(List.of("noun", "foreign"), morphemes.get(1).features()); + } + + /** + * Verifies that bounding unknown-word candidates by the category run does not under + * generate inside the run: a two-kanji unlisted run followed by a Latin letter still + * offers the length-two KANJI candidate, which wins over two single-kanji morphemes. + */ + @Test + void testUnknownRunStillOffersWithinCategoryLengths() { + final String text = "\u5CE0\u9053a"; + Assertions.assertArrayEquals(new String[] {"\u5CE0\u9053", "a"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] {new Span(0, 2), new Span(2, 3)}, + tokenizer.tokenizePos(text)); + } + + @Test + void testWhitespaceSeparatesAndIsNeverAMorpheme() { + final String text = "\u6771\u4EAC \u306B \u884C\u304F"; + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 2), new Span(3, 4), new Span(5, 7)}, + tokenizer.tokenizePos(text)); + Assertions.assertEquals(0, tokenizer.analyze(" ").size()); + Assertions.assertEquals(0, tokenizer.analyze("").size()); + } + + /** + * Verifies that empty input yields empty results from every view of the tokenizer. + */ + @Test + void testEmptyInputYieldsEmptyResults() { + Assertions.assertArrayEquals(new String[0], tokenizer.tokenize("")); + Assertions.assertArrayEquals(new Span[0], tokenizer.tokenizePos("")); + } + + /** + * Verifies single-character input for a listed surface and for an unlisted kanji: + * both come back as exactly one morpheme covering {@code [0, 1)}, and only the + * unlisted one is marked unknown. + */ + @Test + void testSingleCharacterInput() { + Assertions.assertArrayEquals(new String[] {"\u306B"}, tokenizer.tokenize("\u306B")); + Assertions.assertArrayEquals(new Span[] {new Span(0, 1)}, tokenizer.tokenizePos("\u306B")); + Assertions.assertFalse(tokenizer.analyze("\u306B").get(0).unknown()); + + final List unknown = tokenizer.analyze("\u5CE0"); + Assertions.assertEquals(1, unknown.size()); + Assertions.assertEquals("\u5CE0", unknown.get(0).surface()); + Assertions.assertEquals(new Span(0, 1), unknown.get(0).span()); + Assertions.assertTrue(unknown.get(0).unknown()); + } + + /** + * Verifies input made entirely of characters absent from both the lexicon and the + * {@code char.def} mappings: they fall into the DEFAULT category, whose grouping + * setting joins the whole same-category run into one unknown morpheme carrying the + * DEFAULT template's features. + */ + @Test + void testEntirelyUnknownInputGroupsIntoOneDefaultMorpheme() { + final List morphemes = tokenizer.analyze("\u2460\u2461\u2462"); + Assertions.assertEquals(1, morphemes.size()); + Assertions.assertEquals("\u2460\u2461\u2462", morphemes.get(0).surface()); + Assertions.assertEquals(new Span(0, 3), morphemes.get(0).span()); + Assertions.assertTrue(morphemes.get(0).unknown()); + Assertions.assertEquals(List.of("symbol", "unknown"), morphemes.get(0).features()); + } + + /** + * Verifies a mixed run of known and unknown text: the lexicon words around an + * unmapped character are kept intact, the unmapped character becomes its own + * unknown morpheme, and every span stays in original text coordinates. + */ + @Test + void testMixedKnownAndUnknownRuns() { + final String text = "\u6771\u4EAC\u2460\u306B\u884C\u304F"; + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u2460", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 2), new Span(2, 3), new Span(3, 4), new Span(4, 6)}, + tokenizer.tokenizePos(text)); + final List morphemes = tokenizer.analyze(text); + Assertions.assertFalse(morphemes.get(0).unknown()); + Assertions.assertTrue(morphemes.get(1).unknown()); + Assertions.assertFalse(morphemes.get(2).unknown()); + } + + /** + * Verifies that spans keep original text coordinates when the interesting content + * does not start at position zero because of leading whitespace. + */ + @Test + void testSpansStayOriginalAfterLeadingWhitespace() { + final String text = " \u6771\u4EAC\u90FD\u306B\u884C\u304F"; + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u90FD", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(2, 4), new Span(4, 5), new Span(5, 6), new Span(6, 8)}, + tokenizer.tokenizePos(text)); + } + + /** + * Verifies that a lexicon row with fewer than the four mandatory columns is + * rejected at load time. + */ + @Test + void testShortLexiconRowFailsLoud(@TempDir Path broken) throws IOException { + // The rest of the dictionary is well formed, so the short row is what load rejects. + writeUnitMatrixDictionary(broken); + write(broken, LEXICON_CSV, "\u6771,0,0\n"); + Assertions.assertThrows(IOException.class, () -> MecabDictionary.load(broken)); + } + + /** + * Verifies that a non-numeric cost column in a lexicon row is rejected at load + * time. + */ + @Test + void testNonNumericLexiconCostFailsLoud(@TempDir Path broken) throws IOException { + // The rest of the dictionary is well formed, so the cost column is what load rejects. + writeUnitMatrixDictionary(broken); + write(broken, LEXICON_CSV, "\u6771,0,0,abc,noun\n"); + Assertions.assertThrows(IOException.class, () -> MecabDictionary.load(broken)); + } + + /** + * Verifies that a {@code matrix.def} data line with the wrong number of fields is + * rejected at load time. + */ + @Test + void testMalformedMatrixLineFailsLoud(@TempDir Path broken) throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(broken, MATRIX_DEF, "1 1\n0 0\n"); + Assertions.assertThrows(IOException.class, () -> MecabDictionary.load(broken)); + } + + /** + * Verifies that a {@code char.def} code point mapping without a category name is + * rejected at load time. + */ + @Test + void testCharDefMappingWithoutCategoryFailsLoud(@TempDir Path broken) + throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(broken, MATRIX_DEF, UNIT_MATRIX); + write(broken, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\n0x4E00..0x9FFF\n"); + Assertions.assertThrows(IOException.class, () -> MecabDictionary.load(broken)); + } + + /** + * Verifies the fail-loud path when a loadable dictionary cannot cover the input: the + * {@code unk.def} has no DEFAULT template, so a character with neither a lexicon + * entry nor a category template stops segmentation with an exception instead of + * being dropped silently. + */ + @Test + void testMissingDefaultTemplateFailsLoudAtTokenizeTime(@TempDir Path partial) + throws IOException { + write(partial, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(partial, MATRIX_DEF, UNIT_MATRIX); + write(partial, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\nKANJI 0 0 2\n0x4E00..0x9FFF KANJI\n"); + write(partial, UNK_DEF, "KANJI,0,0,8000,noun\n"); + final LatticeTokenizer limited = + new LatticeTokenizer(MecabDictionary.load(partial)); + Assertions.assertThrows(IllegalStateException.class, () -> limited.analyze("\u2460")); + } + + /** + * Verifies that a directory holding a lexicon but none of the definition files is + * rejected at load time, naming the first file that is missing. + */ + @Test + void testMissingDefinitionFileFailsLoud(@TempDir Path broken) throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("required dictionary file is missing: " + + broken.resolve(MATRIX_DEF), e.getMessage()); + } + + /** + * Verifies that a {@code char.def} without the mandatory DEFAULT category is rejected + * at load time rather than leaving unmapped code points without a fallback. + */ + @Test + void testCharDefWithoutDefaultCategoryFailsLoud(@TempDir Path broken) throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(broken, MATRIX_DEF, UNIT_MATRIX); + write(broken, CHAR_DEF, "KANJI 0 0 2\n0x4E00..0x9FFF KANJI\n"); + write(broken, UNK_DEF, "KANJI,0,0,8000,noun\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("char.def defines no DEFAULT category: " + + broken.resolve(CHAR_DEF), e.getMessage()); + } + + /** + * Verifies that a directory with the definition files but no lexicon entry at all is + * rejected at load time, since no text could be segmented against it. + */ + @Test + void testDictionaryWithoutLexiconEntriesFailsLoud(@TempDir Path empty) throws IOException { + writeUnitMatrixDictionary(empty); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(empty)); + Assertions.assertEquals("no lexicon entries found under " + empty, e.getMessage()); + } + + /** + * Verifies that an empty {@code matrix.def} is reported as such instead of as a + * malformed header with nothing to show. + */ + @Test + void testEmptyMatrixDefFailsLoud(@TempDir Path broken) throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(broken, MATRIX_DEF, ""); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("empty matrix.def under " + broken, e.getMessage()); + } + + /** + * Verifies the {@code char.def} fail-loud paths that a malformed line can take: a + * descending code point range, a code point outside the Unicode range, a code point + * field that is not hexadecimal, and a category line missing its length column. + * + * @param charDef The {@code char.def} content under test. + * @param broken The directory the fixture dictionary is written into. + * @throws IOException Thrown if writing the fixture fails. + */ + @ParameterizedTest(name = "[{index}] char.def {0}") + @ValueSource(strings = { + DEFAULT_CATEGORY_LINE + "\n0x0110..0x0100 LATIN\n", + DEFAULT_CATEGORY_LINE + "\n0x110000 LATIN\n", + DEFAULT_CATEGORY_LINE + "\n0xZZ LATIN\n", + "DEFAULT 0 1\n", + "DEFAULT 2 1 0\n", + "DEFAULT 0 true 0\n", + "DEFAULT 0 1 -1\n"}) + void testMalformedCharDefFailsLoud(String charDef, @TempDir Path broken) + throws IOException { + write(broken, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(broken, MATRIX_DEF, UNIT_MATRIX); + write(broken, CHAR_DEF, charDef); + write(broken, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + Assertions.assertThrows(IOException.class, () -> MecabDictionary.load(broken)); + } + + /** + * Verifies that a MeCab-style quoted CSV field may contain a comma, with {@code ""} + * escaping a literal quote, and that the loaded features keep both intact. + */ + @Test + void testQuotedCsvFieldWithCommaLoads(@TempDir Path quoted) throws IOException { + write(quoted, LEXICON_CSV, + "\u6771,0,0,3000,\"noun,common\",\"say \"\"hi\"\"\"\n"); + write(quoted, MATRIX_DEF, UNIT_MATRIX); + write(quoted, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\n"); + write(quoted, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final List morphemes = + new LatticeTokenizer(MecabDictionary.load(quoted)).analyze("\u6771"); + Assertions.assertEquals(1, morphemes.size()); + Assertions.assertEquals(List.of("noun,common", "say \"hi\""), + morphemes.get(0).features()); + } + + /** + * Verifies that an {@code unk.def} template naming a category {@code char.def} never + * defined fails at load with {@link IOException}. + */ + @Test + void testUnkDefUndefinedCategoryFailsLoud(@TempDir Path ghost) throws IOException { + write(ghost, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(ghost, MATRIX_DEF, UNIT_MATRIX); + write(ghost, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\n"); + write(ghost, UNK_DEF, "GHOST,0,0,8000,noun\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(ghost)); + Assertions.assertEquals("unk.def names the undefined category GHOST: " + + ghost.resolve(UNK_DEF), e.getMessage()); + } + + /** + * Writes a miniature dictionary whose {@code char.def} maps a supplementary plane + * range, the shape a UniDic-style distribution uses for the CJK extension blocks. + * + * @param target The directory to write the dictionary files into. Must not be + * {@code null} and must exist. + * @throws IOException Thrown if writing any of the files fails. + */ + private static void writeSupplementaryDictionary(Path target) throws IOException { + write(target, LEXICON_CSV, "\u6771,0,0,6000,noun,common\n"); + write(target, MATRIX_DEF, UNIT_MATRIX); + write(target, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "KANJI 0 0 2", + "LATIN 1 1 0", + "", + "0x4E00..0x9FFF KANJI", + "0x20000..0x2A6DF KANJI", + "0x0061..0x007A LATIN", + "")); + write(target, UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "KANJI,0,0,8000,noun,unknown", + "LATIN,0,0,4000,noun,foreign", + "")); + } + + /** + * Verifies that a {@code char.def} range above U+FFFF is honored rather than + * discarded: a supplementary plane ideograph inside the mapped range takes the + * category the range names, while a supplementary code point outside every mapped + * range still falls back to DEFAULT. + */ + @Test + void testSupplementaryCharDefRangeIsHonored(@TempDir Path supplementary) + throws IOException { + writeSupplementaryDictionary(supplementary); + final MecabDictionary dictionary = MecabDictionary.load(supplementary); + // U+20BB7 is a CJK extension B ideograph inside the mapped range. + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x20BB7).name()); + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x6771).name()); + Assertions.assertEquals("DEFAULT", dictionary.categoryOf(0x2460).name()); + Assertions.assertEquals("DEFAULT", dictionary.categoryOf(0x2A6E0).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf('a').name()); + } + + /** + * Verifies that a supplementary plane ideograph is analyzed as the single character + * it is: one morpheme whose span covers both code units and which carries the + * features of the category its {@code char.def} range names, never one morpheme per + * surrogate. The second case shows the category's length templates count characters, + * not code units, so a run of two supplementary ideographs is still reachable by the + * length-two template. + */ + @Test + void testSupplementaryIdeographIsOneMorpheme(@TempDir Path supplementary) + throws IOException { + writeSupplementaryDictionary(supplementary); + final LatticeTokenizer supplementaryTokenizer = + new LatticeTokenizer(MecabDictionary.load(supplementary)); + // U+20BB7 written as its surrogate pair, per this file's ASCII-only convention. + final String text = "\uD842\uDFB7"; + final List morphemes = supplementaryTokenizer.analyze(text); + Assertions.assertEquals(1, morphemes.size()); + Assertions.assertEquals(text, morphemes.get(0).surface()); + Assertions.assertEquals(new Span(0, 2), morphemes.get(0).span()); + Assertions.assertEquals(List.of("noun", "unknown"), morphemes.get(0).features()); + + final List pair = supplementaryTokenizer.analyze(text + text); + Assertions.assertEquals(1, pair.size()); + Assertions.assertEquals(new Span(0, 4), pair.get(0).span()); + Assertions.assertEquals(List.of("noun", "unknown"), pair.get(0).features()); + } + + /** + * Verifies that a supplementary plane ideograph does not absorb neighbouring text of + * another category: the ideograph and an unmapped symbol beside it stay two + * morphemes, each span covering whole characters. + */ + @Test + void testSupplementaryIdeographDoesNotAbsorbItsNeighbour(@TempDir Path supplementary) + throws IOException { + writeSupplementaryDictionary(supplementary); + final LatticeTokenizer supplementaryTokenizer = + new LatticeTokenizer(MecabDictionary.load(supplementary)); + final String text = "\uD842\uDFB7\u2460"; + Assertions.assertArrayEquals(new String[] {"\uD842\uDFB7", "\u2460"}, + supplementaryTokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] {new Span(0, 2), new Span(2, 3)}, + supplementaryTokenizer.tokenizePos(text)); + } + + /** + * Writes every dictionary file except the lexicon, so a test can supply a lexicon of + * its own against a one by one connection matrix. + * + * @param target The directory to write the dictionary files into. Must not be + * {@code null} and must exist. + * @throws IOException Thrown if writing any of the files fails. + */ + private static void writeUnitMatrixDictionary(Path target) throws IOException { + write(target, MATRIX_DEF, UNIT_MATRIX); + write(target, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\n"); + write(target, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + } + + /** + * Verifies that a lexicon row whose right context id is outside the + * {@code matrix.def} dimensions is rejected at load time, naming the file, the line, + * and the offending id, rather than reaching the cost matrix with an out of range + * index during segmentation. + */ + @Test + void testRightContextIdBeyondMatrixFailsLoudAtLoad(@TempDir Path mismatched) + throws IOException { + writeUnitMatrixDictionary(mismatched); + write(mismatched, LEXICON_CSV, "\u6771,0,5,3000,noun\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(mismatched)); + Assertions.assertEquals("malformed entry at " + mismatched.resolve(LEXICON_CSV) + + " line 1: right context id 5 is outside the matrix.def dimensions 1 1", + e.getMessage()); + } + + /** + * Verifies that a lexicon row whose left context id is outside the {@code matrix.def} + * dimensions is rejected at load time, naming the file, the line, and the offending + * id. + */ + @Test + void testLeftContextIdBeyondMatrixFailsLoudAtLoad(@TempDir Path mismatched) + throws IOException { + writeUnitMatrixDictionary(mismatched); + write(mismatched, LEXICON_CSV, "\u6771,0,0,3000,noun\n\u90FD,7,0,3000,noun\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(mismatched)); + Assertions.assertEquals("malformed entry at " + mismatched.resolve(LEXICON_CSV) + + " line 2: left context id 7 is outside the matrix.def dimensions 1 1", + e.getMessage()); + } + + @Test + void testInvalidArguments() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> new LatticeTokenizer(null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> MecabDictionary.load(null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> MecabDictionary.load(null, StandardCharsets.UTF_8)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> MecabDictionary.load(directory, null)); + Assertions.assertThrows(IllegalArgumentException.class, () -> tokenizer.analyze(null)); + Assertions.assertThrows(IllegalArgumentException.class, () -> tokenizer.tokenize(null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> tokenizer.tokenizePos(null)); + } + + /** + * Verifies the {@link Morpheme} contract every segmentation result is built from: a + * {@code null} span, a {@code null} or empty surface, and {@code null} features are + * all rejected, and the feature list is copied so a later change to the caller's list + * cannot be seen through the morpheme. + */ + @Test + void testMorphemeRejectsInvalidArguments() { + final Span span = new Span(0, 1); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new Morpheme(null, "\u6771", List.of("noun"), false)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new Morpheme(span, null, List.of("noun"), false)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new Morpheme(span, "", List.of("noun"), false)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new Morpheme(span, "\u6771", null, false)); + + final List features = new ArrayList<>(List.of("noun")); + final Morpheme morpheme = new Morpheme(span, "\u6771", features, false); + features.add("proper"); + Assertions.assertEquals(List.of("noun"), morpheme.features()); + } + + /** + * Verifies the supplementary range table's interval cutting and precedence: a later + * {@code char.def} mapping strictly inside an earlier one wins exactly on its own + * stretch, and the earlier category resumes after it, so the cut produces three + * intervals from two overlapping ranges. + */ + @Test + void testLaterSupplementaryMappingWinsInsideAnEarlierRange(@TempDir Path overlapped) + throws IOException { + write(overlapped, LEXICON_CSV, "\u6771,0,0,6000,noun\n"); + write(overlapped, MATRIX_DEF, UNIT_MATRIX); + write(overlapped, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "KANJI 0 0 2", + "LATIN 1 1 0", + "", + "0x20000..0x2FFFF KANJI", + "0x24000..0x25000 LATIN", + "")); + write(overlapped, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final MecabDictionary dictionary = MecabDictionary.load(overlapped); + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x20000).name()); + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x23FFF).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0x24000).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0x25000).name()); + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x25001).name()); + Assertions.assertEquals("KANJI", dictionary.categoryOf(0x2FFFF).name()); + Assertions.assertEquals("DEFAULT", dictionary.categoryOf(0x30000).name()); + } + + /** + * Verifies a {@code char.def} range straddling the BMP boundary: the part up to + * U+FFFF lands in the directly indexed table and the rest in the range table, and + * both halves answer the same category with no gap at the seam. + */ + @Test + void testCharDefRangeStraddlingTheBmpBoundary(@TempDir Path straddling) + throws IOException { + write(straddling, LEXICON_CSV, "\u6771,0,0,6000,noun\n"); + write(straddling, MATRIX_DEF, UNIT_MATRIX); + write(straddling, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "LATIN 1 1 0", + "", + "0xFF00..0x10040 LATIN", + "")); + write(straddling, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final MecabDictionary dictionary = MecabDictionary.load(straddling); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0xFF00).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0xFFFF).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0x10000).name()); + Assertions.assertEquals("LATIN", dictionary.categoryOf(0x10040).name()); + Assertions.assertEquals("DEFAULT", dictionary.categoryOf(0x10041).name()); + Assertions.assertEquals("DEFAULT", dictionary.categoryOf(0xFEFF).name()); + } + + /** + * Verifies that a {@code char.def} mapping to a category its category section never + * defined fails at load, naming the code point and the ghost category, instead of + * silently falling back to DEFAULT at lookup time. + */ + @Test + void testMappingToUndefinedCategoryFailsLoud(@TempDir Path ghost) throws IOException { + write(ghost, LEXICON_CSV, "\u6771,0,0,6000,noun\n"); + write(ghost, MATRIX_DEF, UNIT_MATRIX); + write(ghost, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "", + "0x0100..0x0110 GHOST", + "")); + write(ghost, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(ghost)); + Assertions.assertEquals("char.def declaration at U+0100 names the undefined" + + " category GHOST", e.getMessage()); + } + + /** + * Verifies that a connection cost outside the 16-bit range the binary matrix format + * defines is rejected at load instead of being truncated by the narrowing cast into + * a silently different cost. + */ + @Test + void testMatrixCostOutsideShortRangeFailsLoud(@TempDir Path broken) throws IOException { + writeUnitMatrixDictionary(broken); + write(broken, MATRIX_DEF, "1 1\n0 0 40000\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("malformed matrix.def line 2: connection cost 40000 is" + + " outside the 16-bit range the format defines", e.getMessage()); + } + + /** + * Verifies that {@code matrix.def} dimensions whose product exceeds the addressable + * array size fail loud at the header instead of overflowing the int multiplication + * into a negative or wrapped allocation size. + */ + @Test + void testMatrixDimensionProductBeyondIntRangeFailsLoud(@TempDir Path broken) + throws IOException { + writeUnitMatrixDictionary(broken); + write(broken, MATRIX_DEF, "70000 70000\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("matrix.def dimensions 70000 x 70000 overflow the" + + " addressable connection matrix", e.getMessage()); + } + + /** + * Verifies that a single matrix dimension above {@link ResourceLimits#MAX_ENTRIES} + * is rejected before the connection-cost array is allocated. + */ + @Test + void testMatrixDimensionAboveMaxEntriesFailsLoud(@TempDir Path broken) throws IOException { + writeUnitMatrixDictionary(broken); + final int over = ResourceLimits.MAX_ENTRIES + 1; + write(broken, MATRIX_DEF, over + " 1\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertTrue(e.getMessage().contains("exceed safe limit of " + + ResourceLimits.MAX_ENTRIES), e.getMessage()); + } + + /** + * Verifies that a matrix whose cell count is above + * {@link ResourceLimits#MAX_MATRIX_CELLS} but still below {@link Integer#MAX_VALUE} + * is rejected. Without that bound, a header such as {@code 46340 46340} would + * allocate about 4 GiB of shorts. + */ + @Test + void testMatrixCellCountAboveMaxCellsFailsLoud(@TempDir Path broken) throws IOException { + writeUnitMatrixDictionary(broken); + // 11600 x 11600 = 134_560_000 cells, above the default MAX_MATRIX_CELLS of 2^27. + write(broken, MATRIX_DEF, "11600 11600\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("matrix.def dimensions 11600 x 11600 exceed safe limit of " + + ResourceLimits.MAX_MATRIX_CELLS, e.getMessage()); + } + + /** + * Verifies that the dimensions of a real published distribution pass the header + * bound. mecab-ko-dic 2.1.1 declares {@code 3822 2693}, which is 10,292,646 cells: + * above {@link ResourceLimits#MAX_ENTRIES} but a legitimate 20 MB cost matrix, so + * the cell bound must be sized to cells rather than reusing the entry bound. The + * load still fails on the truncated body, but with the incomplete-matrix message, + * not the safe-limit one. + */ + @Test + void testKoDicSizedMatrixDimensionsPassTheHeaderBound(@TempDir Path koDic) + throws IOException { + writeUnitMatrixDictionary(koDic); + write(koDic, MATRIX_DEF, "3822 2693\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(koDic)); + Assertions.assertEquals("matrix.def declares 3822 x 2693 connection costs but only 0" + + " pairs are listed", e.getMessage()); + } + + /** + * Verifies that a truncated {@code matrix.def} fails loud. Unlisted pairs must not + * keep the short-array default of cost zero, the cheapest connection. + */ + @Test + void testIncompleteMatrixFailsLoud(@TempDir Path broken) throws IOException { + writeUnitMatrixDictionary(broken); + write(broken, MATRIX_DEF, "2 2\n0 0 1\n0 1 2\n1 0 3\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("matrix.def declares 2 x 2 connection costs but only 3" + + " pairs are listed", e.getMessage()); + } + + /** + * Verifies that a {@code matrix.def} data row naming context ids outside the + * declared dimensions is rejected at load with the offending line and ids. + */ + @Test + void testMatrixRowContextIdsOutsideDimensionsFailLoud(@TempDir Path broken) + throws IOException { + writeUnitMatrixDictionary(broken); + write(broken, MATRIX_DEF, "1 1\n2 0 5\n"); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(broken)); + Assertions.assertEquals("malformed matrix.def line 2: context ids 2 0 are outside" + + " the declared dimensions 1 1", e.getMessage()); + } + + @Test + void testZeroLengthCategoryUsesItsUnknownTemplate(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "SINGLE 1 0 0", + "", + "0x2460 SINGLE", + "")); + write(dictionary, UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "SINGLE,0,0,1000,symbol,single", + "")); + + final Morpheme morpheme = new LatticeTokenizer(MecabDictionary.load(dictionary)) + .analyze("\u2460").get(0); + + Assertions.assertEquals(List.of("symbol", "single"), morpheme.features()); + } + + @Test + void testRejectsEmptyLexiconSurface(@TempDir Path dictionary) throws IOException { + writeUnitMatrixDictionary(dictionary); + write(dictionary, LEXICON_CSV, String.join("\n", + "\u6771,0,0,3000,noun", + ",0,0,3000,noun", + "")); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("malformed entry at " + dictionary.resolve(LEXICON_CSV) + + " line 2: surface must not be empty", e.getMessage()); + } + + @ParameterizedTest(name = "word cost {0}") + @ValueSource(ints = {-32769, 32768}) + void testRejectsLexiconCostOutsideShortRange(int cost, @TempDir Path dictionary) + throws IOException { + writeUnitMatrixDictionary(dictionary); + write(dictionary, LEXICON_CSV, "\u6771,0,0," + cost + ",noun\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("malformed entry at " + dictionary.resolve(LEXICON_CSV) + + " line 1: word cost " + cost + + " is outside the 16-bit range the format defines", e.getMessage()); + } + + @Test + void testRejectsDuplicateCharacterCategory(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "\u6771,0,0,3000,noun\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, DEFAULT_CATEGORY_LINE + "\nDEFAULT 1 0 2\n"); + write(dictionary, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("duplicate char.def category DEFAULT at line 2", + e.getMessage()); + } + + @Test + void testRejectsDuplicateMatrixEntry(@TempDir Path dictionary) throws IOException { + writeUnitMatrixDictionary(dictionary); + write(dictionary, MATRIX_DEF, "1 1\n0 0 1\n0 0 2\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("duplicate matrix.def entry 0 0 at line 3", e.getMessage()); + } + + @Test + void testRejectsMalformedDictionaryEncoding(@TempDir Path dictionary) + throws IOException { + writeUnitMatrixDictionary(dictionary); + Files.write(dictionary.resolve(LEXICON_CSV), + new byte[] {TRUNCATED_UTF8_LEAD_BYTE, ',', '0', ',', '0', ',', '1', '\n'}); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertInstanceOf(MalformedInputException.class, e); + Assertions.assertEquals("Input length = 1", e.getMessage()); + } + + @Test + void testLongLexiconSurfaceLoads(@TempDir Path dictionary) throws IOException { + writeUnitMatrixDictionary(dictionary); + final String surface = "a".repeat(20_000); + write(dictionary, LEXICON_CSV, surface + ",0,0,3000,fixture\n"); + + final LatticeTokenizer longSurfaceTokenizer = + new LatticeTokenizer(MecabDictionary.load(dictionary)); + + Assertions.assertArrayEquals(new String[] {surface}, + longSurfaceTokenizer.tokenize(surface)); + } + + @Test + void testSecondaryCharacterCategoryExtendsUnknownRun(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "\u6771,0,0,6000,noun\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "KANJI 1 0 1", + "KANJINUMERIC 1 1 0", + "", + "0x4E00 KANJINUMERIC KANJI", + "0x5C71 KANJI", + "")); + write(dictionary, UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "KANJI,0,0,5000,noun,unknown", + "KANJINUMERIC,0,0,1000,number,unknown", + "")); + + final List morphemes = new LatticeTokenizer( + MecabDictionary.load(dictionary)).analyze("\u4E00\u5C71"); + + Assertions.assertEquals(1, morphemes.size()); + Assertions.assertEquals("\u4E00\u5C71", morphemes.get(0).surface()); + Assertions.assertEquals(List.of("number", "unknown"), morphemes.get(0).features()); + } + + /** + * Verifies MeCab's category-chain grouping. Character A has X and Y, B has X, and C + * has Y. A and B intersect, while B and C do not, so the initial unknown word covers + * {@code ab}. + */ + @Test + void testMultipleCategoriesUsePairwiseGrouping(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "z,0,0,6000,fixture\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "X 1 1 0", + "Y 1 1 0", + "", + "0x0061 X Y", + "0x0062 X", + "0x0063 Y", + "")); + write(dictionary, UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "X,0,0,1000,x,unknown", + "Y,0,0,1000,y,unknown", + "")); + + final LatticeTokenizer groupingTokenizer = + new LatticeTokenizer(MecabDictionary.load(dictionary)); + + Assertions.assertArrayEquals(new String[] {"ab", "c"}, + groupingTokenizer.tokenize("abc")); + } + + /** + * Verifies that an intermediate multi-category character can connect a run. + * Character A has X, B has X and Y, and C has Y. MeCab advances the active + * assignment at each position, allowing B to connect both portions. + */ + @Test + void testCategoryOverlapCanConnectRun(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "z,0,0,6000,fixture\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "X 1 1 0", + "Y 1 1 0", + "", + "0x0061 X", + "0x0062 X Y", + "0x0063 Y", + "")); + write(dictionary, UNK_DEF, String.join("\n", + DEFAULT_UNKNOWN_TEMPLATE, + "X,0,0,1000,x,unknown", + "Y,0,0,1000,y,unknown", + "")); + + final LatticeTokenizer groupingTokenizer = + new LatticeTokenizer(MecabDictionary.load(dictionary)); + + Assertions.assertArrayEquals(new String[] {"abc"}, + groupingTokenizer.tokenize("abc")); + } + + @Test + void testRejectsTooManyCharacterCategories(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "z,0,0,6000,fixture\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + final StringBuilder charDef = new StringBuilder(DEFAULT_CATEGORY_LINE).append('\n'); + for (int i = 1; i < 18; i++) { + charDef.append('C').append(i).append(" 0 0 1\n"); + } + write(dictionary, CHAR_DEF, charDef.toString()); + write(dictionary, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("char.def defines 18 categories; MeCab supports at most 17", + e.getMessage()); + } + + @Test + void testRejectsCharacterCategoryLengthAboveMecabLimit(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "z,0,0,6000,fixture\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, "DEFAULT 0 1 16\n"); + write(dictionary, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("category LENGTH must be between 0 and 15 at " + + dictionary.resolve(CHAR_DEF) + " line 1", e.getMessage()); + } + + @Test + void testRejectsUndefinedSecondaryCharacterCategory(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "\u6771,0,0,6000,noun\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "KANJINUMERIC 1 1 0", + "", + "0x4E00 KANJINUMERIC GHOST", + "")); + write(dictionary, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("char.def declaration at U+4E00 names the undefined" + + " category GHOST", e.getMessage()); + } + + @Test + void testRejectsUndefinedCategoryOnShadowedMapping(@TempDir Path dictionary) + throws IOException { + write(dictionary, LEXICON_CSV, "z,0,0,6000,fixture\n"); + write(dictionary, MATRIX_DEF, UNIT_MATRIX); + write(dictionary, CHAR_DEF, String.join("\n", + DEFAULT_CATEGORY_LINE, + "LATIN 1 1 0", + "", + "0x0061 GHOST", + "0x0061 LATIN", + "")); + write(dictionary, UNK_DEF, DEFAULT_UNKNOWN_TEMPLATE + "\n"); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionary.load(dictionary)); + + Assertions.assertEquals("char.def declaration at U+0061 names the undefined" + + " category GHOST", e.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeUsageExampleTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeUsageExampleTest.java new file mode 100644 index 0000000000..7e3f3e5a88 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/LatticeUsageExampleTest.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.tools.tokenize.lattice; + +import java.io.IOException; +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import opennlp.tools.util.Span; +import opennlp.tools.util.archive.TarArchives; + +/** + * Demonstrates the intended end-to-end usage of this package with miniature, + * project-authored data: a MeCab-format dictionary archive is installed with + * {@link MecabDictionaryInstaller}, loaded as a {@link MecabDictionary}, and segmented + * with a {@link LatticeTokenizer}; a plain frequency lexicon is loaded and segmented + * with a {@link UnigramSegmenter}. Everything is written to a temporary directory by + * the test itself; no external dictionary or lexicon data and no network access are + * involved. + * + *

    Source strings are written as Unicode escapes to keep this file ASCII-only. The + * Japanese fixture words are Tokyo (U+6771 U+4EAC), Kyoto (U+4EAC U+90FD), east + * (U+6771), the metropolis suffix (U+90FD), the case particle ni (U+306B), and the + * verb iku, to go (U+884C U+304F); the Chinese fixture words are wo, I (U+6211), + * laidao, arrive (U+6765 U+5230), Beijing (U+5317 U+4EAC), and Tiananmen + * (U+5929 U+5B89 U+95E8).

    + */ +public class LatticeUsageExampleTest { + + /** + * Walks the full MeCab-format flow: package a miniature Japanese dictionary as a + * {@code tar.gz} archive, install it from a file URI, load it, and tokenize. The + * segmentation must pick the cheaper path (Tokyo plus the metropolis suffix) over + * the competing reading (east plus Kyoto), the spans must be in original text + * coordinates, and the morphemes must carry the dictionary's feature columns. + */ + @Test + void testInstallLoadAndTokenizeAMecabFormatDictionary(@TempDir Path work) + throws IOException { + // A minimal but complete dictionary: one lexicon file plus the three definition + // files every MeCab-format distribution contains, wrapped like a release archive. + final byte[] archive = TarArchives.gzippedTar(new String[][] { + {"mini-dict-0.1/lexicon.csv", String.join("\n", + "\u6771\u4EAC,0,0,3000,noun,proper", + "\u4EAC\u90FD,0,0,3000,noun,proper", + "\u6771,0,0,6000,noun,common", + "\u90FD,0,0,4000,noun,suffix", + "\u306B,0,0,1000,particle,case", + "\u884C\u304F,0,0,3000,verb,base", + "")}, + {"mini-dict-0.1/matrix.def", "1 1\n0 0 0\n"}, + {"mini-dict-0.1/char.def", String.join("\n", + "DEFAULT 0 1 0", + "KANJI 0 0 2", + "HIRAGANA 0 1 0", + "", + "0x3041..0x3096 HIRAGANA", + "0x4E00..0x9FFF KANJI", + "")}, + {"mini-dict-0.1/unk.def", String.join("\n", + "DEFAULT,0,0,10000,symbol,unknown", + "KANJI,0,0,8000,noun,unknown", + "HIRAGANA,0,0,9000,particle,unknown", + "")}, + {"mini-dict-0.1/README", "not a dictionary payload file"}}); + final Path archiveFile = work.resolve("mini-dict-0.1.tar.gz"); + Files.write(archiveFile, archive); + + // Install the local archive and unpack the payload. + final Path dictionaryDirectory = work.resolve("dictionary"); + final int extracted = MecabDictionaryInstaller.install( + archiveFile.toUri(), dictionaryDirectory); + Assertions.assertEquals(4, extracted); + + // Load and tokenize; both views must agree and stay in original coordinates. + final LatticeTokenizer tokenizer = + new LatticeTokenizer(MecabDictionary.load(dictionaryDirectory)); + final String text = "\u6771\u4EAC\u90FD\u306B\u884C\u304F"; + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u90FD", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 2), new Span(2, 3), new Span(3, 4), new Span(4, 6)}, + tokenizer.tokenizePos(text)); + + // The analyze view adds the dictionary's feature columns to every morpheme. + final List morphemes = tokenizer.analyze(text); + Assertions.assertEquals(4, morphemes.size()); + Assertions.assertEquals("\u6771\u4EAC", morphemes.get(0).surface()); + Assertions.assertEquals(List.of("noun", "proper"), morphemes.get(0).features()); + Assertions.assertFalse(morphemes.get(0).unknown()); + } + + /** + * Walks the frequency-lexicon flow: write a miniature word-count lexicon to a file, + * load it, and segment. The segmentation must recover the listed multi-character + * words with spans in original text coordinates. + */ + @Test + void testLoadAndSegmentWithAFrequencyLexicon(@TempDir Path work) throws IOException { + // One word, its count, and an optional tag per line, whitespace separated. + final Path lexicon = work.resolve("words.txt"); + Files.write(lexicon, String.join("\n", + "\u6211 5000 r", + "\u6765\u5230 2000 v", + "\u5317\u4EAC 3000 ns", + "\u5929\u5B89\u95E8 1200 ns", + "").getBytes(StandardCharsets.UTF_8)); + + final UnigramSegmenter segmenter = UnigramSegmenter.load(lexicon); + final String text = "\u6211\u6765\u5230\u5317\u4EAC\u5929\u5B89\u95E8"; + Assertions.assertArrayEquals( + new String[] {"\u6211", "\u6765\u5230", "\u5317\u4EAC", "\u5929\u5B89\u95E8"}, + segmenter.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 1), new Span(1, 3), new Span(3, 5), new Span(5, 8)}, + segmenter.tokenizePos(text)); + } + + /** + * Walks the non-UTF-8 flow that widely used Japanese distributions require: the same + * miniature dictionary is written to disk encoded in EUC-JP and loaded through the + * charset-taking overload. The segmentation must match the UTF-8 run exactly, which + * shows the encoding is a property of loading, not of tokenization. + */ + @Test + void testLoadAnEucJpEncodedDictionary(@TempDir Path work) throws IOException { + final Charset eucJp = Charset.forName("EUC-JP"); + Files.write(work.resolve("lexicon.csv"), String.join("\n", + "\u6771\u4EAC,0,0,3000,noun,proper", + "\u4EAC\u90FD,0,0,3000,noun,proper", + "\u6771,0,0,6000,noun,common", + "\u90FD,0,0,4000,noun,suffix", + "\u306B,0,0,1000,particle,case", + "\u884C\u304F,0,0,3000,verb,base", + "").getBytes(eucJp)); + Files.write(work.resolve("matrix.def"), "1 1\n0 0 0\n".getBytes(eucJp)); + Files.write(work.resolve("char.def"), String.join("\n", + "DEFAULT 0 1 0", + "KANJI 0 0 2", + "HIRAGANA 0 1 0", + "", + "0x3041..0x3096 HIRAGANA", + "0x4E00..0x9FFF KANJI", + "").getBytes(eucJp)); + Files.write(work.resolve("unk.def"), String.join("\n", + "DEFAULT,0,0,10000,symbol,unknown", + "KANJI,0,0,8000,noun,unknown", + "HIRAGANA,0,0,9000,particle,unknown", + "").getBytes(eucJp)); + + final LatticeTokenizer tokenizer = + new LatticeTokenizer(MecabDictionary.load(work, eucJp)); + Assertions.assertArrayEquals( + new String[] {"\u6771\u4EAC", "\u90FD", "\u306B", "\u884C\u304F"}, + tokenizer.tokenize("\u6771\u4EAC\u90FD\u306B\u884C\u304F")); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstallerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstallerTest.java new file mode 100644 index 0000000000..cbbfe4915b --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/MecabDictionaryInstallerTest.java @@ -0,0 +1,333 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.nio.file.FileSystem; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.DictionaryCatalog; +import opennlp.tools.util.DigestTestUtil; +import opennlp.tools.util.archive.TarArchives; + +/** + * Tests the installer against project-authored, in-memory archives; no external + * dictionary data and no network access are involved. Fetch, verification, and + * unpacking limits are exercised in {@code opennlp.tools.util.ResourceInstallerTest}, + * which tests the shared installation path this installer delegates to. + */ +public class MecabDictionaryInstallerTest { + + @Test + void testInstallsDictionaryFilesAndFlattensPaths(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"dict-1.0/lexicon.csv", "cat,0,0,100,noun\n"}, + {"dict-1.0/matrix.def", "1 1\n0 0 0\n"}, + {"dict-1.0/char.def", "DEFAULT 0 1 0\n"}, + {"dict-1.0/unk.def", "DEFAULT,0,0,10000,unknown\n"}, + {"dict-1.0/README", "not a dictionary file"}, + {"dict-1.0/dicrc", "config"}}); + + final int installed = + MecabDictionaryInstaller.install(archiveFile.toUri(), target); + + Assertions.assertEquals(5, installed); + Assertions.assertTrue(Files.exists(target.resolve("lexicon.csv"))); + Assertions.assertTrue(Files.exists(target.resolve("matrix.def"))); + Assertions.assertTrue(Files.exists(target.resolve("char.def"))); + Assertions.assertTrue(Files.exists(target.resolve("unk.def"))); + Assertions.assertTrue(Files.exists(target.resolve("dicrc"))); + Assertions.assertTrue(Files.notExists(target.resolve("README"))); + Assertions.assertEquals("cat,0,0,100,noun\n", + Files.readString(target.resolve("lexicon.csv"))); + } + + /** + * Verifies that only files at the archive root count as dictionary payload. + * mecab-ko-dic ships template user dictionaries under {@code user-dic/} whose + * numeric fields are empty. They are input for {@code mecab-dict-index}, not loadable + * lexicon data. Flattening them next to the real lexicon fails the subsequent load, + * and on a case-insensitive file system a template can silently overwrite a real + * lexicon file of the same base name. + */ + @Test + void testNestedTemplateFilesAreNotInstalled(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"dict-1.0/NNP.csv", "cat,1786,3546,2953,noun\n"}, + {"dict-1.0/matrix.def", "1 1\n0 0 0\n"}, + {"dict-1.0/user-dic/person.csv", "template,,,,noun\n"}}); + + final int installed = + MecabDictionaryInstaller.install(archiveFile.toUri(), target); + + Assertions.assertEquals(2, installed); + Assertions.assertTrue(Files.exists(target.resolve("NNP.csv"))); + Assertions.assertTrue(Files.notExists(target.resolve("person.csv"))); + } + + /** + * Verifies that two payload entries that flatten to the same base name are rejected, + * because keeping either one silently would install an ambiguous dictionary. + */ + @Test + void testEntriesFlatteningToTheSameNameAreRejected(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"words.csv", "cat,0,0,100,noun\n"}, + {"d/words.csv", "dog,0,0,100,noun\n"}}); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + Assertions.assertEquals( + "the archive flattens two entries to the same name: words.csv", + e.getMessage()); + Assertions.assertTrue(Files.notExists(target.resolve("words.csv"))); + } + + @Test + void testInstallReadsAFileUri(@TempDir Path source, @TempDir Path target) + throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"d/words.csv", "cat,0,0,100,noun\n"}, + {"d/matrix.def", "1 1\n0 0 0\n"}}); + + final int installed = + MecabDictionaryInstaller.install(archiveFile.toUri(), target); + + Assertions.assertEquals(2, installed); + Assertions.assertTrue(Files.exists(target.resolve("words.csv"))); + } + + /** + * Checks installation when the target uses a different filesystem provider. + * + * @param source The directory containing the fixture archive. + * @param scratch The directory containing the ZIP filesystem. + * @throws IOException Thrown if creating or installing the fixture fails. + */ + @Test + void testInstallIntoANonDefaultFileSystem(@TempDir Path source, @TempDir Path scratch) + throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"d/words.csv", "cat,0,0,100,noun\n"}, + {"d/matrix.def", "1 1\n0 0 0\n"}}); + final URI zip = URI.create("jar:" + scratch.resolve("target.zip").toUri()); + + try (FileSystem targetFileSystem = + FileSystems.newFileSystem(zip, Map.of("create", "true"))) { + final Path target = targetFileSystem.getPath("/dictionary"); + + Assertions.assertEquals(2, + MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + Assertions.assertEquals("cat,0,0,100,noun\n", + Files.readString(target.resolve("words.csv"))); + } + } + + /** + * Verifies that a pax long-name entry installs under its real name. The common + * distributions ship plain ustar today, but {@code bsdtar} and + * {@code tar --format=posix} write pax archives, whose over-100-byte names live in an + * extension header instead of the header name field. + */ + @Test + void testPaxLongNamedEntryInstallsUnderItsRealName(@TempDir Path source, + @TempDir Path target) throws IOException { + final String longBaseName = "a".repeat(110) + ".csv"; + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + TarArchives.entry(tar, "PaxHeaders.0/lexicon", + TarArchives.paxRecord("path", "dict-1.0/" + longBaseName), 'x'); + TarArchives.entry(tar, "dict-1.0/" + "a".repeat(88), + "cat,0,0,100,noun\n".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TarArchives.TERMINATOR_SIZE]); + final Path archiveFile = source.resolve("dict.tar.gz"); + Files.write(archiveFile, TarArchives.gzip(tar.toByteArray())); + + final int installed = + MecabDictionaryInstaller.install(archiveFile.toUri(), target); + + Assertions.assertEquals(1, installed); + Assertions.assertEquals("cat,0,0,100,noun\n", + Files.readString(target.resolve(longBaseName))); + } + + /** + * Verifies that installing into a target that already holds a dictionary file of the + * same name is rejected, leaving the first installation in place. Refreshing a + * dictionary means removing its old files first. + */ + @Test + void testReinstallOverAnExistingDictionaryIsRejected(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path archiveFile = archive(source, new String[][] { + {"d/words.csv", "cat,0,0,100,noun\n"}}); + Assertions.assertEquals(1, + MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + Assertions.assertTrue(e.getMessage().contains("target already contains: ")); + Assertions.assertEquals("cat,0,0,100,noun\n", + Files.readString(target.resolve("words.csv"))); + } + + @Test + void testRemoteInstallWithoutDigestIsRejected(@TempDir Path target) { + final IllegalArgumentException e = + Assertions.assertThrows(IllegalArgumentException.class, + () -> MecabDictionaryInstaller.install( + URI.create("https://example.invalid/dict.tar.gz"), target)); + Assertions.assertTrue( + e.getMessage().contains("checksum must be given for an http or https source")); + } + + @Test + void testInstallVerifiesDigest(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path archiveFile = archive(source, new String[][] { + {"d/words.csv", "cat,0,0,100,noun\n"}, + {"d/matrix.def", "1 1\n0 0 0\n"}}); + final byte[] archive = Files.readAllBytes(archiveFile); + + final int installed = MecabDictionaryInstaller.install( + archiveFile.toUri(), target, DigestTestUtil.sha512(archive)); + Assertions.assertEquals(2, installed); + + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionaryInstaller.install(archiveFile.toUri(), + target.resolve("other"), DigestTestUtil.sha512(new byte[] {1}))); + Assertions.assertTrue(e.getMessage().contains("checksum mismatch: expected")); + } + + @Test + void testInstallFromCatalogRequiresRemoteProperty(@TempDir Path target) + throws IOException { + final DictionaryCatalog catalog = emptyCatalog(); + final String previous = + System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + try { + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionaryInstaller.installFromCatalog( + catalog, "mecab.ipadic", target)); + Assertions.assertTrue( + e.getMessage().contains(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY)); + } finally { + if (previous == null) { + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + } else { + System.setProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY, previous); + } + } + } + + @Test + void testArchivesWithoutDictionaryFilesAreRejected(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path archiveFile = + archive(source, new String[][] {{"readme.txt", "nothing here"}}); + final IOException e = Assertions.assertThrows(IOException.class, + () -> MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + Assertions.assertEquals("the archive contains no dictionary file", e.getMessage()); + } + + /** + * Checks each public installer parameter independently. + * + * @param argument The invalid parameter. + * @param target A scratch directory managed by the test framework. + * @throws IOException Thrown if the empty catalog cannot be loaded. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"archive", "targetDirectory", "catalog", "dictionaryId", + "catalog targetDirectory"}) + void testInvalidArguments(String argument, @TempDir Path target) throws IOException { + final DictionaryCatalog catalog = emptyCatalog(); + final Executable install = switch (argument) { + case "archive" -> () -> MecabDictionaryInstaller.install(null, target); + case "targetDirectory" -> () -> + MecabDictionaryInstaller.install(target.toUri(), null); + case "catalog" -> () -> + MecabDictionaryInstaller.installFromCatalog(null, "mecab.ipadic", target); + case "dictionaryId" -> () -> + MecabDictionaryInstaller.installFromCatalog(catalog, null, target); + case "catalog targetDirectory" -> () -> + MecabDictionaryInstaller.installFromCatalog(catalog, "mecab.ipadic", null); + default -> throw new IllegalArgumentException("unknown argument: " + argument); + }; + + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, install); + final String parameter = argument.startsWith("catalog ") + ? argument.substring("catalog ".length()) : argument; + Assertions.assertEquals(parameter + " must not be null", thrown.getMessage()); + } + + private static DictionaryCatalog emptyCatalog() throws IOException { + return DictionaryCatalog.load(new ByteArrayInputStream(new byte[0])); + } + + /** + * Writes a gzip-compressed tar archive of the given entries to a file. + * + * @param directory The directory to write the archive into. + * @param entries The entries as {@code {name, content}} pairs. + * @return The archive file. Never {@code null}. + * @throws IOException Thrown if writing fails. + */ + private static Path archive(Path directory, String[][] entries) throws IOException { + final Path archiveFile = directory.resolve("dict.tar.gz"); + Files.write(archiveFile, TarArchives.gzippedTar(entries)); + return archiveFile; + } + + @Test + void testStaleScratchOfAKilledInstallIsRemoved(@TempDir Path source, + @TempDir Path target) throws IOException { + final Path stale = Files.createDirectories(target.resolve(".mecab-dict-OLD")); + Files.writeString(stale.resolve("words.csv"), "half"); + final Path archiveFile = archive(source, new String[][] { + {"d/words.csv", "cat,0,0,100,noun\n"}}); + + Assertions.assertEquals(1, + MecabDictionaryInstaller.install(archiveFile.toUri(), target)); + + try (Stream entries = Files.list(target)) { + Assertions.assertEquals(List.of("words.csv"), + entries.map(path -> path.getFileName().toString()).sorted().toList()); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/UnigramSegmenterTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/UnigramSegmenterTest.java new file mode 100644 index 0000000000..c6e00cf1a0 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/lattice/UnigramSegmenterTest.java @@ -0,0 +1,270 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.tools.tokenize.lattice; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.MalformedInputException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.Span; + +/** + * Tests the frequency-driven segmenter against a project-authored miniature lexicon; + * no external lexicon data is involved. + * + *

    Source strings are written as Unicode escapes to keep this file ASCII-only; the + * class works over the same miniature Chinese frequency lexicon as the sibling usage + * example, and its Javadoc spells out each fixture word.

    + */ +public class UnigramSegmenterTest { + + /** UTF-8 lead byte with the required continuation byte omitted. */ + private static final byte TRUNCATED_UTF8_LEAD_BYTE = (byte) 0xC3; + + private static final String LEXICON = String.join("\n", + "\u6211 5000 r", + "\u6765\u5230 2000 v", + "\u5317\u4EAC 3000 ns", + "\u6E05\u534E\u5927\u5B66 800 nt", + "\u6E05\u534E 400 ns", + "\u534E\u5927 100 ns", + "\u5927\u5B66 1500 n", + "\u7684 9000 uj", + ""); + + private static UnigramSegmenter segmenter; + + @BeforeAll + static void loadLexicon() throws IOException { + segmenter = UnigramSegmenter.load( + new ByteArrayInputStream(LEXICON.getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8); + } + + @Test + void testPrefersWholeWordsOverFragments() { + Assertions.assertArrayEquals( + new String[] {"\u6211", "\u6765\u5230", "\u5317\u4EAC", "\u6E05\u534E\u5927\u5B66"}, + segmenter.tokenize("\u6211\u6765\u5230\u5317\u4EAC\u6E05\u534E\u5927\u5B66")); + } + + @Test + void testSpansStayInOriginalCoordinates() { + Assertions.assertArrayEquals(new Span[] { + new Span(0, 1), new Span(1, 3), new Span(3, 5), new Span(5, 9)}, + segmenter.tokenizePos("\u6211\u6765\u5230\u5317\u4EAC\u6E05\u534E\u5927\u5B66")); + } + + @Test + void testUnknownCharactersFallBackToSingles() { + Assertions.assertArrayEquals( + new String[] {"\u6211", "\u7231", "\u5317\u4EAC"}, + segmenter.tokenize("\u6211\u7231\u5317\u4EAC")); + } + + @Test + void testWhitespaceSeparates() { + Assertions.assertArrayEquals( + new String[] {"\u5317\u4EAC", "\u5927\u5B66"}, + segmenter.tokenize("\u5317\u4EAC \u5927\u5B66")); + Assertions.assertEquals(0, segmenter.tokenizePos(" ").length); + } + + /** + * Verifies that empty input yields empty results from both views of the segmenter. + */ + @Test + void testEmptyInputYieldsEmptyResults() { + Assertions.assertArrayEquals(new String[0], segmenter.tokenize("")); + Assertions.assertArrayEquals(new Span[0], segmenter.tokenizePos("")); + } + + /** + * Verifies single-character input for a listed word and for a character the + * lexicon does not know: both come back as exactly one token covering + * {@code [0, 1)}. + */ + @Test + void testSingleCharacterInput() { + Assertions.assertArrayEquals(new String[] {"\u6211"}, segmenter.tokenize("\u6211")); + Assertions.assertArrayEquals(new Span[] {new Span(0, 1)}, segmenter.tokenizePos("\u6211")); + Assertions.assertArrayEquals(new String[] {"\u7231"}, segmenter.tokenize("\u7231")); + Assertions.assertArrayEquals(new Span[] {new Span(0, 1)}, segmenter.tokenizePos("\u7231")); + } + + /** + * Verifies input made entirely of characters absent from the lexicon: every + * character becomes its own single-character token, since only the unknown + * fallback is available. + */ + @Test + void testEntirelyUnknownInputFallsBackToSingleCharacters() { + Assertions.assertArrayEquals( + new String[] {"x", "y", "z"}, + segmenter.tokenize("xyz")); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 1), new Span(1, 2), new Span(2, 3)}, + segmenter.tokenizePos("xyz")); + } + + /** + * Verifies a mixed run of known and unknown text inside one whitespace-free + * stretch: the unknown character becomes a single token while the listed words + * around it, including the longest listed compound, stay intact. + */ + @Test + void testMixedKnownAndUnknownRuns() { + Assertions.assertArrayEquals( + new String[] {"\u6211", "\u7231", "\u6E05\u534E\u5927\u5B66"}, + segmenter.tokenize("\u6211\u7231\u6E05\u534E\u5927\u5B66")); + Assertions.assertArrayEquals(new Span[] { + new Span(0, 1), new Span(1, 2), new Span(2, 6)}, + segmenter.tokenizePos("\u6211\u7231\u6E05\u534E\u5927\u5B66")); + } + + /** + * Verifies that spans keep original text coordinates when the content does not + * start at position zero because of leading whitespace. + */ + @Test + void testSpansStayOriginalAfterLeadingWhitespace() { + final String text = " \u6211\u6765\u5230\u5317\u4EAC"; + Assertions.assertArrayEquals( + new String[] {"\u6211", "\u6765\u5230", "\u5317\u4EAC"}, + segmenter.tokenize(text)); + Assertions.assertArrayEquals(new Span[] { + new Span(2, 3), new Span(3, 5), new Span(5, 7)}, + segmenter.tokenizePos(text)); + } + + @ParameterizedTest(name = "lexicon content \"{0}\"") + @ValueSource(strings = {"word\n", "word abc\n", "word 0\n", "\n\n"}) + void testMalformedLexiconsFailLoud(String lexicon) { + Assertions.assertThrows(IOException.class, () -> UnigramSegmenter.load( + new ByteArrayInputStream(lexicon.getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8)); + } + + @Test + void testEntryLimit() { + final String lexicon = "first 1\nsecond 1\n"; + final IOException e = Assertions.assertThrows(IOException.class, + () -> UnigramSegmenter.load( + new ByteArrayInputStream(lexicon.getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8, 1)); + Assertions.assertEquals("lexicon entry count exceeds safe limit of 1", e.getMessage()); + } + + @Test + void testCountTotalOverflow() { + final String lexicon = "first " + Long.MAX_VALUE + "\nsecond 1\n"; + final IOException e = Assertions.assertThrows(IOException.class, + () -> UnigramSegmenter.load( + new ByteArrayInputStream(lexicon.getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8)); + Assertions.assertEquals("lexicon count total overflows at line 2", e.getMessage()); + } + + @Test + void testInvalidArguments() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> UnigramSegmenter.load((Path) null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> UnigramSegmenter.load((Path) null, StandardCharsets.UTF_8)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> UnigramSegmenter.load(Path.of("words.txt"), null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> UnigramSegmenter.load((InputStream) null, StandardCharsets.UTF_8)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> UnigramSegmenter.load(new ByteArrayInputStream(new byte[0]), null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> segmenter.tokenize(null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> segmenter.tokenizePos(null)); + } + + /** + * Verifies that the unknown-character fallback advances one code point instead of one + * code unit: a supplementary character absent from the lexicon comes back as one span + * over its surrogate pair, and no span boundary occurs inside it. + */ + @Test + void testUnknownSupplementaryCharacterIsNeverSplit() { + // U+20BB7, a CJK extension B ideograph, written as its surrogate pair + final String text = "\uD842\uDFB7\uD842\uDFB7"; + final Span[] spans = segmenter.tokenizePos(text); + for (final Span span : spans) { + Assertions.assertEquals(0, span.getStart() % 2, + "span must start on a code point boundary: " + span); + Assertions.assertEquals(0, span.getEnd() % 2, + "span must end on a code point boundary: " + span); + } + int covered = 0; + for (final Span span : spans) { + covered += span.length(); + } + Assertions.assertEquals(text.length(), covered); + } + + /** + * Pins Unicode-whitespace trimming of lexicon lines: a leading ideographic space + * (U+3000), common in hand-edited CJK text files, is stripped like ASCII whitespace, + * so the entry loads rather than failing as a malformed count. + */ + @Test + void testLeadingIdeographicSpaceIsTrimmed() throws IOException { + // U+3000 ideographic space, then the fixture word U+6211 and its count + final String lexicon = "\u3000\u6211 5000 r\n"; + final UnigramSegmenter loaded = UnigramSegmenter.load( + new ByteArrayInputStream(lexicon.getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8); + Assertions.assertArrayEquals(new String[] {"\u6211"}, loaded.tokenize("\u6211")); + } + + @Test + void testRejectsMalformedLexiconEncoding() { + final byte[] malformed = {'w', TRUNCATED_UTF8_LEAD_BYTE, ' ', '1', '\n'}; + + final IOException e = Assertions.assertThrows(IOException.class, + () -> UnigramSegmenter.load( + new ByteArrayInputStream(malformed), StandardCharsets.UTF_8)); + + Assertions.assertInstanceOf(MalformedInputException.class, e); + Assertions.assertEquals("Input length = 1", e.getMessage()); + } + + @Test + void testLongLexiconWordLoads() throws IOException { + final String word = "a".repeat(20_000); + + final UnigramSegmenter loaded = UnigramSegmenter.load( + new ByteArrayInputStream((word + " 1\n").getBytes(StandardCharsets.UTF_8)), + StandardCharsets.UTF_8); + + Assertions.assertArrayEquals(new String[] {word}, loaded.tokenize(word)); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DictionaryCatalogTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DictionaryCatalogTest.java new file mode 100644 index 0000000000..cbb65dfb3c --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DictionaryCatalogTest.java @@ -0,0 +1,263 @@ +/* + * 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.util; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.stemmer.hunspell.HunspellDictionary; + +/** + * Tests the opt-in dictionary catalog against an in-memory properties file and a + * local file URI so no network access is required. + */ +public class DictionaryCatalogTest { + + private static final String SHA_512 = "0".repeat(128); + + /** Checks a catalog entry with all required fields. */ + @Test + void testEntryAcceptsItsDocumentedFields() { + final URI absolute = URI.create("https://example.invalid/dictionary.tar.gz"); + + Assertions.assertDoesNotThrow(() -> + new DictionaryCatalog.Entry("demo", absolute, "A".repeat(128), null)); + } + + /** + * Checks each invalid field of the public catalog entry record. + * + * @param field The invalid record field. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"id", "uri", "relative uri", "sha512", "invalid sha512", + "empty filename", "path filename"}) + void testEntryRejectsInvalidFields(String field) { + final URI absolute = URI.create("https://example.invalid/dictionary.tar.gz"); + final Executable construction = switch (field) { + case "id" -> () -> new DictionaryCatalog.Entry(null, absolute, SHA_512, null); + case "uri" -> () -> new DictionaryCatalog.Entry("demo", null, SHA_512, null); + case "relative uri" -> () -> new DictionaryCatalog.Entry( + "demo", URI.create("dictionary.tar.gz"), SHA_512, null); + case "sha512" -> () -> new DictionaryCatalog.Entry("demo", absolute, null, null); + case "invalid sha512" -> + () -> new DictionaryCatalog.Entry("demo", absolute, "not-a-digest", null); + case "empty filename" -> + () -> new DictionaryCatalog.Entry("demo", absolute, SHA_512, ""); + case "path filename" -> + () -> new DictionaryCatalog.Entry("demo", absolute, SHA_512, "../dict.bin"); + default -> throw new IllegalArgumentException("unknown field: " + field); + }; + + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, construction); + final String expectedMessage = switch (field) { + case "id" -> "id must not be null"; + case "uri" -> "uri must not be null"; + case "relative uri" -> "uri must be absolute"; + case "sha512" -> "sha512 must not be null"; + case "invalid sha512" -> "sha512 must be 128 hex digits"; + case "empty filename", "path filename" -> "filename must be a file name"; + default -> throw new IllegalArgumentException("unknown field: " + field); + }; + Assertions.assertEquals(expectedMessage, thrown.getMessage()); + } + + /** + * Checks that invalid property values are reported as catalog I/O errors. + * + * @throws IOException Thrown if a fixture catalog cannot be loaded. + */ + @Test + void testGetRejectsInvalidCatalogValues() throws IOException { + final DictionaryCatalog relativeUri = DictionaryCatalog.load(new ByteArrayInputStream( + ("demo.url=dictionary.tar.gz\ndemo.sha512=" + SHA_512 + "\n") + .getBytes(StandardCharsets.UTF_8))); + final DictionaryCatalog invalidDigest = DictionaryCatalog.load(new ByteArrayInputStream( + "demo.url=https://example.invalid/dictionary.tar.gz\ndemo.sha512=invalid\n" + .getBytes(StandardCharsets.UTF_8))); + + Assertions.assertAll( + () -> Assertions.assertThrows(IOException.class, () -> relativeUri.get("demo")), + () -> Assertions.assertThrows(IOException.class, () -> invalidDigest.get("demo"))); + } + + /** + * Checks argument validation before the opt-in remote setting is read. + * + * @param argument The invalid method parameter. + * @param dir A scratch directory managed by the test framework. + * @throws IOException Thrown if the fixture catalog cannot be loaded. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"id", "targetDirectory"}) + void testInstallValidatesArgumentsBeforeTheRemoteSetting(String argument, + @TempDir Path dir) + throws IOException { + final DictionaryCatalog catalog = DictionaryCatalog.load( + new ByteArrayInputStream(new byte[0])); + final String previous = System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + try { + final Executable install = switch (argument) { + case "id" -> () -> catalog.install(null, dir); + case "targetDirectory" -> () -> catalog.install("demo", null); + default -> throw new IllegalArgumentException("unknown argument: " + argument); + }; + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, install); + Assertions.assertEquals(argument + " must not be null", thrown.getMessage()); + } finally { + restore(previous); + } + } + + /** + * Verifies that a catalog install without the remote-download property fails with + * the property name in the message, before anything is fetched or created. + * + * @param dir A scratch directory managed by the test framework. + * @throws Exception Thrown if the fixture catalog cannot be prepared. + */ + @Test + void testInstallRequiresRemoteProperty(@TempDir Path dir) throws Exception { + final byte[] payload = "payload".getBytes(StandardCharsets.UTF_8); + final DictionaryCatalog loaded = demoCatalog(dir, payload); + final Path target = dir.resolve("out"); + + final String previous = + System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + try { + final IOException e = Assertions.assertThrows(IOException.class, + () -> loaded.install("demo", target)); + Assertions.assertTrue( + e.getMessage().contains(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY)); + Assertions.assertTrue(Files.notExists(target)); + } finally { + restore(previous); + } + } + + /** + * Verifies that an enabled catalog install fetches the entry and stores its + * digest-verified bytes under the source name in the target directory. + * + * @param dir A scratch directory managed by the test framework. + * @throws Exception Thrown if the fixture catalog cannot be prepared or fetched. + */ + @Test + void testInstallStoresTheEntryUnderItsSourceName(@TempDir Path dir) throws Exception { + final byte[] payload = "payload".getBytes(StandardCharsets.UTF_8); + final DictionaryCatalog loaded = demoCatalog(dir, payload); + final Path target = dir.resolve("out"); + + final String previous = + System.getProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + System.setProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY, "true"); + try { + loaded.install("demo", target); + Assertions.assertArrayEquals(payload, + Files.readAllBytes(target.resolve("dict.bin"))); + } finally { + restore(previous); + } + } + + /** + * Verifies that the example catalog holds the MeCab and Hunspell entries, each with + * a full-length SHA-512 digest. Applications load their own catalog through the same + * {@link DictionaryCatalog#load(InputStream)} entry point. + * + * @throws IOException Thrown if the example catalog fails to load. + */ + @Test + void testExampleCatalogContainsMecabAndHunspellEntries() throws IOException { + final DictionaryCatalog catalog = loadExampleCatalog(); + Assertions.assertTrue(catalog.ids().contains("mecab.ipadic")); + Assertions.assertTrue(catalog.ids().contains("mecab.ko-dic")); + Assertions.assertTrue(catalog.ids().contains( + "hunspell.en_US" + HunspellDictionary.AFFIX_FILE_SUFFIX)); + Assertions.assertEquals(128, catalog.get("mecab.ipadic").sha512().length()); + Assertions.assertEquals(128, catalog.get( + "hunspell.en_US" + + HunspellDictionary.DICTIONARY_FILE_SUFFIX) + .sha512().length()); + } + + /** + * Loads the example catalog from test resources rather than a production classpath + * default. + * + * @return The example catalog. Never {@code null}. + * @throws IOException Thrown if the resource is absent or cannot be read. + */ + private static DictionaryCatalog loadExampleCatalog() throws IOException { + try (InputStream in = DictionaryCatalogTest.class.getResourceAsStream( + "/opennlp/tools/util/dictionary-catalog.properties")) { + if (in == null) { + throw new IOException("missing example dictionary catalog"); + } + return DictionaryCatalog.load(in); + } + } + + /** + * Builds a one-entry catalog whose URL is a local file holding {@code payload}, so + * downloads need no network. + * + * @param dir The directory to write the payload file into. + * @param payload The bytes the catalog entry points at. + * @return The loaded catalog. Never {@code null}. + * @throws IOException Thrown if the payload file cannot be written. + */ + private static DictionaryCatalog demoCatalog(Path dir, byte[] payload) + throws IOException { + final Path source = dir.resolve("dict.bin"); + Files.write(source, payload); + final String catalog = "demo.url=" + source.toUri() + "\n" + + "demo.sha512=" + DigestTestUtil.sha512(payload) + "\n"; + return DictionaryCatalog.load( + new ByteArrayInputStream(catalog.getBytes(StandardCharsets.UTF_8))); + } + + /** + * Restores the remote-download property to its value before the test. + * + * @param previous The saved value, or {@code null} when the property was unset. + */ + private static void restore(String previous) { + if (previous == null) { + System.clearProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY); + } else { + System.setProperty(DictionaryCatalog.REMOTE_DOWNLOAD_PROPERTY, previous); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DigestTestUtil.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DigestTestUtil.java new file mode 100644 index 0000000000..20c9e08cc6 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DigestTestUtil.java @@ -0,0 +1,64 @@ +/* + * 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.util; + +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.HexFormat; + +/** + * Computes digests for test fixtures. + */ +public final class DigestTestUtil { + + private DigestTestUtil() { + } + + /** + * {@return the SHA-256 digest of {@code bytes} as 64 lowercase hex digits} + * + * @param bytes The content to digest. Must not be {@code null}. + */ + public static String sha256(byte[] bytes) { + return digest("SHA-256", bytes); + } + + /** + * {@return the SHA-512 digest of {@code bytes} as 128 lowercase hex digits} + * + * @param bytes The content to digest. Must not be {@code null}. + */ + public static String sha512(byte[] bytes) { + return digest("SHA-512", bytes); + } + + /** + * {@return the digest of {@code bytes} as lowercase hex digits} + * + * @param algorithm The digest algorithm name. + * @param bytes The content to digest. Must not be {@code null}. + */ + private static String digest(String algorithm, byte[] bytes) { + try { + return HexFormat.of().formatHex( + MessageDigest.getInstance(algorithm).digest(bytes)); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException(e); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DownloadUtilCacheIntegrityTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DownloadUtilCacheIntegrityTest.java new file mode 100644 index 0000000000..81556dc6a4 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/DownloadUtilCacheIntegrityTest.java @@ -0,0 +1,204 @@ +/* + * 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.util; + +import java.io.IOException; +import java.io.InputStream; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.HexFormat; + +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 opennlp.tools.chunker.ChunkerModel; + +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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies that models served from the local download home are checked against their + * SHA-512 checksum, and that doing so does not require network access. + *

    + * The "remote" side is a {@code file:} URL inside the test's temporary directory, so these + * tests never contact the CDN. + */ +public class DownloadUtilCacheIntegrityTest { + + private static final String DOWNLOAD_HOME_PROPERTY = "OPENNLP_DOWNLOAD_HOME"; + + private static final String MODEL_FILENAME = "opennlp-test-chunker.bin"; + + /** + * The model that gets published, downloaded, and whose checksum is the authoritative one. + */ + private static final String PUBLISHED_MODEL = "/opennlp/tools/chunker/chunker170default.bin"; + + /** + * A different, but equally loadable, model. Used to replace a cached file so that the only + * thing distinguishing it from the expected model is its checksum - not its parsability. + */ + private static final String SUBSTITUTE_MODEL = "/opennlp/tools/chunker/chunker180custom.bin"; + + @TempDir + Path tempDir; + + private Path downloadHome; + private Path remoteModel; + private Path remoteChecksum; + private URL modelUrl; + private String previousDownloadHome; + + @BeforeEach + void setUp() throws IOException { + previousDownloadHome = System.getProperty(DOWNLOAD_HOME_PROPERTY); + System.setProperty(DOWNLOAD_HOME_PROPERTY, tempDir.toString()); + downloadHome = tempDir.resolve(".opennlp"); + + final Path remoteDir = Files.createDirectories(tempDir.resolve("remote")); + remoteModel = remoteDir.resolve(MODEL_FILENAME); + copyResource(PUBLISHED_MODEL, remoteModel); + + remoteChecksum = remoteDir.resolve(MODEL_FILENAME + ".sha512"); + Files.writeString(remoteChecksum, sha512(remoteModel) + " " + MODEL_FILENAME, + StandardCharsets.UTF_8); + + modelUrl = remoteModel.toUri().toURL(); + } + + @AfterEach + void tearDown() { + if (previousDownloadHome == null) { + System.clearProperty(DOWNLOAD_HOME_PROPERTY); + } else { + System.setProperty(DOWNLOAD_HOME_PROPERTY, previousDownloadHome); + } + } + + /** + * Sanity check: a first download validates and populates the cache, storing the published + * checksum file next to the model so that later loads have something to check against. + */ + @Test + void testFirstDownloadPopulatesCache() throws IOException { + assertNotNull(DownloadUtil.downloadModel(modelUrl, ChunkerModel.class)); + assertTrue(Files.exists(downloadHome.resolve(MODEL_FILENAME)), + "The model should have been cached in the download home"); + + final Path sidecar = downloadHome.resolve(MODEL_FILENAME + ".sha512"); + assertTrue(Files.exists(sidecar), "The published checksum should have been stored"); + assertEquals(Files.readString(remoteChecksum, StandardCharsets.UTF_8).strip(), + Files.readString(sidecar, StandardCharsets.UTF_8).strip(), + "The stored checksum should be interchangeable with the published one"); + } + + /** + * The actual defect: once a model is cached, its contents are never re-checked. The cached + * file is replaced with a different but perfectly loadable model, so that a passing + * result cannot be explained by the model parser rejecting garbage. + */ + @Test + void testTamperedCachedModelIsRejected() throws IOException { + assertNotNull(DownloadUtil.downloadModel(modelUrl, ChunkerModel.class)); + + copyResource(SUBSTITUTE_MODEL, downloadHome.resolve(MODEL_FILENAME)); + + final IOException e = assertThrows(IOException.class, + () -> DownloadUtil.downloadModel(modelUrl, ChunkerModel.class), + "A cached model that no longer matches its published checksum must be rejected"); + assertTrue(e.getMessage().contains("SHA512"), + "Expected a checksum failure, but got: " + e.getMessage()); + } + + /** + * Guards the constraint called out in OPENNLP-1902: verifying the cache must not turn every + * cached load into a network request. Both published artifacts are removed after the first + * download, so any attempt to reach the "remote" side would fail. + */ + @Test + void testValidCachedModelLoadsWithoutNetworkAccess() throws IOException { + assertNotNull(DownloadUtil.downloadModel(modelUrl, ChunkerModel.class)); + + Files.delete(remoteChecksum); + Files.delete(remoteModel); + + assertNotNull(DownloadUtil.downloadModel(modelUrl, ChunkerModel.class), + "A valid cached model must still load when the CDN is unreachable"); + } + + /** + * A cache populated by an older OpenNLP release has no checksum sidecar. When the published + * checksum is still reachable it must be used, so that pre-existing caches are covered too. + */ + @Test + void testLegacyCacheWithoutSidecarIsVerified() throws IOException { + Files.createDirectories(downloadHome); + copyResource(SUBSTITUTE_MODEL, downloadHome.resolve(MODEL_FILENAME)); + + final IOException e = assertThrows(IOException.class, + () -> DownloadUtil.downloadModel(modelUrl, ChunkerModel.class), + "A legacy cached model must be verified against the published checksum"); + assertTrue(e.getMessage().contains("SHA512"), + "Expected a checksum failure, but got: " + e.getMessage()); + } + + /** + * The other half of the legacy-cache decision: when there is no sidecar and the published + * checksum cannot be reached either, the model is loaded rather than refused, so that an + * offline download home populated by an older release keeps working. Nothing is stored in + * that case, so the verification is retried the next time the CDN is available. + */ + @Test + void testLegacyCacheIsLoadedWhenChecksumIsUnreachable() throws IOException { + Files.createDirectories(downloadHome); + copyResource(PUBLISHED_MODEL, downloadHome.resolve(MODEL_FILENAME)); + Files.delete(remoteChecksum); + + assertNotNull(DownloadUtil.downloadModel(modelUrl, ChunkerModel.class), + "A legacy cached model must still load when the published checksum is unreachable"); + assertFalse(Files.exists(downloadHome.resolve(MODEL_FILENAME + ".sha512")), + "Nothing should be stored when the published checksum could not be retrieved"); + } + + private void copyResource(String resource, Path target) throws IOException { + try (InputStream in = DownloadUtilCacheIntegrityTest.class.getResourceAsStream(resource)) { + assertNotNull(in, "Missing test resource: " + resource); + Files.createDirectories(target.getParent()); + Files.copy(in, target, StandardCopyOption.REPLACE_EXISTING); + } + } + + private static String sha512(Path file) throws IOException { + try { + final MessageDigest digest = MessageDigest.getInstance("SHA-512"); + return HexFormat.of().formatHex(digest.digest(Files.readAllBytes(file))); + } catch (NoSuchAlgorithmException e) { + throw new IOException("SHA-512 algorithm not found", e); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/InstallerTestSupport.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/InstallerTestSupport.java new file mode 100644 index 0000000000..a341bc9c20 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/InstallerTestSupport.java @@ -0,0 +1,120 @@ +/* + * 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.util; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.stream.Stream; + +import opennlp.tools.util.archive.TarArchives; + +/** + * Shared fixtures for the {@code ResourceInstaller} test classes: tar archive building, + * gzip compression, digest computation, and installed-file listing. + */ +final class InstallerTestSupport { + + static final int BLOCK = TarArchives.BLOCK; + static final int TERMINATOR_SIZE = TarArchives.TERMINATOR_SIZE; + + /** One kibibyte, a convenient small ceiling for limit tests. */ + static final long KIBIBYTE = 1024; + + /** One mebibyte, a convenient generous ceiling for tests that do not exercise it. */ + static final long MEBIBYTE = 1024 * KIBIBYTE; + + private InstallerTestSupport() { + } + + /** + * Writes one regular-file tar entry into the given buffer. + * + * @param tar The buffer receiving the entry bytes. Must not be {@code null}. + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param content The entry content. Must not be {@code null}. + * @throws IOException Thrown if writing to the buffer fails. + * @throws IllegalArgumentException Thrown if the name exceeds the tar name field. + */ + static void tarEntry(ByteArrayOutputStream tar, String name, byte[] content) + throws IOException { + TarArchives.entry(tar, name, content); + } + + /** + * Builds a gzip-compressed tar archive from name and content pairs. + * + * @param entries Pairs of entry name and UTF-8 text content. Must not be {@code null}. + * @return The archive bytes. Never {@code null}. + * @throws IOException Thrown if assembling the archive fails. + */ + static byte[] tarGz(String[][] entries) throws IOException { + return TarArchives.gzippedTar(entries); + } + + /** + * Compresses the given bytes with gzip. + * + * @param content The bytes to compress. Must not be {@code null}. + * @return The gzip-compressed bytes. Never {@code null}. + * @throws IOException Thrown if compressing fails. + */ + static byte[] gzip(byte[] content) throws IOException { + return TarArchives.gzip(content); + } + + /** + * Computes the SHA-256 of the given bytes as a lowercase hex string. + * + * @param content The bytes to digest. Must not be {@code null}. + * @return The 64-character lowercase hex digest. Never {@code null}. + */ + static String sha256(byte[] content) { + return DigestTestUtil.sha256(content); + } + + /** + * Computes the SHA-512 of the given bytes as a lowercase hex string. + * + * @param content The bytes to digest. Must not be {@code null}. + * @return The 128-character lowercase hex digest. Never {@code null}. + */ + static String sha512(byte[] content) { + return DigestTestUtil.sha512(content); + } + + /** + * Lists every regular file below the given directory as relative paths with forward + * slashes, sorted lexicographically, so tests can assert the exact installed file + * set. + * + * @param root The directory to walk. Must not be {@code null}. + * @return The sorted relative paths. Never {@code null}. + * @throws IOException Thrown if walking the directory fails. + */ + static List installedFiles(Path root) throws IOException { + try (Stream walk = Files.walk(root)) { + return walk.filter(Files::isRegularFile) + .map(file -> root.relativize(file).toString().replace('\\', '/')) + .sorted() + .toList(); + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerHttpTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerHttpTest.java new file mode 100644 index 0000000000..7e27bb528b --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerHttpTest.java @@ -0,0 +1,635 @@ +/* + * 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.util; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.net.SocketTimeoutException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.function.Executable; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import static opennlp.tools.util.InstallerTestSupport.KIBIBYTE; +import static opennlp.tools.util.InstallerTestSupport.MEBIBYTE; +import static opennlp.tools.util.InstallerTestSupport.installedFiles; +import static opennlp.tools.util.InstallerTestSupport.sha256; +import static opennlp.tools.util.InstallerTestSupport.tarGz; + +/** + * Exercises {@link ResourceInstaller} against a local scripted HTTP server: happy + * downloads, redirect handling and its policy, error statuses, stalled responses + * against the read timeout, and download limits against incorrect or oversized bodies. + */ +public class ResourceInstallerHttpTest { + + private static final Duration GENEROUS = Duration.ofSeconds(10); + + /** The limits these tests never exercise, kept at their defaults. */ + private static final long DEFAULT_ENTRIES = + ResourceInstaller.Limits.DEFAULT.maxEntries(); + private static final long DEFAULT_RATIO = + ResourceInstaller.Limits.DEFAULT.maxExpansionRatio(); + + /** Long enough that a stalled route outlives any timeout a test configures. */ + private static final Duration STALL = Duration.ofSeconds(30); + + /** + * A well-formed digest for fetches that fail before verification runs; its value is + * never compared. + */ + private static final String UNREACHED_CHECKSUM = "0".repeat(64); + + private StubServer server; + + @BeforeEach + void startServer() throws IOException { + server = new StubServer(); + } + + @AfterEach + void stopServer() throws IOException { + server.close(); + } + + /** + * Builds installation limits with the given read timeout and otherwise generous + * values, so timeout tests state only the value they exercise. + * + * @param readTimeout The read timeout to apply. + * @return The limits. Never {@code null}. + */ + private static ResourceInstaller.Limits withReadTimeout(Duration readTimeout) { + return new ResourceInstaller.Limits(GENEROUS, readTimeout, 5, MEBIBYTE, MEBIBYTE, + DEFAULT_ENTRIES, DEFAULT_RATIO); + } + + /** + * Builds installation limits with the given redirect allowance and otherwise + * generous values. + * + * @param maxRedirects The number of redirects to follow. + * @return The limits. Never {@code null}. + */ + private static ResourceInstaller.Limits withMaxRedirects(int maxRedirects) { + return new ResourceInstaller.Limits(GENEROUS, GENEROUS, maxRedirects, + MEBIBYTE, MEBIBYTE, DEFAULT_ENTRIES, DEFAULT_RATIO); + } + + /** + * Builds installation limits with the given download limit and otherwise generous + * values. + * + * @param maxDownloadBytes The download limit in bytes. + * @return The limits. Never {@code null}. + */ + private static ResourceInstaller.Limits withDownloadLimit(long maxDownloadBytes) { + return new ResourceInstaller.Limits(GENEROUS, GENEROUS, 5, maxDownloadBytes, + MEBIBYTE, DEFAULT_ENTRIES, DEFAULT_RATIO); + } + + /** + * Asserts that the given call is rejected as an argument error demanding a checksum + * for the remote source. + * + * @param source The remote source the message must name. + * @param call The call under test. + */ + private static void assertChecksumRequired(URI source, Executable call) { + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, call); + Assertions.assertEquals( + "checksum must be given for an http or https source: " + source, + thrown.getMessage()); + } + + @Test + void testHttpSourceWithoutChecksumIsRejectedBeforeFetching(@TempDir Path target) + throws Exception { + final AtomicBoolean fetched = new AtomicBoolean(); + server.route("/corpus.tar.gz", out -> { + fetched.set(true); + StubServer.ok(out, tarGz(new String[][] {{"corpus/data.txt", "unverified"}})); + }); + + final URI source = server.uri("/corpus.tar.gz"); + Assertions.assertAll( + () -> assertChecksumRequired(source, + () -> ResourceInstaller.install(source, target)), + () -> assertChecksumRequired(source, + () -> ResourceInstaller.install(source, target, null)), + () -> assertChecksumRequired(source, + () -> ResourceInstaller.install(source, target, null, withMaxRedirects(5)))); + Assertions.assertFalse(fetched.get()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testHttpsSourceWithoutChecksumIsRejectedBeforeCreatingTheTarget( + @TempDir Path parent) { + final Path target = parent.resolve("not-created-yet"); + final URI source = URI.create("https://example.invalid/corpus.tar.gz"); + + assertChecksumRequired(source, () -> ResourceInstaller.install(source, target)); + Assertions.assertTrue(Files.notExists(target)); + } + + @Test + void testHttpDownloadInstallsArchive(@TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "over http"}}); + server.route("/corpus.tar.gz", out -> StubServer.ok(out, archive)); + + ResourceInstaller.install(server.uri("/corpus.tar.gz"), target, sha256(archive)); + + Assertions.assertEquals("over http", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testInvalidSourceNameIsRejectedBeforeFetching(@TempDir Path parent) + throws Exception { + final AtomicBoolean fetched = new AtomicBoolean(); + final byte[] content = "dictionary".getBytes(StandardCharsets.UTF_8); + server.route("/bad%00name.dat", out -> { + fetched.set(true); + StubServer.ok(out, content); + }); + final Path target = parent.resolve("not-created-yet"); + + final IllegalArgumentException thrown = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ResourceInstaller.install(server.uri("/bad%00name.dat"), target, + sha256(content))); + + Assertions.assertEquals("name must be a file name", thrown.getMessage()); + Assertions.assertFalse(fetched.get()); + Assertions.assertTrue(Files.notExists(target)); + } + + @Test + void testAbsoluteRedirectIsFollowed(@TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "moved"}}); + server.route("/old.tar.gz", out -> StubServer.redirect(out, + server.uri("/new.tar.gz").toString())); + server.route("/new.tar.gz", out -> StubServer.ok(out, archive)); + + ResourceInstaller.install(server.uri("/old.tar.gz"), target, sha256(archive)); + + Assertions.assertEquals("moved", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testRelativeRedirectIsResolvedAgainstTheSource(@TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "relative"}}); + server.route("/mirror/old.tar.gz", + out -> StubServer.redirect(out, "new.tar.gz")); + server.route("/mirror/new.tar.gz", out -> StubServer.ok(out, archive)); + + ResourceInstaller.install(server.uri("/mirror/old.tar.gz"), target, sha256(archive)); + + Assertions.assertEquals("relative", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"301 Moved Permanently", "302 Found", "303 See Other", + "307 Temporary Redirect", "308 Permanent Redirect"}) + void testEveryRedirectStatusIsFollowed(String status, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "followed"}}); + server.route("/old.tar.gz", out -> StubServer.redirect(out, status, "/new.tar.gz")); + server.route("/new.tar.gz", out -> StubServer.ok(out, archive)); + + ResourceInstaller.install(server.uri("/old.tar.gz"), target, sha256(archive)); + + Assertions.assertEquals("followed", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testZeroRedirectAllowanceRejectsTheFirstRedirect(@TempDir Path target) + throws Exception { + server.route("/once", out -> StubServer.redirect(out, + server.uri("/anywhere.tar.gz").toString())); + + final URI source = server.uri("/once"); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(source, target, UNREACHED_CHECKSUM, withMaxRedirects(0))); + Assertions.assertEquals("more than 0 redirects: " + source, thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testMalformedRedirectLocationFails() { + final URI from = URI.create("http://example.invalid/archive.tar.gz"); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.resolveRedirect(from, "http://mirror.invalid/bad path")); + Assertions.assertEquals("redirect from " + from + + " contains a malformed Location: http://mirror.invalid/bad path", + thrown.getMessage()); + } + + @Test + void testRedirectChainBeyondLimitFails(@TempDir Path target) throws Exception { + server.route("/loop", out -> StubServer.redirect(out, "/loop")); + + final URI source = server.uri("/loop"); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(source, target, UNREACHED_CHECKSUM, withMaxRedirects(3))); + Assertions.assertEquals("more than 3 redirects: " + source, thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testRedirectWithoutLocationFails(@TempDir Path target) throws Exception { + server.route("/broken", out -> StubServer.head(out, "302 Found", + "Content-Length: 4", "", "gone")); + + final URI source = server.uri("/broken"); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(source, target, UNREACHED_CHECKSUM)); + Assertions.assertEquals("redirect from " + source + " contains no Location header", + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testRedirectToNonHttpSchemeFails(@TempDir Path target) throws Exception { + server.route("/non-http", out -> StubServer.redirect(out, "file:///etc/passwd")); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/non-http"), target, + UNREACHED_CHECKSUM)); + Assertions.assertEquals( + "redirect target is not an http or https location: file:///etc/passwd", + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testHttpsToHttpDowngradeIsRejected() { + final URI https = URI.create("https://example.invalid/archive.tar.gz"); + + // Cover the rejected downgrade and both accepted upgrade and same-scheme cases. + Assertions.assertAll( + () -> { + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.resolveRedirect(https, + "http://example.invalid/archive.tar.gz")); + Assertions.assertEquals( + "redirect downgrades https to http: http://example.invalid/archive.tar.gz", + thrown.getMessage()); + }, + () -> Assertions.assertDoesNotThrow(() -> ResourceInstaller.resolveRedirect(https, + "https://mirror.invalid/archive.tar.gz")), + () -> Assertions.assertDoesNotThrow(() -> ResourceInstaller.resolveRedirect( + URI.create("http://example.invalid/archive.tar.gz"), + "https://mirror.invalid/archive.tar.gz"))); + } + + @Test + void testHttpErrorStatusFails(@TempDir Path target) throws Exception { + server.route("/missing.tar.gz", out -> StubServer.head(out, "404 Not Found", + "Content-Length: 0", "")); + + final URI source = server.uri("/missing.tar.gz"); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(source, target, UNREACHED_CHECKSUM)); + Assertions.assertEquals("download failed with HTTP status 404: " + source, + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testStalledResponseHitsReadTimeout(@TempDir Path target) { + server.route("/stall", out -> StubServer.sleep(STALL)); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/stall"), target, + UNREACHED_CHECKSUM, withReadTimeout(Duration.ofMillis(250)))); + Assertions.assertInstanceOf(SocketTimeoutException.class, thrown); + } + + @Test + void testStalledBodyHitsReadTimeout(@TempDir Path target) throws Exception { + server.route("/drip", out -> { + StubServer.head(out, "200 OK", "Content-Length: 100000", ""); + out.write("just a few bytes".getBytes(StandardCharsets.UTF_8)); + out.flush(); + StubServer.sleep(STALL); + }); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/drip"), target, + UNREACHED_CHECKSUM, withReadTimeout(Duration.ofMillis(250)))); + Assertions.assertInstanceOf(SocketTimeoutException.class, thrown); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testDeclaredContentLengthBeyondLimitFailsFast(@TempDir Path target) + throws Exception { + server.route("/liar.tar.gz", out -> StubServer.head(out, "200 OK", + "Content-Length: 10000000", "")); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/liar.tar.gz"), target, + UNREACHED_CHECKSUM, withDownloadLimit(KIBIBYTE))); + Assertions.assertEquals( + "declared content length 10000000 exceeds the download limit of 1024 bytes", + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testStreamedBodyBeyondLimitAborts(@TempDir Path target) throws Exception { + server.route("/endless", out -> { + StubServer.head(out, "200 OK", ""); + out.write(new byte[64 * 1024]); + out.flush(); + }); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/endless"), target, + UNREACHED_CHECKSUM, withDownloadLimit(KIBIBYTE))); + Assertions.assertEquals("download exceeds the limit of 1024 bytes", + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + @Timeout(value = 15, threadMode = Timeout.ThreadMode.SEPARATE_THREAD) + void testSubMillisecondReadTimeoutStillTimesOut(@TempDir Path target) { + server.route("/stall", out -> StubServer.sleep(STALL)); + final ResourceInstaller.Limits limits = ResourceInstaller.Limits.builder() + .readTimeout(Duration.ofNanos(1)) + .build(); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(server.uri("/stall"), target, + UNREACHED_CHECKSUM, limits)); + + Assertions.assertInstanceOf(SocketTimeoutException.class, thrown); + } + + + @Test + void testTimeoutBeyondTheMillisecondRangeIsCapped(@TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"payload/data.txt", "content"}}); + server.route("/payload.tar.gz", out -> StubServer.ok(out, archive)); + final Duration beyondMillis = Duration.ofSeconds(Long.MAX_VALUE / 1000 + 1); + final ResourceInstaller.Limits limits = ResourceInstaller.Limits.builder() + .connectTimeout(beyondMillis) + .readTimeout(beyondMillis) + .build(); + + ResourceInstaller.install(server.uri("/payload.tar.gz"), target, sha256(archive), + limits); + + Assertions.assertEquals("content", + Files.readString(target.resolve("payload/data.txt"))); + } + + /** + * A scripted HTTP server on a loopback socket. Each registered route writes a raw + * response for redirect, timeout, malformed-response, and size-limit tests. + */ + private static final class StubServer implements AutoCloseable { + + /** Writes a raw HTTP response to the connected client. */ + @FunctionalInterface + interface Responder { + + /** + * Writes the raw response bytes for one request. + * + * @param out The response stream. + * @throws IOException Thrown if writing fails. + */ + void respond(OutputStream out) throws IOException; + } + + private final ServerSocket socket; + private final Map routes = new ConcurrentHashMap<>(); + private final List connections = new ArrayList<>(); + + StubServer() throws IOException { + socket = new ServerSocket(0, 50, InetAddress.getLoopbackAddress()); + final Thread acceptor = new Thread(this::acceptLoop, "stub-http-acceptor"); + acceptor.setDaemon(true); + acceptor.start(); + } + + /** + * Registers the responder serving the given absolute request path. + * + * @param path The absolute request path, starting with {@code /}. + * @param responder The script producing the raw response. + */ + void route(String path, Responder responder) { + routes.put(path, responder); + } + + /** + * Builds the http URI of the given absolute path on this server. + * + * @param path The absolute request path, starting with {@code /}. + * @return The URI. Never {@code null}. + */ + URI uri(String path) { + return URI.create("http://127.0.0.1:" + socket.getLocalPort() + path); + } + + /** + * Writes the HTTP/1.0 status line followed by the given lines, each terminated by + * CRLF. The caller supplies the header lines, then an empty string for the blank + * line separating head from body, then optional body text. + * + * @param out The response stream. + * @param status The status line content after the protocol, such as {@code 200 OK}. + * @param lines Header lines, then an empty string separator, then optional body + * text, each written with a trailing CRLF. + * @throws IOException Thrown if writing fails. + */ + static void head(OutputStream out, String status, String... lines) + throws IOException { + final StringBuilder response = new StringBuilder("HTTP/1.0 ").append(status) + .append("\r\n"); + for (final String line : lines) { + response.append(line).append("\r\n"); + } + out.write(response.toString().getBytes(StandardCharsets.US_ASCII)); + out.flush(); + } + + /** + * Writes a complete 200 response carrying the given body with its exact length. + * + * @param out The response stream. + * @param body The response body bytes. + * @throws IOException Thrown if writing fails. + */ + static void ok(OutputStream out, byte[] body) throws IOException { + head(out, "200 OK", "Content-Length: " + body.length, ""); + out.write(body); + out.flush(); + } + + /** + * Writes a 302 redirect to the given location. + * + * @param out The response stream. + * @param location The Location header value, absolute or relative. + * @throws IOException Thrown if writing fails. + */ + static void redirect(OutputStream out, String location) throws IOException { + redirect(out, "302 Found", location); + } + + /** + * Writes a redirect with the given status to the given location. + * + * @param out The response stream. + * @param status The status line content after the protocol, such as + * {@code 301 Moved Permanently}. + * @param location The Location header value, absolute or relative. + * @throws IOException Thrown if writing fails. + */ + static void redirect(OutputStream out, String status, String location) + throws IOException { + head(out, status, "Location: " + location, "Content-Length: 0", ""); + } + + /** + * Blocks the handler thread, simulating a stalled server. Interruption during + * teardown ends the sleep early. + * + * @param duration How long to stall. + */ + static void sleep(Duration duration) { + try { + Thread.sleep(duration.toMillis()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + private void acceptLoop() { + while (!socket.isClosed()) { + final Socket connection; + try { + connection = socket.accept(); + } catch (IOException e) { + return; + } + synchronized (connections) { + connections.add(connection); + } + final Thread handler = new Thread(() -> handle(connection), + "stub-http-handler"); + handler.setDaemon(true); + handler.start(); + } + } + + /** + * Reads one request, dispatches it to the registered responder, and closes the + * connection. Write failures from clients that abort mid-transfer are expected + * and ignored. + * + * @param connection The accepted client connection. + */ + private void handle(Socket connection) { + try (connection) { + final String path = readRequestPath(connection.getInputStream()); + final Responder responder = routes.get(path); + if (responder == null) { + head(connection.getOutputStream(), "404 Not Found", "Content-Length: 0", ""); + return; + } + responder.respond(connection.getOutputStream()); + } catch (IOException e) { + // The client hung up or the server is shutting down; both are test-normal. + } + } + + /** + * Reads the request head up to its terminating blank line and returns the path + * from the request line. + * + * @param in The request stream. + * @return The request path. Never {@code null}. + * @throws IOException Thrown if the request head is malformed or truncated. + */ + private static String readRequestPath(InputStream in) throws IOException { + final ByteArrayOutputStream headBytes = new ByteArrayOutputStream(); + int last4 = 0; + int b; + while ((b = in.read()) >= 0) { + headBytes.write(b); + last4 = (last4 << 8) | b; + if (last4 == 0x0D0A0D0A) { + break; + } + } + final String head = headBytes.toString(StandardCharsets.US_ASCII); + final int firstLineEnd = head.indexOf("\r\n"); + final String requestLine = firstLineEnd < 0 ? head : head.substring(0, firstLineEnd); + final String[] parts = requestLine.split(" "); + if (parts.length < 2) { + throw new IOException("malformed request line: " + requestLine); + } + return parts[1]; + } + + @Override + public void close() throws IOException { + socket.close(); + synchronized (connections) { + for (final Socket connection : connections) { + connection.close(); + } + } + } + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerTest.java new file mode 100644 index 0000000000..f336234142 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/ResourceInstallerTest.java @@ -0,0 +1,1412 @@ +/* + * 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.util; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.nio.file.FileAlreadyExistsException; +import java.nio.file.FileSystem; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.Path; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.stream.Stream; +import java.util.zip.GZIPOutputStream; +import java.util.zip.ZipEntry; +import java.util.zip.ZipOutputStream; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.DisabledOnOs; +import org.junit.jupiter.api.condition.OS; +import org.junit.jupiter.api.function.Executable; +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 org.junit.jupiter.params.provider.ValueSource; + +import opennlp.tools.util.archive.TarArchives; + +import static opennlp.tools.util.InstallerTestSupport.BLOCK; +import static opennlp.tools.util.InstallerTestSupport.KIBIBYTE; +import static opennlp.tools.util.InstallerTestSupport.MEBIBYTE; +import static opennlp.tools.util.InstallerTestSupport.TERMINATOR_SIZE; +import static opennlp.tools.util.InstallerTestSupport.gzip; +import static opennlp.tools.util.InstallerTestSupport.installedFiles; +import static opennlp.tools.util.InstallerTestSupport.sha256; +import static opennlp.tools.util.InstallerTestSupport.sha512; +import static opennlp.tools.util.InstallerTestSupport.tarEntry; +import static opennlp.tools.util.InstallerTestSupport.tarGz; + +public class ResourceInstallerTest { + + private static final String CHECKSUM_ARGUMENT_ERROR = + "checksum must be 64 (SHA-256) or 128 (SHA-512) hex characters; pass null to skip"; + private static final String ESCAPE_ERROR = + "archive entry escapes the target directory: "; + private static final String EXPANSION_LIMIT_ERROR = + "expanded content exceeds the limit of " + KIBIBYTE + " bytes"; + private static final String ENTRY_LIMIT_ERROR = + "archive entry count exceeds the limit of 2 entries"; + private static final String COLLISION_ERROR = "target already contains: "; + private static final String DUPLICATE_ENTRY_ERROR = + "archive contains duplicate file entry: "; + private static final String RATIO_ERROR = + "content expands beyond 100 times its compressed size"; + private static final String ZIP_MISMATCH_ERROR = + "zip local headers and central directory list different files"; + + /** The property name used by the parser tests; never read by the installer. */ + private static final String TEST_LIMIT_PROPERTY = "opennlp.test.limit"; + + /** + * Installs the given file under the default limits and asserts that it fails with the + * expected message, leaving the target directory without a single installed file. + * + * @param file The source file to install. + * @param target The target directory, which must have been empty before the attempt. + * @param message The exact failure message expected. + * @throws IOException Thrown if listing the target directory fails. + */ + private static void assertInstallFails(Path file, Path target, String message) + throws IOException { + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + Assertions.assertEquals(message, thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + /** + * Installs the given file under the given limits and asserts that it fails with the + * expected message, leaving the target directory without a single installed file. + * + * @param file The source file to install. + * @param target The target directory, which must have been empty before the attempt. + * @param limits The limits to install under. + * @param message The exact failure message expected. + * @throws IOException Thrown if listing the target directory fails. + */ + private static void assertInstallFails(Path file, Path target, + ResourceInstaller.Limits limits, String message) throws IOException { + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target, null, limits)); + Assertions.assertEquals(message, thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + /** + * Computes the SHA-256 of the given bytes as an uppercase hex string, so tests can + * prove that checksum comparison does not depend on the hex letter case. + * + * @param content The bytes to digest. Must not be {@code null}. + * @return The 64-character uppercase hex digest. Never {@code null}. + * @throws NoSuchAlgorithmException Thrown if the digest algorithm is unavailable. + */ + private static String sha256UpperCase(byte[] content) throws NoSuchAlgorithmException { + return sha256(content).toUpperCase(Locale.ROOT); + } + + /** + * Builds installation limits with generous timeouts and redirect allowance but the + * given size limits, so limit tests state only the values they exercise. + * + * @param maxDownloadBytes The download limit in bytes. + * @param maxExpandedBytes The expansion limit in bytes. + * @return The limits. Never {@code null}. + */ + private static ResourceInstaller.Limits limits(long maxDownloadBytes, + long maxExpandedBytes) { + return new ResourceInstaller.Limits(Duration.ofSeconds(10), Duration.ofSeconds(10), + 5, maxDownloadBytes, maxExpandedBytes, + ResourceInstaller.Limits.DEFAULT.maxEntries(), + ResourceInstaller.Limits.DEFAULT.maxExpansionRatio()); + } + + @Test + void testInstallEndToEndUsageExample(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"corpus/README", "A tiny example corpus.\n"}, + {"corpus/tokens.txt", "the\ncat\n"}, + {"corpus/pos/tags.tsv", "the\tDET\ncat\tNOUN\n"}}); + final Path file = source.resolve("corpus.tar.gz"); + Files.write(file, archive); + + final Path result = ResourceInstaller.install(file.toUri(), target, sha256(archive)); + + Assertions.assertEquals(target, result); + Assertions.assertEquals( + List.of("corpus/README", "corpus/pos/tags.tsv", "corpus/tokens.txt"), + installedFiles(target)); + Assertions.assertEquals("A tiny example corpus.\n", + Files.readString(target.resolve("corpus/README"))); + Assertions.assertEquals("the\ncat\n", + Files.readString(target.resolve("corpus/tokens.txt"))); + Assertions.assertEquals("the\tDET\ncat\tNOUN\n", + Files.readString(target.resolve("corpus/pos/tags.tsv"))); + } + + @Test + void testTarGzUnpacksWithStructure(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"corpus-1.0/train.conllu", "# sent_id = 1\n"}, + {"corpus-1.0/sub/readme.txt", "hello"}}); + final Path file = source.resolve("corpus.tgz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target, sha256(archive)); + + Assertions.assertEquals("# sent_id = 1\n", + Files.readString(target.resolve("corpus-1.0/train.conllu"))); + Assertions.assertEquals("hello", + Files.readString(target.resolve("corpus-1.0/sub/readme.txt"))); + } + + @Test + void testChecksumMismatchFailsBeforeUnpacking(@TempDir Path source, + @TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"a/file.txt", "content"}}); + final Path file = source.resolve("archive.tar.gz"); + Files.write(file, archive); + + final String wrong = "0".repeat(64); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target, wrong)); + Assertions.assertEquals("checksum mismatch: expected " + wrong + + " but downloaded " + sha256(archive), thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testChecksumComparisonIgnoresHexLetterCase(@TempDir Path source, + @TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"data/entry.txt", "payload"}}); + final Path file = source.resolve("cased.tar.gz"); + Files.write(file, archive); + // The uppercase digest must differ textually from the lowercase one, otherwise + // this test would not exercise the case handling at all. + Assertions.assertNotEquals(sha256(archive), sha256UpperCase(archive)); + + ResourceInstaller.install(file.toUri(), target, sha256UpperCase(archive)); + + Assertions.assertEquals("payload", + Files.readString(target.resolve("data/entry.txt"))); + } + + @Test + void testChecksumIgnoresUnicodeWhitespace(@TempDir Path source, + @TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "verified"}}); + final Path file = source.resolve("unicode-space.tar.gz"); + Files.write(file, archive); + final String emSpace = Character.toString(0x2003); + + ResourceInstaller.install(file.toUri(), target, + emSpace + sha256(archive) + emSpace); + + Assertions.assertEquals("verified", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testEscapingEntriesAreRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"../escape.txt", "bad"}}); + final Path file = source.resolve("evil.tar.gz"); + Files.write(file, archive); + + assertInstallFails(file, target, ESCAPE_ERROR + "../escape.txt"); + Assertions.assertTrue(Files.notExists(target.getParent().resolve("escape.txt"))); + } + + @Test + void testAbsoluteTarEntryIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"/absolute-escape-attempt/evil.txt", "bad"}}); + final Path file = source.resolve("absolute.tar.gz"); + Files.write(file, archive); + + assertInstallFails(file, target, + ESCAPE_ERROR + "/absolute-escape-attempt/evil.txt"); + Assertions.assertTrue(Files.notExists(Path.of("/absolute-escape-attempt"))); + } + + @Test + void testEscapingTarDirectoryEntryIsRejected(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + TarArchives.entry(tar, "../outside/", new byte[0], TarArchives.TYPE_DIRECTORY); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("directory-escape.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, ESCAPE_ERROR + "../outside/"); + } + + @Test + void testZipEntryWithTraversalIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("../zip-escape.txt")); + zip.write("bad".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("evil.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, ESCAPE_ERROR + "../zip-escape.txt"); + Assertions.assertTrue(Files.notExists(target.getParent().resolve("zip-escape.txt"))); + } + + @Test + void testEscapingZipDirectoryEntryIsRejected(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("../outside/")); + zip.closeEntry(); + } + final Path file = source.resolve("directory-escape.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, ESCAPE_ERROR + "../outside/"); + } + + @Test + void testZipUnpacks(@TempDir Path source, @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("lexicon/words.txt")); + zip.write("cat 100\n".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("lexicon.zip"); + Files.write(file, out.toByteArray()); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals("cat 100\n", + Files.readString(target.resolve("lexicon/words.txt"))); + } + + /** + * Checks that ZIP validation works when the target uses a non-default file system. + * + * @param source The directory containing the source ZIP file. + * @param scratch The directory containing the target file system. + * @throws IOException Thrown if the fixture or installation cannot be read or written. + */ + @Test + void testZipUnpacksIntoANonDefaultFileSystem(@TempDir Path source, + @TempDir Path scratch) throws IOException { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(bytes)) { + zip.putNextEntry(new ZipEntry("payload/data.txt")); + zip.write("content".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path archive = source.resolve("payload.zip"); + Files.write(archive, bytes.toByteArray()); + final URI targetUri = URI.create("jar:" + scratch.resolve("target.zip").toUri()); + + try (FileSystem fileSystem = + FileSystems.newFileSystem(targetUri, Map.of("create", "true"))) { + final Path target = fileSystem.getPath("/installed"); + + ResourceInstaller.install(archive.toUri(), target); + + Assertions.assertEquals("content", + Files.readString(target.resolve("payload/data.txt"))); + } + } + + @Test + void testModelBinIsStoredPacked(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("manifest.properties")); + zip.write("OpenNLP-Version: 0.0.0\n".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("en-ner-person.bin"); + Files.write(file, out.toByteArray()); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals(List.of("en-ner-person.bin"), installedFiles(target)); + Assertions.assertArrayEquals(out.toByteArray(), + Files.readAllBytes(target.resolve("en-ner-person.bin"))); + } + + @Test + void testPlainGzipDecompressesToTheSourceName(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (GZIPOutputStream gzip = new GZIPOutputStream(out)) { + gzip.write("word\tlemma\n".getBytes(StandardCharsets.UTF_8)); + } + final Path file = source.resolve("lexicon.tsv.gz"); + Files.write(file, out.toByteArray()); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals("word\tlemma\n", + Files.readString(target.resolve("lexicon.tsv"))); + } + + @Test + void testPlainGzipWithoutABaseNameUsesResource(@TempDir Path source, + @TempDir Path target) throws Exception { + final Path file = source.resolve(".gz"); + Files.write(file, gzip("word\tlemma\n".getBytes(StandardCharsets.UTF_8))); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals("word\tlemma\n", + Files.readString(target.resolve("resource"))); + } + + @Test + void testPlainFilesAreStoredUnderTheirSourceName(@TempDir Path source, + @TempDir Path target) throws Exception { + final Path file = source.resolve("frequencies.txt"); + Files.writeString(file, "cat 100"); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals("cat 100", Files.readString(target.resolve("frequencies.txt"))); + } + + /** + * Checks each required installation parameter. + * + * @param argument The invalid method parameter. + * @param target A scratch directory managed by the test framework. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(strings = {"source", "targetDirectory", "limits"}) + void testInvalidArguments(String argument, @TempDir Path target) { + final Executable call = switch (argument) { + case "source" -> () -> ResourceInstaller.install(null, target); + case "targetDirectory" -> () -> ResourceInstaller.install(target.toUri(), null); + case "limits" -> + () -> ResourceInstaller.install(target.toUri(), target, null, null); + default -> throw new IllegalArgumentException("unknown argument: " + argument); + }; + + assertArgumentError(argument + " must not be null", call); + } + + /** + * Supplies checksum arguments that are neither a valid SHA-256 nor a valid SHA-512 + * digest. + * + * @return One case per rejected digest. Never {@code null}. + */ + static Stream rejectedChecksums() { + return Stream.of( + Arguments.of("blank", " "), + Arguments.of("too short", "abc123"), + Arguments.of("right length, non-hex characters", "g".repeat(64)), + Arguments.of("between the two supported lengths", "a".repeat(96)), + Arguments.of("longer than SHA-512", "a".repeat(129))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("rejectedChecksums") + void testInvalidChecksumIsRejected(String label, String checksum, @TempDir Path target) { + assertArgumentError(CHECKSUM_ARGUMENT_ERROR, + () -> ResourceInstaller.install(target.toUri(), target, checksum)); + } + + /** + * Asserts that the given call fails as an argument error carrying the exact message. + * + * @param message The exact failure message expected. + * @param call The call under test. + */ + private static void assertArgumentError(String message, Executable call) { + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, call); + Assertions.assertEquals(message, thrown.getMessage()); + } + + @Test + void testSha512ChecksumVerifies(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"data/entry.txt", "payload"}}); + final Path file = source.resolve("checked.tar.gz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target, sha512(archive)); + + Assertions.assertEquals("payload", + Files.readString(target.resolve("data/entry.txt"))); + } + + @Test + void testSha512ChecksumMismatchFailsBeforeUnpacking(@TempDir Path source, + @TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"a/file.txt", "content"}}); + final Path file = source.resolve("archive.tar.gz"); + Files.write(file, archive); + + final String wrong = "0".repeat(128); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target, wrong)); + Assertions.assertEquals("checksum mismatch: expected " + wrong + + " but downloaded " + sha512(archive), thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testFailedTarUnpackLeavesTargetEmpty(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"good.txt", "fine"}, + {"../escape.txt", "bad"}}); + final Path file = source.resolve("partial.tar.gz"); + Files.write(file, archive); + + assertInstallFails(file, target, ESCAPE_ERROR + "../escape.txt"); + } + + @Test + void testFailedZipUnpackLeavesTargetEmpty(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("good.txt")); + zip.write("fine".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + zip.putNextEntry(new ZipEntry("../zip-escape.txt")); + zip.write("bad".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("partial.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, ESCAPE_ERROR + "../zip-escape.txt"); + } + + @Test + void testDuplicateTarFileEntryIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "corpus/data.txt", "first".getBytes(StandardCharsets.UTF_8)); + tarEntry(tar, "corpus/data.txt", "second".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("duplicate.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, DUPLICATE_ENTRY_ERROR + "corpus/data.txt"); + } + + @Test + void testEquivalentZipFileEntriesAreRejected(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("corpus/data.txt")); + zip.write("first".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + zip.putNextEntry(new ZipEntry("corpus/./data.txt")); + zip.write("second".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("duplicate.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, DUPLICATE_ENTRY_ERROR + "corpus/./data.txt"); + } + + @Test + void testPlainFileBeginningWithPkIsStored(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] content = "PK plain dictionary".getBytes(StandardCharsets.UTF_8); + final Path file = source.resolve("dictionary.dat"); + Files.write(file, content); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertArrayEquals(content, + Files.readAllBytes(target.resolve("dictionary.dat"))); + } + + @Test + void testTruncatedZipHeaderIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path file = source.resolve("truncated.zip"); + Files.write(file, new byte[] {'P', 'K', 3, 4}); + + assertInstallFails(file, target, "malformed zip archive"); + } + + @Test + void testZipWithoutCentralDirectoryIsRejected(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("corpus/data.txt")); + zip.write("content".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final byte[] complete = out.toByteArray(); + int centralDirectory = -1; + for (int i = 0; i <= complete.length - 4; i++) { + if (complete[i] == 'P' && complete[i + 1] == 'K' + && complete[i + 2] == 1 && complete[i + 3] == 2) { + centralDirectory = i; + break; + } + } + Assertions.assertTrue(centralDirectory > 0); + final Path file = source.resolve("missing-central-directory.zip"); + Files.write(file, Arrays.copyOf(complete, centralDirectory)); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + Assertions.assertEquals("malformed zip archive", thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testTruncatedEmptyZipHeaderIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path file = source.resolve("truncated-empty.zip"); + Files.write(file, new byte[] {'P', 'K', 5, 6}); + + assertInstallFails(file, target, "malformed zip archive"); + } + + @Test + void testValidEmptyZipInstallsNothing(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] endHeader = new byte[22]; + endHeader[0] = 'P'; + endHeader[1] = 'K'; + endHeader[2] = 5; + endHeader[3] = 6; + final Path file = source.resolve("empty.zip"); + Files.write(file, endHeader); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testValidEmptyTarGzInstallsNothing(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path file = source.resolve("empty.tar.gz"); + Files.write(file, gzip(new byte[TERMINATOR_SIZE])); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testInvalidZipEntryPathLeavesNoStagingFiles(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("good.txt")); + zip.write("good".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + zip.putNextEntry(new ZipEntry("bad\0name.txt")); + zip.write("bad".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("invalid-path.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, + "archive entry has an invalid path: bad\0name.txt"); + } + + @Test + void testTruncatedTarLeavesTargetEmpty(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "first.txt", "complete".getBytes(StandardCharsets.UTF_8)); + tarEntry(tar, "second.txt", "never finished".getBytes(StandardCharsets.UTF_8)); + final byte[] whole = tar.toByteArray(); + // Cut inside the second entry's header: first entry occupies two 512-byte blocks. + final byte[] truncated = Arrays.copyOf(whole, 2 * BLOCK + 100); + final Path file = source.resolve("truncated.tar.gz"); + Files.write(file, gzip(truncated)); + + Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testFailedInstallKeepsPreexistingTargetContent(@TempDir Path source, + @TempDir Path target) throws Exception { + Files.writeString(target.resolve("existing.txt"), "keep"); + final byte[] archive = tarGz(new String[][] { + {"good.txt", "fine"}, + {"../escape.txt", "bad"}}); + final Path file = source.resolve("partial.tar.gz"); + Files.write(file, archive); + + Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + Assertions.assertEquals(List.of("existing.txt"), installedFiles(target)); + Assertions.assertEquals("keep", Files.readString(target.resolve("existing.txt"))); + } + + @Test + void testInstallationLeavesNoStagingResidue(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "content"}}); + final Path file = source.resolve("clean.tar.gz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target); + + try (Stream walk = Files.walk(target)) { + final List hidden = walk + .filter(path -> !path.equals(target)) + .map(path -> path.getFileName().toString()) + .filter(fileName -> fileName.startsWith(".")) + .toList(); + Assertions.assertEquals(List.of(), hidden); + } + Assertions.assertEquals(List.of("corpus/data.txt"), installedFiles(target)); + } + + /** + * Supplies invalid {@code Limits} constructions with the argument error each one + * must raise. + * + * @return One case per invalid argument. Never {@code null}. + */ + static Stream invalidLimits() { + final Duration valid = Duration.ofSeconds(10); + return Stream.of( + Arguments.of("null connectTimeout", (Executable) + () -> new ResourceInstaller.Limits(null, valid, 5, 1024, 1024, 10, 100), + "connectTimeout must not be null"), + Arguments.of("zero connectTimeout", (Executable) + () -> new ResourceInstaller.Limits(Duration.ZERO, valid, 5, 1024, 1024, 10, 100), + "connectTimeout must be positive"), + Arguments.of("null readTimeout", (Executable) + () -> new ResourceInstaller.Limits(valid, null, 5, 1024, 1024, 10, 100), + "readTimeout must not be null"), + Arguments.of("negative readTimeout", (Executable) + () -> new ResourceInstaller.Limits(valid, Duration.ofSeconds(-1), + 5, 1024, 1024, 10, 100), + "readTimeout must be positive"), + Arguments.of("negative maxRedirects", (Executable) + () -> new ResourceInstaller.Limits(valid, valid, -1, 1024, 1024, 10, 100), + "maxRedirects must not be negative"), + Arguments.of("zero maxDownloadBytes", (Executable) + () -> new ResourceInstaller.Limits(valid, valid, 5, 0, 1024, 10, 100), + "maxDownloadBytes must be positive"), + Arguments.of("zero maxExpandedBytes", (Executable) + () -> new ResourceInstaller.Limits(valid, valid, 5, 1024, 0, 10, 100), + "maxExpandedBytes must be positive"), + Arguments.of("zero maxEntries", (Executable) + () -> new ResourceInstaller.Limits(valid, valid, 5, 1024, 1024, 0, 100), + "maxEntries must be positive"), + Arguments.of("zero maxExpansionRatio", (Executable) + () -> new ResourceInstaller.Limits(valid, valid, 5, 1024, 1024, 10, 0), + "maxExpansionRatio must be positive"), + Arguments.of("builder zero connectTimeout", (Executable) () -> + ResourceInstaller.Limits.builder().connectTimeout(Duration.ZERO).build(), + "connectTimeout must be positive"), + Arguments.of("builder null readTimeout", (Executable) () -> + ResourceInstaller.Limits.builder().readTimeout(null).build(), + "readTimeout must not be null"), + Arguments.of("builder negative maxRedirects", (Executable) () -> + ResourceInstaller.Limits.builder().maxRedirects(-1).build(), + "maxRedirects must not be negative"), + Arguments.of("builder zero maxDownloadBytes", (Executable) () -> + ResourceInstaller.Limits.builder().maxDownloadBytes(0).build(), + "maxDownloadBytes must be positive"), + Arguments.of("builder negative maxExpandedBytes", (Executable) () -> + ResourceInstaller.Limits.builder().maxExpandedBytes(-1).build(), + "maxExpandedBytes must be positive"), + Arguments.of("builder zero maxEntries", (Executable) () -> + ResourceInstaller.Limits.builder().maxEntries(0).build(), + "maxEntries must be positive")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("invalidLimits") + void testLimitsRejectInvalidValues(String label, Executable construction, + String message) { + final IllegalArgumentException thrown = + Assertions.assertThrows(IllegalArgumentException.class, construction); + Assertions.assertEquals(message, thrown.getMessage()); + } + + @Test + void testDownloadLimitRejectsOversizedSource(@TempDir Path source, + @TempDir Path target) throws Exception { + final Path file = source.resolve("large.txt"); + Files.write(file, new byte[8192]); + + assertInstallFails(file, target, limits(KIBIBYTE, MEBIBYTE), + "download exceeds the limit of " + KIBIBYTE + " bytes"); + } + + @Test + void testTarExpansionLimitRejectsArchive(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "bomb/zeros.bin", new byte[64 * 1024]); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("bomb.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testTarExpansionLimitCountsBytesAfterTheTerminator(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "corpus/data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + tar.write(new byte[64 * 1024]); + final Path file = source.resolve("trailing-data.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, limits(MEBIBYTE, 16 * KIBIBYTE), + "expanded content exceeds the limit of " + 16 * KIBIBYTE + " bytes"); + } + + @Test + void testTarGzipTrailerIsVerified(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "corpus/data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + tar.write(new byte[64 * 1024]); + final byte[] archive = gzip(tar.toByteArray()); + archive[archive.length - 8] ^= 1; + final Path file = source.resolve("bad-trailer.tar.gz"); + Files.write(file, archive); + + Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testTarMetadataCountsTowardExpansionLimit(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + TarArchives.entry(tar, "pax_global_header", + TarArchives.paxRecord("comment", "a".repeat(64 * 1024)), 'g'); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("metadata-bomb.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testZipExpansionLimitRejectsArchive(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("bomb/zeros.bin")); + zip.write(new byte[64 * 1024]); + zip.closeEntry(); + } + final Path file = source.resolve("bomb.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testZipDirectoryContentCountsTowardExpansionLimit(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("payload/")); + zip.write(new byte[64 * 1024]); + zip.closeEntry(); + } + final Path file = source.resolve("directory-content.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testPlainGzipExpansionLimitRejectsFile(@TempDir Path source, + @TempDir Path target) throws Exception { + final Path file = source.resolve("zeros.bin.gz"); + Files.write(file, gzip(new byte[64 * 1024])); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testInstallWithinCustomLimitsSucceeds(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "small"}}); + final Path file = source.resolve("small.tar.gz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target, sha256(archive), + limits(MEBIBYTE, MEBIBYTE)); + + Assertions.assertEquals("small", Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testDefaultLimitsArePinned() { + final ResourceInstaller.Limits defaults = ResourceInstaller.Limits.DEFAULT; + Assertions.assertEquals(Duration.ofSeconds(20), defaults.connectTimeout()); + Assertions.assertEquals(Duration.ofSeconds(60), defaults.readTimeout()); + Assertions.assertEquals(5, defaults.maxRedirects()); + Assertions.assertEquals(1L << 30, defaults.maxDownloadBytes()); + Assertions.assertEquals(4L << 30, defaults.maxExpandedBytes()); + Assertions.assertEquals(100_000L, defaults.maxEntries()); + Assertions.assertEquals(100L, defaults.maxExpansionRatio()); + } + + @Test + void testLimitPropertyNames() { + Assertions.assertEquals("opennlp.download.max.bytes", + ResourceInstaller.Limits.MAX_DOWNLOAD_BYTES_PROPERTY); + Assertions.assertEquals("opennlp.install.max.total.bytes", + ResourceInstaller.Limits.MAX_EXPANDED_BYTES_PROPERTY); + Assertions.assertEquals("opennlp.install.max.entries", + ResourceInstaller.Limits.MAX_ENTRIES_PROPERTY); + Assertions.assertEquals("opennlp.install.max.expansion.ratio", + ResourceInstaller.Limits.MAX_EXPANSION_RATIO_PROPERTY); + } + + @Test + void testLimitPropertyOverrideIsRead() { + System.setProperty(TEST_LIMIT_PROPERTY, " 123 "); + try { + Assertions.assertEquals(123L, + ResourceInstaller.Limits.longProperty(TEST_LIMIT_PROPERTY, 7L)); + } finally { + System.clearProperty(TEST_LIMIT_PROPERTY); + } + } + + /** + * Supplies property values that must fall back to the built-in default: absent, + * not a number, zero, negative, and empty. + * + * @return One case per unusable value. Never {@code null}. + */ + static Stream unusableLimitProperties() { + return Stream.of( + Arguments.of("absent", null), + Arguments.of("not a number", "abc"), + Arguments.of("zero", "0"), + Arguments.of("negative", "-5"), + Arguments.of("empty", "")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("unusableLimitProperties") + void testLimitPropertyFallsBackOnUnusableValues(String label, String value) { + if (value == null) { + System.clearProperty(TEST_LIMIT_PROPERTY); + } else { + System.setProperty(TEST_LIMIT_PROPERTY, value); + } + try { + Assertions.assertEquals(7L, + ResourceInstaller.Limits.longProperty(TEST_LIMIT_PROPERTY, 7L)); + } finally { + System.clearProperty(TEST_LIMIT_PROPERTY); + } + } + + /** + * Builds installation limits with the given entry limit and otherwise default + * values, so entry-count tests state only the value they exercise. + * + * @param maxEntries The entry limit. + * @return The limits. Never {@code null}. + */ + private static ResourceInstaller.Limits entryLimit(long maxEntries) { + return ResourceInstaller.Limits.builder().maxEntries(maxEntries).build(); + } + + @Test + void testTarEntryCountLimitRejectsArchive(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"corpus/one.txt", "1"}, + {"corpus/two.txt", "2"}, + {"corpus/three.txt", "3"}}); + final Path file = source.resolve("many.tar.gz"); + Files.write(file, archive); + + assertInstallFails(file, target, entryLimit(2), + ENTRY_LIMIT_ERROR); + } + + /** + * Checks that tar extension headers count toward the archive entry limit. + * + * @param source A scratch directory for the source archive. + * @param target A scratch installation directory. + * @throws Exception Thrown if the fixture cannot be created or installed. + */ + @Test + void testTarEntryCountLimitIncludesMetadata(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + TarArchives.entry(tar, "PaxHeaders/one", + TarArchives.paxRecord("comment", "first"), 'x'); + TarArchives.entry(tar, "PaxHeaders/b", + TarArchives.paxRecord("comment", "right"), 'x'); + tarEntry(tar, "payload/data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("metadata-entries.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, entryLimit(2), ENTRY_LIMIT_ERROR); + } + + @Test + void testZipEntryCountLimitRejectsArchive(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + for (int i = 0; i < 3; i++) { + zip.putNextEntry(new ZipEntry("corpus/entry-" + i + ".txt")); + zip.write("x".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + } + final Path file = source.resolve("many.zip"); + Files.write(file, out.toByteArray()); + + assertInstallFails(file, target, entryLimit(2), + ENTRY_LIMIT_ERROR); + } + + @Test + void testEntryCountExactlyAtLimitSucceeds(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] archive = tarGz(new String[][] { + {"corpus/one.txt", "1"}, + {"corpus/two.txt", "2"}}); + final Path file = source.resolve("exact.tar.gz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target, sha256(archive), entryLimit(2)); + + Assertions.assertEquals(List.of("corpus/one.txt", "corpus/two.txt"), + installedFiles(target)); + } + + @Test + void testSha512ChecksumComparisonIgnoresHexLetterCase(@TempDir Path source, + @TempDir Path target) throws Exception { + final byte[] archive = tarGz(new String[][] {{"data/entry.txt", "payload"}}); + final Path file = source.resolve("cased512.tar.gz"); + Files.write(file, archive); + final String upperCase = sha512(archive).toUpperCase(Locale.ROOT); + Assertions.assertNotEquals(sha512(archive), upperCase); + + ResourceInstaller.install(file.toUri(), target, upperCase); + + Assertions.assertEquals("payload", + Files.readString(target.resolve("data/entry.txt"))); + } + + @Test + void testDownloadExactlyAtLimitSucceeds(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path file = source.resolve("exact.dat"); + Files.write(file, new byte[1024]); + + ResourceInstaller.install(file.toUri(), target, null, limits(KIBIBYTE, MEBIBYTE)); + + Assertions.assertEquals(List.of("exact.dat"), installedFiles(target)); + Assertions.assertEquals(1024, Files.size(target.resolve("exact.dat"))); + } + + @Test + void testExpansionExactlyAtLimitSucceeds(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "corpus/exact.bin", new byte[1024]); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("exact.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + ResourceInstaller.install(file.toUri(), target, null, + limits(MEBIBYTE, tar.size())); + + Assertions.assertEquals(List.of("corpus/exact.bin"), installedFiles(target)); + Assertions.assertEquals(1024, Files.size(target.resolve("corpus/exact.bin"))); + } + + @Test + void testCumulativeExpansionAcrossEntriesHitsLimit(@TempDir Path source, + @TempDir Path target) throws Exception { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tarEntry(tar, "corpus/first.bin", new byte[768]); + tarEntry(tar, "corpus/second.bin", new byte[768]); + tar.write(new byte[TERMINATOR_SIZE]); + final Path file = source.resolve("cumulative.tar.gz"); + Files.write(file, gzip(tar.toByteArray())); + + assertInstallFails(file, target, limits(MEBIBYTE, KIBIBYTE), + EXPANSION_LIMIT_ERROR); + } + + @Test + void testReinstallOverAnExistingFileIsRejected(@TempDir Path source, @TempDir Path target) + throws Exception { + final byte[] first = tarGz(new String[][] {{"corpus/data.txt", "version one"}}); + final byte[] second = tarGz(new String[][] {{"corpus/data.txt", "version two"}}); + final Path firstFile = source.resolve("first.tar.gz"); + final Path secondFile = source.resolve("second.tar.gz"); + Files.write(firstFile, first); + Files.write(secondFile, second); + ResourceInstaller.install(firstFile.toUri(), target, sha256(first)); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(secondFile.toUri(), target, sha256(second))); + + Assertions.assertEquals( + COLLISION_ERROR + target.resolve("corpus/data.txt"), + thrown.getMessage()); + Assertions.assertEquals("version one", + Files.readString(target.resolve("corpus/data.txt"))); + + Files.delete(target.resolve("corpus/data.txt")); + ResourceInstaller.install(secondFile.toUri(), target, sha256(second)); + Assertions.assertEquals("version two", + Files.readString(target.resolve("corpus/data.txt"))); + } + + @Test + void testCollidingInstallPromotesNothing(@TempDir Path source, @TempDir Path target) + throws Exception { + Files.createDirectories(target.resolve("corpus")); + Files.writeString(target.resolve("corpus/data.txt"), "keep"); + final byte[] archive = tarGz(new String[][] { + {"corpus/fresh.txt", "new"}, + {"corpus/data.txt", "replacement"}}); + final Path file = source.resolve("colliding.tar.gz"); + Files.write(file, archive); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target, sha256(archive))); + + Assertions.assertEquals( + COLLISION_ERROR + target.resolve("corpus/data.txt"), + thrown.getMessage()); + Assertions.assertEquals("keep", Files.readString(target.resolve("corpus/data.txt"))); + Assertions.assertTrue(Files.notExists(target.resolve("corpus/fresh.txt"))); + Assertions.assertEquals(List.of("corpus/data.txt"), installedFiles(target)); + } + + @Test + @DisabledOnOs(OS.WINDOWS) + void testPromotionRejectsToFollowASymlinkedDirectory(@TempDir Path source, + @TempDir Path target, @TempDir Path outside) throws Exception { + final Path link = target.resolve("link"); + Files.createSymbolicLink(link, outside); + final byte[] archive = tarGz(new String[][] {{"link/planted.txt", "escaped"}}); + final Path file = source.resolve("symlink.tar.gz"); + Files.write(file, archive); + + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> ResourceInstaller.install(file.toUri(), target)); + + Assertions.assertEquals("installation path crosses a symbolic link: " + link, + thrown.getMessage()); + Assertions.assertEquals(List.of(), installedFiles(outside)); + Assertions.assertTrue(Files.notExists(outside.resolve("planted.txt"))); + } + + @Test + void testNestedDirectoryThatIsNotASymlinkStillInstalls(@TempDir Path source, + @TempDir Path target) throws Exception { + Files.createDirectory(target.resolve("link")); + final byte[] archive = tarGz(new String[][] {{"link/planted.txt", "fine"}}); + final Path file = source.resolve("nested.tar.gz"); + Files.write(file, archive); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals("fine", + Files.readString(target.resolve("link/planted.txt"))); + } + + /** + * Supplies source schemes the installer rejects because they are outside the bounded + * HTTP and local-file paths. + * + * @return One case per rejected scheme. Never {@code null}. + */ + static Stream rejectedSourceSchemes() { + return Stream.of( + Arguments.of("ftp", URI.create("ftp://example.invalid/corpus.tar.gz")), + Arguments.of("jar", URI.create("jar:file:/tmp/a.jar!/corpus.tar.gz")), + Arguments.of("mailto", URI.create("mailto:someone@example.invalid")), + Arguments.of("no scheme at all", URI.create("corpus.tar.gz"))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("rejectedSourceSchemes") + void testUnsupportedSourceSchemeIsRejected(String label, URI source, + @TempDir Path target) throws IOException { + assertArgumentError("source scheme must be http, https, or file, but was: " + source, + () -> ResourceInstaller.install(source, target)); + Assertions.assertEquals(List.of(), installedFiles(target)); + } + + @Test + void testUnsupportedSourceSchemeIsRejectedBeforeCreatingTheTarget(@TempDir Path parent) { + final Path target = parent.resolve("not-created-yet"); + + assertArgumentError( + "source scheme must be http, https, or file, but was: ftp://example.invalid/c.gz", + () -> ResourceInstaller.install(URI.create("ftp://example.invalid/c.gz"), target)); + Assertions.assertTrue(Files.notExists(target)); + } + + + + + @Test + void testBuilderWithoutOverridesEqualsTheDefaults() { + Assertions.assertEquals(ResourceInstaller.Limits.DEFAULT, + ResourceInstaller.Limits.builder().build()); + } + + @Test + void testBuilderOverridesOnlyWhatIsSet() { + final ResourceInstaller.Limits limits = ResourceInstaller.Limits.builder() + .maxDownloadBytes(MEBIBYTE) + .maxExpandedBytes(2 * MEBIBYTE) + .build(); + + Assertions.assertEquals(MEBIBYTE, limits.maxDownloadBytes()); + Assertions.assertEquals(2 * MEBIBYTE, limits.maxExpandedBytes()); + Assertions.assertEquals(ResourceInstaller.Limits.DEFAULT.connectTimeout(), + limits.connectTimeout()); + Assertions.assertEquals(ResourceInstaller.Limits.DEFAULT.readTimeout(), + limits.readTimeout()); + Assertions.assertEquals(ResourceInstaller.Limits.DEFAULT.maxRedirects(), + limits.maxRedirects()); + Assertions.assertEquals(ResourceInstaller.Limits.DEFAULT.maxEntries(), + limits.maxEntries()); + } + + @Test + void testDownloadFileIsCreatedInTheTargetDirectory(@TempDir Path target) + throws Exception { + final Path downloaded = ResourceInstaller.createDownloadFile(target); + try { + Assertions.assertEquals(target, downloaded.getParent()); + Assertions.assertTrue( + downloaded.getFileName().toString().startsWith(".opennlp-download"), + downloaded.getFileName().toString()); + } finally { + Files.deleteIfExists(downloaded); + } + } + + static Stream expansionBombs() throws IOException { + // Four mebibytes of repeated text compress to a few kibibytes in either format: far + // beyond the ratio ceiling, far below the absolute expansion limit. Repeated text + // rather than zeros, because a gzip stream of zeros reads as an empty tar archive. + final byte[] repetitive = + "opennlp ".repeat(512 * 1024).getBytes(StandardCharsets.UTF_8); + return Stream.of( + Arguments.of("corpus.txt.gz", gzip(repetitive)), + Arguments.of("corpus.zip", zipOf("corpus.txt", repetitive))); + } + + @ParameterizedTest + @MethodSource("expansionBombs") + void testExpansionRatioIsBounded(String name, byte[] content, @TempDir Path source, + @TempDir Path target) throws Exception { + final Path file = source.resolve(name); + Files.write(file, content); + + assertInstallFails(file, target, RATIO_ERROR); + } + + @ParameterizedTest + @MethodSource("expansionBombs") + void testRaisingTheExpansionRatioAcceptsWhatTheDefaultRejects(String name, + byte[] content, @TempDir Path source, @TempDir Path target) throws Exception { + final Path file = source.resolve(name); + Files.write(file, content); + + // Raising the byte limit alone cannot lift the ratio: the tighter budget wins. + final ResourceInstaller.Limits raised = ResourceInstaller.Limits.builder() + .maxExpansionRatio(2000) + .build(); + ResourceInstaller.install(file.toUri(), target, null, raised); + + Assertions.assertEquals(4 << 20, Files.size(target.resolve("corpus.txt"))); + } + + @Test + void testMoveIntoPlaceDoesNotReplaceAnExistingFile(@TempDir Path directory) + throws Exception { + final Path staged = directory.resolve("staged.txt"); + final Path destination = directory.resolve("installed.txt"); + Files.writeString(staged, "new"); + Files.writeString(destination, "old"); + + Assertions.assertThrows(FileAlreadyExistsException.class, + () -> ResourceInstaller.moveIntoPlace(staged, destination)); + + Assertions.assertEquals("old", Files.readString(destination)); + Assertions.assertEquals("new", Files.readString(staged)); + } + + @Test + void testModelSuffixIsMatchedIgnoringCase(@TempDir Path source, @TempDir Path target) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry("manifest.properties")); + zip.write("OpenNLP-Version: 0.0.0\n".getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + final Path file = source.resolve("en-ner-person.BIN"); + Files.write(file, out.toByteArray()); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals(List.of("en-ner-person.BIN"), installedFiles(target)); + } + + @Test + void testGzipSuffixIsStrippedIgnoringCase(@TempDir Path source, @TempDir Path target) + throws Exception { + final Path file = source.resolve("lexicon.tsv.GZ"); + Files.write(file, gzip("word\tlemma\n".getBytes(StandardCharsets.UTF_8))); + + ResourceInstaller.install(file.toUri(), target); + + Assertions.assertEquals(List.of("lexicon.tsv"), installedFiles(target)); + } + + @Test + void testStaleWorkFilesOfAKilledInstallAreRemoved(@TempDir Path source, + @TempDir Path target) throws Exception { + final Path staleStaging = Files.createDirectory(target.resolve(".opennlp-stagingOLD")); + Files.writeString(staleStaging.resolve("partial.txt"), "half"); + Files.writeString(target.resolve(".opennlp-downloadOLD.part"), "half"); + final Path file = source.resolve("corpus.tar.gz"); + Files.write(file, tarGz(new String[][] {{"corpus/data.txt", "content"}})); + + ResourceInstaller.install(file.toUri(), target); + + try (Stream entries = Files.list(target)) { + Assertions.assertEquals(List.of("corpus"), + entries.map(path -> path.getFileName().toString()).sorted().toList()); + } + } + + @Test + void testFailedInstallRemovesTheTargetDirectoryItCreated(@TempDir Path source, + @TempDir Path parent) throws Exception { + final byte[] archive = tarGz(new String[][] {{"corpus/data.txt", "content"}}); + final Path file = source.resolve("corpus.tar.gz"); + Files.write(file, archive); + final Path target = parent.resolve("fresh"); + + Assertions.assertThrows(IOException.class, () -> ResourceInstaller.install( + file.toUri(), target, sha256("other".getBytes(StandardCharsets.UTF_8)))); + + Assertions.assertTrue(Files.notExists(target)); + } + + @Test + void testZipLocalHeadersMustMatchTheCentralDirectory(@TempDir Path source, + @TempDir Path target) throws Exception { + // The same content under two names, so the local file sections are the same size and + // one archive's central directory fits the other's local headers. + final byte[] first = zipOf("a.txt", "content"); + final byte[] second = zipOf("b.txt", "content"); + final int centralFirst = centralDirectoryStart(first); + final int centralSecond = centralDirectoryStart(second); + final byte[] hybrid = new byte[centralFirst + second.length - centralSecond]; + System.arraycopy(first, 0, hybrid, 0, centralFirst); + System.arraycopy(second, centralSecond, hybrid, centralFirst, + second.length - centralSecond); + final Path file = source.resolve("hybrid.zip"); + Files.write(file, hybrid); + + assertInstallFails(file, target, ZIP_MISMATCH_ERROR); + } + + /** Builds a zip archive holding one text entry. */ + private static byte[] zipOf(String name, String content) throws IOException { + return zipOf(name, content.getBytes(StandardCharsets.UTF_8)); + } + + /** Builds a zip archive holding one entry with the given bytes. */ + private static byte[] zipOf(String name, byte[] content) throws IOException { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(out)) { + zip.putNextEntry(new ZipEntry(name)); + zip.write(content); + zip.closeEntry(); + } + return out.toByteArray(); + } + + /** Finds the first central directory header signature in a zip archive. */ + private static int centralDirectoryStart(byte[] zip) { + for (int i = 0; i + 3 < zip.length; i++) { + if (zip[i] == 'P' && zip[i + 1] == 'K' && zip[i + 2] == 1 && zip[i + 3] == 2) { + return i; + } + } + throw new AssertionError("no central directory"); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarArchives.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarArchives.java new file mode 100644 index 0000000000..6646cd1732 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarArchives.java @@ -0,0 +1,345 @@ +/* + * 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.util.archive; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.zip.GZIPOutputStream; + +/** + * Builds tar archives byte by byte, so tests can assemble well-formed, boundary, and + * malformed archives without an extra library. Each header has the ustar + * magic and a valid checksum, as a real archive does; {@link #reseal(byte[])} restores + * the checksum after a test has corrupted some other field on purpose. + */ +public final class TarArchives { + + /** The tar block size; headers, contents, and terminators are multiples of it. */ + public static final int BLOCK = 512; + + /** The size of the two all-zero blocks that terminate a tar archive. */ + public static final int TERMINATOR_SIZE = 2 * BLOCK; + + /** The type flag of a regular file. */ + public static final char TYPE_REGULAR_FILE = '0'; + + /** The type flag an old-style archive leaves NUL for a regular file. */ + public static final char TYPE_REGULAR_FILE_CLASSIC = '\0'; + + /** The type flag of a directory. */ + public static final char TYPE_DIRECTORY = '5'; + + /** The offset of the 155-byte ustar name prefix field. */ + public static final int PREFIX_OFFSET = 345; + + private static final int NAME_LENGTH = 100; + private static final int PREFIX_LENGTH = 155; + private static final int MODE_OFFSET = 100; + private static final int SIZE_OFFSET = 124; + private static final int SIZE_LENGTH = 12; + private static final int BASE_256_MARKER = 0x80; + private static final int BASE_256_NEGATIVE = 0x40; + private static final int CHECKSUM_OFFSET = 148; + private static final int CHECKSUM_LENGTH = 8; + private static final int TYPE_OFFSET = 156; + private static final int MAGIC_OFFSET = 257; + private static final int BLANK = ' '; + private static final String SIZE_FORMAT = "%011o"; + private static final String CHECKSUM_FORMAT = "%06o"; + private static final String MODE = "0000644 "; + + /** The ustar magic and version fields: {@code "ustar"}, NUL, then {@code "00"}. */ + private static final byte[] USTAR_MAGIC = {'u', 's', 't', 'a', 'r', 0, '0', '0'}; + + /** The GNU magic and version fields: {@code "ustar"}, two blanks, then NUL. */ + private static final byte[] GNU_MAGIC = {'u', 's', 't', 'a', 'r', ' ', ' ', 0}; + + private TarArchives() { + } + + /** + * Builds one 512-byte tar header block with the given name, declared content size, + * and type flag. The mode field is populated with a non-zero octal value so that a + * name filling the whole 100-byte name field is followed by non-zero bytes, which + * makes name-boundary tests meaningful. + * + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param size The content size to declare in the octal size field; must not be + * negative. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if the name does not fit the name field or + * the size is negative. + */ + public static byte[] header(String name, long size, char typeFlag) { + return header(name, "", size, typeFlag, true); + } + + /** + * Builds one 512-byte classic v7 header block, which carries no ustar magic and no + * name prefix field. Real archives in this format still exist, and a reader that keys + * off the magic rather than the checksum will not see them. + * + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param size The content size to declare in the octal size field; must not be + * negative. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if the name does not fit the name field or + * the size is negative. + */ + public static byte[] classicHeader(String name, long size, char typeFlag) { + return header(name, "", size, typeFlag, false); + } + + /** + * Builds one 512-byte GNU header block. GNU writes {@code "ustar"} followed by two + * blanks and a NUL where ustar writes {@code "ustar"}, a NUL, and the version, and it + * uses the offset ustar gives to the name prefix for {@code atime} instead. + * + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param size The content size to declare in the octal size field; must not be + * negative. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @param atime The octal {@code atime} to write at {@link #PREFIX_OFFSET}, as a GNU + * incremental archive does, or empty to leave it NUL. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if the name does not fit the name field or + * the size is negative. + */ + public static byte[] gnuHeader(String name, long size, char typeFlag, String atime) { + final byte[] block = header(name, "", size, typeFlag, false); + System.arraycopy(GNU_MAGIC, 0, block, MAGIC_OFFSET, GNU_MAGIC.length); + write(block, PREFIX_OFFSET, atime); + return reseal(block); + } + + /** + * Builds one 512-byte header block whose size field uses the base-256 encoding, which + * GNU writes when a length does not fit the eleven octal digits the field holds. + * + *

    The leading bit marks the encoding, the next bit carries the sign, and the + * remaining bits of that byte followed by every later byte form a big-endian two's + * complement number.

    + * + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param size The content size to encode. May be negative, which no real writer emits + * for a size, so that the reader's rejection of it can be exercised. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if the name does not fit the name field. + */ + public static byte[] base256Header(String name, long size, char typeFlag) { + final byte[] block = header(name, 0, typeFlag); + long remaining = size; + for (int i = SIZE_OFFSET + SIZE_LENGTH - 1; i >= SIZE_OFFSET; i--) { + block[i] = (byte) (remaining & 0xFF); + remaining >>= Byte.SIZE; + } + block[SIZE_OFFSET] = (byte) (size < 0 + ? block[SIZE_OFFSET] | BASE_256_MARKER | BASE_256_NEGATIVE + : block[SIZE_OFFSET] & ~BASE_256_NEGATIVE | BASE_256_MARKER); + return reseal(block); + } + + /** + * Encodes one pax extended header record, {@code " =\n"}. The + * length includes all record bytes: digits, the blank, and the newline. + * + * @param keyword The pax keyword. + * @param value The keyword's value. + * @return The encoded record. Never {@code null}. + */ + public static byte[] paxRecord(String keyword, String value) { + final byte[] body = (keyword + "=" + value + "\n").getBytes(StandardCharsets.UTF_8); + int digits = 1; + while (Integer.toString(body.length + 1 + digits).length() > digits) { + digits++; + } + final byte[] prefix = ((body.length + 1 + digits) + " ") + .getBytes(StandardCharsets.US_ASCII); + final byte[] record = new byte[prefix.length + body.length]; + System.arraycopy(prefix, 0, record, 0, prefix.length); + System.arraycopy(body, 0, record, prefix.length, body.length); + return record; + } + + /** + * Builds one 512-byte ustar header block whose entry name is split across the name + * prefix field and the name field, which is how a real archive stores a name longer + * than 100 bytes. + * + * @param name The name field content; at most 100 bytes when encoded as UTF-8. + * @param prefix The name prefix field content; at most 155 bytes when encoded as + * UTF-8. Empty for a header without a prefix. + * @param size The content size to declare in the octal size field; must not be + * negative. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if either name part does not fit its field + * or the size is negative. + */ + public static byte[] header(String name, String prefix, long size, char typeFlag) { + return header(name, prefix, size, typeFlag, true); + } + + /** + * Builds one 512-byte header block, in either the ustar or the classic v7 format. + * + * @param name The name field content; at most 100 bytes when encoded as UTF-8. + * @param prefix The name prefix field content; at most 155 bytes when encoded as + * UTF-8. Empty for a header without a prefix, and required to be empty + * for a classic header, which has no prefix field. + * @param size The content size to declare in the octal size field; must not be + * negative. + * @param typeFlag The tar type flag, for example {@link #TYPE_REGULAR_FILE}. + * @param ustar Whether to write the ustar magic and version fields. + * @return The header block. Never {@code null}. + * @throws IllegalArgumentException Thrown if either name part does not fit its field, + * the size is negative, or a classic header is asked for with a prefix. + */ + private static byte[] header(String name, String prefix, long size, char typeFlag, + boolean ustar) { + final byte[] nameBytes = name.getBytes(StandardCharsets.UTF_8); + if (nameBytes.length > NAME_LENGTH) { + throw new IllegalArgumentException( + "entry name exceeds " + NAME_LENGTH + " bytes: " + name); + } + final byte[] prefixBytes = prefix.getBytes(StandardCharsets.UTF_8); + if (prefixBytes.length > PREFIX_LENGTH) { + throw new IllegalArgumentException( + "entry name prefix exceeds " + PREFIX_LENGTH + " bytes: " + prefix); + } + if (size < 0) { + throw new IllegalArgumentException("size must not be negative: " + size); + } + if (!ustar && prefixBytes.length > 0) { + throw new IllegalArgumentException("a classic header has no name prefix field"); + } + final byte[] block = new byte[BLOCK]; + System.arraycopy(nameBytes, 0, block, 0, nameBytes.length); + System.arraycopy(prefixBytes, 0, block, PREFIX_OFFSET, prefixBytes.length); + write(block, MODE_OFFSET, MODE); + write(block, SIZE_OFFSET, String.format(SIZE_FORMAT, size)); + if (ustar) { + System.arraycopy(USTAR_MAGIC, 0, block, MAGIC_OFFSET, USTAR_MAGIC.length); + } + block[TYPE_OFFSET] = (byte) typeFlag; + return reseal(block); + } + + /** + * Recomputes and writes the header checksum of the given block, so a test can corrupt + * a field on purpose and still hand the reader an otherwise well-formed header. + * + * @param block The header block to seal. Must be 512 bytes long. + * @return The same block, with its checksum field filled in. Never {@code null}. + */ + public static byte[] reseal(byte[] block) { + for (int i = CHECKSUM_OFFSET; i < CHECKSUM_OFFSET + CHECKSUM_LENGTH; i++) { + block[i] = BLANK; + } + int sum = 0; + for (final byte b : block) { + sum += b & 0xFF; + } + write(block, CHECKSUM_OFFSET, String.format(CHECKSUM_FORMAT, sum)); + block[CHECKSUM_OFFSET + CHECKSUM_LENGTH - 1] = BLANK; + return block; + } + + /** + * Writes US-ASCII text into a header field. + * + * @param block The header block to write into. + * @param offset The field's offset in the block. + * @param text The text to write. + */ + private static void write(byte[] block, int offset, String text) { + final byte[] bytes = text.getBytes(StandardCharsets.US_ASCII); + System.arraycopy(bytes, 0, block, offset, bytes.length); + } + + /** + * Writes one complete tar entry into the given buffer: the header block declaring the + * content's actual length, the content itself, and zero padding up to the next + * 512-byte block boundary. + * + * @param tar The buffer receiving the entry bytes. Must not be {@code null}. + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param content The entry content. Must not be {@code null}. + * @param typeFlag The tar type flag for the header. + * @throws IOException Thrown if writing to the buffer fails. + * @throws IllegalArgumentException Thrown if the name does not fit the name field. + */ + public static void entry(ByteArrayOutputStream tar, String name, byte[] content, + char typeFlag) throws IOException { + tar.write(header(name, content.length, typeFlag)); + tar.write(content); + tar.write(new byte[(BLOCK - content.length % BLOCK) % BLOCK]); + } + + /** + * Writes one complete regular-file tar entry into the given buffer. + * + * @param tar The buffer receiving the entry bytes. Must not be {@code null}. + * @param name The entry name; at most 100 bytes when encoded as UTF-8. + * @param content The entry content. Must not be {@code null}. + * @throws IOException Thrown if writing to the buffer fails. + * @throws IllegalArgumentException Thrown if the name does not fit the name field. + */ + public static void entry(ByteArrayOutputStream tar, String name, byte[] content) + throws IOException { + entry(tar, name, content, TYPE_REGULAR_FILE); + } + + /** + * Builds a gzip-compressed tar archive from name and content pairs, terminated by + * the two all-zero blocks that end a tar archive. + * + * @param entries The entries as {@code {name, content}} pairs of UTF-8 text. Must not + * be {@code null}. + * @return The compressed archive bytes. Never {@code null}. + * @throws IOException Thrown if writing to the in-memory streams fails. + */ + public static byte[] gzippedTar(String[][] entries) throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + for (final String[] entry : entries) { + entry(tar, entry[0], entry[1].getBytes(StandardCharsets.UTF_8)); + } + tar.write(new byte[TERMINATOR_SIZE]); + return gzip(tar.toByteArray()); + } + + /** + * Compresses raw tar bytes the way a {@code .tar.gz} distribution is shipped, so + * tests can wrap hand-built or deliberately truncated tar content. + * + * @param content The raw tar bytes. Must not be {@code null}. + * @return The gzip-compressed bytes. Never {@code null}. + * @throws IOException Thrown if writing to the in-memory stream fails. + */ + public static byte[] gzip(byte[] content) throws IOException { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (GZIPOutputStream compressed = new GZIPOutputStream(out)) { + compressed.write(content); + } + return out.toByteArray(); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarStreamTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarStreamTest.java new file mode 100644 index 0000000000..52b4bd96c1 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/archive/TarStreamTest.java @@ -0,0 +1,783 @@ +/* + * 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.util.archive; + +import java.io.BufferedInputStream; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.function.Supplier; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +import static opennlp.tools.util.archive.TarArchives.BLOCK; +import static opennlp.tools.util.archive.TarArchives.TERMINATOR_SIZE; +import static opennlp.tools.util.archive.TarArchives.TYPE_DIRECTORY; +import static opennlp.tools.util.archive.TarArchives.TYPE_REGULAR_FILE; +import static opennlp.tools.util.archive.TarArchives.TYPE_REGULAR_FILE_CLASSIC; +import static opennlp.tools.util.archive.TarArchives.entry; +import static opennlp.tools.util.archive.TarArchives.header; + +/** + * Exercises {@link TarStream} against archives built byte by byte in the test, covering + * regular traversal as well as boundary and corruption cases. + */ +public class TarStreamTest { + + private static final int NAME_LENGTH = 100; + private static final int SIZE_OFFSET = 124; + private static final int CHECKSUM_OFFSET = 148; + + /** The largest length eleven octal digits can hold, one byte short of 8 GiB. */ + private static final long LARGEST_OCTAL_SIZE = (1L << 33) - 1; + + private static final String MALFORMED_RECORD = "malformed pax extended header record"; + private static final String SPARSE_REJECTED = "sparse tar entries are not supported: " + + "the archived bytes describe file holes, not contiguous content"; + + /** The metadata records GNU tar and bsdtar write ahead of an ordinary entry. */ + private static final String[][] METADATA = { + {"mtime", "1786557589.505896519"}, + {"atime", "1786557589.496852337"}, + {"ctime", "1786557589.505896519"}, + {"uid", "1000"}, + {"gname", "krickert"}, + {"SCHILY.dev", "66306"}, + {"hdrcharset", "BINARY"}}; + + /** + * Supplies blocks that must not be mistaken for a tar header, each with a description + * naming the reason. + * + * @return The rejection cases. Never {@code null}. + */ + private static Stream nonHeaderContent() { + final byte[] filler = new byte[BLOCK]; + Arrays.fill(filler, (byte) 'x'); + return Stream.of( + Arguments.of("fewer bytes than one block", + "too short to be a tar header".getBytes(StandardCharsets.US_ASCII)), + Arguments.of("no ustar magic and a non-octal size field", filler), + Arguments.of("an all-zero block", new byte[BLOCK])); + } + + @Test + void testEmptyStreamHasNoEntries() throws IOException { + final TarStream stream = new TarStream(new ByteArrayInputStream(new byte[0])); + Assertions.assertFalse(stream.next()); + } + + @Test + void testTerminatorOnlyArchiveHasNoEntries() throws IOException { + final TarStream stream = + new TarStream(new ByteArrayInputStream(new byte[TERMINATOR_SIZE])); + Assertions.assertFalse(stream.next()); + } + + /** + * Checks that repeated reads after the terminator remain at end of archive. + * + * @throws IOException Thrown if the fixture archive cannot be read. + */ + @Test + void testNextRemainsAtEndAfterAnUnalignedEntry() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "data.txt", "x".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertArrayEquals("x".getBytes(StandardCharsets.UTF_8), + stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testReadsEntriesSizesTypesAndContent() throws IOException { + final byte[] blockSized = new byte[BLOCK]; + for (int i = 0; i < blockSized.length; i++) { + blockSized[i] = (byte) (i % 251); + } + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "data/", new byte[0], TYPE_DIRECTORY); + entry(tar, "data/skip.bin", "0123456789".getBytes(StandardCharsets.US_ASCII), + TYPE_REGULAR_FILE); + entry(tar, "data/alpha.txt", "alpha\n".getBytes(StandardCharsets.UTF_8), + TYPE_REGULAR_FILE_CLASSIC); + entry(tar, "block.bin", blockSized, TYPE_REGULAR_FILE); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data/", stream.name()); + Assertions.assertEquals(0, stream.size()); + Assertions.assertFalse(stream.isFile()); + + // The next call must skip the unread content and padding of this entry. + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data/skip.bin", stream.name()); + Assertions.assertEquals(10, stream.size()); + Assertions.assertTrue(stream.isFile()); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data/alpha.txt", stream.name()); + Assertions.assertEquals(6, stream.size()); + Assertions.assertTrue(stream.isFile()); + Assertions.assertArrayEquals("alpha\n".getBytes(StandardCharsets.UTF_8), + stream.entryStream().readAllBytes()); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("block.bin", stream.name()); + Assertions.assertEquals(BLOCK, stream.size()); + Assertions.assertTrue(stream.isFile()); + Assertions.assertArrayEquals(blockSized, stream.entryStream().readAllBytes()); + + Assertions.assertFalse(stream.next()); + } + + @Test + void testNameFillsFullHundredByteField() throws IOException { + final StringBuilder longName = new StringBuilder("d/"); + while (longName.length() < NAME_LENGTH) { + longName.append('x'); + } + final String name = longName.toString(); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, name, "n".getBytes(StandardCharsets.US_ASCII), TYPE_REGULAR_FILE); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(NAME_LENGTH, stream.name().length()); + Assertions.assertEquals(name, stream.name()); + Assertions.assertArrayEquals("n".getBytes(StandardCharsets.US_ASCII), + stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testTruncatedHeaderReportsError() { + final byte[] partial = + Arrays.copyOf(header("cut.bin", 0, TYPE_REGULAR_FILE), BLOCK / 2); + final TarStream stream = new TarStream(new ByteArrayInputStream(partial)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("truncated tar header", thrown.getMessage()); + } + + @Test + void testTruncatedEntryContentReportsError() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tar.write(header("data.bin", 10, TYPE_REGULAR_FILE)); + tar.write("1234".getBytes(StandardCharsets.US_ASCII)); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + final IOException thrown = Assertions.assertThrows(IOException.class, + () -> stream.entryStream().readAllBytes()); + Assertions.assertEquals("truncated tar entry: data.bin", thrown.getMessage()); + } + + @Test + void testTruncatedArchiveWhenSkippingReportsError() throws IOException { + final TarStream stream = new TarStream( + new ByteArrayInputStream(header("gone.bin", 600, TYPE_REGULAR_FILE))); + + Assertions.assertTrue(stream.next()); + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("truncated tar archive", thrown.getMessage()); + } + + @Test + void testMalformedSizeFieldReportsError() { + final byte[] block = header("bad.bin", 0, TYPE_REGULAR_FILE); + block[SIZE_OFFSET] = '9'; + // Reseal, so the reader reaches the size field instead of stopping at the checksum. + TarArchives.reseal(block); + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("malformed tar size field in entry header", + thrown.getMessage()); + } + + @Test + void testSizeFieldWithEmbeddedPaddingIsRejected() { + final byte[] block = header("bad.bin", 0, TYPE_REGULAR_FILE); + Arrays.fill(block, SIZE_OFFSET, SIZE_OFFSET + 12, (byte) 0); + block[SIZE_OFFSET] = '1'; + block[SIZE_OFFSET + 1] = ' '; + block[SIZE_OFFSET + 2] = '2'; + TarArchives.reseal(block); + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("malformed tar size field in entry header", + thrown.getMessage()); + } + + @Test + void testStartsWithHeaderDetectsTarAndKeepsPosition() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "data/alpha.txt", "alpha\n".getBytes(StandardCharsets.UTF_8), + TYPE_REGULAR_FILE); + tar.write(new byte[TERMINATOR_SIZE]); + final InputStream in = + new BufferedInputStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(TarStream.startsWithHeader(in)); + + final TarStream stream = new TarStream(in); + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data/alpha.txt", stream.name()); + Assertions.assertArrayEquals("alpha\n".getBytes(StandardCharsets.UTF_8), + stream.entryStream().readAllBytes()); + } + + @ParameterizedTest + @MethodSource("nonHeaderContent") + void testStartsWithHeaderRejectsNonTarContent(String description, byte[] content) + throws IOException { + Assertions.assertFalse( + TarStream.startsWithHeader(new ByteArrayInputStream(content)), description); + } + + @Test + void testStartsWithHeaderRejectsUnusableStreams() { + final InputStream notMarkable = new InputStream() { + @Override + public int read() { + return -1; + } + }; + // assertAll so a missing check on one argument does not hide the other. + Assertions.assertAll( + () -> Assertions.assertThrows(IllegalArgumentException.class, + () -> TarStream.startsWithHeader(null)), + () -> Assertions.assertThrows(IllegalArgumentException.class, + () -> TarStream.startsWithHeader(notMarkable))); + } + + @Test + void testNullStreamIsRejected() { + Assertions.assertThrows(IllegalArgumentException.class, () -> new TarStream(null)); + } + + /** + * @param maxEntries The invalid entry limit. + */ + @ParameterizedTest(name = "{0}") + @ValueSource(longs = {0, -1}) + void testInvalidEntryLimitIsRejected(long maxEntries) { + final IllegalArgumentException thrown = Assertions.assertThrows( + IllegalArgumentException.class, + () -> new TarStream(InputStream.nullInputStream(), maxEntries)); + Assertions.assertEquals("maxEntries must be positive", thrown.getMessage()); + } + + @Test + void testHeaderWithWrongChecksumIsRejected() { + final byte[] block = header("tampered.bin", 10, TYPE_REGULAR_FILE); + block[0] = 'X'; + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("malformed tar header checksum", thrown.getMessage()); + } + + @Test + void testChecksumWithEmbeddedPaddingIsRejected() { + final byte[] block = header("bad-checksum.bin", 0, TYPE_REGULAR_FILE); + System.arraycopy(block, CHECKSUM_OFFSET + 1, block, CHECKSUM_OFFSET + 2, 5); + block[CHECKSUM_OFFSET + 1] = ' '; + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("malformed tar header checksum", thrown.getMessage()); + } + + @Test + void testStartsWithHeaderRejectsUstarMagicWithoutAChecksum() throws IOException { + final byte[] block = header("looks-real.bin", 0, TYPE_REGULAR_FILE); + Arrays.fill(block, CHECKSUM_OFFSET, CHECKSUM_OFFSET + 8, (byte) '0'); + + Assertions.assertFalse(TarStream.startsWithHeader(new ByteArrayInputStream(block))); + } + + @Test + void testUstarPrefixIsJoinedToTheName() throws IOException { + final String prefix = "corpus-1.0/" + "d".repeat(120); + final String tail = "annotations/train.conllu"; + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + final byte[] content = "# sent_id = 1\n".getBytes(StandardCharsets.UTF_8); + tar.write(TarArchives.header(tail, prefix, content.length, TYPE_REGULAR_FILE)); + tar.write(content); + tar.write(new byte[BLOCK - content.length]); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(prefix + "/" + tail, stream.name()); + Assertions.assertTrue(stream.name().length() > 100); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + } + + @Test + void testEmptyUstarPrefixLeavesTheNameAlone() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "plain.txt", "x".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("plain.txt", stream.name()); + } + + @Test + void testPaxGlobalHeaderWithOnlyACommentIsConsumed() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "pax_global_header", + TarArchives.paxRecord("comment", "0123456789abcdef"), 'g'); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data.txt", stream.name()); + Assertions.assertTrue(stream.isFile()); + Assertions.assertArrayEquals("content".getBytes(StandardCharsets.UTF_8), + stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + } + + /** + * Supplies pax global headers that must be rejected: the two keywords that would change + * every entry after them, the sparse family whose content cannot be unpacked, and + * payloads that are not pax records at all. + * + * @return One case per rejected global header. Never {@code null}. + */ + private static Stream rejectedGlobalHeaders() { + return Stream.of( + Arguments.of("path redirects every following entry", + TarArchives.paxRecord("path", "../escape.txt"), + "pax global header contains path, which would change every entry after it"), + Arguments.of("size resizes every following entry", + TarArchives.paxRecord("size", "999999"), + "pax global header contains size, which would change every entry after it"), + Arguments.of("sparse map cannot be unpacked", + TarArchives.paxRecord("GNU.sparse.name", "holes.bin"), + SPARSE_REJECTED), + Arguments.of("no length prefix", + "comment=0\n".getBytes(StandardCharsets.UTF_8), MALFORMED_RECORD), + Arguments.of("length longer than the payload", + "99 comment=0\n".getBytes(StandardCharsets.UTF_8), MALFORMED_RECORD), + Arguments.of("length shorter than its own prefix", + "1 comment=0\n".getBytes(StandardCharsets.UTF_8), MALFORMED_RECORD), + Arguments.of("no keyword before the equals sign", + "12 =value\n".getBytes(StandardCharsets.UTF_8), MALFORMED_RECORD), + Arguments.of("no equals sign at all", + "12 comment0\n".getBytes(StandardCharsets.UTF_8), MALFORMED_RECORD)); + } + @ParameterizedTest(name = "{0}") + @MethodSource("rejectedGlobalHeaders") + void testRejectedPaxGlobalHeader(String description, byte[] payload, String message) + throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "pax_global_header", payload, 'g'); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals(message, thrown.getMessage()); + } + + @Test + void testClassicHeaderWithoutUstarMagicIsRead() throws IOException { + final byte[] content = "classic\n".getBytes(StandardCharsets.UTF_8); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tar.write(TarArchives.classicHeader("v7/data.txt", content.length, + TYPE_REGULAR_FILE)); + tar.write(content); + tar.write(new byte[BLOCK - content.length]); + tar.write(new byte[TERMINATOR_SIZE]); + final byte[] archive = tar.toByteArray(); + + Assertions.assertTrue(TarStream.startsWithHeader( + new BufferedInputStream(new ByteArrayInputStream(archive)))); + final TarStream stream = new TarStream(new ByteArrayInputStream(archive)); + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("v7/data.txt", stream.name()); + Assertions.assertTrue(stream.isFile()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testHeaderWithAnEmptyNameIsRejected() throws IOException { + final byte[] block = TarArchives.reseal(header("", 0, TYPE_REGULAR_FILE)); + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + Assertions.assertAll( + () -> { + final IOException thrown = + Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("tar entry header contains an empty name", + thrown.getMessage()); + }, + () -> Assertions.assertFalse( + TarStream.startsWithHeader(new ByteArrayInputStream(block)))); + } + + @Test + void testHeaderNameWithMalformedUtf8IsRejected() { + final byte[] block = header("data.txt", 0, TYPE_REGULAR_FILE); + block[0] = (byte) 0xC3; + block[1] = 0; + TarArchives.reseal(block); + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("tar entry name is not valid UTF-8", thrown.getMessage()); + } + + @Test + void testEntryStreamRejectsInvalidReadRanges() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + Assertions.assertTrue(stream.next()); + final InputStream content = stream.entryStream(); + final byte[] buffer = new byte[8]; + + Assertions.assertAll( + () -> Assertions.assertThrows(NullPointerException.class, + () -> content.read(null, 0, 1)), + () -> Assertions.assertThrows(NullPointerException.class, + () -> content.read(null, 0, 0)), + () -> Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> content.read(buffer, -1, 1)), + () -> Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> content.read(buffer, 0, -1)), + () -> Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> content.read(buffer, 0, buffer.length + 1)), + () -> Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> content.read(buffer, buffer.length, 1)), + // Zero length at a valid offset is legal and must not be mistaken for an error. + () -> Assertions.assertEquals(0, content.read(buffer, buffer.length, 0))); + } + + @Test + void testZeroLengthReadReturnsZero() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + Assertions.assertTrue(stream.next()); + final InputStream content = stream.entryStream(); + final byte[] buffer = new byte[8]; + + Assertions.assertEquals(0, content.read(buffer, 0, 0)); + Assertions.assertArrayEquals("content".getBytes(StandardCharsets.UTF_8), + content.readAllBytes()); + // Exhausted: a zero-length read still reports zero, not end of stream. + Assertions.assertEquals(0, content.read(buffer, 0, 0)); + Assertions.assertEquals(-1, content.read(buffer, 0, buffer.length)); + // Exhausted, and the range is still checked before the end-of-stream answer. + Assertions.assertThrows(IndexOutOfBoundsException.class, + () -> content.read(buffer, 0, buffer.length + 1)); + } + + /** + * Concatenates byte arrays, so a test can assemble one extension header payload from + * several records. + * + * @param parts The pieces to join. Must not be {@code null}. + * @return The concatenation. Never {@code null}. + */ + private static byte[] concat(byte[]... parts) { + final ByteArrayOutputStream joined = new ByteArrayOutputStream(); + for (final byte[] part : parts) { + joined.writeBytes(part); + } + return joined.toByteArray(); + } + + /** + * Encodes metadata records emitted by real archive writers, so tests cover their output. + * + * @return The encoded records. Never {@code null}. + */ + private static byte[] metadataRecords() { + final ByteArrayOutputStream records = new ByteArrayOutputStream(); + for (final String[] record : METADATA) { + records.writeBytes(TarArchives.paxRecord(record[0], record[1])); + } + return records.toByteArray(); + } + + @Test + void testPaxExtendedHeaderSuppliesTheEntryName() throws IOException { + final String path = "./corpus-1.0/" + "d".repeat(120) + "/annotations/train.conllu"; + final byte[] content = "content\n".getBytes(StandardCharsets.UTF_8); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "./PaxHeaders/train.conllu", + concat(TarArchives.paxRecord("path", path), metadataRecords()), 'x'); + // GNU tar stores the truncated name in the entry header. + entry(tar, path.substring(0, 100), content); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(path, stream.name()); + Assertions.assertTrue(stream.isFile()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testPaxPathWithMalformedUtf8IsRejected() throws IOException { + final byte[] malformedPath = { + '1', '0', ' ', 'p', 'a', 't', 'h', '=', (byte) 0xC3, '\n'}; + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "./PaxHeaders/data.txt", malformedPath, 'x'); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("pax record is not valid UTF-8", thrown.getMessage()); + } + + @Test + void testPaxMetadataOnlyHeaderLeavesTheEntryAlone() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "./PaxHeaders/short.txt", metadataRecords(), 'x'); + entry(tar, "./short.txt", "short\n".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("./short.txt", stream.name()); + Assertions.assertEquals(6, stream.size()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testPaxExtendedHeaderAppliesOnlyToTheNextEntry() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "./PaxHeaders/first", + TarArchives.paxRecord("path", "overridden/first.txt"), 'x'); + entry(tar, "truncated-first", "one".getBytes(StandardCharsets.UTF_8)); + entry(tar, "second.txt", "two".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("overridden/first.txt", stream.name()); + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("second.txt", stream.name()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testPaxExtendedHeaderSuppliesTheEntrySize() throws IOException { + final byte[] content = "0123456789".getBytes(StandardCharsets.US_ASCII); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "./PaxHeaders/data.bin", TarArchives.paxRecord("size", "10"), 'x'); + tar.write(header("data.bin", 0, TYPE_REGULAR_FILE)); + tar.write(content); + tar.write(new byte[BLOCK - content.length]); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(10, stream.size()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + } + + @Test + void testGnuLongNameHeaderSuppliesTheEntryName() throws IOException { + final String path = "./corpus-1.0/" + "d".repeat(120) + "/annotations/train.conllu"; + final byte[] content = "content\n".getBytes(StandardCharsets.UTF_8); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "././@LongLink", + (path + "\0").getBytes(StandardCharsets.UTF_8), 'L'); + entry(tar, path.substring(0, 100), content); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(path, stream.name()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + Assertions.assertFalse(stream.next()); + } + + @Test + void testGnuLongNameWithMalformedUtf8IsRejected() throws IOException { + final byte[] malformedName = {'b', 'a', 'd', (byte) 0xC3, 0}; + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "././@LongLink", malformedName, 'L'); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("GNU long name is not valid UTF-8", thrown.getMessage()); + } + + @Test + void testGnuLongLinkHeaderIsConsumed() throws IOException { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + entry(tar, "././@LongLink", + ("../" + "t".repeat(120) + "/target\0").getBytes(StandardCharsets.UTF_8), 'K'); + entry(tar, "data.txt", "content".getBytes(StandardCharsets.UTF_8)); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("data.txt", stream.name()); + Assertions.assertArrayEquals("content".getBytes(StandardCharsets.UTF_8), + stream.entryStream().readAllBytes()); + } + + @Test + void testGnuHeaderDoesNotReadItsAtimeAsANamePrefix() throws IOException { + final byte[] content = "short\n".getBytes(StandardCharsets.UTF_8); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tar.write(TarArchives.gnuHeader("./short.txt", content.length, TYPE_REGULAR_FILE, + "15237132225")); + tar.write(content); + tar.write(new byte[BLOCK - content.length]); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("./short.txt", stream.name()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + } + + /** + * Supplies the two sparse entry declarations. Their archived bytes describe file holes, + * so copying those bytes would produce incorrect content. + * + * @return One case per sparse declaration. Never {@code null}. + */ + private static Stream sparseEntries() { + return Stream.of( + Arguments.of("GNU sparse type flag", (Supplier) () -> { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tar.writeBytes(header("holes.bin", 0, 'S')); + return tar.toByteArray(); + }), + Arguments.of("pax sparse records", (Supplier) () -> { + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + try { + entry(tar, "./PaxHeaders/holes.bin", + concat(TarArchives.paxRecord("GNU.sparse.major", "1"), + TarArchives.paxRecord("GNU.sparse.name", "holes.bin")), 'x'); + entry(tar, "holes.bin", "not the content".getBytes(StandardCharsets.UTF_8)); + } catch (IOException e) { + throw new IllegalStateException(e); + } + return tar.toByteArray(); + })); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("sparseEntries") + void testSparseEntriesAreRejected(String description, Supplier archive) { + final TarStream stream = new TarStream(new ByteArrayInputStream(archive.get())); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals(SPARSE_REJECTED, thrown.getMessage()); + } + + @Test + void testBase256SizeFieldIsRead() throws IOException { + final byte[] content = "0123456789".getBytes(StandardCharsets.US_ASCII); + final ByteArrayOutputStream tar = new ByteArrayOutputStream(); + tar.write(TarArchives.base256Header("big.bin", content.length, TYPE_REGULAR_FILE)); + tar.write(content); + tar.write(new byte[BLOCK - content.length]); + tar.write(new byte[TERMINATOR_SIZE]); + final TarStream stream = new TarStream(new ByteArrayInputStream(tar.toByteArray())); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals("big.bin", stream.name()); + Assertions.assertEquals(content.length, stream.size()); + Assertions.assertArrayEquals(content, stream.entryStream().readAllBytes()); + } + + @Test + void testBase256SizeFieldCarriesLengthsBeyondTheOctalRange() throws IOException { + final long beyondOctal = 8L * 1024 * 1024 * 1024; + Assertions.assertTrue(beyondOctal > LARGEST_OCTAL_SIZE); + final TarStream stream = new TarStream(new ByteArrayInputStream( + TarArchives.base256Header("huge.bin", beyondOctal, TYPE_REGULAR_FILE))); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(beyondOctal, stream.size()); + } + + @Test + void testBase256SizeFieldAcceptsTheLargestRepresentableLength() throws IOException { + final TarStream stream = new TarStream(new ByteArrayInputStream( + TarArchives.base256Header("max.bin", Long.MAX_VALUE, TYPE_REGULAR_FILE))); + + Assertions.assertTrue(stream.next()); + Assertions.assertEquals(Long.MAX_VALUE, stream.size()); + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("truncated tar archive", thrown.getMessage()); + } + + @Test + void testBase256SizeFieldBeyondTheLongRangeIsRejected() { + final byte[] block = TarArchives.base256Header("huge.bin", 0, TYPE_REGULAR_FILE); + // One bit above the 63 a long can hold, left of everything the encoder can write. + block[SIZE_OFFSET + 1] = 1; + TarArchives.reseal(block); + final TarStream stream = new TarStream(new ByteArrayInputStream(block)); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals( + "tar size field exceeds the largest length this reader can represent", + thrown.getMessage()); + } + + @Test + void testNegativeBase256SizeFieldIsRejected() { + final TarStream stream = new TarStream(new ByteArrayInputStream( + TarArchives.base256Header("negative.bin", -1, TYPE_REGULAR_FILE))); + + final IOException thrown = Assertions.assertThrows(IOException.class, stream::next); + Assertions.assertEquals("tar size field is negative", thrown.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizerTest.java new file mode 100644 index 0000000000..bb5a0267c2 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AggregateCharSequenceNormalizerTest.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.tools.util.normalizer; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +class AggregateCharSequenceNormalizerTest { + + @Test + void testAppliesNormalizersInConstructionOrder() { + CharSequenceNormalizer upper = text -> text.toString().replace('a', 'B'); + CharSequenceNormalizer strip = text -> text.toString().replace("B", ""); + assertEquals("cc", new AggregateCharSequenceNormalizer(upper, strip) + .normalize("accB").toString()); + assertEquals("Bcc", new AggregateCharSequenceNormalizer(strip, upper) + .normalize("accB").toString()); + } + + @Test + void testRejectsNullNormalizersLoudly() { + IllegalArgumentException nullArray = assertThrows(IllegalArgumentException.class, + () -> new AggregateCharSequenceNormalizer((CharSequenceNormalizer[]) null)); + assertEquals("The normalizers must not be null.", nullArray.getMessage()); + + IllegalArgumentException nullElement = assertThrows(IllegalArgumentException.class, + () -> new AggregateCharSequenceNormalizer( + NfcCharSequenceNormalizer.getInstance(), null)); + assertEquals("The normalizers must not contain null.", nullElement.getMessage()); + } + + @Test + void testChangingTheCallerArrayDoesNotReachTheAggregate() { + CharSequenceNormalizer[] normalizers = {NfcCharSequenceNormalizer.getInstance()}; + AggregateCharSequenceNormalizer aggregate = + new AggregateCharSequenceNormalizer(normalizers); + normalizers[0] = text -> "changed"; + assertEquals("abc", aggregate.normalize("abc").toString()); + } + + @Test + void testRejectsNullTextLoudly() { + AggregateCharSequenceNormalizer aggregate = new AggregateCharSequenceNormalizer(); + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, + () -> aggregate.normalize(null)); + assertEquals("The text must not be null.", e.getMessage()); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AlignedNormalizerPipelineTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AlignedNormalizerPipelineTest.java index 07098dfbc8..a4da8ed564 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AlignedNormalizerPipelineTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/AlignedNormalizerPipelineTest.java @@ -28,9 +28,9 @@ /** * Exercises {@link OffsetAwareNormalizer} and {@code TextNormalizer.Builder.buildAligned()}: the - * cursor-based rungs report alignments, an aligned pipeline composes them with + * cursor-based normalizers report alignments, an aligned pipeline composes them with * {@link Alignment#andThen(Alignment)} so a span found in the fully normalized text maps back to the - * original input, and a non-alignable rung is rejected loudly. + * original input, and a non-alignable normalizer is rejected loudly. */ public class AlignedNormalizerPipelineTest { @@ -50,10 +50,11 @@ private static String covered(AlignedText aligned, int normalizedStart, int norm // The aligned form must always reproduce exactly what the plain form produces. @Test - void alignedNormalizedTextMatchesPlainForEveryRung() { - final OffsetAwareNormalizer[] rungs = { + void alignedNormalizedTextMatchesPlainForEveryNormalizer() { + final OffsetAwareNormalizer[] normalizers = { WhitespaceCharSequenceNormalizer.getInstance(), LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance(), + ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance(), DashCharSequenceNormalizer.getInstance(), InvisibleCharSequenceNormalizer.getInstance(), QuoteCharSequenceNormalizer.getInstance(), @@ -74,10 +75,10 @@ void alignedNormalizedTextMatchesPlainForEveryRung() { cp(0x201C) + "don" + cp(0x2019) + "t " + cp(0x2026) + " Stra" + cp(0x00DF) + "e " + cp(0x2022) + " " + cp(0xFF15) + cp(MATH_BOLD_DIGIT_ZERO) }; - for (final OffsetAwareNormalizer rung : rungs) { + for (final OffsetAwareNormalizer normalizer : normalizers) { for (final String input : inputs) { - assertEquals(rung.normalize(input).toString(), rung.normalizeAligned(input).normalized(), - rung.getClass().getSimpleName() + " on [" + input + "]"); + assertEquals(normalizer.normalize(input).toString(), normalizer.normalizeAligned(input).normalized(), + normalizer.getClass().getSimpleName() + " on [" + input + "]"); } } } @@ -137,7 +138,7 @@ void emptyAlignedPipelineIsIdentity() { } @Test - void buildAlignedRejectsNonAlignableRungLoudly() { + void buildAlignedRejectsNonAlignableNormalizerLoudly() { final IllegalStateException ex = assertThrows(IllegalStateException.class, () -> TextNormalizer.builder().nfc().whitespace().buildAligned()); assertTrue(ex.getMessage().contains("Nfc"), ex.getMessage()); @@ -145,7 +146,7 @@ void buildAlignedRejectsNonAlignableRungLoudly() { } @Test - void buildAlignedReportsTheOffendingRungIndexWhenItIsNotFirst() { + void buildAlignedReportsTheOffendingNormalizerIndexWhenItIsNotFirst() { // A non-alignable step after several offset-aware ones must still be rejected, and the message // must name its 0-based position (index 2) and type so the failure points at the right fold. final IllegalStateException ex = assertThrows(IllegalStateException.class, @@ -155,7 +156,7 @@ void buildAlignedReportsTheOffendingRungIndexWhenItIsNotFirst() { } @Test - void buildAlignedRejectsEachKindOfNonAlignableRung() { + void buildAlignedRejectsEachKindOfNonAlignableNormalizer() { // Every fold that routes through java.text.Normalizer or JDK case mapping is rejected, named. assertThrows(IllegalStateException.class, () -> TextNormalizer.builder().nfkc().buildAligned()); @@ -180,6 +181,8 @@ void capabilityIsDetectableByInstanceOf() { assertTrue(BulletCharSequenceNormalizer.getInstance() instanceof OffsetAwareNormalizer); assertTrue(GermanUmlautCharSequenceNormalizer.getInstance() instanceof OffsetAwareNormalizer); assertTrue(FullCaseFoldCharSequenceNormalizer.getInstance() instanceof OffsetAwareNormalizer); + assertTrue(ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance() + instanceof OffsetAwareNormalizer); // The folds that route through java.text.Normalizer or JDK case mapping cannot, by design. assertFalse(NfkcCharSequenceNormalizer.getInstance() instanceof OffsetAwareNormalizer); assertFalse(CaseFoldCharSequenceNormalizer.getInstance() instanceof OffsetAwareNormalizer); @@ -203,17 +206,17 @@ void roundTripOfAFullySpanningMatchReturnsTheWholeOriginal() { @Test void lineBreakPreservingCollapsesHorizontalRunsButKeepsBreaks() { - final LineBreakPreservingWhitespaceCharSequenceNormalizer rung = + final LineBreakPreservingWhitespaceCharSequenceNormalizer normalizer = LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance(); final String original = "Hello world\n\n\tfoo bar"; - assertEquals("Hello world\nfoo bar", rung.normalize(original).toString()); + assertEquals("Hello world\nfoo bar", normalizer.normalize(original).toString()); - // The plain whitespace rung instead flattens the blank line into a single space. + // The plain whitespace normalizer instead flattens the blank line into a single space. assertEquals("Hello world foo bar", WhitespaceCharSequenceNormalizer.getInstance().normalize(original).toString()); - final AlignedText aligned = rung.normalizeAligned(original); - assertEquals(rung.normalize(original).toString(), aligned.normalized()); + final AlignedText aligned = normalizer.normalizeAligned(original); + assertEquals(normalizer.normalize(original).toString(), aligned.normalized()); // "bar" sits at [16, 19) in the collapsed form and at [21, 24) in the original. assertEquals(original.indexOf("bar"), aligned.toOriginalSpan(16, 19).getStart()); assertEquals("bar", covered(aligned, 16, 19)); @@ -223,10 +226,10 @@ void lineBreakPreservingCollapsesHorizontalRunsButKeepsBreaks() { @Test void lineBreakPreservingTrimsLeadingAndTrailingBreaks() { - final LineBreakPreservingWhitespaceCharSequenceNormalizer rung = + final LineBreakPreservingWhitespaceCharSequenceNormalizer normalizer = LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance(); final String original = "\n\nHello\n\n"; - final AlignedText aligned = rung.normalizeAligned(original); + final AlignedText aligned = normalizer.normalizeAligned(original); assertEquals("Hello", aligned.normalized()); assertEquals("Hello", covered(aligned, 0, 5)); assertEquals(original.indexOf("Hello"), aligned.toOriginalSpan(0, 5).getStart()); @@ -261,18 +264,49 @@ void pipelineMapsAnOriginalSpanForwardToTheNormalizedText() { @Test void lineBreakPreservingNormalizesCrLfAndUnicodeSeparators() { - final LineBreakPreservingWhitespaceCharSequenceNormalizer rung = + final LineBreakPreservingWhitespaceCharSequenceNormalizer normalizer = LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance(); - assertEquals("a\nb", rung.normalize("a\r\nb").toString()); // CRLF -> one newline - assertEquals("a\nb", rung.normalize("a\n\n\n\nb").toString()); // blank lines -> one newline - assertEquals("x\ny", rung.normalize("x" + cp(0x2028) + "y").toString()); // line separator - assertEquals("p\nq", rung.normalize("p" + cp(0x2029) + "q").toString()); // paragraph separator + assertEquals("a\nb", normalizer.normalize("a\r\nb").toString()); // CRLF -> one newline + assertEquals("a\nb", normalizer.normalize("a\n\n\n\nb").toString()); // blank lines -> one newline + assertEquals("x\ny", normalizer.normalize("x" + cp(0x2028) + "y").toString()); // line separator + assertEquals("p\nq", normalizer.normalize("p" + cp(0x2029) + "q").toString()); // paragraph separator // A horizontal run still collapses to a space even when mixed with a break-bearing run. - assertEquals("a b\nc", rung.normalize("a b \n c").toString()); + assertEquals("a b\nc", normalizer.normalize("a b \n c").toString()); + } + + @Test + void paragraphPreservingUnwrapsHardWrapsButKeepsBlankLines() { + final ParagraphPreservingWhitespaceCharSequenceNormalizer normalizer = + ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance(); + final String original = "Hello world\n\n\tfoo bar"; + assertEquals("Hello world\nfoo bar", normalizer.normalize(original).toString()); + assertEquals("on the bank", normalizer.normalize("on the\nbank").toString()); + assertEquals("a\nb", LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance() + .normalize("a\nb").toString()); + + final AlignedText aligned = normalizer.normalizeAligned(original); + assertEquals(normalizer.normalize(original).toString(), aligned.normalized()); + assertEquals(original.indexOf("bar"), aligned.toOriginalSpan(16, 19).getStart()); + assertEquals("bar", covered(aligned, 16, 19)); + assertEquals("\n\n\t", covered(aligned, 11, 12)); + } + + @Test + void paragraphPreservingComposesInAnAlignedPipeline() { + assertTrue(ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance() + instanceof OffsetAwareNormalizer); + final String original = "a" + cp(ZERO_WIDTH_SPACE) + " b\n\nc" + cp(EM_DASH) + "d"; + final OffsetAwareNormalizer pipeline = TextNormalizer.builder() + .stripInvisible().whitespacePreservingParagraphs().dashes().buildAligned(); + + final AlignedText aligned = pipeline.normalizeAligned(original); + assertEquals("a b\nc-d", aligned.normalized()); + assertEquals(pipeline.normalize(original).toString(), aligned.normalized()); + assertEquals("c" + cp(EM_DASH) + "d", covered(aligned, 4, 7)); } @Test - void whitespaceRungCollapsesAllWhitespaceToEmptyWithAValidSpan() { + void whitespaceNormalizerCollapsesAllWhitespaceToEmptyWithAValidSpan() { final AlignedText aligned = WhitespaceCharSequenceNormalizer.getInstance().normalizeAligned(" "); assertEquals("", aligned.normalized()); diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/CharSequenceNormalizerContractTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/CharSequenceNormalizerContractTest.java index 40e537d496..240548336b 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/CharSequenceNormalizerContractTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/CharSequenceNormalizerContractTest.java @@ -55,6 +55,8 @@ static Stream allNormalizers() { Arguments.of("invisible", InvisibleCharSequenceNormalizer.getInstance()), Arguments.of("lineBreakPreservingWhitespace", LineBreakPreservingWhitespaceCharSequenceNormalizer.getInstance()), + Arguments.of("paragraphPreservingWhitespace", + ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance()), Arguments.of("quote", QuoteCharSequenceNormalizer.getInstance()), Arguments.of("whitespace", WhitespaceCharSequenceNormalizer.getInstance())); } diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizerTest.java new file mode 100644 index 0000000000..250ad75c19 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/ParagraphPreservingWhitespaceCharSequenceNormalizerTest.java @@ -0,0 +1,166 @@ +/* + * 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.util.normalizer; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +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.util.Span; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class ParagraphPreservingWhitespaceCharSequenceNormalizerTest { + + private static ParagraphPreservingWhitespaceCharSequenceNormalizer norm() { + return ParagraphPreservingWhitespaceCharSequenceNormalizer.getInstance(); + } + + private static String cp(int codePoint) { + return new String(Character.toChars(codePoint)); + } + + static Stream unicodeLineBreakCodePoints() { + return UnicodeWhitespace.lineBreaks().stream() + .map(ws -> Arguments.of(ws.codePoint(), ws.abbreviation())); + } + + @Test + void getInstanceReturnsTheSharedSingleton() { + assertSame(norm(), norm()); + } + + @Test + void deserializationReturnsTheSharedSingleton() throws Exception { + final ByteArrayOutputStream bos = new ByteArrayOutputStream(); + try (ObjectOutputStream oos = new ObjectOutputStream(bos)) { + oos.writeObject(norm()); + } + try (ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()))) { + assertSame(norm(), ois.readObject()); + } + } + + @Test + void horizontalRunsCollapseToASingleSpace() { + assertEquals("a b c", norm().normalize("a \t b" + cp(0x00A0) + cp(0x3000) + "c").toString()); + } + + @Test + void singleLineBreakInRunCollapsesToSpace() { + assertEquals("a b", norm().normalize("a \n b").toString()); + assertEquals("on the bank", norm().normalize("on the\nbank").toString()); + } + + @ParameterizedTest(name = "single {1} break unwraps to a space") + @MethodSource("unicodeLineBreakCodePoints") + void singleUnicodeLineBreakInRunCollapsesToSpace(int codePoint, String abbreviation) { + if (codePoint == 0x000D) { + assertEquals("a b", norm().normalize("a\rb").toString()); + assertEquals("a b", norm().normalize("a\r\nb").toString()); + } else { + assertEquals("a b", norm().normalize("a" + cp(codePoint) + "b").toString()); + } + } + + @Test + void twoOrMoreLineBreaksInRunCollapseToOneNewline() { + assertEquals("a\nb", norm().normalize("a\n\nb").toString()); + assertEquals("a\nb", norm().normalize("a\n\n\n\nb").toString()); + assertEquals("Hello world\nfoo bar", norm().normalize("Hello world\n\n\tfoo bar").toString()); + } + + @ParameterizedTest(name = "doubled {1} break keeps a paragraph boundary") + @MethodSource("unicodeLineBreakCodePoints") + void twoUnicodeLineBreaksInRunCollapseToNewline(int codePoint, String abbreviation) { + if (codePoint == 0x000D) { + assertEquals("a\nb", norm().normalize("a\r\n\r\nb").toString()); + } else { + assertEquals("a\nb", norm().normalize("a" + cp(codePoint) + cp(codePoint) + "b").toString()); + } + } + + @Test + void carriageReturnLineFeedCountsAsOneBreak() { + assertEquals("a b", norm().normalize("a\r\nb").toString()); + assertEquals("a\nb", norm().normalize("a\r\n\r\nb").toString()); + } + + @Test + void gutenbergHardWrapUnwrapsButKeepsParagraphBreak() { + final String input = """ + Alice was beginning to get very tired of sitting by her sister on the + bank, and of having nothing to do: once or twice she had peeped into + the book her sister was reading, but it had no pictures or + conversations in it, "and what is the use of a book," thought Alice + "without pictures or conversations?" + + So she was considering in her own mind (as well as she could, for the\ + """; + final String normalized = norm().normalize(input).toString(); + assertEquals( + "Alice was beginning to get very tired of sitting by her sister on the bank, " + + "and of having nothing to do: once or twice she had peeped into " + + "the book her sister was reading, but it had no pictures or " + + "conversations in it, \"and what is the use of a book,\" thought Alice " + + "\"without pictures or conversations?\"\n" + + "So she was considering in her own mind (as well as she could, for the", + normalized); + } + + @Test + void edgesAreTrimmed() { + assertEquals("a", norm().normalize(" a ").toString()); + assertEquals("a", norm().normalize("\n a \n").toString()); + } + + @Test + void whitespaceOnlyInputNormalizesToEmpty() { + assertEquals("", norm().normalize(" \n ").toString()); + assertEquals("", norm().normalize("").toString()); + } + + @Test + void alignedNormalizedMatchesNormalize() { + final String in = " one \t two\r\n\r\nthree " + cp(0x2028) + " four "; + assertEquals("one two\nthree four", norm().normalize(in).toString()); + assertEquals(norm().normalize(in).toString(), norm().normalizeAligned(in).normalizedString()); + } + + @Test + void alignmentMapsThroughCollapseAndTrim() { + final AlignedText at = norm().normalizeAligned(" x\n\ny "); + assertEquals("x\ny", at.normalizedString()); + assertEquals(new Span(2, 3), at.toOriginalSpan(0, 1)); // "x" + assertEquals(new Span(5, 6), at.toOriginalSpan(2, 3)); // "y" + } + + @Test + void nullTextIsRejected() { + assertThrows(IllegalArgumentException.class, () -> norm().normalize(null)); + assertThrows(IllegalArgumentException.class, () -> norm().normalizeAligned(null)); + } +} diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/TextNormalizerTest.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/TextNormalizerTest.java index 6391a311a7..7845b12238 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/TextNormalizerTest.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/util/normalizer/TextNormalizerTest.java @@ -30,7 +30,7 @@ private static String cp(int codePoint) { } @Test - void testRungsApplyInOrder() { + void testNormalizersApplyInOrder() { final CharSequenceNormalizer n = TextNormalizer.builder().caseFold().accentFold().build(); assertEquals("cafe", n.normalize("CAF" + cp(0x00C9)).toString()); // CAFE-acute -> cafe } @@ -66,7 +66,7 @@ void testDefaultChainCleansMessyInput() { } @Test - void testEveryRungIsInvokable() { + void testEveryNormalizerIsInvokable() { final CharSequenceNormalizer n = TextNormalizer.builder() .stripInvisible().nfc().nfkc().whitespace().quotes().dashes().digits().ellipsis().bullets() .fullCaseFold().accentFold().build(); @@ -76,14 +76,14 @@ void testEveryRungIsInvokable() { } @Test - void testCaseFoldRungIsInvokable() { + void testCaseFoldNormalizerIsInvokable() { final CharSequenceNormalizer n = TextNormalizer.builder().caseFold().build(); assertEquals("cafe", n.normalize("CAFE").toString()); } @Test void testCaseFoldWithFullCaseFoldIsPermittedButRedundant() { - // This builder composes rungs freely and does not enforce the exclusion the TermAnalyzer + // This builder composes normalizers freely and does not enforce the exclusion the TermAnalyzer // layer does; the combination is documented as redundant, and this pins that it changes // nothing over full case folding alone. final String input = "STRA" + cp(0x00DF) + "E"; diff --git a/opennlp-core/opennlp-runtime/src/test/resources/opennlp/tools/util/dictionary-catalog.properties b/opennlp-core/opennlp-runtime/src/test/resources/opennlp/tools/util/dictionary-catalog.properties new file mode 100644 index 0000000000..748258fe69 --- /dev/null +++ b/opennlp-core/opennlp-runtime/src/test/resources/opennlp/tools/util/dictionary-catalog.properties @@ -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. +# + +# Pinned remote dictionary files used as catalog parsing fixtures. Tests never fetch +# these URLs. Applications provide and maintain their own catalogs. + +# MeCab IPADIC 2.7.0 (EUC-JP). Upstream: MeCab project on SourceForge. +mecab.ipadic.url=https://downloads.sourceforge.net/project/mecab/\ +mecab-ipadic/2.7.0-20070801/mecab-ipadic-2.7.0-20070801.tar.gz +mecab.ipadic.sha512=35ea662cb62f1967849f7ed5781bd6dafef0fe20d63e88d9\ +a0057666e57ed23d5a0e6fb8d0701a0cc4da43a1050c1b0246\ +3bb862decc71c36b7fc2acdc158d86 +mecab.ipadic.filename=mecab-ipadic-2.7.0-20070801.tar.gz + +# mecab-ko-dic 2.1.1 (UTF-8). Upstream: eunjeon/mecab-ko-dic on Bitbucket. +mecab.ko-dic.url=https://bitbucket.org/eunjeon/mecab-ko-dic/downloads/\ +mecab-ko-dic-2.1.1-20180720.tar.gz +mecab.ko-dic.sha512=986f8f9c66c53accd296756bf632c979d2d44b695ada33f3\ +6c662f210dba34cd95d67b61dd8c84a1f7d59f80ee6bc22eb1\ +e9afb5dc6a7f9b6b75b4fbf2f8164f +mecab.ko-dic.filename=mecab-ko-dic-2.1.1-20180720.tar.gz + +# LibreOffice en_US Hunspell pair, pinned to dictionaries commit 208a9fd8. +hunspell.en_US.aff.url=https://raw.githubusercontent.com/LibreOffice/\ +dictionaries/208a9fd80b2a182fe20f224cd615119c6323ae2e/en/en_US.aff +hunspell.en_US.aff.sha512=2b4448dfdff03caf300914415f4642f8d2ba5b650c5f024a\ +12355b420a279ffc12146649fce092ba591504476634a3d6\ +fd4c079335a27085b396fa76bfd28b74 +hunspell.en_US.aff.filename=en_US.aff + +hunspell.en_US.dic.url=https://raw.githubusercontent.com/LibreOffice/\ +dictionaries/208a9fd80b2a182fe20f224cd615119c6323ae2e/en/en_US.dic +hunspell.en_US.dic.sha512=4be737249a8a436d20a02be575dcf6cf2f06f5f2abb840ea\ +5ec0ef0ac73a71fa0e4669e527c703d5c6b50ef61713a674\ +1b55bc136d559a54bcdeebcd62027988 +hunspell.en_US.dic.filename=en_US.dic + +hunspell.en_US.readme.url=https://raw.githubusercontent.com/LibreOffice/\ +dictionaries/208a9fd80b2a182fe20f224cd615119c6323ae2e/en/README_en_US.txt +hunspell.en_US.readme.sha512=aa23ebc8adc0649b540264c7bf98cef5b6e383fec0e4a1a7\ +dd49d1c887cfeefd8edf6a568afc8a651521a3864c5b1ab5\ +0748ad16230d386809080b5b09135082 +hunspell.en_US.readme.filename=README_en_US.txt diff --git a/opennlp-distr/pom.xml b/opennlp-distr/pom.xml index 7e8501bca9..c1cf8e841c 100644 --- a/opennlp-distr/pom.xml +++ b/opennlp-distr/pom.xml @@ -96,6 +96,11 @@ opennlp-subword + + org.apache.opennlp + opennlp-embeddings + + org.apache.opennlp diff --git a/opennlp-distr/src/main/assembly/bin.xml b/opennlp-distr/src/main/assembly/bin.xml index b861de7caf..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 diff --git a/opennlp-distr/src/main/assembly/src.xml b/opennlp-distr/src/main/assembly/src.xml index 1dea96bd62..6287486289 100644 --- a/opennlp-distr/src/main/assembly/src.xml +++ b/opennlp-distr/src/main/assembly/src.xml @@ -34,5 +34,12 @@ **/release.properties + + + ../.mvn + .mvn + \ No newline at end of file diff --git a/opennlp-distr/src/main/readme/NOTICE b/opennlp-distr/src/main/readme/NOTICE index 4ae212e882..5de2119011 100644 --- a/opennlp-distr/src/main/readme/NOTICE +++ b/opennlp-distr/src/main/readme/NOTICE @@ -162,10 +162,10 @@ List of third-party dependencies grouped by their license type. MIT License - * onnx-runtime (com.microsoft.onnxruntime:onnxruntime_gpu:1.28.0 - https://microsoft.github.io/onnxruntime/) - * onnx-runtime (com.microsoft.onnxruntime:onnxruntime:1.28.0 - https://microsoft.github.io/onnxruntime/) + * onnx-runtime (com.microsoft.onnxruntime:onnxruntime_gpu:1.29.0 - https://microsoft.github.io/onnxruntime/) + * onnx-runtime (com.microsoft.onnxruntime:onnxruntime:1.29.0 - https://microsoft.github.io/onnxruntime/) * SLF4J API Module (org.slf4j:slf4j-api:2.0.18 - http://www.slf4j.org) The MIT License (MIT) - * ClassGraph (io.github.classgraph:classgraph:4.8.186 - https://github.com/classgraph/classgraph) + * ClassGraph (io.github.classgraph:classgraph:4.8.194 - https://github.com/classgraph/classgraph) diff --git a/opennlp-docs/src/docbkx/document.xml b/opennlp-docs/src/docbkx/document.xml new file mode 100644 index 0000000000..169a34c2fd --- /dev/null +++ b/opennlp-docs/src/docbkx/document.xml @@ -0,0 +1,394 @@ + + + + + + + Document Annotation Container + +
    + Introduction + + The task APIs of OpenNLP each return one kind of result: a tokenizer returns + spans, a tagger returns tags, a name finder returns names. An application that + runs several of them over one text has to keep those results aligned itself, + and structured results such as coreference chains or dependency arcs, which + refer to other results, have no place to live at all. The document container + gives them one. It adds a small API next to the task APIs and replaces none of + them: a task that needs only one component keeps calling it directly. + + + The package opennlp.tools.document provides an immutable container + that carries the original text of one document together with any number of + annotation layers over it. A layer is a list of Annotation values, + each pairing a Span with a value. A span's offsets count Java + chars (UTF-16 units, so a supplementary-plane character counts as two) from + the beginning of the text exactly as the caller supplied it, never a + normalized or derived form, so every annotation can be highlighted in + the source text. The value's Java type is part of the layer's identity: a token + layer reads back as List<Annotation<String>>, for + example the token The paired with the span [0..3). + + + A document with three annotation layers + +text: The dog barks. It naps. + 0 1 2 + 01234567890123456789012 + +sentences: [0..14) "The dog barks." [15..23) "It naps." +tokens: [0..3) "The" [4..7) "dog" [8..14) "barks." [15..17) "It" [18..23) "naps." +pos-tags: [0..3) "DT" [4..7) "NN" [8..14) "VBZ" [15..17) "PRP" [18..23) "VBZ" + + + Each layer in the figure above is identified by a LayerKey that + names the layer and declares the type of its values: + Layers.TOKENS is a LayerKey<String>, so + document.get(Layers.TOKENS) returns the token annotations with + string values. Two keys are equal when their id, their value type, and their + scope are equal, so independently created constants for the same layer + interoperate. The keys of the core linguistic layers live in + Layers; a capability-specific layer's key lives on the annotator + that provides it, and adding a capability never changes the container. + + + Key ids are namespaced. Keys OpenNLP defines start with the + opennlp: prefix; a user extension must use its own prefix. A + bare id, like the token-lengths key below, is legal for an + application-local layer. Toolkit keys are created through + Layers.key and Layers.documentKey, which apply the + prefix. + + + A document is never modified in place: adding a layer returns a new document + that shares the unchanged layers with its ancestor. Documents created through + Document.of capture their text at construction and are safe to + share between threads. Annotations reference other annotations by index within + their layer, never by object identity. Three invariants make those index + references sound: a layer keeps its insertion order, a layer is immutable once + added and detached from the caller's input list, and providing a layer that + already exists is rejected, so an index reference stays valid for the lifetime + of the document. + + + A key declares its scope. Keys are positional by default: every annotation + carries a span, and consumers never null-check it. A document-scoped key + carries whole-document values without spans, which is where a language id, + a category distribution, or provenance belongs: + + LANGUAGE = LayerKey.document("app:language", String.class); + +Document tagged = document.with(LANGUAGE, List.of(Annotation.of("eng"))); +String language = tagged.get(LANGUAGE).get(0).value(); // "eng", span is null]]> + + + A corpus may carry a hand-annotated version of a layer beside a produced one. + The convention is a gold: id prefix on the same key scheme, for + example gold:opennlp:tokens beside opennlp:tokens; + since adding a layer is once-only, competing versions always live under + distinct keys and never replace each other. + + + Each rule above is pinned by a test in DocumentContractTest: + testKeysWithSameIdButDifferentTypesAreDifferentLayers, + testToolkitKeysCarryTheNamespacePrefix, + testDocumentScopedLayersCarrySpanlessValues, + testLayerPreservesInsertionOrder, + testDuplicateLayerIsRejectedWithExactMessage, + testGoldLayerLivesBesideThePredictedLayer, and + testCustomLayerNeedsNoContainerChange. + + + None of this is new architecture, and that is intended. Standoff annotations + over an unchanged original text are how the UIMA CAS works; typed keys that let + any module add a layer without the container knowing about it are the pattern of + Stanford CoreNLP's annotation map; and annotators that declare what they require + and provide, checked when the pipeline is assembled, follow spaCy's pipeline + analysis. The container takes the lightweight end of each: string ids with a + namespace instead of class identity, and no type system descriptors. + +
    + +
    + What the layers contain + + A layer has one value type and one annotation per unit, anchored on the span + of that unit in the original text. The toolkit's own adapters provide these + layers: + + + Layers provided by the toolkit's adapters + + + + Key + Value + One annotation per + Provided by + + + + + opennlp:sentences + String, the sentence text + sentence + SentenceDetectorAnnotator + + + opennlp:tokens + String, the token + token + TokenizerAnnotator + + + opennlp:pos + String, the tag + token, on the token's span + POSTaggerAnnotator + + + opennlp:lemmas + String, the lemma + token, on the token's span + LemmatizerAnnotator + + + opennlp:stems + String, the stem + token, on the token's span + StemmerAnnotator + + + opennlp:entities + String, the entity type + name, on the name's span + NameFinderAnnotator + + + opennlp:chunks + String, the chunk type + chunk + ChunkerAnnotator + + + opennlp:phrases + Phrase: the label and the head span + constituent + ParserAnnotator + + + +
    + + The keys of the first four layers in the table live in Layers; + the others live on the annotator that provides them, which is also where an + extension puts its key. Layers that refer to other layers do so by index into + them, which the invariants above keep valid: a coreference layer, for example, + puts one annotation on each mention's span with the chain it belongs to, and a + dependency layer puts one annotation on each token's span naming the head token + by its index in opennlp:tokens. Adding such a layer requires no change + in this package; testCustomLayerNeedsNoContainerChange pins that, + and the section below shows how it is done. + +
    + +
    + Building a pipeline + + A pipeline step implements DocumentAnnotator: it reads layers from a + document and returns a new document with its own layers added. An annotator + declares the layers it requires and provides, and a + DocumentAnalyzer validates those declarations when the pipeline is + assembled: every required layer must be provided by an earlier annotator, and no + two annotators may provide the same layer, so a misordered pipeline fails at + build time rather than midway through a document. + + + Adapters for the toolkit's own components are provided: + SentenceDetectorAnnotator, TokenizerAnnotator, + POSTaggerAnnotator, NameFinderAnnotator, + ChunkerAnnotator, ParserAnnotator, + LemmatizerAnnotator, and StemmerAnnotator. Each wraps + one of the existing component APIs; an application that needs only one task + can keep using that component's API directly. The following pipeline combines + three adapters with one custom annotator and analyzes the text from the + introduction's figure; the components behind the adapters are any + SentenceDetector, Tokenizer, and + POSTagger, for example the ME implementations loaded from models: + + + + + DocumentPipelineExampleTest asserts the pipeline and layer + round-trip shown here; DocumentAnalyzerTest covers the build-time + checks, for example testMisorderedPipelineFailsAtBuildTime. + + + The resulting document carries exactly the four layers the pipeline provides: + the figure's three plus the custom token-lengths layer. The sentence + layer holds two annotations, [0..14) covering + The dog barks. and [15..23) covering + It naps.. The token layer holds five tokens whose spans refer to + the document text even inside the second sentence, so + It is [15..17) and naps. is + [18..23). Layers produced per token stay aligned with the token + layer by position: + + > tokens = document.get(Layers.TOKENS); +List> tags = document.get(Layers.POS_TAGS); +for (int i = 0; i < tags.size(); i++) { + // each tag sits on its token's span, e.g. "DT" on [0..3) for "The" + Span span = tags.get(i).span(); +} + +// every span refers to the original text, so covered text round-trips +for (Annotation token : tokens) { + CharSequence covered = token.span().getCoveredText(document.text()); +}]]> + + + Because a document is immutable, independent pipelines can process the same + text in parallel, each starting from its own Document.of(text), + and their results can be joined afterwards: merge returns a new + document carrying the layers of both. The texts must match, and by default a + layer key present on both documents is rejected, so two branches that both + run a tokenizer collide on opennlp:tokens. When the branches + rebuild a shared prefix identically, passing + DuplicateLayerPolicy.KEEP_EQUAL keeps one copy of each agreeing + layer; copies whose contents differ are still rejected. + +
    + +
    + Writing a custom annotator + + The annotator below defines its own key in its own code, reads the token + layer, and provides one integer annotation per token; the analyzer's + build-time validation guarantees a tokenizer ran earlier: + + TOKEN_LENGTHS = + LayerKey.of("token-lengths", Integer.class); + +class TokenLengthAnnotator implements DocumentAnnotator { + + @Override + public Document annotate(Document document) { + DocumentAnnotators.requireLayers(document, Layers.TOKENS); + List> tokens = document.get(Layers.TOKENS); + List> lengths = new ArrayList<>(tokens.size()); + for (Annotation token : tokens) { + lengths.add(new Annotation<>(token.span(), token.value().length())); + } + return document.with(TOKEN_LENGTHS, lengths); + } + + @Override + public Set> requires() { + return Set.of(Layers.TOKENS); + } + + @Override + public Set> provides() { + return Set.of(TOKEN_LENGTHS); + } +}]]> + + + Because the key declares the value type, the values come back as numbers + without a cast; for the text above the five values are + 3, 3, 6, 2, 5, each on its token's span: + + > lengths = document.get(TOKEN_LENGTHS); +int firstTokenLength = lengths.get(0).value(); // 3, for "The" at [0..3)]]> + + + Required layers must be present, but they may be empty: an empty token layer + yields the annotator's provided layers present but empty, so a pipeline degrades + gracefully on documents without content. An absent required layer is rejected + with an IllegalArgumentException naming the layer, because a + missing pipeline stage is an assembly error, not an empty document. + DocumentAnnotators.requireLayers performs exactly that rejection, + including the null check on the document, and is what the toolkit's own + adapters use. + +
    + +
    + Term vectors + + TermVectorAnnotator rolls the token layer up into a document-scoped + term vector layer for index consumers: one TermVector per distinct + term, carrying the term string, its occurrence count, and the occurrence spans in + original text coordinates. Term identity comes from the annotator's inputs: without + a normalizer the term is the token's covered text as-is; with a plain + CharSequenceNormalizer, the general path, each token's covered text is + normalized on its own to become the term; with an + OffsetAwareNormalizer the document text is normalized once and each + token's covered normalized text becomes the term. On every path tokens differing + only by a normalization fold group together while their spans keep pointing into + the original text. A token whose normalized form is empty is omitted from the + layer. The layer preserves first-occurrence order. + TermVectorPipelineTest#testTokenizerAndTermVectorPipeline asserts the + behavior shown here. + > vectors = document.get(TermVectorAnnotator.TERM_VECTORS); +// vectors.get(1).value() is ("dog", 2, [[4..7), [19..22)])]]> + + + + The normalized-term workflow passes any shipped + CharSequenceNormalizer to the plain-normalizer constructor: the + normalizer is applied to each token's covered text to produce the term, so the + case, NFC, and accent folds that buildAligned() rejects all work + here, and every occurrence span stays the token's own span in the original text. + TermVectorNormalizedExampleTest#testCaseFoldedTermsKeepOriginalSpans + asserts the behavior shown here. + > vectors = document.get(TermVectorAnnotator.TERM_VECTORS); +// vectors.get(0).value() is ("word", 3, [[0..4), [5..9), [10..14)])]]> + + + + Scoring-only consumers that never read offsets can skip storing them: + new TermVectorAnnotator(TermVectorAnnotator.Mode.SCORING_ONLY) emits + vectors that carry counts but no spans, and + TermVector.count("dog", 2) builds the matching expectation in tests. + +
    +     diff --git a/opennlp-docs/src/docbkx/embeddings.xml b/opennlp-docs/src/docbkx/embeddings.xml new file mode 100644 index 0000000000..74dc3f03f0 --- /dev/null +++ b/opennlp-docs/src/docbkx/embeddings.xml @@ -0,0 +1,261 @@ + + + + + + + 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 carries vocab.txt, + model.safetensors, config.json, and + tokenizer_config.json; a Model2Vec Unigram model carries + tokenizer.json, model.safetensors, and + config.json. Its JSON file contains the Unigram vocabulary, scores, + normalizer, and pre-tokenizer. 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 carries the model's own 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. + +
    +
    + +
    + 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/model-loading.xml b/opennlp-docs/src/docbkx/model-loading.xml index eb2356a80d..8a8d6cb394 100644 --- a/opennlp-docs/src/docbkx/model-loading.xml +++ b/opennlp-docs/src/docbkx/model-loading.xml @@ -164,4 +164,57 @@ for(ClassPathModelEntry entry : models) { - \ No newline at end of file + +
    + Installing third-party resources + + ResourceInstaller copies a supplied http, + https, or file URI into a target directory. Http and https + sources require a SHA-256 or SHA-512 checksum. File sources may omit it. + Gzip-compressed tar and zip archives are unpacked, plain gzip files are decompressed, + and other content is stored under its source name. OpenNLP *.bin model + files remain packed. + + + + + + Content is unpacked in a staging directory on the target filesystem. Archive paths + cannot escape that directory, duplicate file paths are rejected, and existing target + files are not replaced. A checksum, download, or unpacking error promotes no files. + Promotion also rejects symbolic links below the target. Concurrent changes to the + target directory are outside this guarantee. + + + Http and https requests use connection and read timeouts, follow up to five + redirects by default, and reject an https-to-http downgrade. Default limits are 1 GiB + downloaded, 4 GiB expanded, and 100000 archive entries; compressed content may + expand to at most 100 times its compressed size, with a 1 MiB floor for small + sources. These defaults can be set at + JVM startup with + opennlp.download.max.bytes, opennlp.install.max.total.bytes, + opennlp.install.max.entries, and + opennlp.install.max.expansion.ratio. Use + ResourceInstaller.Limits to set limits for one call. + + + The expansion ratio and the expansion limit are enforced independently and the + tighter one applies, so raising opennlp.install.max.total.bytes alone + does not admit an archive that exceeds the ratio. A resource that legitimately + compresses better than 100 to 1, such as a corpus of highly repetitive text, needs + opennlp.install.max.expansion.ratio raised as well. + + + + +
    + diff --git a/opennlp-docs/src/docbkx/normalizer.xml b/opennlp-docs/src/docbkx/normalizer.xml index d955e0366a..1844dd0f61 100644 --- a/opennlp-docs/src/docbkx/normalizer.xml +++ b/opennlp-docs/src/docbkx/normalizer.xml @@ -106,6 +106,14 @@ run containing a line break as a single newline, so paragraph structure survives; trims the edges. + + ParagraphPreservingWhitespaceCharSequenceNormalizer + Unwraps hard-wrapped prose (Gutenberg-style fixed-width lines): a + whitespace run with at most one line break collapses to a space, two or + more collapse to a single newline, and CRLF counts as one + break; trims the edges. Use before sentence detection on plain text, not + on raw markdown with lists or code blocks. + DashCharSequenceNormalizer Maps every Unicode dash to the ASCII hyphen-minus. The mathematical @@ -198,7 +206,7 @@ String hyphen = DashCharSequenceNormalizer.getInstance()
    Composing a pipeline - TextNormalizer is a fluent builder that composes the rungs, in the order + TextNormalizer is a fluent builder that composes the normalizers, in the order they are added, into a single CharSequenceNormalizer: @@ -259,7 +267,7 @@ Span hit = aligned.toOriginalSpan(5, 14); // "the-match" in the normalized tex CharSequenceNormalizer and adds normalizeAligned, so a caller tests for it with a plain instanceof, the same pattern the name finder uses for OffsetMappingNameFinder. Every per-code-point fold implements it: whitespace, the - line-break-preserving whitespace rung, dashes, invisible-control stripping, quotes, digits, + line-break-preserving whitespace normalizer, dashes, invisible-control stripping, quotes, digits, ellipsis, bullets, the German umlaut transliteration, Unicode full case folding (fullCaseFold(), whose expansions come from a bundled table with known lengths, so it reports its edits), and the emoji/emoticon folds (emojiToEmoticon() and @@ -419,6 +427,15 @@ dash.normalize("a—b"); // "a-b"]]> you squish horizontal whitespace while keeping line breaks. + + + collapseParagraphPreserving(text, lineBreaks, paragraphReplacement) + collapses runs but emits paragraphReplacement only when a run contains + two or more logical line breaks (a single break, including CRLF, becomes + the usual replacement). This is the hard-wrap unwrap used by + ParagraphPreservingWhitespaceCharSequenceNormalizer. + + So the replacement is your choice and the method picks the behavior. Folding tabs and @@ -434,7 +451,8 @@ ws.collapsePreserving(text, CodePointSet.of('\n'), '\n'); // squish spaces, ke When you need the normalized form together with a map back to the original, the normalizeAligned, collapseAligned, - collapsePreservingAligned, trimAligned, and + collapsePreservingAligned, collapseParagraphPreservingAligned, + trimAligned, and removeAllAligned variants return an AlignedText that carries an Alignment. The alignment maps spans between the two forms with toOriginalSpan and toNormalizedSpan, staying correct across diff --git a/opennlp-docs/src/docbkx/opennlp.xml b/opennlp-docs/src/docbkx/opennlp.xml index 36641c2c89..93a9045434 100644 --- a/opennlp-docs/src/docbkx/opennlp.xml +++ b/opennlp-docs/src/docbkx/opennlp.xml @@ -109,6 +109,7 @@ under the License. + @@ -120,6 +121,7 @@ under the License. + diff --git a/opennlp-docs/src/docbkx/stemmer.xml b/opennlp-docs/src/docbkx/stemmer.xml index 248b310c15..d2367684b5 100644 --- a/opennlp-docs/src/docbkx/stemmer.xml +++ b/opennlp-docs/src/docbkx/stemmer.xml @@ -69,4 +69,56 @@ new CachingStemmer(factory).stem("running"); // "run"]]> longer uses a sharing or caching stemmer.
    + +
    + Hunspell dictionaries + + opennlp.tools.stemmer.hunspell implements the documented + Hunspell dictionary format: a user-supplied + .aff affix file and its .dic word list. OpenNLP + bundles no dictionary data; dictionaries are downloaded separately, and + each states its own license. The dictionary is immutable and safe to share; + HunspellStemmerFactory creates a fresh stemmer per call. + HunspellManualExampleTest asserts the behavior shown here. + + + The stems above are those of the project-authored miniature dictionary the + test loads, which lists work with an agentive and a plural + suffix; the test asserts the same stem for worker. Which stem a + published dictionary yields for a given form is decided by that dictionary. + Acquisition helpers and the supported affix feature set live in + dev/README-hunspell-dictionaries.md. + An opt-in catalog download accepts an application-supplied + DictionaryCatalog, needs + -Dopennlp.download.remote=true, and fetches through the + digest-verified ResourceInstaller path. A file that already exists + in the target is not replaced. Remove old files before refreshing a dictionary. + Directives outside the supported affix-stemming subset are skipped, so + published dictionaries can still use their supported rules. Conversion, + suggestion, and advanced compound behavior from skipped directives is not + applied to the returned stems. + This includes ICONV, OCONV, + COMPLEXPREFIXES, COMPOUNDRULE, + IGNORE, and KEEPCASE. Results can differ from + Hunspell for words that need these rules. + FLAG and AF declarations apply to the complete + affix file, including rules listed before those declarations. Parsing + rejects malformed text in parsed rules, invalid counts, numeric flags + outside the range 1 through 65000, and COMPOUNDMIN values + that cannot be doubled safely. Comments and unused metadata may retain a + legacy encoding. Default and long flag modes preserve raw + one-byte flag values in UTF-8 files. Compound length and boundary checks + count Unicode code points. + A rule that strips a whole stem applies only when the affix file + declares FULLSTRIP, as in Hunspell itself. + Each affix or dictionary stream is rejected when it exceeds + HunspellDictionary.MAX_STREAM_BYTES (64 MiB). + +
    diff --git a/opennlp-docs/src/docbkx/tokenizer.xml b/opennlp-docs/src/docbkx/tokenizer.xml index 6733728d3a..e9a9370a2a 100644 --- a/opennlp-docs/src/docbkx/tokenizer.xml +++ b/opennlp-docs/src/docbkx/tokenizer.xml @@ -538,67 +538,82 @@ wordTokenizer.tokenize("The quick brown fox.", (start, end, type) -> { - -
    - Subword Tokenization - - Neural models usually operate on subword units drawn from a fixed vocabulary rather than - on words. The SubwordTokenizer interface in - opennlp.tools.tokenize covers this case: encode splits text into - pieces, and every returned SubwordPiece carries the piece string in the - model's normalized form, its vocabulary id, and the exact span of the original text it - came from. Because the segmentation runs after the model's own normalization, a piece is - generally not a substring of the input; the spans always refer to the caller's original - text, so annotations computed over the pieces can be mapped back without guesswork. The - encodeToIds and encodeToPieces methods return just the ids or - the piece strings when the spans are not needed. Thread safety is implementation specific; - both implementations described below are immutable and safe for concurrent use by - multiple threads. - -
    - SentencePiece +
    + Lattice tokenization for CJK - SentencePieceTokenizer in the opennlp-subword artifact - (package opennlp.subword.sentencepiece) runs a trained SentencePiece - .model file purely in Java, with no native library. The file itself - carries everything the class needs: the vocabulary with piece scores and types, the - segmentation algorithm (unigram language model or byte-pair encoding), and the text - normalizer the model was trained with. Loading is a one-time cost and the resulting - instance is immutable, so one tokenizer can be shared by any number of threads. + Languages written without spaces need a dictionary-backed segmenter. + LatticeTokenizer scores paths over a MeCab-format dictionary and + emits the cheapest segmentation with spans in original text coordinates. + UnigramSegmenter does the same from a plain frequency lexicon. + Install a dictionary archive with MecabDictionaryInstaller, load it + as a MecabDictionary, and tokenize. Remote archives require an + expected SHA-256 or SHA-512 digest; an application-supplied opt-in catalog + (installFromCatalog) also needs + -Dopennlp.download.remote=true. + matrix.def must list a cost for every declared cell; matrix + dimensions plus lexicon size are bounded by the shared + ResourceLimits.MAX_ENTRIES limit and the matrix cell count + by ResourceLimits.MAX_MATRIX_CELLS. + Fetching, verification, and unpacking use ResourceInstaller, + described in the model loading chapter. Default limits are a 1 GiB + download, 4 GiB unpacked, and 100000 archive entries. Set higher limits at + JVM startup via + opennlp.download.max.bytes, + opennlp.install.max.total.bytes, + opennlp.install.max.entries, and + opennlp.install.max.expansion.ratio for larger dictionaries such as + UniDic. Installed files are flattened to their base names. Existing target + files are not replaced and must be removed before a refresh. + Dictionary loading rejects empty lexicon surfaces, duplicate + matrix.def entries, duplicate char.def + categories, costs outside the signed 16-bit range, malformed text in the + selected encoding, and unk.def templates for undefined + categories. MeCab-quoted CSV fields are accepted. + A char.def may define up to 17 categories, and category + lengths range from 0 through 15, matching the MeCab character-property + compiler. + If a char.def mapping lists multiple categories, the first + provides unknown-word settings. Grouping continues while successive + assignments overlap. Unknown words use their category's unk.def + template first, with DEFAULT used when no category template exists. + LatticeUsageExampleTest asserts the install-load-tokenize and + lexicon flows shown here. - id " + piece.id() - + ", original text [" + piece.start() + ", " + piece.end() + ")"); -} +String text = "東京に行く"; +String[] tokens = tokenizer.tokenize(text); +Span[] spans = tokenizer.tokenizePos(text); +List morphemes = tokenizer.analyze(text); -int[] ids = tokenizer.encodeToIds("Ready for the embedding layer.");]]> +UnigramSegmenter segmenter = UnigramSegmenter.load(Path.of("chinese-words.txt")); +String[] words = segmenter.tokenize("我来到北京清华大学");]]> - SentencePieceUsageExampleTest asserts the load-and-encode workflow - shown here. The vocabulary can be inspected through - vocabularySize, idToPiece, pieceToId, and - score, and the algorithm method reports whether the - model uses the unigram or the byte-pair encoding algorithm. - - - The model's own normalizer is also exposed directly: - SentencePieceTokenizer implements OffsetAwareNormalizer, so - normalize applies the model's normalization rules to arbitrary text and - normalizeAligned additionally returns the character alignment described - in . This is useful when other processing must see - text exactly as the subword model does. - - - Output parity with the reference implementation is pinned by test fixtures whose - expected pieces, ids, and spans are produced by the reference - sentencepiece - package rather than by the Java code under test. The fixtures README at - opennlp-extensions/opennlp-subword/src/test/resources/opennlp/subword/sentencepiece/README.md - documents how the bundled test models were created and walks through regenerating - them and validating the implementation, including against real published models.
    +
    +
    + Subword Tokenization + + Neural models commonly use subword units from a fixed vocabulary. The + SubwordTokenizer interface encodes text as SubwordPiece values. + Each value contains the normalized piece string, vocabulary id, and UTF-16 span in the + original text. The span remains valid when model normalization changes the piece text. + encodeToIds and encodeToPieces return only the ids or strings. + An implementation can include model control pieces with empty spans, such as the + classification and separator pieces used by BERT. Thread safety is implementation specific. + + + The interface, SubwordPiece, and vocabulary-based + WordpieceEncoder are in opennlp-api. Implementations that read + trained model formats can implement the same contract. + + + A piece without source text has an empty span. When normalization reorders characters, the + range can include adjacent source characters. +
    WordPiece @@ -606,26 +621,28 @@ int[] ids = tokenizer.encodeToIds("Ready for the embedding layer.");]]> BERT tokenization pipeline: control character removal, whitespace normalization, CJK isolation, optional lower casing with accent stripping, punctuation isolation, and greedy longest-match wordpiece segmentation. It is constructed from the model's - vocabulary, conventionally the lines of a vocab.txt file, where a piece's - id is its line number. + vocabulary, conventionally the lines of a vocab.txt file, where the line + number is the id. vocabulary = Files.readAllLines(Path.of("vocab.txt")); WordpieceEncoder encoder = new WordpieceEncoder(vocabulary); // uncased model WordpieceEncoder cased = new WordpieceEncoder(vocabulary, false); // cased model -List pieces = encoder.encode("OpenNLP encodes text for BERT models."); -int[] ids = encoder.encodeToIds("OpenNLP encodes text for BERT models.");]]> +String text = "Alice was beginning to get very tired."; +List pieces = encoder.encode(text); +for (SubwordPiece piece : pieces) { + CharSequence source = text.subSequence(piece.start(), piece.end()); + System.out.println(piece.id() + "\t" + piece.piece() + "\t" + source); +} +int[] ids = encoder.encodeToIds(text);]]> - Every encoding is framed by the classification and separator pieces, which carry empty - spans at the text's boundaries, and words the vocabulary cannot cover become the - unknown piece. Models with other special tokens or with non-contiguous ids are + The classification and separator pieces appear at the beginning and end with empty + spans. Words the vocabulary cannot cover, and words exceeding 100 normalized Unicode + code points, become the unknown piece. A constructor parameter can set a different maximum. + Models with other special tokens or non-contiguous ids are supported through the constructors taking explicit special tokens or a piece-to-id - map. The older BertTokenizer is deprecated and scheduled for removal: - its tokenize method returns the same piece sequence that - encodeToPieces returns. Migrating means passing the vocabulary as a - List or a piece-to-id map instead of a Set; the output - is unchanged. + map.
    diff --git a/opennlp-extensions/opennlp-embeddings/README.md b/opennlp-extensions/opennlp-embeddings/README.md new file mode 100644 index 0000000000..3b06b15083 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/README.md @@ -0,0 +1,231 @@ + + +# 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. + +## 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..091f11cde6 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/TRAINING.md @@ -0,0 +1,81 @@ + + +# 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. + +## 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/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/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..2242e502e3 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/ModelFileNames.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package opennlp.embeddings; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +/** + * The file names of a static embedding model directory, shared by + * {@link StaticEmbeddingModel}'s loader and {@link ModelAssembler}. A WordPiece directory holds + * {@link #SAFETENSORS}, {@link #CONFIG}, {@link #VOCABULARY}, and {@link #TOKENIZER_CONFIG}; a + * Unigram directory holds {@link #SAFETENSORS}, {@link #CONFIG}, and + * {@link #TOKENIZER_JSON}. A separate-file SentencePiece directory additionally holds 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 tokenizer description whose Unigram {@code model.vocab} order names the matrix rows. */ + static final String TOKENIZER_JSON = "tokenizer.json"; + + /** The model configuration carrying the {@code normalize} pooling switch. */ + static final String CONFIG = "config.json"; + + /** The BERT-style vocabulary of a WordPiece model, one token per line in row order. */ + static final String VOCABULARY = "vocab.txt"; + + /** The tokenizer configuration carrying the WordPiece {@code do_lower_case} switch. */ + static final String TOKENIZER_CONFIG = "tokenizer_config.json"; + + /** The optional term rows of the matrix, one normalized term per line in row order. */ + static final String TERMS = "terms.txt"; + + /** The file names SentencePiece models ship their trained {@code .model} under, in try order. */ + static final List SENTENCEPIECE_MODELS = + List.of("sentencepiece.bpe.model", "spiece.model", "tokenizer.model"); + + /** The ONNX graph of a teacher, relative to the teacher directory's root. */ + static final String ONNX_MODEL = "onnx/model.onnx"; + + /** The external weights an ONNX export splits out of {@link #ONNX_MODEL}, if it splits them. */ + static final String ONNX_MODEL_DATA = "onnx/model.onnx_data"; + + /** Not instantiable. */ + private ModelFileNames() { + } + + /** + * {@return the first of the given file names that exists as a regular file in the directory, + * or {@code null} when none does} + * + * @param directory The directory to look in. + * @param names The file names to try, in order. + */ + static Path firstRegularFile(Path directory, List names) { + for (final String name : names) { + final Path file = directory.resolve(name); + if (Files.isRegularFile(file)) { + return file; + } + } + return null; + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/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/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..619a8056ca --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/StaticEmbeddingModel.java @@ -0,0 +1,1244 @@ +/* + * 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.

    + * + *

    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; + private 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 float[] embeddings; + 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; + // Per-row L2 norms and special-token mask, precomputed at load time for the neighbor scan. + private final double[] rowNorms; + 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(float[] embeddings, float[] weights, int dimension, + EmbeddingVocabulary vocabulary, SubwordTokenizer tokenizer, + IntPredicate skipPieceId, boolean normalize, double[] rowNorms, + boolean[] specialRows, TermTable terms) { + this.embeddings = embeddings; + this.weights = weights; + this.dimension = dimension; + this.vocabulary = vocabulary; + this.tokenizer = tokenizer; + this.skipPieceId = skipPieceId; + this.normalize = normalize; + this.rowNorms = rowNorms; + this.specialRows = specialRows; + this.terms = terms; + } + + /** + * 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.

    + * + * @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)) { + return loadSentencePiece(sentencePieceModelFile, tokenizerJsonFile, + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)), + termLines, termsFile.toString()); + } + if (Files.isRegularFile(tokenizerJsonFile)) { + return loadModel2VecUnigram(tokenizerJsonFile, + requiredFile(modelDirectory, ModelFileNames.SAFETENSORS), + requiredNormalize(requiredFile(modelDirectory, ModelFileNames.CONFIG)), + termLines, termsFile.toString()); + } + throw new InvalidFormatException("Model directory " + modelDirectory + " has neither a " + + ModelFileNames.VOCABULARY + " (WordPiece layout) nor a " + + ModelFileNames.TOKENIZER_JSON + " (Unigram layout)"); + } + + /** Loads a self-contained Model2Vec Unigram directory. */ + private static StaticEmbeddingModel loadModel2VecUnigram( + Path tokenizerJsonFile, Path safetensorsFile, Normalization normalization, + List termLines, String termsSourceName) throws IOException { + final EmbeddingVocabulary vocabulary = + EmbeddingVocabulary.fromTokenizerJson(tokenizerJsonFile); + final TermTable terms = TermTable.of(termLines, vocabulary.size(), termsSourceName); + final Model2VecUnigramTokenizer tokenizer; + try { + tokenizer = Model2VecUnigramTokenizer.load(tokenizerJsonFile); + } catch (InvalidFormatException e) { + throw new InvalidFormatException("Unigram model needs either a self-contained " + + "tokenizer.json or a trained SentencePiece .model file: " + e.getMessage(), e); + } + requireVocabularyCoverage(tokenizer, vocabulary, tokenizerJsonFile); + final Matrix matrix = readMatrix(vocabulary, terms.size(), safetensorsFile, + tokenizerJsonFile.toString()); + final IntPredicate skipPieceId = + id -> tokenizer.isUnknown(id) || tokenizer.isControl(id); + return new StaticEmbeddingModel(matrix.embeddings(), matrix.weights(), matrix.dimension(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + rowNorms(matrix.embeddings(), matrix.dimension(), vocabulary.size() + terms.size()), + specialRows(vocabulary, SENTENCEPIECE_SPECIAL_TOKENS, + vocabulary.size() + terms.size(), 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 safetensorsFile = requiredFile(modelDirectory, ModelFileNames.SAFETENSORS); + 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"); + } + return loadWordpiece(vocabularyFile, safetensorsFile, + lowerCase ? Casing.UNCASED : Casing.CASED, normalization, termLines, termsSourceName); + } + + /** + * 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 int unknownId = vocabulary.id(WordpieceTokenizer.BERT_UNK_TOKEN); + if (unknownId < 0) { + throw new InvalidFormatException("Vocabulary " + vocabularyFile + " 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(matrix.embeddings(), matrix.weights(), matrix.dimension(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + rowNorms(matrix.embeddings(), matrix.dimension(), vocabulary.size() + terms.size()), + specialRows(vocabulary, WORDPIECE_SPECIAL_TOKENS, vocabulary.size() + terms.size()), + terms); + } + + /** + * Builds the WordPiece encoder, mapping {@code [CLS]} and {@code [SEP]} onto the unknown row + * when the distilled vocabulary dropped them. A static embedding table mean-pools its content + * pieces and never pools the delimiters, so distillers routinely remove + * {@code [CLS]}/{@code [SEP]} from the table; the encoder still wraps every encoding in them + * and needs an id for each, and pooling skips them regardless of their ids, so pointing the + * absent delimiter tokens at the unknown row makes the model loadable without changing which + * pieces are pooled. + * + * @param vocabulary The matrix row vocabulary; must contain the unknown token. + * @param lowerCase Whether the tokenizer lower-cases and strips accents. + * @param unknownId The unknown token's row, reused as the id of {@code [CLS]} or + * {@code [SEP]} when that token is absent. + * @return The encoder. + */ + private static WordpieceEncoder wordpieceEncoder(EmbeddingVocabulary vocabulary, + boolean lowerCase, int unknownId) { + if (vocabulary.id(WordpieceTokenizer.BERT_CLS_TOKEN) >= 0 + && vocabulary.id(WordpieceTokenizer.BERT_SEP_TOKEN) >= 0) { + return new WordpieceEncoder(vocabulary.orderedTokens(), lowerCase); + } + final List tokens = vocabulary.orderedTokens(); + final Map ids = new HashMap<>(tokens.size() * 2); + for (int id = 0; id < tokens.size(); id++) { + ids.put(tokens.get(id), id); + } + ids.putIfAbsent(WordpieceTokenizer.BERT_CLS_TOKEN, unknownId); + ids.putIfAbsent(WordpieceTokenizer.BERT_SEP_TOKEN, unknownId); + return new WordpieceEncoder(ids, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN, + WordpieceTokenizer.BERT_SEP_TOKEN, WordpieceTokenizer.BERT_UNK_TOKEN); + } + + /** + * Loads a SentencePiece static embedding model from a trained SentencePiece {@code .model} + * file, the Unigram {@code tokenizer.json} naming the matrix rows, and a safetensors weight + * file. No model is bundled with this module; the caller supplies the files. + * + *

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

    + * + * @param sentencePieceModelFile The trained SentencePiece {@code .model} file. Must not be + * {@code null} and must exist. + * @param tokenizerJsonFile The Unigram {@code tokenizer.json} file; its + * {@code model.vocab} list order is the matrix row order, with + * {@code added_tokens} overlaid. Must not be {@code null} and + * must exist. + * @param safetensorsFile The {@code model.safetensors} file. Must not be {@code null} + * and must exist, and must contain exactly one 2-D float tensor + * (the embedding matrix) whose row count matches the vocabulary + * size. An optional 1-D 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 IntPredicate skipPieceId = + id -> tokenizer.isUnknown(id) || tokenizer.isControl(id); + return new StaticEmbeddingModel(matrix.embeddings(), matrix.weights(), matrix.dimension(), + vocabulary, tokenizer, skipPieceId, normalization == Normalization.L2, + rowNorms(matrix.embeddings(), matrix.dimension(), vocabulary.size() + terms.size()), + specialRows(vocabulary, SENTENCEPIECE_SPECIAL_TOKENS, + vocabulary.size() + terms.size(), 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 L2 norm of every matrix row, precomputed for the neighbor scan} + * + * @param embeddings The flat row-major matrix. + * @param dimension The row width. + * @param rowCount The number of rows. + */ + private static double[] rowNorms(float[] embeddings, int dimension, int rowCount) { + final double[] 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 = embeddings[base + d]; + sumOfSquares += (double) value * value; + } + rowNorms[row] = Math.sqrt(sumOfSquares); + } + return rowNorms; + } + + /** + * {@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"); + } + final double[] sum = new double[dimension]; + // IntConsumer needs a mutable counter for the pooled rows. + final int[] pooled = new int[1]; + forEachPooledRow(text, row -> { + final int base = row * dimension; + if (weights == null) { + for (int d = 0; d < dimension; d++) { + sum[d] += embeddings[base + d]; + } + } else { + final float weight = weights[row]; + for (int d = 0; d < dimension; d++) { + sum[d] += (double) embeddings[base + d] * weight; + } + } + pooled[0]++; + }); + final int denominator = Math.max(pooled[0], 1); + final float[] pooledVector = new float[dimension]; + for (int d = 0; d < dimension; d++) { + sum[d] /= denominator; + } + if (normalize) { + double sumOfSquares = 0; + for (final double value : sum) { + sumOfSquares += value * value; + } + final double norm = Math.max(Math.sqrt(sumOfSquares), NORMALIZE_EPSILON); + for (int d = 0; d < dimension; d++) { + pooledVector[d] = (float) (sum[d] / norm); + } + } else { + for (int d = 0; d < dimension; d++) { + pooledVector[d] = finiteFloat(sum[d]); + } + } + return pooledVector; + } + + /** Converts a finite double to the nearest finite float 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; + } + + /** + * 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(); + } + final int rowCount = rowNorms.length; + // The result cannot contain more entries than the matrix has rows. + 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 = rowNorms[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; + } + final int base = row * dimension; + // Keep four independent sums in the dot product. + double dot0 = 0; + double dot1 = 0; + double dot2 = 0; + double dot3 = 0; + int d = 0; + for (final int limit = dimension - 3; d < limit; d += 4) { + dot0 += query[d] * embeddings[base + d]; + dot1 += query[d + 1] * embeddings[base + d + 1]; + dot2 += query[d + 2] * embeddings[base + d + 2]; + dot3 += query[d + 3] * embeddings[base + d + 3]; + } + double dot = dot0 + dot1 + dot2 + dot3; + for (; d < dimension; d++) { + dot += query[d] * embeddings[base + d]; + } + best.offer(row, boundedCosine(dot / (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..50d0c42189 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/main/java/opennlp/embeddings/cmdline/CLI.java @@ -0,0 +1,140 @@ +/* + * 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()); + + 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/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/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/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/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/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..dad82c8b85 --- /dev/null +++ b/opennlp-extensions/opennlp-embeddings/src/test/java/opennlp/embeddings/cmdline/CLITest.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package opennlp.embeddings.cmdline; + +import java.util.Set; +import java.util.stream.Stream; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import opennlp.tools.cmdline.BasicCmdLineTool; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The command names the dispatcher offers and the help every tool produces. The names are the + * module's public surface (TRAINING.md and the manual quote them), so a class rename that changes + * a command has to fail here rather than in a user's shell. + */ +class CLITest { + + /** {@return the tools the dispatcher registers, as parameterized-test arguments} */ + static Stream tools() { + return Stream.of(new AssembleModelTool(), new DistillModelTool()); + } + + @Test + void testOffersExactlyTheDistillationCommands() { + assertEquals(Set.of("AssembleModel", "DistillModel"), CLI.getToolNames()); + } + + @Test + void testTheToolNamesCannotBeModifiedByACaller() { + final Set names = CLI.getToolNames(); + + assertThrows(UnsupportedOperationException.class, () -> names.add("Other")); + } + + @ParameterizedTest + @MethodSource("tools") + void testEveryRegisteredToolDescribesItself(BasicCmdLineTool tool) { + assertTrue(CLI.getToolNames().contains(tool.getName()), + tool.getName() + " must be registered with the dispatcher"); + assertFalse(tool.getShortDescription().isBlank(), + tool.getName() + " must have a short description for the usage listing"); + assertTrue(tool.getHelp().contains(tool.getName()), tool.getHelp()); + } + + @Test + void testDistillHelpNamesEveryParameter() { + final String help = new DistillModelTool().getHelp(); + + assertTrue(help.contains("-teacher hf-id-or-path"), help); + assertTrue(help.contains("-out dir"), help); + // The optional parameters are bracketed, so a user can see they may be omitted. + assertTrue(help.contains("[-pcaDims "), help); + assertTrue(help.contains("[-terms "), help); + } + + @Test + void testAssembleHelpNamesItsParameter() { + final String help = new AssembleModelTool().getHelp(); + + assertTrue(help.contains("-modelDir dir"), help); + } +} diff --git a/opennlp-extensions/opennlp-embeddings/src/test/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-morfologik/src/main/java/opennlp/morfologik/lemmatizer/MorfologikLemmatizer.java b/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/lemmatizer/MorfologikLemmatizer.java index 42d4ecb4eb..6665a01a99 100644 --- a/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/lemmatizer/MorfologikLemmatizer.java +++ b/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/lemmatizer/MorfologikLemmatizer.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Objects; import java.util.Set; @@ -63,7 +64,7 @@ public MorfologikLemmatizer(Dictionary dictionary) { } private List lemmatize(String word, String postag) { - List dictMap = new DictionaryLookup(dictionary).lookup(word.toLowerCase()); + List dictMap = new DictionaryLookup(dictionary).lookup(word.toLowerCase(Locale.ROOT)); Set lemmas = new HashSet<>(); for (WordData wordData : dictMap) { if (Objects.equals(postag, asString(wordData.getTag()))) { diff --git a/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/tagdict/MorfologikTagDictionary.java b/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/tagdict/MorfologikTagDictionary.java index ac7d42b22a..6a63a7b007 100644 --- a/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/tagdict/MorfologikTagDictionary.java +++ b/opennlp-extensions/opennlp-morfologik/src/main/java/opennlp/morfologik/tagdict/MorfologikTagDictionary.java @@ -19,6 +19,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Locale; import morfologik.stemming.Dictionary; import morfologik.stemming.DictionaryLookup; @@ -64,7 +65,7 @@ public MorfologikTagDictionary(Dictionary dict, boolean caseSensitive) @Override public String[] getTags(String word) { if (!isCaseSensitive) { - word = word.toLowerCase(); + word = word.toLowerCase(Locale.ROOT); } List data = dictLookup.lookup(word); diff --git a/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/AbstractMorfologikTest.java b/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/AbstractMorfologikTest.java index ce3fe06316..7ebeb30027 100644 --- a/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/AbstractMorfologikTest.java +++ b/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/AbstractMorfologikTest.java @@ -31,14 +31,26 @@ public abstract class AbstractMorfologikTest { protected static Path createMorfologikDictionary() throws Exception { + return createMorfologikDictionary("dictionaryWithLemma"); + } + + /** + * Builds a Morfologik FSA dictionary from the {@code .txt} and {@code .info} test + * resources sharing the given base name. + * + * @param resourceBaseName The base name of the tab separated dictionary resource pair. + * + * @return The {@link Path} of the built FSA dictionary. + */ + protected static Path createMorfologikDictionary(String resourceBaseName) throws Exception { Path tabFilePath = File.createTempFile(AbstractMorfologikTest.class.getName(), ".txt").toPath(); tabFilePath.toFile().deleteOnExit(); Path infoFilePath = DictionaryMetadata.getExpectedMetadataLocation(tabFilePath); infoFilePath.toFile().deleteOnExit(); - Files.copy(getResourceStream("/dictionaryWithLemma.txt"), tabFilePath, + Files.copy(getResourceStream("/" + resourceBaseName + ".txt"), tabFilePath, StandardCopyOption.REPLACE_EXISTING); - Files.copy(getResourceStream("/dictionaryWithLemma.info"), infoFilePath, + Files.copy(getResourceStream("/" + resourceBaseName + ".info"), infoFilePath, StandardCopyOption.REPLACE_EXISTING); MorfologikDictionaryBuilder builder = new MorfologikDictionaryBuilder(); diff --git a/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/MorfologikLocaleTest.java b/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/MorfologikLocaleTest.java new file mode 100644 index 0000000000..a5ba15c632 --- /dev/null +++ b/opennlp-extensions/opennlp-morfologik/src/test/java/opennlp/morfologik/MorfologikLocaleTest.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package opennlp.morfologik; + +import java.nio.file.Path; +import java.util.Locale; + +import morfologik.stemming.Dictionary; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import opennlp.morfologik.lemmatizer.MorfologikLemmatizer; +import opennlp.morfologik.tagdict.MorfologikTagDictionary; + +/** + * Tests that the case insensitive lookups into a prebuilt Morfologik FSA dictionary do not + * depend on the JVM's default {@link Locale}. + *

    + * The FSA is built once and shipped; folding the query with the default locale makes the + * same dictionary answer differently on different JVMs. + */ +public class MorfologikLocaleTest extends AbstractMorfologikTest { + + /** + * Turkish folds {@code 'I'} to the dotless {@code 'ı'} (U+0131) instead of {@code 'i'}. + */ + private static final Locale TURKISH = Locale.of("tr", "TR"); + + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + + @Test + public void testLemmatizeIsIndependentOfDefaultLocale() throws Exception { + final MorfologikLemmatizer lemmatizer = new MorfologikLemmatizer(createLocaleDictionary()); + + Locale.setDefault(TURKISH); + + final String[] lemmas = lemmatizer.lemmatize( + new String[] {"Illinois", "INDICES"}, new String[] {"PROP", "NOUN"}); + + Assertions.assertArrayEquals(new String[] {"Illinois", "index"}, lemmas); + } + + @Test + public void testGetTagsIsIndependentOfDefaultLocale() throws Exception { + final MorfologikTagDictionary tagDictionary = + new MorfologikTagDictionary(Dictionary.read(createLocaleDictionary()), false); + + Locale.setDefault(TURKISH); + + Assertions.assertArrayEquals(new String[] {"PROP"}, tagDictionary.getTags("Illinois")); + Assertions.assertArrayEquals(new String[] {"NOUN"}, tagDictionary.getTags("INDICES")); + } + + private static Path createLocaleDictionary() throws Exception { + final Path output = createMorfologikDictionary("dictionaryLocaleSafety"); + output.toFile().deleteOnExit(); + return output; + } +} diff --git a/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.info b/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.info new file mode 100644 index 0000000000..02f8003500 --- /dev/null +++ b/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.info @@ -0,0 +1,15 @@ +# +# REQUIRED PROPERTIES +# + +# Column (lemma, inflected, tag) separator. This must be a single byte in the target encoding. +fsa.dict.separator=, + +# The charset in which the input is encoded. UTF-8 is strongly recommended. +fsa.dict.encoding=UTF-8 + +# The type of lemma-inflected form encoding compression that precedes automaton +# construction. Allowed values: [suffix, infix, prefix, none]. +# Details are in Daciuk's paper and in the code. +# Leave at 'prefix' if not sure. +fsa.dict.encoder=prefix diff --git a/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.txt b/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.txt new file mode 100644 index 0000000000..181ff2acc0 --- /dev/null +++ b/opennlp-extensions/opennlp-morfologik/src/test/resources/dictionaryLocaleSafety.txt @@ -0,0 +1,2 @@ +Illinois,illinois,PROP +index,indices,NOUN 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 index 1c813d5314..78a434ebd8 100644 --- 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 @@ -26,7 +26,7 @@ /** * Byte-pair-encoding segmentation: the normalized text starts as single characters (or - * user-defined symbols, which are frozen whole) and adjacent pairs merge greedily, highest piece + * 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 @@ -34,11 +34,11 @@ */ final class BpeEncoder implements Serializable { - private static final long serialVersionUID = -57799941356582785L; + private static final long serialVersionUID = 112585252536688886L; private static final int MAX_RESEGMENT_DEPTH = 100; - private final Map pieces; + private final HashMap pieces; private final float[] scores; private final boolean[] unused; private final boolean[] reserved; @@ -57,7 +57,7 @@ final class BpeEncoder implements Serializable { * @param userDefinedMatcher Longest-match trie over user-defined symbols, or null when the * model defines none. */ - BpeEncoder(Map pieces, float[] scores, boolean[] unused, boolean[] reserved, + BpeEncoder(HashMap pieces, float[] scores, boolean[] unused, boolean[] reserved, int unkId, PieceTrie userDefinedMatcher) { this.pieces = pieces; this.scores = scores; @@ -91,36 +91,36 @@ List encode(byte[] normalized, int size) { } // The symbol list as index-linked ranges of the normalized bytes; merged-away symbols - // become empty ranges. Freeze flags travel as 0/1 bytes parallel to the ranges. + // 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 freezeB = new ByteBuilder(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 frozen = matched > 0; - final int length = frozen ? matched + 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); - freezeB.append(frozen ? (byte) 1 : (byte) 0); + protectedB.append(protectedSymbol ? (byte) 1 : (byte) 0); position += length; } - final int symbolCount = freezeB.length(); + final int symbolCount = protectedB.length(); final int[] from = fromB.toArray(); final int[] to = toB.toArray(); - final byte[] frozenFlags = freezeB.array(); + final byte[] protectedFlags = protectedB.array(); final int[] prev = new int[symbolCount]; final int[] next = new int[symbolCount]; - final boolean[] freeze = new boolean[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; - freeze[i] = frozenFlags[i] != 0; + protectedSymbols[i] = protectedFlags[i] != 0; } // Higher score first; equal scores break towards the leftmost pair. @@ -132,7 +132,7 @@ List encode(byte[] normalized, int size) { final Map revMerge = new HashMap<>(); for (int left = 0; left + 1 < symbolCount; left++) { - maybeAddPair(normalized, from, to, freeze, left, left + 1, agenda, revMerge); + maybeAddPair(normalized, from, to, protectedSymbols, left, left + 1, agenda, revMerge); } while (!agenda.isEmpty()) { @@ -152,8 +152,10 @@ List encode(byte[] normalized, int size) { } from[top.right()] = to[top.right()]; - maybeAddPair(normalized, from, to, freeze, prev[top.left()], top.left(), agenda, revMerge); - maybeAddPair(normalized, from, to, freeze, top.left(), next[top.left()], agenda, revMerge); + 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); @@ -174,16 +176,16 @@ List encode(byte[] normalized, int size) { * @param normalized The buffer holding the normalized UTF-8 bytes. * @param from Per symbol, the inclusive start offset in {@code normalized}. * @param to Per symbol, the exclusive end offset in {@code normalized}. - * @param freeze Per symbol, whether it is a user-defined symbol excluded from merging. + * @param 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[] freeze, + private void maybeAddPair(byte[] normalized, int[] from, int[] to, boolean[] protectedSymbols, int left, int right, PriorityQueue agenda, Map revMerge) { - if (left == -1 || right == -1 || freeze[left] || freeze[right]) { + if (left == -1 || right == -1 || protectedSymbols[left] || protectedSymbols[right]) { return; } final String piece = 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 index 5bed8a6852..4cbd594482 100644 --- 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 @@ -25,7 +25,7 @@ *

    Each unit is one little-endian 32-bit word encoding a label, an offset to the unit's * children, and a leaf flag; traversal XORs the offset with the next key byte. Only the longest * prefix match is needed here, so this walks the byte key once and remembers the last accepting - * state. Out-of-range unit references, which a well-formed trie never produces, fail loudly + * 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 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 index e9f4cb0bbd..64b54aca1e 100644 --- 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 @@ -16,6 +16,9 @@ */ 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; @@ -29,7 +32,7 @@ * ({@code sentencepiece_model.proto}, Apache License 2.0). This reader walks the tag stream * directly and keeps only the fields inference needs: the pieces with scores and types, the * normalizer spec, the trainer-spec fields that change runtime behavior, and the embedded - * self-test samples. Unknown fields are skipped, and malformed input fails loudly.

    + * self-test samples. Unknown fields are skipped, and malformed input is rejected.

    * * @see @@ -98,13 +101,16 @@ static RawModel read(byte[] data) throws InvalidFormatException { final ModelProtoReader reader = new ModelProtoReader(data); final RawModel model = new RawModel(); while (reader.pos < data.length) { - final long tag = reader.varint(); + final long tag = reader.readTag(data.length); switch (fieldOf(tag)) { - case FIELD_MODEL_PIECES -> reader.piece(model, reader.lenPayload(tag)); - case FIELD_MODEL_TRAINER_SPEC -> reader.trainerSpec(model, reader.lenPayload(tag)); - case FIELD_MODEL_NORMALIZER_SPEC -> reader.normalizerSpec(model, reader.lenPayload(tag)); - case FIELD_MODEL_SELF_TEST_DATA -> reader.selfTestData(model, reader.lenPayload(tag)); - default -> reader.skip(tag); + 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()) { @@ -146,12 +152,12 @@ private void piece(RawModel model, int end) throws InvalidFormatException { float score = 0; int type = RawModel.TYPE_NORMAL; while (pos < end) { - final long tag = varint(); + final long tag = readTag(end); switch (fieldOf(tag)) { - case FIELD_PIECE_PIECE -> piece = utf8(lenPayload(tag)); - case FIELD_PIECE_SCORE -> score = fixed32Float(tag); - case FIELD_PIECE_TYPE -> type = (int) varintOf(tag); - default -> skip(tag); + 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()) { @@ -175,13 +181,13 @@ private void piece(RawModel model, int end) throws InvalidFormatException { */ private void trainerSpec(RawModel model, int end) throws InvalidFormatException { while (pos < end) { - final long tag = varint(); + final long tag = readTag(end); switch (fieldOf(tag)) { - case FIELD_TRAINER_MODEL_TYPE -> model.modelType = (int) varintOf(tag); + case FIELD_TRAINER_MODEL_TYPE -> model.modelType = (int) varintOf(tag, end); case FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX -> - model.treatWhitespaceAsSuffix = varintOf(tag) != 0; - case FIELD_TRAINER_BYTE_FALLBACK -> model.byteFallback = varintOf(tag) != 0; - default -> skip(tag); + model.treatWhitespaceAsSuffix = varintOf(tag, end) != 0; + case FIELD_TRAINER_BYTE_FALLBACK -> model.byteFallback = varintOf(tag, end) != 0; + default -> skip(tag, end); } } } @@ -196,15 +202,17 @@ private void trainerSpec(RawModel model, int end) throws InvalidFormatException */ private void normalizerSpec(RawModel model, int end) throws InvalidFormatException { while (pos < end) { - final long tag = varint(); + final long tag = readTag(end); switch (fieldOf(tag)) { case FIELD_NORMALIZER_PRECOMPILED_CHARSMAP -> - model.precompiledCharsMap = bytes(lenPayload(tag)); - case FIELD_NORMALIZER_ADD_DUMMY_PREFIX -> model.addDummyPrefix = varintOf(tag) != 0; + 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) != 0; - case FIELD_NORMALIZER_ESCAPE_WHITESPACES -> model.escapeWhitespaces = varintOf(tag) != 0; - default -> skip(tag); + model.removeExtraWhitespaces = varintOf(tag, end) != 0; + case FIELD_NORMALIZER_ESCAPE_WHITESPACES -> + model.escapeWhitespaces = varintOf(tag, end) != 0; + default -> skip(tag, end); } } } @@ -219,17 +227,17 @@ private void normalizerSpec(RawModel model, int end) throws InvalidFormatExcepti */ private void selfTestData(RawModel model, int end) throws InvalidFormatException { while (pos < end) { - final long tag = varint(); + final long tag = readTag(end); if (fieldOf(tag) == FIELD_SELF_TEST_SAMPLES) { - final int sampleEnd = lenPayload(tag); + final int sampleEnd = lenPayload(tag, end); String input = null; String expected = null; while (pos < sampleEnd) { - final long sampleTag = varint(); + final long sampleTag = readTag(sampleEnd); switch (fieldOf(sampleTag)) { - case FIELD_SAMPLE_INPUT -> input = utf8(lenPayload(sampleTag)); - case FIELD_SAMPLE_EXPECTED -> expected = utf8(lenPayload(sampleTag)); - default -> skip(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) { @@ -237,7 +245,7 @@ private void selfTestData(RawModel model, int end) throws InvalidFormatException model.selfTestExpected.add(expected); } } else { - skip(tag); + skip(tag, end); } } } @@ -246,18 +254,19 @@ private void selfTestData(RawModel model, int end) throws InvalidFormatException * 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 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) throws InvalidFormatException { + 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(); - if (length < 0 || pos + length > data.length) { - throw malformed("length " + length + " exceeds the remaining input"); + final long length = varint(limit); + if (length < 0 || length > limit - pos) { + throw malformed("length " + length + " crosses its message boundary"); } return pos + (int) length; } @@ -265,29 +274,31 @@ private int lenPayload(long tag) throws InvalidFormatException { /** * Reads the varint value of a field after checking its wire type. * - * @param tag The field tag, whose wire type must be varint. + * @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) throws InvalidFormatException { + 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(); + 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 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) throws InvalidFormatException { + 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 (pos + 4 > data.length) { + if (limit - pos < 4) { throw malformed("truncated 32-bit value"); } final int bits = (data[pos] & 0xFF) | (data[pos + 1] & 0xFF) << 8 @@ -301,11 +312,20 @@ private float fixed32Float(long tag) throws InvalidFormatException { * * @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) { - final String s = new String(data, pos, end - pos, StandardCharsets.UTF_8); - pos = end; - return s; + 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"); + } } /** @@ -322,19 +342,39 @@ private byte[] bytes(int end) { } /** - * Reads a base-128 varint from the current position, advancing past it. + * 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() throws InvalidFormatException { + private long varint(int limit) throws InvalidFormatException { long value = 0; for (int shift = 0; shift < 64; shift += 7) { - if (pos >= data.length) { + 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; @@ -346,16 +386,17 @@ private long varint() throws InvalidFormatException { /** * Skips the value of an unrecognized field according to its wire type. * - * @param tag The field tag. + * @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) throws InvalidFormatException { + private void skip(long tag, int limit) throws InvalidFormatException { switch (wireTypeOf(tag)) { - case WIRE_VARINT -> varint(); - case WIRE_FIXED64 -> advance(8); - case WIRE_LEN -> pos = lenPayload(tag); - case WIRE_FIXED32 -> advance(4); + 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)); } } @@ -364,10 +405,11 @@ private void skip(long tag) throws InvalidFormatException { * 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) throws InvalidFormatException { - if (pos + count > data.length) { + private void advance(int count, int limit) throws InvalidFormatException { + if (limit - pos < count) { throw malformed("truncated field"); } pos += count; 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 index 42477251ab..5b58d225b0 100644 --- 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 @@ -103,8 +103,7 @@ static PieceTrie build(byte[][] pieces, int[] ids) { } Arrays.sort(order, Comparator.comparing(i -> pieces[i], Arrays::compareUnsigned)); - // First pass counts nodes and edges, second pass fills the packed arrays; both walk the - // sorted keys with the same recursion, so the shapes agree by construction. + // 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(); 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 index 74caac0683..cf2df762db 100644 --- 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 @@ -18,6 +18,7 @@ import java.io.IOException; import java.io.InputStream; +import java.io.InvalidClassException; import java.io.ObjectInputFilter; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; @@ -53,11 +54,9 @@ * *

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

    * - *

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

    + *

    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), @@ -68,7 +67,7 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwareNormalizer { // Serializable through the OffsetAwareNormalizer contract. - private static final long serialVersionUID = -4472058014098085134L; + private static final long serialVersionUID = 7751888381608757475L; /** The segmentation algorithm a model was trained with. */ public enum Algorithm { @@ -94,14 +93,14 @@ public enum Algorithm { private final int[] types; private final int unkId; private final boolean byteFallback; - private final Map mainPieces; - private final Map reservedPieces; + 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 List selfTestInputs; - private final List selfTestExpected; + private final String[] selfTestInputs; + private final String[] selfTestExpected; /** * Validates a parsed model and derives the runtime structures: the piece maps, the byte-piece @@ -229,8 +228,8 @@ private SentencePieceTokenizer(ModelProtoReader.RawModel model) throws InvalidFo userDefinedMatcher); } - selfTestInputs = List.copyOf(model.selfTestInputs); - selfTestExpected = List.copyOf(model.selfTestExpected); + selfTestInputs = model.selfTestInputs.toArray(String[]::new); + selfTestExpected = model.selfTestExpected.toArray(String[]::new); } /** @@ -310,8 +309,8 @@ public void serialize(OutputStream out) throws IOException { * {@link java.io.InvalidClassException} before {@link ObjectInputStream#readObject()} * returns.

    * - *

    Callers should still treat this method as defense-in-depth: only invoke it on streams from - * trusted sources. If the default limits are too tight for an unusually large model, use + *

    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.

    @@ -357,7 +356,13 @@ public static SentencePieceTokenizer deserialize(InputStream in, Deserialization } try (ObjectInputStream ois = new ObjectInputStream(in)) { ois.setObjectInputFilter(buildFilter(limits)); - return (SentencePieceTokenizer) ois.readObject(); + 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; } } @@ -365,10 +370,8 @@ public static SentencePieceTokenizer deserialize(InputStream in, Deserialization * Resource limits applied by the {@link ObjectInputFilter} used by * {@link SentencePieceTokenizer#deserialize(InputStream, DeserializationLimits)}. * - *

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

    + *

    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. @@ -429,14 +432,7 @@ public record DeserializationLimits(long maxDepth, long maxRefs, long maxArrayLe "java.util.HashMap", // HashMap.readObject() requests permission to allocate a Map.Entry[] before reading // entries; the array type itself never appears as a value in the stream. - "java.util.Map$Entry", - // The unmodifiable lists created by List.copyOf serialize through the CollSer proxy, - // which requests an Object[] allocation for the elements, and the filter is also invoked - // for the concrete list class the proxy resolves to. - "java.util.CollSer", - "java.util.ImmutableCollections$List12", - "java.util.ImmutableCollections$ListN", - "java.lang.Object" + "java.util.Map$Entry" ); /** @@ -732,12 +728,12 @@ private void checkId(int id) { /** {@return the embedded self-test input samples} */ List selfTestInputs() { - return selfTestInputs; + return List.of(selfTestInputs); } /** {@return the embedded self-test expected segmentations} */ List selfTestExpected() { - return selfTestExpected; + return List.of(selfTestExpected); } // The prefix of a byte-fallback piece string; a full piece has the form "<0xAB>". @@ -763,8 +759,22 @@ private static int parseBytePiece(String piece) { if (piece.length() != 6 || !piece.startsWith(BYTE_PIECE_PREFIX) || piece.charAt(5) != '>') { return -1; } - final int high = Character.digit(piece.charAt(3), 16); - final int low = Character.digit(piece.charAt(4), 16); + 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/test/java/opennlp/subword/sentencepiece/PieceTrieTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/PieceTrieTest.java index adaf2a292e..189f3fb25f 100644 --- 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 @@ -110,7 +110,7 @@ void testWideRootDispatchesAllByteValues() { } @Test - void testDuplicatePiecesFailLoudly() { + 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 index 7ac28e67ba..16c40e9e2d 100644 --- 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 @@ -137,7 +137,7 @@ private static String cp(int codePoint) { } @Test - void testNullInputsFailLoudly() { + 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/SentencePieceModelValidationTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java index 5f340eda3a..f2c859eb7e 100644 --- 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 @@ -40,14 +40,13 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -/** - * Fail-loud behavior on malformed models, plus the concurrency guarantee: one loaded tokenizer - * must produce identical results from many threads. - */ +/** Validates malformed-model rejection and concurrent tokenizer use. */ class SentencePieceModelValidationTest { + private static final String HEX_DIGITS = "0123456789ABCDEF"; + @Test - void testNullAndEmptyInputFailLoudly() { + void testRejectsNullAndEmptyInput() { assertThrows(IllegalArgumentException.class, () -> SentencePieceTokenizer.load((Path) null)); assertThrows(IllegalArgumentException.class, @@ -57,14 +56,14 @@ void testNullAndEmptyInputFailLoudly() { } @Test - void testGarbageBytesFailLoudly() { + 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 testTruncatedModelFailsLoudly() throws IOException { + void testRejectsTruncatedModel() throws IOException { final byte[] whole = readModel(); final byte[] truncated = Arrays.copyOf(whole, whole.length / 3); assertThrows(InvalidFormatException.class, @@ -72,7 +71,49 @@ void testTruncatedModelFailsLoudly() throws IOException { } @Test - void testUnsupportedModelTypeFailsLoudly() { + 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, @@ -81,7 +122,7 @@ void testUnsupportedModelTypeFailsLoudly() { } @Test - void testMissingUnknownPieceFailsLoudly() { + void testRejectsMissingUnknownPiece() { final byte[] model = minimalModelWithoutUnk(); final InvalidFormatException e = assertThrows(InvalidFormatException.class, () -> SentencePieceTokenizer.load(new ByteArrayInputStream(model))); @@ -89,7 +130,7 @@ void testMissingUnknownPieceFailsLoudly() { } @Test - void testMalformedPrecompiledCharsMapFailsLoudly() { + 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(); @@ -193,6 +234,38 @@ void testBytePiecesExistOnlyInByteFallbackModels() { 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")) { 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 index e5bb42d148..eacd8a94f4 100644 --- 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 @@ -121,6 +121,30 @@ void testForeignPayloadIsRejected() throws IOException { 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. 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 index b1bd4cc9ee..bcd16eb53c 100644 --- 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 @@ -33,9 +33,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue; /** - * Pins the cookbook path documented in {@code tokenizer.xml}: load a - * {@link SentencePieceTokenizer} from a {@code .model} file, encode text to pieces with - * original offsets, and obtain id arrays. + * Demonstrates loading a {@link SentencePieceTokenizer}, retaining original-text offsets, and + * obtaining token ids. */ class SentencePieceUsageExampleTest { 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 index 6b8cbc39db..faa795ccd6 100644 --- 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 @@ -42,9 +42,9 @@ 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) +- `.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. diff --git a/opennlp-extensions/pom.xml b/opennlp-extensions/pom.xml index 2b0e87273d..4fed4c405e 100644 --- a/opennlp-extensions/pom.xml +++ b/opennlp-extensions/pom.xml @@ -38,6 +38,7 @@ + opennlp-embeddings opennlp-morfologik opennlp-spellcheck opennlp-subword diff --git a/opennlp-tools/src/test/java/opennlp/tools/chunker/ChunkerDetailedFMeasureListenerTest.java b/opennlp-tools/src/test/java/opennlp/tools/chunker/ChunkerDetailedFMeasureListenerTest.java index 00031b41c4..de790cd71a 100644 --- a/opennlp-tools/src/test/java/opennlp/tools/chunker/ChunkerDetailedFMeasureListenerTest.java +++ b/opennlp-tools/src/test/java/opennlp/tools/chunker/ChunkerDetailedFMeasureListenerTest.java @@ -23,6 +23,7 @@ import java.nio.charset.StandardCharsets; import java.util.Locale; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -32,6 +33,13 @@ public class ChunkerDetailedFMeasureListenerTest { + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + @Test void testEvaluator() throws IOException { @@ -68,4 +76,33 @@ void testEvaluator() throws IOException { Assertions.assertEquals(expected.toString().trim(), listener.createReport(Locale.ENGLISH).trim()); } + + /** + * The no-argument report has to render the same figures everywhere, so a JVM defaulting to + * a locale with a comma decimal separator must not change the numbers in the output. + */ + @Test + void testReportIsIndependentOfDefaultLocale() throws IOException { + ChunkerDetailedFMeasureListener listener = new ChunkerDetailedFMeasureListener(); + evaluate(listener); + + Locale.setDefault(Locale.GERMANY); + + Assertions.assertEquals(listener.createReport(Locale.ENGLISH).trim(), + listener.createReport().trim()); + } + + private static void evaluate(ChunkerDetailedFMeasureListener listener) throws IOException { + ResourceAsStreamFactory inPredicted = new ResourceAsStreamFactory( + ChunkerDetailedFMeasureListenerTest.class, "/opennlp/tools/chunker/output.txt"); + ResourceAsStreamFactory inExpected = new ResourceAsStreamFactory( + ChunkerDetailedFMeasureListenerTest.class, "/opennlp/tools/chunker/output.txt"); + + DummyChunkSampleStream predictedSample = new DummyChunkSampleStream( + new PlainTextByLineStream(inPredicted, StandardCharsets.UTF_8), true); + DummyChunkSampleStream expectedSample = new DummyChunkSampleStream( + new PlainTextByLineStream(inExpected, StandardCharsets.UTF_8), false); + + new ChunkerEvaluator(new DummyChunker(predictedSample), listener).evaluate(expectedSample); + } } diff --git a/opennlp-tools/src/test/java/opennlp/tools/util/StringUtilTest.java b/opennlp-tools/src/test/java/opennlp/tools/util/StringUtilTest.java index a47306d3d4..1951c114a0 100644 --- a/opennlp-tools/src/test/java/opennlp/tools/util/StringUtilTest.java +++ b/opennlp-tools/src/test/java/opennlp/tools/util/StringUtilTest.java @@ -679,4 +679,31 @@ void testLowercaseBeyondBMP() { String lc = StringUtil.toLowerCase(input); Assertions.assertArrayEquals(expectedCodePoints, lc.codePoints().toArray()); } + + /** + * Verifies the accepting side of the blank check against the toolkit's whitespace + * definition: empty and JDK-whitespace values are blank, and so are the no-break + * spaces U+00A0 and U+2007, which {@link String#isBlank()} does not cover. + */ + @ParameterizedTest + @ValueSource(strings = {"", " \t\n", "\u00A0", " \u00A0\u2007 "}) + void testIsBlankAcceptsWhitespaceOnlyValues(String input) { + Assertions.assertTrue(StringUtil.isBlank(input)); + } + + /** + * Verifies the rejecting side of the blank check: any non-whitespace code point + * makes a value non-blank, including the supplementary-plane letter U+10428, which + * must be read as one code point rather than two chars. + */ + @ParameterizedTest + @ValueSource(strings = {"a", " a ", "\uD801\uDC28"}) + void testIsBlankRejectsValuesWithContent(String input) { + Assertions.assertFalse(StringUtil.isBlank(input)); + } + + @Test + void testIsBlankWithNullString() { + Assertions.assertThrows(NullPointerException.class, () -> StringUtil.isBlank(null)); + } } diff --git a/opennlp-tools/src/test/java/opennlp/tools/util/TrainingParametersTest.java b/opennlp-tools/src/test/java/opennlp/tools/util/TrainingParametersTest.java index a68547feb9..99505e9c46 100644 --- a/opennlp-tools/src/test/java/opennlp/tools/util/TrainingParametersTest.java +++ b/opennlp-tools/src/test/java/opennlp/tools/util/TrainingParametersTest.java @@ -19,8 +19,10 @@ import java.io.ByteArrayInputStream; import java.util.HashMap; +import java.util.Locale; import java.util.Map; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -28,6 +30,18 @@ public class TrainingParametersTest { + /** + * Turkish folds {@code 'I'} to the dotless {@code 'ı'} (U+0131) instead of {@code 'i'}. + */ + private static final Locale TURKISH = Locale.of("tr", "TR"); + + private final Locale defaultLocale = Locale.getDefault(); + + @AfterEach + void restoreDefaultLocale() { + Locale.setDefault(defaultLocale); + } + @Test void testConstructors() throws Exception { TrainingParameters tp1 = @@ -124,6 +138,28 @@ public void testSetParamsWithoutIterationsCLIParams() { Parameters.CUTOFF_DEFAULT_VALUE)); } + /** + * The switch names are derived by folding the parameter constants, which start with a + * capital {@code 'I'} respectively {@code 'C'}. Folding them with the default locale makes + * a Turkish JVM look for {@code "-\u0131terations"} and silently train with the defaults. + */ + @Test + public void testSetParamsIsIndependentOfDefaultLocale() { + String[] args = + { "-model" , "en-token-test.bin" , "-lang" , "en" , "-data" , + "en-token.train" , "-encoding" , "UTF-8" , "-cutoff" , "10" , "-iterations" , "50" }; + + Locale.setDefault(TURKISH); + TrainingParameters tr = TrainingParameters.setParams(args); + + Assertions.assertEquals(50 , + tr.getIntParameter(Parameters.ITERATIONS_PARAM , + Parameters.ITERATIONS_DEFAULT_VALUE)); + Assertions.assertEquals(10 , + tr.getIntParameter(Parameters.CUTOFF_PARAM , + Parameters.CUTOFF_DEFAULT_VALUE)); + } + @Test void testGetAlgorithm() { TrainingParameters tp = build("Algorithm=Perceptron,n1.Algorithm=SVM"); diff --git a/pom.xml b/pom.xml index 3f77763325..dec360a7c1 100644 --- a/pom.xml +++ b/pom.xml @@ -204,6 +204,12 @@ test-jar + + opennlp-embeddings + ${project.groupId} + ${project.version} + + opennlp-morfologik ${project.groupId} @@ -253,13 +259,13 @@ https://opennlp.sourceforge.net/models-1.5/ - 6.1.2 + 6.1.3 2.0.3 - 1.28.0 - 2.0.18 + 1.29.0 + 2.0.19 2.26.1 - 2.12.6 - 4.8.186 + 2.12.7 + 4.8.194 3.0.0 1.37 @@ -267,7 +273,7 @@ 4.3.0 0.8.15 3.8.0 - 3.5.6 + 3.6.0 3.12.0 3.10 2.7.1 diff --git a/rat-excludes b/rat-excludes index a505ab2d60..90b6a69803 100644 --- a/rat-excludes +++ b/rat-excludes @@ -71,7 +71,13 @@ 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